|
import requests |
|
import json |
|
from typing import List, Dict, Optional |
|
|
|
|
|
class ResponseWrapper: |
|
def __init__(self, response_data): |
|
""" |
|
Wrap the response data to support both dict-like and attribute-like access |
|
|
|
:param response_data: The raw response dictionary from OpenRouter |
|
""" |
|
self._data = response_data |
|
|
|
def __getattr__(self, name): |
|
""" |
|
Allow attribute-style access to the response data |
|
|
|
:param name: Attribute name to access |
|
:return: Corresponding value from the response data |
|
""" |
|
if name in self._data: |
|
value = self._data[name] |
|
return self._wrap(value) |
|
raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'") |
|
|
|
def __getitem__(self, key): |
|
""" |
|
Allow dictionary-style access to the response data |
|
|
|
:param key: Key to access |
|
:return: Corresponding value from the response data |
|
""" |
|
value = self._data[key] |
|
return self._wrap(value) |
|
|
|
def _wrap(self, value): |
|
""" |
|
Recursively wrap dictionaries and lists to support attribute access |
|
|
|
:param value: Value to wrap |
|
:return: Wrapped value |
|
""" |
|
if isinstance(value, dict): |
|
return ResponseWrapper(value) |
|
elif isinstance(value, list): |
|
return [self._wrap(item) for item in value] |
|
return value |
|
|
|
def __iter__(self): |
|
""" |
|
Allow iteration over the wrapped dictionary |
|
""" |
|
return iter(self._data) |
|
|
|
def get(self, key, default=None): |
|
""" |
|
Provide a get method similar to dictionary |
|
""" |
|
return self._wrap(self._data.get(key, default)) |
|
|
|
def keys(self): |
|
""" |
|
Return dictionary keys |
|
""" |
|
return self._data.keys() |
|
|
|
def items(self): |
|
""" |
|
Return dictionary items |
|
""" |
|
return [(k, self._wrap(v)) for k, v in self._data.items()] |
|
|
|
def __str__(self): |
|
""" |
|
Return a JSON string representation of the response data |
|
|
|
:return: JSON-formatted string of the response |
|
""" |
|
return json.dumps(self._data, indent=2) |
|
|
|
def __repr__(self): |
|
""" |
|
Return a string representation for debugging |
|
|
|
:return: Representation of the ResponseWrapper |
|
""" |
|
return f"ResponseWrapper({json.dumps(self._data, indent=2)})" |
|
|
|
|
|
class OpenRouter: |
|
def __init__(self, api_key: str, base_url: str = "https://openrouter.ai/api/v1"): |
|
""" |
|
Initialize OpenRouter client |
|
|
|
:param api_key: API key for OpenRouter |
|
:param base_url: Base URL for OpenRouter API (default is standard endpoint) |
|
""" |
|
self.api_key = api_key |
|
self.base_url = base_url |
|
self.chat = self.ChatNamespace(self) |
|
|
|
class ChatNamespace: |
|
def __init__(self, client): |
|
self._client = client |
|
self.completions = self.CompletionsNamespace(client) |
|
|
|
class CompletionsNamespace: |
|
def __init__(self, client): |
|
self._client = client |
|
|
|
def create( |
|
self, |
|
model: str, |
|
messages: List[Dict[str, str]], |
|
temperature: float = 0.7, |
|
max_tokens: Optional[int] = None, |
|
**kwargs |
|
): |
|
""" |
|
Create a chat completion request |
|
|
|
:param model: Model to use |
|
:param messages: List of message dictionaries |
|
:param temperature: Sampling temperature |
|
:param max_tokens: Maximum number of tokens to generate |
|
:return: Wrapped response object |
|
""" |
|
headers = { |
|
"Authorization": f"Bearer {self._client.api_key}", |
|
"Content-Type": "application/json", |
|
"HTTP-Referer": kwargs.get("http_referer", "https://your-app-domain.com"), |
|
"X-Title": kwargs.get("x_title", "AI Ad Generator") |
|
} |
|
|
|
payload = { |
|
"model": model, |
|
"messages": messages, |
|
"temperature": temperature, |
|
} |
|
|
|
if model.startswith("deepseek"): |
|
payload["provider"] = { |
|
"order": [ |
|
"DeepSeek", |
|
"DeepInfra", |
|
"Fireworks", |
|
], |
|
"allow_fallbacks": False |
|
} |
|
|
|
if max_tokens is not None: |
|
payload["max_tokens"] = max_tokens |
|
|
|
|
|
payload.update({k: v for k, v in kwargs.items() |
|
if k not in ["http_referer", "x_title"]}) |
|
|
|
try: |
|
response = requests.post( |
|
f"{self._client.base_url}/chat/completions", |
|
headers=headers, |
|
data=json.dumps(payload) |
|
) |
|
|
|
response.raise_for_status() |
|
|
|
|
|
return ResponseWrapper(response.json()) |
|
|
|
except requests.RequestException as e: |
|
raise Exception(f"OpenRouter API request failed: {e}") |