| """Model wrapper for LiteLLM""" |
|
|
| import os |
| import json |
| 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 |
| |
| |
| if "groq" in model_id.lower(): |
| if not os.getenv("GROQ_API_KEY"): |
| print("⚠️ GROQ_API_KEY not set in environment") |
| raise RuntimeError("GROQ_API_KEY not set. Please add it to your Space secrets.") |
| |
| 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: |
| formatted_tools = [ |
| { |
| "type": "function", |
| "function": { |
| "name": tool.name, |
| "description": tool.description, |
| "parameters": tool.parameters |
| } |
| } |
| for tool in tools |
| ] |
| |
| |
| if "groq" in self.model_id.lower(): |
| api_key = os.getenv("GROQ_API_KEY") |
| if not api_key: |
| raise RuntimeError("GROQ_API_KEY not set in environment") |
| |
| print(f"DEBUG: Using Groq model: {self.model_id}") |
| |
| response = litellm.completion( |
| model=self.model_id, |
| api_key=api_key, |
| messages=messages, |
| tools=formatted_tools, |
| temperature=0.1 |
| ) |
| else: |
| |
| 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 "" |
| } |
| |
| if hasattr(message, 'tool_calls') and message.tool_calls: |
| result["tool_calls"] = [] |
| for tc in message.tool_calls: |
| |
| args = tc.function.arguments |
| if isinstance(args, str): |
| try: |
| args = json.loads(args) |
| except: |
| args = {} |
| |
| result["tool_calls"].append({ |
| "id": tc.id if hasattr(tc, 'id') else f"call_{tc.function.name}", |
| "name": tc.function.name, |
| "arguments": args |
| }) |
| |
| 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}") |
