2.1 (#2)
Final release of version 2.1. Headlights: ### Core Features - 🤖 Interactive chat with 300+ AI models via OpenRouter - 🔍 Model selection with search and filtering - 💾 Conversation save/load/export (Markdown, JSON, HTML) - 📎 File attachments (images, PDFs, code files) - 💰 Real-time cost tracking and credit monitoring - 🎨 Rich terminal UI with syntax highlighting - 📝 Persistent command history with search (Ctrl+R) - 🌐 Online mode (web search capabilities) - 🧠 Conversation memory toggle ### MCP Integration - 🔧 **File Mode**: AI can read, search, and list local files - Automatic .gitignore filtering - Virtual environment exclusion - Large file handling (auto-truncates >50KB) - ✍️ **Write Mode**: AI can modify files with permission - Create, edit, delete files - Move, copy, organize files - Always requires explicit opt-in - 🗄️ **Database Mode**: AI can query SQLite databases - Read-only access (safe) - Schema inspection - Full SQL query support Reviewed-on: #2 Co-authored-by: Rune Olsen <rune@rune.pm> Co-committed-by: Rune Olsen <rune@rune.pm>
This commit was merged in pull request #2.
This commit is contained in:
@@ -0,0 +1,32 @@
|
||||
"""
|
||||
Provider abstraction for oAI.
|
||||
|
||||
This module provides a unified interface for AI model providers,
|
||||
enabling easy extension to support additional providers beyond OpenRouter.
|
||||
"""
|
||||
|
||||
from oai.providers.base import (
|
||||
AIProvider,
|
||||
ChatMessage,
|
||||
ChatResponse,
|
||||
ToolCall,
|
||||
ToolFunction,
|
||||
UsageStats,
|
||||
ModelInfo,
|
||||
ProviderCapabilities,
|
||||
)
|
||||
from oai.providers.openrouter import OpenRouterProvider
|
||||
|
||||
__all__ = [
|
||||
# Base classes and types
|
||||
"AIProvider",
|
||||
"ChatMessage",
|
||||
"ChatResponse",
|
||||
"ToolCall",
|
||||
"ToolFunction",
|
||||
"UsageStats",
|
||||
"ModelInfo",
|
||||
"ProviderCapabilities",
|
||||
# Provider implementations
|
||||
"OpenRouterProvider",
|
||||
]
|
||||
@@ -0,0 +1,413 @@
|
||||
"""
|
||||
Abstract base provider for AI model integration.
|
||||
|
||||
This module defines the interface that all AI providers must implement,
|
||||
along with common data structures for requests and responses.
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import (
|
||||
Any,
|
||||
AsyncIterator,
|
||||
Dict,
|
||||
Iterator,
|
||||
List,
|
||||
Optional,
|
||||
Union,
|
||||
)
|
||||
|
||||
|
||||
class MessageRole(str, Enum):
|
||||
"""Message roles in a conversation."""
|
||||
|
||||
SYSTEM = "system"
|
||||
USER = "user"
|
||||
ASSISTANT = "assistant"
|
||||
TOOL = "tool"
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolFunction:
|
||||
"""
|
||||
Represents a function within a tool call.
|
||||
|
||||
Attributes:
|
||||
name: The function name
|
||||
arguments: JSON string of function arguments
|
||||
"""
|
||||
|
||||
name: str
|
||||
arguments: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolCall:
|
||||
"""
|
||||
Represents a tool/function call requested by the model.
|
||||
|
||||
Attributes:
|
||||
id: Unique identifier for this tool call
|
||||
type: Type of tool call (usually "function")
|
||||
function: The function being called
|
||||
"""
|
||||
|
||||
id: str
|
||||
type: str
|
||||
function: ToolFunction
|
||||
|
||||
|
||||
@dataclass
|
||||
class UsageStats:
|
||||
"""
|
||||
Token usage statistics from an API response.
|
||||
|
||||
Attributes:
|
||||
prompt_tokens: Number of tokens in the prompt
|
||||
completion_tokens: Number of tokens in the completion
|
||||
total_tokens: Total tokens used
|
||||
total_cost_usd: Cost in USD (if available from API)
|
||||
"""
|
||||
|
||||
prompt_tokens: int = 0
|
||||
completion_tokens: int = 0
|
||||
total_tokens: int = 0
|
||||
total_cost_usd: Optional[float] = None
|
||||
|
||||
@property
|
||||
def input_tokens(self) -> int:
|
||||
"""Alias for prompt_tokens."""
|
||||
return self.prompt_tokens
|
||||
|
||||
@property
|
||||
def output_tokens(self) -> int:
|
||||
"""Alias for completion_tokens."""
|
||||
return self.completion_tokens
|
||||
|
||||
|
||||
@dataclass
|
||||
class ChatMessage:
|
||||
"""
|
||||
A single message in a chat conversation.
|
||||
|
||||
Attributes:
|
||||
role: The role of the message sender
|
||||
content: Message content (text or structured content blocks)
|
||||
name: Optional name for the sender
|
||||
tool_calls: List of tool calls (for assistant messages)
|
||||
tool_call_id: Tool call ID this message responds to (for tool messages)
|
||||
"""
|
||||
|
||||
role: str
|
||||
content: Union[str, List[Dict[str, Any]], None] = None
|
||||
name: Optional[str] = None
|
||||
tool_calls: Optional[List[ToolCall]] = None
|
||||
tool_call_id: Optional[str] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert to dictionary format for API requests."""
|
||||
result: Dict[str, Any] = {"role": self.role}
|
||||
|
||||
if self.content is not None:
|
||||
result["content"] = self.content
|
||||
|
||||
if self.name:
|
||||
result["name"] = self.name
|
||||
|
||||
if self.tool_calls:
|
||||
result["tool_calls"] = [
|
||||
{
|
||||
"id": tc.id,
|
||||
"type": tc.type,
|
||||
"function": {
|
||||
"name": tc.function.name,
|
||||
"arguments": tc.function.arguments,
|
||||
},
|
||||
}
|
||||
for tc in self.tool_calls
|
||||
]
|
||||
|
||||
if self.tool_call_id:
|
||||
result["tool_call_id"] = self.tool_call_id
|
||||
|
||||
return result
|
||||
|
||||
|
||||
@dataclass
|
||||
class ChatResponseChoice:
|
||||
"""
|
||||
A single choice in a chat response.
|
||||
|
||||
Attributes:
|
||||
index: Index of this choice
|
||||
message: The response message
|
||||
finish_reason: Why the response ended
|
||||
"""
|
||||
|
||||
index: int
|
||||
message: ChatMessage
|
||||
finish_reason: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ChatResponse:
|
||||
"""
|
||||
Response from a chat completion request.
|
||||
|
||||
Attributes:
|
||||
id: Unique identifier for this response
|
||||
choices: List of response choices
|
||||
usage: Token usage statistics
|
||||
model: Model that generated this response
|
||||
created: Unix timestamp of creation
|
||||
"""
|
||||
|
||||
id: str
|
||||
choices: List[ChatResponseChoice]
|
||||
usage: Optional[UsageStats] = None
|
||||
model: Optional[str] = None
|
||||
created: Optional[int] = None
|
||||
|
||||
@property
|
||||
def message(self) -> Optional[ChatMessage]:
|
||||
"""Get the first choice's message."""
|
||||
if self.choices:
|
||||
return self.choices[0].message
|
||||
return None
|
||||
|
||||
@property
|
||||
def content(self) -> Optional[str]:
|
||||
"""Get the text content of the first choice."""
|
||||
msg = self.message
|
||||
if msg and isinstance(msg.content, str):
|
||||
return msg.content
|
||||
return None
|
||||
|
||||
@property
|
||||
def tool_calls(self) -> Optional[List[ToolCall]]:
|
||||
"""Get tool calls from the first choice."""
|
||||
msg = self.message
|
||||
if msg:
|
||||
return msg.tool_calls
|
||||
return None
|
||||
|
||||
|
||||
@dataclass
|
||||
class StreamChunk:
|
||||
"""
|
||||
A single chunk from a streaming response.
|
||||
|
||||
Attributes:
|
||||
id: Response ID
|
||||
delta_content: New content in this chunk
|
||||
finish_reason: Finish reason (if this is the last chunk)
|
||||
usage: Usage stats (usually in the last chunk)
|
||||
error: Error message if something went wrong
|
||||
"""
|
||||
|
||||
id: str
|
||||
delta_content: Optional[str] = None
|
||||
finish_reason: Optional[str] = None
|
||||
usage: Optional[UsageStats] = None
|
||||
error: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelInfo:
|
||||
"""
|
||||
Information about an AI model.
|
||||
|
||||
Attributes:
|
||||
id: Unique model identifier
|
||||
name: Display name
|
||||
description: Model description
|
||||
context_length: Maximum context window size
|
||||
pricing: Pricing info (input/output per million tokens)
|
||||
supported_parameters: List of supported API parameters
|
||||
input_modalities: Supported input types (text, image, etc.)
|
||||
output_modalities: Supported output types
|
||||
"""
|
||||
|
||||
id: str
|
||||
name: str
|
||||
description: str = ""
|
||||
context_length: int = 0
|
||||
pricing: Dict[str, float] = field(default_factory=dict)
|
||||
supported_parameters: List[str] = field(default_factory=list)
|
||||
input_modalities: List[str] = field(default_factory=lambda: ["text"])
|
||||
output_modalities: List[str] = field(default_factory=lambda: ["text"])
|
||||
|
||||
def supports_images(self) -> bool:
|
||||
"""Check if model supports image input."""
|
||||
return "image" in self.input_modalities
|
||||
|
||||
def supports_tools(self) -> bool:
|
||||
"""Check if model supports function calling/tools."""
|
||||
return "tools" in self.supported_parameters or "functions" in self.supported_parameters
|
||||
|
||||
def supports_streaming(self) -> bool:
|
||||
"""Check if model supports streaming responses."""
|
||||
return "stream" in self.supported_parameters
|
||||
|
||||
def supports_online(self) -> bool:
|
||||
"""Check if model supports web search (online mode)."""
|
||||
return self.supports_tools()
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProviderCapabilities:
|
||||
"""
|
||||
Capabilities supported by a provider.
|
||||
|
||||
Attributes:
|
||||
streaming: Provider supports streaming responses
|
||||
tools: Provider supports function calling
|
||||
images: Provider supports image inputs
|
||||
online: Provider supports web search
|
||||
max_context: Maximum context length across all models
|
||||
"""
|
||||
|
||||
streaming: bool = True
|
||||
tools: bool = True
|
||||
images: bool = True
|
||||
online: bool = False
|
||||
max_context: int = 128000
|
||||
|
||||
|
||||
class AIProvider(ABC):
|
||||
"""
|
||||
Abstract base class for AI model providers.
|
||||
|
||||
All provider implementations must inherit from this class
|
||||
and implement the required abstract methods.
|
||||
"""
|
||||
|
||||
def __init__(self, api_key: str, base_url: Optional[str] = None):
|
||||
"""
|
||||
Initialize the provider.
|
||||
|
||||
Args:
|
||||
api_key: API key for authentication
|
||||
base_url: Optional custom base URL for the API
|
||||
"""
|
||||
self.api_key = api_key
|
||||
self.base_url = base_url
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def name(self) -> str:
|
||||
"""Get the provider name."""
|
||||
pass
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def capabilities(self) -> ProviderCapabilities:
|
||||
"""Get provider capabilities."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def list_models(self) -> List[ModelInfo]:
|
||||
"""
|
||||
Fetch available models from the provider.
|
||||
|
||||
Returns:
|
||||
List of available models with their info
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_model(self, model_id: str) -> Optional[ModelInfo]:
|
||||
"""
|
||||
Get information about a specific model.
|
||||
|
||||
Args:
|
||||
model_id: The model identifier
|
||||
|
||||
Returns:
|
||||
Model information or None if not found
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def chat(
|
||||
self,
|
||||
model: str,
|
||||
messages: List[ChatMessage],
|
||||
stream: bool = False,
|
||||
max_tokens: Optional[int] = None,
|
||||
temperature: Optional[float] = None,
|
||||
tools: Optional[List[Dict[str, Any]]] = None,
|
||||
tool_choice: Optional[str] = None,
|
||||
**kwargs: Any,
|
||||
) -> Union[ChatResponse, Iterator[StreamChunk]]:
|
||||
"""
|
||||
Send a chat completion request.
|
||||
|
||||
Args:
|
||||
model: Model ID to use
|
||||
messages: List of chat messages
|
||||
stream: Whether to stream the response
|
||||
max_tokens: Maximum tokens in response
|
||||
temperature: Sampling temperature
|
||||
tools: List of tool definitions for function calling
|
||||
tool_choice: How to handle tool selection ("auto", "none", etc.)
|
||||
**kwargs: Additional provider-specific parameters
|
||||
|
||||
Returns:
|
||||
ChatResponse for non-streaming, Iterator[StreamChunk] for streaming
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def chat_async(
|
||||
self,
|
||||
model: str,
|
||||
messages: List[ChatMessage],
|
||||
stream: bool = False,
|
||||
max_tokens: Optional[int] = None,
|
||||
temperature: Optional[float] = None,
|
||||
tools: Optional[List[Dict[str, Any]]] = None,
|
||||
tool_choice: Optional[str] = None,
|
||||
**kwargs: Any,
|
||||
) -> Union[ChatResponse, AsyncIterator[StreamChunk]]:
|
||||
"""
|
||||
Send an async chat completion request.
|
||||
|
||||
Args:
|
||||
model: Model ID to use
|
||||
messages: List of chat messages
|
||||
stream: Whether to stream the response
|
||||
max_tokens: Maximum tokens in response
|
||||
temperature: Sampling temperature
|
||||
tools: List of tool definitions for function calling
|
||||
tool_choice: How to handle tool selection
|
||||
**kwargs: Additional provider-specific parameters
|
||||
|
||||
Returns:
|
||||
ChatResponse for non-streaming, AsyncIterator[StreamChunk] for streaming
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_credits(self) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Get account credit/balance information.
|
||||
|
||||
Returns:
|
||||
Dict with credit info or None if not supported
|
||||
"""
|
||||
pass
|
||||
|
||||
def validate_api_key(self) -> bool:
|
||||
"""
|
||||
Validate that the API key is valid.
|
||||
|
||||
Returns:
|
||||
True if API key is valid
|
||||
"""
|
||||
try:
|
||||
self.list_models()
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
@@ -0,0 +1,623 @@
|
||||
"""
|
||||
OpenRouter provider implementation.
|
||||
|
||||
This module implements the AIProvider interface for OpenRouter,
|
||||
supporting chat completions, streaming, and function calling.
|
||||
"""
|
||||
|
||||
from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union
|
||||
|
||||
import requests
|
||||
from openrouter import OpenRouter
|
||||
|
||||
from oai.constants import APP_NAME, APP_URL, DEFAULT_BASE_URL
|
||||
from oai.providers.base import (
|
||||
AIProvider,
|
||||
ChatMessage,
|
||||
ChatResponse,
|
||||
ChatResponseChoice,
|
||||
ModelInfo,
|
||||
ProviderCapabilities,
|
||||
StreamChunk,
|
||||
ToolCall,
|
||||
ToolFunction,
|
||||
UsageStats,
|
||||
)
|
||||
from oai.utils.logging import get_logger
|
||||
|
||||
|
||||
class OpenRouterProvider(AIProvider):
|
||||
"""
|
||||
OpenRouter API provider implementation.
|
||||
|
||||
Provides access to multiple AI models through OpenRouter's unified API,
|
||||
supporting chat completions, streaming responses, and function calling.
|
||||
|
||||
Attributes:
|
||||
client: The underlying OpenRouter client
|
||||
_models_cache: Cached list of available models
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
base_url: Optional[str] = None,
|
||||
app_name: str = APP_NAME,
|
||||
app_url: str = APP_URL,
|
||||
):
|
||||
"""
|
||||
Initialize the OpenRouter provider.
|
||||
|
||||
Args:
|
||||
api_key: OpenRouter API key
|
||||
base_url: Optional custom base URL
|
||||
app_name: Application name for API headers
|
||||
app_url: Application URL for API headers
|
||||
"""
|
||||
super().__init__(api_key, base_url or DEFAULT_BASE_URL)
|
||||
self.app_name = app_name
|
||||
self.app_url = app_url
|
||||
self.client = OpenRouter(api_key=api_key)
|
||||
self._models_cache: Optional[List[ModelInfo]] = None
|
||||
self._raw_models_cache: Optional[List[Dict[str, Any]]] = None
|
||||
|
||||
self.logger = get_logger()
|
||||
self.logger.info(f"OpenRouter provider initialized with base URL: {self.base_url}")
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
"""Get the provider name."""
|
||||
return "OpenRouter"
|
||||
|
||||
@property
|
||||
def capabilities(self) -> ProviderCapabilities:
|
||||
"""Get provider capabilities."""
|
||||
return ProviderCapabilities(
|
||||
streaming=True,
|
||||
tools=True,
|
||||
images=True,
|
||||
online=True,
|
||||
max_context=2000000, # Claude models support up to 200k
|
||||
)
|
||||
|
||||
def _get_headers(self) -> Dict[str, str]:
|
||||
"""Get standard HTTP headers for API requests."""
|
||||
headers = {
|
||||
"HTTP-Referer": self.app_url,
|
||||
"X-Title": self.app_name,
|
||||
}
|
||||
if self.api_key:
|
||||
headers["Authorization"] = f"Bearer {self.api_key}"
|
||||
return headers
|
||||
|
||||
def _parse_model(self, model_data: Dict[str, Any]) -> ModelInfo:
|
||||
"""
|
||||
Parse raw model data into ModelInfo.
|
||||
|
||||
Args:
|
||||
model_data: Raw model data from API
|
||||
|
||||
Returns:
|
||||
Parsed ModelInfo object
|
||||
"""
|
||||
architecture = model_data.get("architecture", {})
|
||||
pricing_data = model_data.get("pricing", {})
|
||||
|
||||
# Parse pricing (convert from string to float if needed)
|
||||
pricing = {}
|
||||
for key in ["prompt", "completion"]:
|
||||
value = pricing_data.get(key)
|
||||
if value is not None:
|
||||
try:
|
||||
# Convert from per-token to per-million-tokens
|
||||
pricing[key] = float(value) * 1_000_000
|
||||
except (ValueError, TypeError):
|
||||
pricing[key] = 0.0
|
||||
|
||||
return ModelInfo(
|
||||
id=model_data.get("id", ""),
|
||||
name=model_data.get("name", model_data.get("id", "")),
|
||||
description=model_data.get("description", ""),
|
||||
context_length=model_data.get("context_length", 0),
|
||||
pricing=pricing,
|
||||
supported_parameters=model_data.get("supported_parameters", []),
|
||||
input_modalities=architecture.get("input_modalities", ["text"]),
|
||||
output_modalities=architecture.get("output_modalities", ["text"]),
|
||||
)
|
||||
|
||||
def list_models(self, filter_text_only: bool = True) -> List[ModelInfo]:
|
||||
"""
|
||||
Fetch available models from OpenRouter.
|
||||
|
||||
Args:
|
||||
filter_text_only: If True, exclude video-only models
|
||||
|
||||
Returns:
|
||||
List of available models
|
||||
|
||||
Raises:
|
||||
Exception: If API request fails
|
||||
"""
|
||||
if self._models_cache is not None:
|
||||
return self._models_cache
|
||||
|
||||
try:
|
||||
response = requests.get(
|
||||
f"{self.base_url}/models",
|
||||
headers=self._get_headers(),
|
||||
timeout=10,
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
raw_models = response.json().get("data", [])
|
||||
self._raw_models_cache = raw_models
|
||||
|
||||
models = []
|
||||
for model_data in raw_models:
|
||||
# Optionally filter out video-only models
|
||||
if filter_text_only:
|
||||
modalities = model_data.get("modalities", [])
|
||||
if modalities and "video" in modalities and "text" not in modalities:
|
||||
continue
|
||||
|
||||
models.append(self._parse_model(model_data))
|
||||
|
||||
self._models_cache = models
|
||||
self.logger.info(f"Fetched {len(models)} models from OpenRouter")
|
||||
return models
|
||||
|
||||
except requests.RequestException as e:
|
||||
self.logger.error(f"Failed to fetch models: {e}")
|
||||
raise
|
||||
|
||||
def get_raw_models(self) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Get raw model data as returned by the API.
|
||||
|
||||
Useful for accessing provider-specific fields not in ModelInfo.
|
||||
|
||||
Returns:
|
||||
List of raw model dictionaries
|
||||
"""
|
||||
if self._raw_models_cache is None:
|
||||
self.list_models()
|
||||
return self._raw_models_cache or []
|
||||
|
||||
def get_model(self, model_id: str) -> Optional[ModelInfo]:
|
||||
"""
|
||||
Get information about a specific model.
|
||||
|
||||
Args:
|
||||
model_id: The model identifier
|
||||
|
||||
Returns:
|
||||
Model information or None if not found
|
||||
"""
|
||||
models = self.list_models()
|
||||
for model in models:
|
||||
if model.id == model_id:
|
||||
return model
|
||||
return None
|
||||
|
||||
def get_raw_model(self, model_id: str) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Get raw model data for a specific model.
|
||||
|
||||
Args:
|
||||
model_id: The model identifier
|
||||
|
||||
Returns:
|
||||
Raw model dictionary or None if not found
|
||||
"""
|
||||
raw_models = self.get_raw_models()
|
||||
for model in raw_models:
|
||||
if model.get("id") == model_id:
|
||||
return model
|
||||
return None
|
||||
|
||||
def _convert_messages(self, messages: List[ChatMessage]) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Convert ChatMessage objects to API format.
|
||||
|
||||
Args:
|
||||
messages: List of ChatMessage objects
|
||||
|
||||
Returns:
|
||||
List of message dictionaries for the API
|
||||
"""
|
||||
return [msg.to_dict() for msg in messages]
|
||||
|
||||
def _parse_usage(self, usage_data: Any) -> Optional[UsageStats]:
|
||||
"""
|
||||
Parse usage data from API response.
|
||||
|
||||
Args:
|
||||
usage_data: Raw usage data from API
|
||||
|
||||
Returns:
|
||||
Parsed UsageStats or None
|
||||
"""
|
||||
if not usage_data:
|
||||
return None
|
||||
|
||||
# Handle both attribute and dict access
|
||||
prompt_tokens = 0
|
||||
completion_tokens = 0
|
||||
total_cost = None
|
||||
|
||||
if hasattr(usage_data, "prompt_tokens"):
|
||||
prompt_tokens = getattr(usage_data, "prompt_tokens", 0) or 0
|
||||
elif isinstance(usage_data, dict):
|
||||
prompt_tokens = usage_data.get("prompt_tokens", 0) or 0
|
||||
|
||||
if hasattr(usage_data, "completion_tokens"):
|
||||
completion_tokens = getattr(usage_data, "completion_tokens", 0) or 0
|
||||
elif isinstance(usage_data, dict):
|
||||
completion_tokens = usage_data.get("completion_tokens", 0) or 0
|
||||
|
||||
# Try alternative naming (input_tokens/output_tokens)
|
||||
if prompt_tokens == 0:
|
||||
if hasattr(usage_data, "input_tokens"):
|
||||
prompt_tokens = getattr(usage_data, "input_tokens", 0) or 0
|
||||
elif isinstance(usage_data, dict):
|
||||
prompt_tokens = usage_data.get("input_tokens", 0) or 0
|
||||
|
||||
if completion_tokens == 0:
|
||||
if hasattr(usage_data, "output_tokens"):
|
||||
completion_tokens = getattr(usage_data, "output_tokens", 0) or 0
|
||||
elif isinstance(usage_data, dict):
|
||||
completion_tokens = usage_data.get("output_tokens", 0) or 0
|
||||
|
||||
# Get cost if available
|
||||
if hasattr(usage_data, "total_cost_usd"):
|
||||
total_cost = getattr(usage_data, "total_cost_usd", None)
|
||||
elif isinstance(usage_data, dict):
|
||||
total_cost = usage_data.get("total_cost_usd")
|
||||
|
||||
return UsageStats(
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
total_tokens=prompt_tokens + completion_tokens,
|
||||
total_cost_usd=float(total_cost) if total_cost else None,
|
||||
)
|
||||
|
||||
def _parse_tool_calls(self, tool_calls_data: Any) -> Optional[List[ToolCall]]:
|
||||
"""
|
||||
Parse tool calls from API response.
|
||||
|
||||
Args:
|
||||
tool_calls_data: Raw tool calls data
|
||||
|
||||
Returns:
|
||||
List of ToolCall objects or None
|
||||
"""
|
||||
if not tool_calls_data:
|
||||
return None
|
||||
|
||||
tool_calls = []
|
||||
for tc in tool_calls_data:
|
||||
# Handle both attribute and dict access
|
||||
if hasattr(tc, "id"):
|
||||
tc_id = tc.id
|
||||
tc_type = getattr(tc, "type", "function")
|
||||
func = tc.function
|
||||
func_name = func.name
|
||||
func_args = func.arguments
|
||||
else:
|
||||
tc_id = tc.get("id", "")
|
||||
tc_type = tc.get("type", "function")
|
||||
func = tc.get("function", {})
|
||||
func_name = func.get("name", "")
|
||||
func_args = func.get("arguments", "{}")
|
||||
|
||||
tool_calls.append(
|
||||
ToolCall(
|
||||
id=tc_id,
|
||||
type=tc_type,
|
||||
function=ToolFunction(name=func_name, arguments=func_args),
|
||||
)
|
||||
)
|
||||
|
||||
return tool_calls if tool_calls else None
|
||||
|
||||
def _parse_response(self, response: Any) -> ChatResponse:
|
||||
"""
|
||||
Parse API response into ChatResponse.
|
||||
|
||||
Args:
|
||||
response: Raw API response
|
||||
|
||||
Returns:
|
||||
Parsed ChatResponse
|
||||
"""
|
||||
choices = []
|
||||
for choice in response.choices:
|
||||
msg = choice.message
|
||||
message = ChatMessage(
|
||||
role=msg.role if hasattr(msg, "role") else "assistant",
|
||||
content=msg.content if hasattr(msg, "content") else None,
|
||||
tool_calls=self._parse_tool_calls(
|
||||
getattr(msg, "tool_calls", None)
|
||||
),
|
||||
)
|
||||
choices.append(
|
||||
ChatResponseChoice(
|
||||
index=choice.index if hasattr(choice, "index") else 0,
|
||||
message=message,
|
||||
finish_reason=getattr(choice, "finish_reason", None),
|
||||
)
|
||||
)
|
||||
|
||||
return ChatResponse(
|
||||
id=response.id if hasattr(response, "id") else "",
|
||||
choices=choices,
|
||||
usage=self._parse_usage(getattr(response, "usage", None)),
|
||||
model=getattr(response, "model", None),
|
||||
created=getattr(response, "created", None),
|
||||
)
|
||||
|
||||
def chat(
|
||||
self,
|
||||
model: str,
|
||||
messages: List[ChatMessage],
|
||||
stream: bool = False,
|
||||
max_tokens: Optional[int] = None,
|
||||
temperature: Optional[float] = None,
|
||||
tools: Optional[List[Dict[str, Any]]] = None,
|
||||
tool_choice: Optional[str] = None,
|
||||
transforms: Optional[List[str]] = None,
|
||||
**kwargs: Any,
|
||||
) -> Union[ChatResponse, Iterator[StreamChunk]]:
|
||||
"""
|
||||
Send a chat completion request to OpenRouter.
|
||||
|
||||
Args:
|
||||
model: Model ID to use
|
||||
messages: List of chat messages
|
||||
stream: Whether to stream the response
|
||||
max_tokens: Maximum tokens in response
|
||||
temperature: Sampling temperature (0-2)
|
||||
tools: List of tool definitions for function calling
|
||||
tool_choice: How to handle tool selection ("auto", "none", etc.)
|
||||
transforms: List of transforms (e.g., ["middle-out"])
|
||||
**kwargs: Additional parameters
|
||||
|
||||
Returns:
|
||||
ChatResponse for non-streaming, Iterator[StreamChunk] for streaming
|
||||
"""
|
||||
# Build request parameters
|
||||
params: Dict[str, Any] = {
|
||||
"model": model,
|
||||
"messages": self._convert_messages(messages),
|
||||
"stream": stream,
|
||||
"http_headers": self._get_headers(),
|
||||
}
|
||||
|
||||
# Request usage stats in streaming responses
|
||||
if stream:
|
||||
params["stream_options"] = {"include_usage": True}
|
||||
|
||||
if max_tokens is not None:
|
||||
params["max_tokens"] = max_tokens
|
||||
|
||||
if temperature is not None:
|
||||
params["temperature"] = temperature
|
||||
|
||||
if tools:
|
||||
params["tools"] = tools
|
||||
params["tool_choice"] = tool_choice or "auto"
|
||||
|
||||
if transforms:
|
||||
params["transforms"] = transforms
|
||||
|
||||
# Add any additional parameters
|
||||
params.update(kwargs)
|
||||
|
||||
self.logger.debug(f"Sending chat request to model {model}")
|
||||
|
||||
try:
|
||||
response = self.client.chat.send(**params)
|
||||
|
||||
if stream:
|
||||
return self._stream_response(response)
|
||||
else:
|
||||
return self._parse_response(response)
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Chat request failed: {e}")
|
||||
raise
|
||||
|
||||
def _stream_response(self, response: Any) -> Iterator[StreamChunk]:
|
||||
"""
|
||||
Process a streaming response.
|
||||
|
||||
Args:
|
||||
response: Streaming response from API
|
||||
|
||||
Yields:
|
||||
StreamChunk objects
|
||||
"""
|
||||
last_usage = None
|
||||
|
||||
try:
|
||||
for chunk in response:
|
||||
# Check for errors
|
||||
if hasattr(chunk, "error") and chunk.error:
|
||||
yield StreamChunk(
|
||||
id=getattr(chunk, "id", ""),
|
||||
error=chunk.error.message if hasattr(chunk.error, "message") else str(chunk.error),
|
||||
)
|
||||
return
|
||||
|
||||
# Extract delta content
|
||||
delta_content = None
|
||||
finish_reason = None
|
||||
|
||||
if hasattr(chunk, "choices") and chunk.choices:
|
||||
choice = chunk.choices[0]
|
||||
if hasattr(choice, "delta"):
|
||||
delta = choice.delta
|
||||
if hasattr(delta, "content") and delta.content:
|
||||
delta_content = delta.content
|
||||
finish_reason = getattr(choice, "finish_reason", None)
|
||||
|
||||
# Track usage from last chunk
|
||||
if hasattr(chunk, "usage") and chunk.usage:
|
||||
last_usage = self._parse_usage(chunk.usage)
|
||||
|
||||
yield StreamChunk(
|
||||
id=getattr(chunk, "id", ""),
|
||||
delta_content=delta_content,
|
||||
finish_reason=finish_reason,
|
||||
usage=last_usage if finish_reason else None,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Stream error: {e}")
|
||||
yield StreamChunk(id="", error=str(e))
|
||||
|
||||
async def chat_async(
|
||||
self,
|
||||
model: str,
|
||||
messages: List[ChatMessage],
|
||||
stream: bool = False,
|
||||
max_tokens: Optional[int] = None,
|
||||
temperature: Optional[float] = None,
|
||||
tools: Optional[List[Dict[str, Any]]] = None,
|
||||
tool_choice: Optional[str] = None,
|
||||
**kwargs: Any,
|
||||
) -> Union[ChatResponse, AsyncIterator[StreamChunk]]:
|
||||
"""
|
||||
Send an async chat completion request.
|
||||
|
||||
Note: Currently wraps the sync implementation.
|
||||
TODO: Implement true async support when OpenRouter SDK supports it.
|
||||
|
||||
Args:
|
||||
model: Model ID to use
|
||||
messages: List of chat messages
|
||||
stream: Whether to stream the response
|
||||
max_tokens: Maximum tokens in response
|
||||
temperature: Sampling temperature
|
||||
tools: List of tool definitions
|
||||
tool_choice: Tool selection mode
|
||||
**kwargs: Additional parameters
|
||||
|
||||
Returns:
|
||||
ChatResponse for non-streaming, AsyncIterator for streaming
|
||||
"""
|
||||
# For now, use sync implementation
|
||||
# TODO: Add true async when SDK supports it
|
||||
result = self.chat(
|
||||
model=model,
|
||||
messages=messages,
|
||||
stream=stream,
|
||||
max_tokens=max_tokens,
|
||||
temperature=temperature,
|
||||
tools=tools,
|
||||
tool_choice=tool_choice,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
if stream and isinstance(result, Iterator):
|
||||
# Convert sync iterator to async
|
||||
async def async_iter() -> AsyncIterator[StreamChunk]:
|
||||
for chunk in result:
|
||||
yield chunk
|
||||
|
||||
return async_iter()
|
||||
|
||||
return result
|
||||
|
||||
def get_credits(self) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Get OpenRouter account credit information.
|
||||
|
||||
Returns:
|
||||
Dict with credit info:
|
||||
- total_credits: Total credits purchased
|
||||
- used_credits: Credits used
|
||||
- credits_left: Remaining credits
|
||||
|
||||
Raises:
|
||||
Exception: If API request fails
|
||||
"""
|
||||
if not self.api_key:
|
||||
return None
|
||||
|
||||
try:
|
||||
response = requests.get(
|
||||
f"{self.base_url}/credits",
|
||||
headers=self._get_headers(),
|
||||
timeout=10,
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
data = response.json().get("data", {})
|
||||
total_credits = float(data.get("total_credits", 0))
|
||||
total_usage = float(data.get("total_usage", 0))
|
||||
credits_left = total_credits - total_usage
|
||||
|
||||
return {
|
||||
"total_credits": total_credits,
|
||||
"used_credits": total_usage,
|
||||
"credits_left": credits_left,
|
||||
"total_credits_formatted": f"${total_credits:.2f}",
|
||||
"used_credits_formatted": f"${total_usage:.2f}",
|
||||
"credits_left_formatted": f"${credits_left:.2f}",
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to fetch credits: {e}")
|
||||
return None
|
||||
|
||||
def clear_cache(self) -> None:
|
||||
"""Clear the models cache to force a refresh."""
|
||||
self._models_cache = None
|
||||
self._raw_models_cache = None
|
||||
self.logger.debug("Models cache cleared")
|
||||
|
||||
def get_effective_model_id(self, model_id: str, online_enabled: bool) -> str:
|
||||
"""
|
||||
Get the effective model ID with online suffix if needed.
|
||||
|
||||
Args:
|
||||
model_id: Base model ID
|
||||
online_enabled: Whether online mode is enabled
|
||||
|
||||
Returns:
|
||||
Model ID with :online suffix if applicable
|
||||
"""
|
||||
if online_enabled and not model_id.endswith(":online"):
|
||||
return f"{model_id}:online"
|
||||
return model_id
|
||||
|
||||
def estimate_cost(
|
||||
self,
|
||||
model_id: str,
|
||||
input_tokens: int,
|
||||
output_tokens: int,
|
||||
) -> float:
|
||||
"""
|
||||
Estimate the cost for a completion.
|
||||
|
||||
Args:
|
||||
model_id: Model ID
|
||||
input_tokens: Number of input tokens
|
||||
output_tokens: Number of output tokens
|
||||
|
||||
Returns:
|
||||
Estimated cost in USD
|
||||
"""
|
||||
model = self.get_model(model_id)
|
||||
if model and model.pricing:
|
||||
input_cost = model.pricing.get("prompt", 0) * input_tokens / 1_000_000
|
||||
output_cost = model.pricing.get("completion", 0) * output_tokens / 1_000_000
|
||||
return input_cost + output_cost
|
||||
|
||||
# Fallback to default pricing if model not found
|
||||
from oai.constants import MODEL_PRICING
|
||||
|
||||
input_cost = MODEL_PRICING["input"] * input_tokens / 1_000_000
|
||||
output_cost = MODEL_PRICING["output"] * output_tokens / 1_000_000
|
||||
return input_cost + output_cost
|
||||
Reference in New Issue
Block a user