File size: 2,785 Bytes
5d286f3
 
04b780e
5d286f3
 
04b780e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5d286f3
04b780e
 
 
 
5d286f3
04b780e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
import random
import httpx
from typing import Optional
from .config import OPENROUTER_API_KEY, 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
        }

        payload = {
            "model": model,
            "messages": [{"role": "user", "content": prompt}]
        }

        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