Update model.py

model.py

@@ -1,6 +1,4 @@
-"""
-Model wrapper for LiteLLM
-"""
+"""Model wrapper for LiteLLM"""
 
 import os
 from typing import List, Dict, Any, Optional
@@ -11,20 +9,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
-
+        
         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:
         if not litellm:
             return {"content": "Unknown - litellm not installed"}
-
+        
         try:
             formatted_tools = None
             if tools:
@@ -39,25 +38,16 @@ class LiteLLMModel:
                     }
                     for tool in tools
                 ]
-
+            
             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"
-
-                # 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"  
-
-                print("DEBUG: GEMINI_API_KEY last4:", api_key[-4:])
-                print(f"DEBUG: Using supported Gemini model id: {supported_model_id}")
-
+                
+                print(f"DEBUG: Using model id: {self.model_id}")
+                
                 response = litellm.completion(
-                    provider=provider,
-                    model=supported_model_id,  # Use supported model id here
+                    model=self.model_id,
                     api_key=api_key,
                     messages=messages,
                     tools=formatted_tools,
@@ -70,12 +60,12 @@ class LiteLLMModel:
                     tools=formatted_tools,
                     temperature=0.1
                 )
-
+            
             message = response.choices[0].message
             result = {
                 "content": message.content or ""
             }
-
+            
             if hasattr(message, 'tool_calls') and message.tool_calls:
                 result["tool_calls"] = [
                     {
@@ -84,15 +74,16 @@ 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):
     if model_type == "LiteLLMModel":
         return LiteLLMModel(model_id)
     else:
-        raise ValueError(f"Unknown model type: {model_type}")
+        raise ValueError(f"Unknown model type: {model_type}")