152 lines
4.8 KiB
Python
152 lines
4.8 KiB
Python
"""LiteLLM wrapper for unified model access."""
|
|
import time
|
|
import json
|
|
from typing import AsyncIterator, Optional, Dict, Any, List
|
|
from litellm import acompletion
|
|
|
|
from ..models import Provider
|
|
from ..services.logger import async_logger, LogEntry
|
|
|
|
|
|
class LiteLLMWrapper:
|
|
"""Wrapper around LiteLLM for proxy functionality."""
|
|
|
|
@staticmethod
|
|
def _build_model_string(provider: Provider, model: str) -> str:
|
|
"""Build LiteLLM model string.
|
|
|
|
Args:
|
|
provider: The provider configuration
|
|
model: The model name
|
|
|
|
Returns:
|
|
LiteLLM-compatible model string
|
|
"""
|
|
if provider.api_type == "anthropic":
|
|
return f"anthropic/{model}"
|
|
else:
|
|
return f"openai/{model}"
|
|
|
|
@staticmethod
|
|
async def chat_completion(
|
|
provider: Provider,
|
|
api_key: str,
|
|
model: str,
|
|
messages: List[Dict[str, Any]],
|
|
stream: bool = False,
|
|
**kwargs
|
|
) -> tuple[Any, int, int]:
|
|
"""Execute a chat completion request.
|
|
|
|
Args:
|
|
provider: The provider configuration
|
|
api_key: The API key to use
|
|
model: The model name
|
|
messages: The chat messages
|
|
stream: Whether to stream the response
|
|
**kwargs: Additional parameters
|
|
|
|
Returns:
|
|
Tuple of (response, prompt_tokens, completion_tokens)
|
|
"""
|
|
model_str = LiteLLMWrapper._build_model_string(provider, model)
|
|
|
|
# Set API key for this request
|
|
api_base = provider.base_url if provider.api_type == "openai" else None
|
|
|
|
start_time = time.time()
|
|
|
|
try:
|
|
response = await acompletion(
|
|
model=model_str,
|
|
messages=messages,
|
|
api_key=api_key,
|
|
api_base=api_base,
|
|
stream=stream,
|
|
**kwargs
|
|
)
|
|
|
|
if stream:
|
|
# For streaming, return the async iterator
|
|
# Token counts will be collected during iteration
|
|
return response, 0, 0
|
|
else:
|
|
# Extract token usage from response
|
|
usage = getattr(response, "usage", None) or {}
|
|
prompt_tokens = getattr(usage, "prompt_tokens", 0) or 0
|
|
completion_tokens = getattr(usage, "completion_tokens", 0) or 0
|
|
return response, prompt_tokens, completion_tokens
|
|
|
|
except Exception as e:
|
|
# Re-raise with context
|
|
raise
|
|
|
|
@staticmethod
|
|
async def stream_response(
|
|
stream_iterator: AsyncIterator,
|
|
client_key_id: int,
|
|
provider_id: int,
|
|
model: str,
|
|
start_time: float
|
|
) -> AsyncIterator[str]:
|
|
"""Process streaming response and yield SSE-formatted data.
|
|
|
|
Args:
|
|
stream_iterator: The LiteLLM stream iterator
|
|
client_key_id: Client key ID for logging
|
|
provider_id: Provider ID for logging
|
|
model: The model name
|
|
start_time: Request start time
|
|
|
|
Yields:
|
|
SSE-formatted strings
|
|
"""
|
|
prompt_tokens = 0
|
|
completion_tokens = 0
|
|
|
|
try:
|
|
async for chunk in stream_iterator:
|
|
# Extract token counts if available
|
|
if hasattr(chunk, "usage") and chunk.usage:
|
|
prompt_tokens = getattr(chunk.usage, "prompt_tokens", 0) or 0
|
|
completion_tokens = getattr(chunk.usage, "completion_tokens", 0) or 0
|
|
|
|
# Convert to SSE format
|
|
if hasattr(chunk, "model_dump"):
|
|
chunk_dict = chunk.model_dump()
|
|
else:
|
|
chunk_dict = chunk
|
|
|
|
yield f"data: {json.dumps(chunk_dict)}\n\n"
|
|
|
|
yield "data: [DONE]\n\n"
|
|
|
|
# Log successful request
|
|
latency_ms = int((time.time() - start_time) * 1000)
|
|
await async_logger.log(LogEntry(
|
|
client_key_id=client_key_id,
|
|
provider_id=provider_id,
|
|
model=model,
|
|
prompt_tokens=prompt_tokens,
|
|
completion_tokens=completion_tokens,
|
|
latency_ms=latency_ms,
|
|
success=True
|
|
))
|
|
|
|
except Exception as e:
|
|
# Log failed request
|
|
latency_ms = int((time.time() - start_time) * 1000)
|
|
await async_logger.log(LogEntry(
|
|
client_key_id=client_key_id,
|
|
provider_id=provider_id,
|
|
model=model,
|
|
prompt_tokens=prompt_tokens,
|
|
completion_tokens=completion_tokens,
|
|
latency_ms=latency_ms,
|
|
success=False,
|
|
error_message=str(e)
|
|
))
|
|
|
|
# Send error as SSE
|
|
yield f"data: {json.dumps({'error': str(e)})}\n\n"
|