Create model.py

model.py

@@ -0,0 +1,103 @@
+"""
+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 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 = [
+                    {
+                        "type": "function",
+                        "function": {
+                            "name": tool.name,
+                            "description": tool.description,
+                            "parameters": tool.parameters
+                        }
+                    }
+                    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"
+                model_name = (
+                    f"google/{self.model_id}" 
+                    if not self.model_id.startswith("google/") 
+                    else self.model_id
+                )
+                response = litellm.completion(
+                    provider=provider,
+                    model=model_name,
+                    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)
+                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 present
+            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}")