multiapi / app /model_router.py
Juna190825's picture
Update app/model_router.py
04b780e verified
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