144 lines
5.5 KiB
Python
144 lines
5.5 KiB
Python
# api/openrouter_api.py
|
|
import aiohttp
|
|
import logging
|
|
import json
|
|
from typing import Optional, Dict, Any, List
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
async def send_openrouter_image_generation_request(
|
|
api_key: str,
|
|
model: str,
|
|
prompt: str,
|
|
n: int = 1,
|
|
timeout: int = 120,
|
|
input_image_base64: Optional[List[str]] = None,
|
|
**kwargs
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Sends a request to OpenRouter's chat completions endpoint for image generation,
|
|
optionally including input images for context.
|
|
"""
|
|
api_url = "https://openrouter.ai/api/v1/chat/completions"
|
|
headers = {
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
# Build the message content
|
|
content_parts = []
|
|
if prompt:
|
|
content_parts.append({"type": "text", "text": prompt})
|
|
|
|
if input_image_base64:
|
|
for b64_image in input_image_base64:
|
|
content_parts.append({
|
|
"type": "image_url",
|
|
"image_url": {
|
|
"url": f"data:image/jpeg;base64,{b64_image}"
|
|
}
|
|
})
|
|
|
|
messages = [{"role": "user", "content": content_parts}]
|
|
|
|
payload = {
|
|
"model": model,
|
|
"messages": messages,
|
|
"modalities": ["image", "text"],
|
|
}
|
|
|
|
# Add 'n' if it's greater than 1, as some models may support it.
|
|
if n > 1:
|
|
payload['n'] = n
|
|
|
|
# Add only SEED from kwargs as it's a documented advanced parameter.
|
|
# 'size' is not a standard for chat completions and likely cause 400 errors.
|
|
if 'seed' in kwargs and kwargs.get('seed') != -1 and kwargs.get('seed') is not None:
|
|
payload['seed'] = kwargs['seed']
|
|
|
|
# Log a sanitized version of the payload for debugging
|
|
log_payload = payload.copy()
|
|
if log_payload.get("messages") and log_payload["messages"][0].get("content"):
|
|
import copy
|
|
log_payload = copy.deepcopy(payload)
|
|
for part in log_payload["messages"][0]["content"]:
|
|
if part.get("type") == "image_url":
|
|
part["image_url"]["url"] = "[base64 data omitted]"
|
|
logger.info(f"Sending image generation request to OpenRouter with payload: {json.dumps(log_payload, indent=2)}")
|
|
|
|
try:
|
|
async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=timeout)) as session:
|
|
async with session.post(api_url, headers=headers, json=payload) as response:
|
|
response.raise_for_status()
|
|
response_json = await response.json()
|
|
|
|
if response_json.get("choices"):
|
|
image_data = []
|
|
for choice in response_json["choices"]:
|
|
message = choice.get("message", {})
|
|
if "images" in message and message["images"]:
|
|
for img in message["images"]:
|
|
base64_string = img.get("image_url", {}).get("url", "")
|
|
if "base64," in base64_string:
|
|
base64_string = base64_string.split("base64,")[1]
|
|
image_data.append({"b64_json": base64_string})
|
|
|
|
if image_data:
|
|
return {"data": image_data}
|
|
|
|
logger.error(f"OpenRouter API response did not contain expected image data. Response: {response_json}")
|
|
return {"error": "No image data in response", "data": []}
|
|
|
|
except aiohttp.ClientResponseError as e:
|
|
logger.error(f"HTTP error from OpenRouter API: {e.status} {e.message}")
|
|
error_body = "Could not read error body."
|
|
try:
|
|
error_json = await e.json()
|
|
error_body = str(error_json)
|
|
logger.error(f"Error Body (JSON): {error_body}")
|
|
except Exception:
|
|
try:
|
|
error_body = await e.text()
|
|
logger.error(f"Error Body (text): {error_body}")
|
|
except Exception:
|
|
pass
|
|
return {"error": f"HTTP error: {e.status} {e.message}. Body: {error_body}", "data": []}
|
|
except Exception as e:
|
|
logger.error(f"Error during OpenRouter API call: {e}", exc_info=True)
|
|
return {"error": str(e), "data": []}
|
|
|
|
async def send_openrouter_request(
|
|
api_url: str,
|
|
model: str,
|
|
messages: List[Dict[str, Any]],
|
|
api_key: str,
|
|
**kwargs,
|
|
) -> Dict[str, Any]:
|
|
"""Generic OpenRouter chat completions handler."""
|
|
headers = {
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
payload = {"model": model, "messages": messages, **kwargs}
|
|
payload = {k: v for k, v in payload.items() if v is not None}
|
|
|
|
# Log sanitized payload
|
|
log_payload = payload.copy()
|
|
logger.info(f"Sending request to OpenRouter with payload: {json.dumps(log_payload, indent=2)}")
|
|
|
|
try:
|
|
async with aiohttp.ClientSession() as session:
|
|
async with session.post(
|
|
api_url, headers=headers, json=payload
|
|
) as response:
|
|
response.raise_for_status()
|
|
return await response.json()
|
|
except aiohttp.ClientResponseError as e:
|
|
logger.error(f"HTTP error from OpenRouter API: {e.status} {e.message}")
|
|
error_body = await response.text()
|
|
logger.error(f"Error body: {error_body}")
|
|
return {"error": f"HTTP error: {e.status} {e.message}. Body: {error_body}"}
|
|
except Exception as e:
|
|
logger.error(f"Error during OpenRouter API call: {e}", exc_info=True)
|
|
return {"error": str(e)}
|