Create model.py

model.py

@@ -0,0 +1,89 @@
+
+"""
+Model wrapper for LiteLLM
+"""
+
+import os
+from typing import List, Dict, Any, Optional
+
+try:
+    import litellm
+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
+        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 = [
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": tool.name,
+                            "description": tool.description,
+                            "parameters": tool.parameters
+                        }
+                    }
+                    for tool in tools
+                ]
+            
+            # Call LiteLLM
+            response = litellm.completion(
+                model=self.model_id,
+                messages=messages,
+                tools=formatted_tools,
+                temperature=0.1
+            )
+            
+            message = response.choices[0].message
+            
+            result = {
+                "content": message.content or ""
+            }
+            
+            # Handle tool calls
+            if hasattr(message, 'tool_calls') and message.tool_calls:
+                result["tool_calls"] = [
+                    {
+                        "name": tc.function.name,
+                        "arguments": eval(tc.function.arguments) if isinstance(tc.function.arguments, str) else tc.function.arguments
+                    }
+                    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":
+        return LiteLLMModel(model_id)
+    else:
+        raise ValueError(f"Unknown model type: {model_type}")