query-sql / clients /openRouter.py
Ashhar
first commit
9f9844d
raw
history blame
5.64 kB
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
# Add any additional parameters
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()
# Wrap the response data
return ResponseWrapper(response.json())
except requests.RequestException as e:
raise Exception(f"OpenRouter API request failed: {e}")