Update model.py

model.py

@@ -11,29 +11,21 @@ except ImportError:
     print("⚠️ litellm not installed. Install with: pip install litellm")
     litellm = None
 
-
 class LiteLLMModel:
     """Wrapper for LiteLLM models"""
 
     def __init__(self, model_id: str):
         self.model_id = model_id
 
-        # Set API keys from environment for Gemini
         if "gemini" in model_id.lower():
             if not os.getenv("GEMINI_API_KEY"):
                 print("⚠️ GEMINI_API_KEY not set in environment")
 
     def generate(self, messages: List[Dict], tools: Optional[List] = None) -> Dict:
-        """
-        Generate response from model.
-        Returns:
-            Dict with 'content' and optionally 'tool_calls'
-        """
         if not litellm:
             return {"content": "Unknown - litellm not installed"}
 
         try:
-            # Convert tools to OpenAI format
             formatted_tools = None
             if tools:
                 formatted_tools = [
@@ -48,45 +40,30 @@ class LiteLLMModel:
                     for tool in tools
                 ]
 
-            # --- Provider Handling ---
             if "gemini" in self.model_id.lower():
                 api_key = os.getenv("GEMINI_API_KEY")
                 if not api_key:
                     raise RuntimeError("GEMINI_API_KEY not set in environment")
+                provider = "google"
 
-                # Try with original model ID
-                model_name_1 = self.model_id
-                # Try adding "google/" prefix if missing
-                model_name_2 = f"google/{self.model_id}" if not self.model_id.startswith("google/") else self.model_id
+                # Use a known supported Gemini model id from LiteLLM and Google docs
+                # Replace MODEL_ID in your app.py with one of these if needed
+                # Example model names: "vertex_ai/gemini-1.0-pro-001", "vertex_ai/gemini-2.5-flash"
+                # Choose one according to what is currently supported
+                supported_model_id = "vertex_ai/gemini-1.0-pro-001"  
 
-                # Debug print for environment details
                 print("DEBUG: GEMINI_API_KEY last4:", api_key[-4:])
+                print(f"DEBUG: Using supported Gemini model id: {supported_model_id}")
 
-                # Try first model_name_1
-                try:
-                    print(f"DEBUG: Trying Gemini with model_name_1 = {model_name_1}")
-                    response = litellm.completion(
-                        provider="google",
-                        model=model_name_1,
-                        api_key=api_key,
-                        messages=messages,
-                        tools=formatted_tools,
-                        temperature=0.1
-                    )
-                except Exception as e1:
-                    print(f"DEBUG: model_name_1 failed: {e1}")
-                    # Try model_name_2
-                    print(f"DEBUG: Trying Gemini with model_name_2 = {model_name_2}")
-                    response = litellm.completion(
-                        provider="google",
-                        model=model_name_2,
-                        api_key=api_key,
-                        messages=messages,
-                        tools=formatted_tools,
-                        temperature=0.1
-                    )
+                response = litellm.completion(
+                    provider=provider,
+                    model=supported_model_id,  # Use supported model id here
+                    api_key=api_key,
+                    messages=messages,
+                    tools=formatted_tools,
+                    temperature=0.1
+                )
             else:
-                # For other models/providers, pass model without provider and api_key
                 response = litellm.completion(
                     model=self.model_id,
                     messages=messages,
@@ -99,7 +76,6 @@ class LiteLLMModel:
                 "content": message.content or ""
             }
 
-            # Handle tool calls if present
             if hasattr(message, 'tool_calls') and message.tool_calls:
                 result["tool_calls"] = [
                     {
@@ -115,9 +91,7 @@ class LiteLLMModel:
             print(f"Model error: {e}")
             return {"content": "Unknown"}
 
-
 def get_model(model_type: str, model_id: str):
-    """Get model instance"""
     if model_type == "LiteLLMModel":
         return LiteLLMModel(model_id)
     else: