Update model.py

model.py

@@ -11,6 +11,7 @@ except ImportError:
     print("⚠️ litellm not installed. Install with: pip install litellm")
     litellm = None
 
+
 class LiteLLMModel:
     """Wrapper for LiteLLM models"""
 
@@ -52,22 +53,40 @@ class LiteLLMModel:
                 api_key = os.getenv("GEMINI_API_KEY")
                 if not api_key:
                     raise RuntimeError("GEMINI_API_KEY not set in environment")
-                provider = "google"
-                model_name = (
-                    f"google/{self.model_id}" 
-                    if not self.model_id.startswith("google/") 
-                    else self.model_id
-                )
-                response = litellm.completion(
-                    provider=provider,           # <-- REQUIRED
-                    model=model_name,
-                    api_key=api_key,             # <-- REQUIRED
-                    messages=messages,
-                    tools=formatted_tools,
-                    temperature=0.1
-                )
+
+                # 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
+
+                # Debug print for environment details
+                print("DEBUG: GEMINI_API_KEY last4:", api_key[-4:])
+
+                # 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
+                    )
             else:
-                # For other models/providers, only pass model (configure as needed for your environment)
+                # For other models/providers, pass model without provider and api_key
                 response = litellm.completion(
                     model=self.model_id,
                     messages=messages,
@@ -89,12 +108,14 @@ class LiteLLMModel:
                     }
                     for tc in message.tool_calls
                 ]
+
             return result
 
         except Exception as e:
             print(f"Model error: {e}")
             return {"content": "Unknown"}
 
+
 def get_model(model_type: str, model_id: str):
     """Get model instance"""
     if model_type == "LiteLLMModel":