Implement developer review recommendations: enhanced session management, improved error handling, and dependency updates

api/chat.py

@@ -3,8 +3,11 @@ from fastapi import APIRouter, HTTPException
 from fastapi.responses import StreamingResponse
 from core.llm import LLMClient
 from core.memory import save_user_state, load_user_state
+from core.session import session_manager
+import logging
 
 router = APIRouter()
+logger = logging.getLogger(__name__)
 
 llm_client = LLMClient(provider="ollama")  # Default to Ollama
 
@@ -12,28 +15,30 @@ llm_client = LLMClient(provider="ollama")  # Default to Ollama
 async def chat(user_id: str, message: str):
     if not message:
         raise HTTPException(status_code=400, detail="Message is required")
-
-    # Load user state from Redis
-    user_state = load_user_state(user_id)
-    conversation_history = json.loads(user_state.get("conversation", "[]")) if user_state else []
-
-    # Add user message to history
-    conversation_history.append({"role": "user", "content": message})
-
-    # Generate AI response
+    
     try:
+        # Use session manager for better session handling
+        session = session_manager.get_session(user_id)
+        conversation_history = session.get("conversation", [])
+        
+        # Add user message to history
+        conversation_history.append({"role": "user", "content": message})
+        
+        # Generate AI response
         full_response = ""
         response_stream = llm_client.generate(prompt=message, stream=True)
-
+        
         # Collect streamed response
         for chunk in response_stream:
             full_response += chunk
-
-        # Save updated conversation
+        
+        # Save updated conversation using session manager
         conversation_history.append({"role": "assistant", "content": full_response})
-        save_user_state(user_id, {"conversation": json.dumps(conversation_history)})
-
+        session_manager.update_session(user_id, {"conversation": conversation_history})
+        
+        logger.info(f"Successfully processed chat for user {user_id}")
         return {"response": full_response}
-
+        
     except Exception as e:
+        logger.error(f"LLM generation failed for user {user_id}: {e}")
         raise HTTPException(status_code=500, detail=f"LLM generation failed: {e}")

api/main.py

@@ -2,8 +2,10 @@ from fastapi import FastAPI
 from api.status import router as status_router
 from api.chat import router as chat_router
 from core.memory import check_redis_health
+import logging
 
 app = FastAPI()
+logger = logging.getLogger(__name__)
 
 # Mount routers
 app.include_router(status_router, prefix="/api")
@@ -11,13 +13,34 @@ app.include_router(chat_router, prefix="/api")
 
 @app.get("/")
 async def root():
+    logger.info("API root endpoint accessed")
     return {"message": "AI Life Coach API is running"}
 
 @app.get("/health")
 async def health_check():
-    """Health check endpoint"""
+    """Comprehensive health check endpoint"""
     redis_healthy = check_redis_health()
-    return {
-        "status": "healthy" if redis_healthy else "degraded",
-        "redis": "healthy" if redis_healthy else "unhealthy"
+    
+    # Additional health checks could be added here
+    overall_healthy = redis_healthy
+    
+    health_status = {
+        "status": "healthy" if overall_healthy else "degraded",
+        "services": {
+            "redis": "healthy" if redis_healthy else "unhealthy"
+        },
+        "timestamp": __import__('time').time()
     }
+    
+    if overall_healthy:
+        logger.info("Health check passed")
+    else:
+        logger.warning("Health check degraded")
+            
+    return health_status
+
+# Add startup event for initialization logging
+@app.on_event("startup")
+async def startup_event():
+    logger.info("AI Life Coach API starting up...")
+    logger.info("Redis health: %s", "healthy" if check_redis_health() else "unhealthy")

app.py

@@ -1,4 +1,4 @@
-# Force redeploy trigger - version 2.2
+# Force redeploy trigger - version 2.3
 import streamlit as st
 from utils.config import config
 import requests
@@ -6,22 +6,29 @@ import json
 import os
 from core.memory import load_user_state, check_redis_health
 
-# Set page config
+# Set page config FIRST (before any other Streamlit commands)
 st.set_page_config(page_title="AI Life Coach", page_icon="🧘", layout="centered")
 
-# Comprehensive session state initialization
-session_keys_defaults = {
-    'ngrok_url': config.ollama_host,
-    'model_status': "checking",
-    'available_models': [],
-    'selected_model': config.local_model_name,
-    'selected_model_index': 0,
-    'user_message_input': ""
-}
+# ROBUST SESSION STATE INITIALIZATION
+# This must happen before ANY widget creation
+def init_session_state():
+    """Initialize all session state variables with proper defaults"""
+    defaults = {
+        'ngrok_url': config.ollama_host,
+        'model_status': "checking",
+        'available_models': [],
+        'selected_model': config.local_model_name,
+        'selected_model_index': 0,
+        'user_message_input': "",
+        'user_selector': "Rob"  # Add missing default
+    }
+    
+    for key, default_value in defaults.items():
+        if key not in st.session_state:
+            st.session_state[key] = default_value
 
-for key, default_value in session_keys_defaults.items():
-    if key not in st.session_state:
-        st.session_state[key] = default_value
+# CALL THIS FIRST
+init_session_state()
 
 # Sidebar for user selection
 st.sidebar.title("🧘 AI Life Coach")
@@ -32,11 +39,12 @@ st.sidebar.markdown("---")
 st.sidebar.subheader("Ollama Connection")
 ngrok_input = st.sidebar.text_input(
     "Ngrok URL", 
-    value=st.session_state.ngrok_url,
+    value=st.session_state.ngrok_url, 
     key="ngrok_url_input"  # Explicit key
 )
 
-if st.sidebar.button("Update Ngrok URL", key="update_ngrok_button"):  # Explicit key
+if st.sidebar.button("Update Ngrok URL", key="update_ngrok_button"):
+    # Explicit key
     st.session_state.ngrok_url = ngrok_input
     st.session_state.model_status = "checking"
     st.session_state.available_models = []
@@ -54,8 +62,8 @@ NGROK_HEADERS = {
 def fetch_available_models(ngrok_url):
     try:
         response = requests.get(
-            f"{ngrok_url}/api/tags",
-            headers=NGROK_HEADERS,
+            f"{ngrok_url}/api/tags", 
+            headers=NGROK_HEADERS, 
             timeout=5
         )
         if response.status_code == 200:
@@ -65,15 +73,14 @@ def fetch_available_models(ngrok_url):
         pass
     return []
 
-# Update available models
-if st.session_state.ngrok_url and st.session_state.model_status != "unreachable":
-    model_names = fetch_available_models(st.session_state.ngrok_url)
-    if model_names:
-        st.session_state.available_models = model_names
-        # If current selected model not in list, select the first one
-        if st.session_state.selected_model not in model_names:
-            st.session_state.selected_model = model_names[0]
-            st.session_state.selected_model_index = 0
+# Update available models if st.session_state.ngrok_url and st.session_state.model_status != "unreachable":
+model_names = fetch_available_models(st.session_state.ngrok_url)
+if model_names:
+    st.session_state.available_models = model_names
+    # If current selected model not in list, select the first one
+    if st.session_state.selected_model not in model_names:
+        st.session_state.selected_model = model_names[0]
+        st.session_state.selected_model_index = 0
 
 # Model selector dropdown - REPLACED ENTIRE SECTION
 st.sidebar.markdown("---")
@@ -108,7 +115,7 @@ if st.session_state.available_models:
 else:
     st.sidebar.warning("No models available - check Ollama connection")
     model_input = st.sidebar.text_input(
-        "Or enter model name", 
+        "Or enter model name",
         value=st.session_state.selected_model,
         key="manual_model_input"  # Explicit key
     )
@@ -132,7 +139,6 @@ def get_ollama_status(ngrok_url):
             models = response.json().get("models", [])
             model_names = [m.get("name") for m in models]
             st.session_state.available_models = model_names
-            
             if models:
                 selected_model_available = st.session_state.selected_model in model_names
                 return {
@@ -236,11 +242,9 @@ use_fallback = not ollama_status.get("running", False) or config.use_fallback
 # Display Ollama status - Enhanced section with Hugging Face scaling behavior info
 if use_fallback:
     st.sidebar.warning("🌐 Using Hugging Face fallback (Ollama not available)")
-    
     # Add special note for Hugging Face scaling behavior
     if config.hf_api_url and "endpoints.huggingface.cloud" in config.hf_api_url:
         st.sidebar.info("ℹ️ HF Endpoint may be initializing (up to 4 min)")
-    
     if "error" in ollama_status:
         st.sidebar.caption(f"Error: {ollama_status['error'][:50]}...")
 else:
@@ -272,18 +276,17 @@ else:
 st.title("🧘 AI Life Coach")
 st.markdown("Talk to your personal development assistant.")
 
-# Show detailed status
-with st.expander("🔍 Connection Status"):
-    st.write("Ollama Status:", ollama_status)
-    st.write("Model Status:", st.session_state.model_status)
-    st.write("Selected Model:", st.session_state.selected_model)
-    st.write("Available Models:", st.session_state.available_models)
-    st.write("Environment Info:")
-    st.write("- Is HF Space:", IS_HF_SPACE)
-    st.write("- Base URL:", BASE_URL or "Not in HF Space")
-    st.write("- Current Ngrok URL:", st.session_state.ngrok_url)
-    st.write("- Using Fallback:", use_fallback)
-    st.write("- Redis Health:", check_redis_health())
+# Show detailed status with st.expander("🔍 Connection Status"):
+st.write("Ollama Status:", ollama_status)
+st.write("Model Status:", st.session_state.model_status)
+st.write("Selected Model:", st.session_state.selected_model)
+st.write("Available Models:", st.session_state.available_models)
+st.write("Environment Info:")
+st.write("- Is HF Space:", IS_HF_SPACE)
+st.write("- Base URL:", BASE_URL or "Not in HF Space")
+st.write("- Current Ngrok URL:", st.session_state.ngrok_url)
+st.write("- Using Fallback:", use_fallback)
+st.write("- Redis Health:", check_redis_health())
 
 # Function to send message to Ollama
 def send_to_ollama(user_input, conversation_history, ngrok_url, model_name):
@@ -299,9 +302,9 @@ def send_to_ollama(user_input, conversation_history, ngrok_url, model_name):
             }
         }
         response = requests.post(
-            f"{ngrok_url}/api/chat",
-            json=payload,
-            headers=NGROK_HEADERS,
+            f"{ngrok_url}/api/chat", 
+            json=payload, 
+            headers=NGROK_HEADERS, 
             timeout=60
         )
         if response.status_code == 200:
@@ -339,7 +342,7 @@ conversation = get_conversation_history(user)
 for msg in conversation:
     role = msg["role"].capitalize()
     content = msg["content"]
-    st.markdown(f"**{role}:** {content}")
+    st.markdown(f"{role}: {content}")
 
 # Chat input - REPLACED SECTION
 user_input = st.text_input(
@@ -357,7 +360,7 @@ if send_button:
         st.warning("Please enter a message.")
     else:
         # Display user message
-        st.markdown(f"**You:** {user_input}")
+        st.markdown(f"You: {user_input}")
         
         # Prepare conversation history
         conversation_history = [{"role": msg["role"], "content": msg["content"]} for msg in conversation[-5:]]
@@ -370,17 +373,17 @@ if send_button:
                 backend_used = "Hugging Face"
             else:
                 ai_response = send_to_ollama(
-                    user_input,
-                    conversation_history,
-                    st.session_state.ngrok_url,
+                    user_input, 
+                    conversation_history, 
+                    st.session_state.ngrok_url, 
                     st.session_state.selected_model
                 )
                 backend_used = "Ollama"
-        
-        if ai_response:
-            st.markdown(f"**AI Coach ({backend_used}):** {ai_response}")
-        else:
-            st.error(f"Failed to get response from {backend_used}.")
+            
+            if ai_response:
+                st.markdown(f"AI Coach ({backend_used}): {ai_response}")
+            else:
+                st.error(f"Failed to get response from {backend_used}.")
         
         # Clear input after sending (this helps prevent duplicate sends)
         st.session_state.user_message_input = ""

core/session.py

@@ -2,6 +2,11 @@ import json
 import time
 from typing import Dict, Any, Optional
 from core.memory import load_user_state, save_user_state
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 class SessionManager:
     """Manages user sessions and conversation context"""
@@ -14,7 +19,7 @@ class SessionManager:
             session_timeout: Session timeout in seconds (default: 1 hour)
         """
         self.session_timeout = session_timeout
-    
+
     def get_session(self, user_id: str) -> Dict[str, Any]:
         """
         Retrieve user session data
@@ -28,18 +33,20 @@ class SessionManager:
         try:
             state = load_user_state(user_id)
             if not state:
+                logger.info(f"Creating new session for user {user_id}")
                 return self._create_new_session()
             
             # Check if session has expired
-            last_activity = state.get('last_activity', 0)
+            last_activity = float(state.get('last_activity', 0))
             if time.time() - last_activity > self.session_timeout:
+                logger.info(f"Session expired for user {user_id}, creating new session")
                 return self._create_new_session()
-                
+            
             return state
         except Exception as e:
-            print(f"Error retrieving session for user {user_id}: {e}")
+            logger.error(f"Error retrieving session for user {user_id}: {e}")
             return self._create_new_session()
-    
+
     def update_session(self, user_id: str, data: Dict[str, Any]) -> bool:
         """
         Update user session data
@@ -60,11 +67,16 @@ class SessionManager:
             session['last_activity'] = time.time()
             
             # Save updated session
-            return save_user_state(user_id, session)
+            result = save_user_state(user_id, session)
+            if result:
+                logger.debug(f"Successfully updated session for user {user_id}")
+            else:
+                logger.warning(f"Failed to save session for user {user_id}")
+            return result
         except Exception as e:
-            print(f"Error updating session for user {user_id}: {e}")
+            logger.error(f"Error updating session for user {user_id}: {e}")
             return False
-    
+
     def clear_session(self, user_id: str) -> bool:
         """
         Clear user session data
@@ -76,11 +88,14 @@ class SessionManager:
             Boolean indicating success
         """
         try:
-            return save_user_state(user_id, {})
+            result = save_user_state(user_id, {})
+            if result:
+                logger.info(f"Cleared session for user {user_id}")
+            return result
         except Exception as e:
-            print(f"Error clearing session for user {user_id}: {e}")
+            logger.error(f"Error clearing session for user {user_id}: {e}")
             return False
-    
+
     def _create_new_session(self) -> Dict[str, Any]:
         """
         Create a new session with default values
@@ -88,12 +103,14 @@ class SessionManager:
         Returns:
             Dictionary containing new session data
         """
-        return {
+        session = {
             'conversation': [],
             'preferences': {},
             'last_activity': time.time(),
             'created_at': time.time()
         }
+        logger.debug("Created new session")
+        return session
 
 # Global session manager instance
 session_manager = SessionManager()

requirements.txt

@@ -8,3 +8,6 @@ tavily-python>=0.1.0,<1.0.0
 requests==2.31.0
 docker==6.1.3
 pygame==2.5.2
+# Add these for enhanced functionality
+pydantic==1.10.7
+typing-extensions>=4.5.0