Implement model override functionality for AI clients and enhance session management in the simplified medical app

src/config/ai_providers_config.py

@@ -97,12 +97,13 @@ AGENT_CONFIGURATIONS = {
         "reasoning": "Gentle triage requires empathy and nuanced understanding"
     },
     
-    # Spiritual Distress Analyzer uses Anthropic for empathy and safety
+    # Spiritual Distress Analyzer uses Google Gemini by default for speed/throughput
+    # (This matches the baseline UI defaults in Model Settings).
     "SpiritualDistressAnalyzer": {
-        "provider": AIProvider.ANTHROPIC,
-        "model": AIModel.CLAUDE_SONNET_4_5,
-        "temperature": 0.2,
-        "reasoning": "Spiritual distress assessment requires empathy, safety, and nuanced understanding"
+        "provider": AIProvider.GEMINI,
+        "model": AIModel.GEMINI_2_5_FLASH,
+        "temperature": 0.1,
+        "reasoning": "Fast classification task, optimized for speed"
     },
     
     # Referral Message Generator uses Anthropic for compassionate communication

src/core/ai_client.py

@@ -230,19 +230,51 @@ class UniversalAIClient:
     based on agent configuration and availability
     """
     
-    def __init__(self, agent_name: str):
+    def __init__(self, agent_name: str, model_override: Optional[str] = None):
         self.agent_name = agent_name
+        self.model_override = model_override
         self.config = get_agent_config(agent_name)
         self.client = None
         self.fallback_client = None
         
         self._initialize_clients()
+
+    @staticmethod
+    def _resolve_override_model(model_override: str) -> tuple[Optional[AIProvider], Optional[AIModel]]:
+        """Resolve a UI-provided model string into provider+AIModel.
+
+        Expected strings (from UI dropdowns):
+        - gemini-2.5-flash / gemini-2.0-flash / gemini-flash-latest
+        - claude-sonnet-4-5-20250929 / ...
+        """
+        if not model_override:
+            return None, None
+        override = model_override.strip()
+        if not override:
+            return None, None
+
+        try:
+            if override.startswith("gemini"):
+                return AIProvider.GEMINI, AIModel(override)
+            if override.startswith("claude"):
+                return AIProvider.ANTHROPIC, AIModel(override)
+        except Exception:
+            return None, None
+
+        return None, None
     
     def _initialize_clients(self):
         """Initialize primary and fallback clients"""
         primary_provider = self.config["provider"]
         primary_model = self.config["model"]
         temperature = self.config.get("temperature", 0.3)
+
+        # Optional: override model/provider (session-level setting from UI)
+        if self.model_override:
+            override_provider, override_model = self._resolve_override_model(self.model_override)
+            if override_provider is not None and override_model is not None:
+                primary_provider = override_provider
+                primary_model = override_model
         
         # Try to initialize primary client
         try:
@@ -343,14 +375,25 @@ class AIClientManager:
     def __init__(self):
         self._clients = {}  # Cache for AI clients
         self.call_counter = 0  # Backward compatibility
+
+        # Optional: allow an owning session/app to attach per-session overrides.
+        # Expected shape: {agent_name: model_string}
+        self.model_overrides: Dict[str, str] = {}
         
         # NEW: Enhanced client management for medical AI optimization
         self.provider_performance_metrics = {}
         self.medical_context_routing = {}
     
     # Enhanced client retrieval with performance tracking
-    def get_client(self, agent_name: str):
-        """Get or create an AI client for the specified agent"""
+    def get_client(self, agent_name: str, model_override: Optional[str] = None):
+        """Get or create an AI client for the specified agent.
+
+        If `model_override` is provided, a new (non-cached) client is returned
+        to avoid cross-session leakage.
+        """
+        if model_override:
+            return create_ai_client(agent_name, model_override=model_override)
+
         if agent_name not in self._clients:
             self._clients[agent_name] = create_ai_client(agent_name)
         return self._clients[agent_name]
@@ -358,7 +401,8 @@ class AIClientManager:
     def generate_response(self, system_prompt: str, user_prompt: str,
                         temperature: float = None, call_type: str = "",
                         agent_name: str = "DefaultAgent",
-                        medical_context: Optional[Dict] = None):
+                        medical_context: Optional[Dict] = None,
+                        model_override: Optional[str] = None):
         """
         Enhanced response generation with medical context awareness
         
@@ -369,7 +413,7 @@ class AIClientManager:
         - Maintain full backward compatibility
         """
         try:
-            client = self.get_client(agent_name)
+            client = self.get_client(agent_name, model_override=model_override)
             response = client.generate_response(
                 system_prompt=system_prompt,
                 user_prompt=user_prompt,
@@ -421,7 +465,8 @@ class AIClientManager:
         return {name: self.get_client_info(name) for name in self._clients}
     
     def call_spiritual_api(self, system_prompt: str, user_prompt: str,
-                          temperature: float = 0.7) -> str:
+                          temperature: float = 0.7,
+                          model_override: Optional[str] = None) -> str:
         """
         Call AI API for spiritual/emotional analysis.
         
@@ -435,16 +480,21 @@ class AIClientManager:
         Returns:
             AI response as string
         """
+        if model_override is None and self.model_overrides:
+            model_override = self.model_overrides.get("SpiritualDistressAnalyzer")
+
         return self.generate_response(
             system_prompt=system_prompt,
             user_prompt=user_prompt,
             temperature=temperature,
             call_type="spiritual_analysis",
-            agent_name="SpiritualDistressAnalyzer"
+            agent_name="SpiritualDistressAnalyzer",
+            model_override=model_override,
         )
     
     def call_medical_api(self, system_prompt: str, user_prompt: str,
-                        temperature: float = 0.3) -> str:
+                        temperature: float = 0.3,
+                        model_override: Optional[str] = None) -> str:
         """
         Call AI API for medical assistance.
         
@@ -458,16 +508,20 @@ class AIClientManager:
         Returns:
             AI response as string
         """
+        if model_override is None and self.model_overrides:
+            model_override = self.model_overrides.get("SoftMedicalTriage")
+
         return self.generate_response(
             system_prompt=system_prompt,
             user_prompt=user_prompt,
             temperature=temperature,
             call_type="medical_assistance",
-            agent_name="SoftMedicalTriage"
+            agent_name="SoftMedicalTriage",
+            model_override=model_override,
         )
 
 # Factory function for easy client creation
-def create_ai_client(agent_name: str) -> UniversalAIClient:
+def create_ai_client(agent_name: str, model_override: Optional[str] = None) -> UniversalAIClient:
     """
     Create an AI client for a specific agent
     
@@ -477,7 +531,7 @@ def create_ai_client(agent_name: str) -> UniversalAIClient:
     Returns:
         Configured UniversalAIClient instance
     """
-    return UniversalAIClient(agent_name)
+    return UniversalAIClient(agent_name, model_override=model_override)
 
 if __name__ == "__main__":
     print(" AI Client Test")

src/core/simplified_medical_app.py

@@ -64,6 +64,8 @@ class SimplifiedMedicalApp:
         
         # AI client
         self.api = AIClientManager()
+        # Optional per-session model overrides (set by the UI Model Settings tab)
+        self.model_overrides = {}
         
         # Medical components
         self.medical_assistant = MedicalAssistant(self.api)
@@ -89,6 +91,29 @@ class SimplifiedMedicalApp:
         )
         
         logger.info("✅ SimplifiedMedicalApp initialized")
+
+    def set_model_overrides(self, overrides: Optional[dict] = None) -> None:
+        """Set per-session model overrides.
+
+        `overrides` is expected to be a mapping of agent_name -> model string.
+        Example keys used by the UI:
+        - SpiritualDistressAnalyzer
+        - SoftSpiritualTriage
+        - TriageResponseEvaluator
+        - MedicalAssistant
+        - SoftMedicalTriage
+        """
+        self.model_overrides = dict(overrides or {})
+
+        # Propagate to AI manager so core components can read overrides.
+        if hasattr(self, "api") and hasattr(self.api, "model_overrides"):
+            self.api.model_overrides = dict(self.model_overrides)
+
+    def _get_model_override(self, agent_name: str) -> Optional[str]:
+        """Return the model override for an agent, if any."""
+        if not getattr(self, "model_overrides", None):
+            return None
+        return self.model_overrides.get(agent_name)
     
     def process_message(
         self,

src/interface/simplified_gradio_app.py

@@ -597,6 +597,13 @@ Changes apply only to your current session.
                 session = SimplifiedSessionData()
             
             session.update_activity()
+
+            # Apply per-session model overrides (if configured in Model Settings)
+            custom_models = getattr(session, 'custom_models', None)
+            if custom_models:
+                session.app_instance.set_model_overrides(custom_models)
+            else:
+                session.app_instance.set_model_overrides({})
             new_history, status = session.app_instance.process_message(message, history)
             
             # Get updated conversation stats