Update app/model_router.py

app/model_router.py

@@ -1,24 +1,72 @@
 import random
 import httpx
+from typing import Optional
 from .config import OPENROUTER_API_KEY, FREE_MODELS
 
-async def query_model(prompt: str, model: str = None):
-    if not model:
-        model = random.choice(FREE_MODELS)
+async def query_model(prompt: str, model: Optional[str] = None) -> str:
+    """
+    Query the OpenRouter API with a prompt and optional model selection.
+    
+    Args:
+        prompt: The user's input message
+        model: Optional model name to use (random free model selected if None)
+    
+    Returns:
+        The generated response content
+    
+    Raises:
+        httpx.HTTPStatusError: If the API request fails
+        ValueError: If the response format is invalid
+        Exception: For other unexpected errors
+    """
+    try:
+        # Select a random free model if none specified
+        model = model or random.choice(FREE_MODELS)
+        
+        headers = {
+            "Authorization": f"Bearer {OPENROUTER_API_KEY}",
+            "Content-Type": "application/json",
+            "HTTP-Referer": "https://github.com/your-repo",  # Recommended by OpenRouter
+            "X-Title": "Your App Name"  # Recommended by OpenRouter
+        }
 
-    headers = {
-        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
-        "Content-Type": "application/json"
-    }
+        payload = {
+            "model": model,
+            "messages": [{"role": "user", "content": prompt}]
+        }
 
-    payload = {
-        "model": model,
-        "messages": [
-            {"role": "user", "content": prompt}
-        ]
-    }
-
-    async with httpx.AsyncClient() as client:
-        response = await client.post("https://openrouter.ai/api/v1/chat/completions", json=payload, headers=headers)
-        response.raise_for_status()
-        return response.json()["choices"][0]["message"]["content"]
+        async with httpx.AsyncClient(timeout=30.0) as client:
+            response = await client.post(
+                "https://openrouter.ai/api/v1/chat/completions",
+                json=payload,
+                headers=headers
+            )
+            response.raise_for_status()
+            
+            # Safely extract the response content
+            response_data = response.json()
+            if not response_data.get("choices"):
+                raise ValueError("No choices in response")
+            
+            first_choice = response_data["choices"][0]
+            if "message" not in first_choice or "content" not in first_choice["message"]:
+                raise ValueError("Invalid message format in response")
+            
+            return first_choice["message"]["content"]
+            
+    except httpx.HTTPStatusError as e:
+        # Include more details in the error message
+        error_info = f"HTTP error {e.response.status_code}"
+        if e.response.text:
+            try:
+                error_detail = e.response.json().get("error", {}).get("message", e.response.text)
+                error_info += f": {error_detail}"
+            except ValueError:
+                error_info += f": {e.response.text[:200]}"  # Truncate long text
+        raise httpx.HTTPStatusError(error_info, request=e.request, response=e.response)
+        
+    except (KeyError, IndexError, ValueError) as e:
+        raise ValueError(f"Failed to parse API response: {str(e)}") from e
+        
+    except Exception as e:
+        raise Exception(f"Unexpected error while querying model: {str(e)}") from e