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