Spaces:
Sleeping
Sleeping
| 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 |