136 lines
4.4 KiB
Python
136 lines
4.4 KiB
Python
"""OpenAI-compatible chat completion endpoint."""
|
|
import time
|
|
from typing import Optional, List, Dict, Any
|
|
from fastapi import APIRouter, Depends, HTTPException
|
|
from fastapi.responses import StreamingResponse
|
|
from pydantic import BaseModel
|
|
|
|
from ..models import ClientKey
|
|
from ..middleware.auth import get_current_client
|
|
from ..services.router import model_router
|
|
from ..services.key_selector import key_selector
|
|
from ..services.logger import async_logger, LogEntry
|
|
from ..providers.litellm_wrapper import LiteLLMWrapper
|
|
|
|
router = APIRouter(prefix="/v1", tags=["chat"])
|
|
|
|
|
|
class ChatMessage(BaseModel):
|
|
"""Chat message."""
|
|
role: str
|
|
content: str
|
|
|
|
|
|
class ChatCompletionRequest(BaseModel):
|
|
"""Chat completion request."""
|
|
model: str
|
|
messages: List[ChatMessage]
|
|
stream: bool = False
|
|
temperature: Optional[float] = None
|
|
max_tokens: Optional[int] = None
|
|
top_p: Optional[float] = None
|
|
|
|
|
|
@router.post("/chat/completions")
|
|
async def chat_completions(
|
|
request: ChatCompletionRequest,
|
|
client: ClientKey = Depends(get_current_client)
|
|
):
|
|
"""OpenAI-compatible chat completions endpoint."""
|
|
start_time = time.time()
|
|
|
|
# Route model to provider
|
|
provider, error = await model_router.route_model(request.model)
|
|
if error:
|
|
raise HTTPException(status_code=400, detail={"error": error})
|
|
|
|
# Get next API key
|
|
provider_key = await key_selector.get_next_key(provider.id)
|
|
if not provider_key:
|
|
raise HTTPException(
|
|
status_code=503,
|
|
detail={"error": f"No available API key for provider: {provider.name}"}
|
|
)
|
|
|
|
# Build messages
|
|
messages = [{"role": m.role, "content": m.content} for m in request.messages]
|
|
|
|
# Build kwargs
|
|
kwargs = {}
|
|
if request.temperature is not None:
|
|
kwargs["temperature"] = request.temperature
|
|
if request.max_tokens is not None:
|
|
kwargs["max_tokens"] = request.max_tokens
|
|
if request.top_p is not None:
|
|
kwargs["top_p"] = request.top_p
|
|
|
|
try:
|
|
response, prompt_tokens, completion_tokens = await LiteLLMWrapper.chat_completion(
|
|
provider=provider,
|
|
api_key=provider_key.key,
|
|
model=request.model,
|
|
messages=messages,
|
|
stream=request.stream,
|
|
**kwargs
|
|
)
|
|
|
|
if request.stream:
|
|
return StreamingResponse(
|
|
LiteLLMWrapper.stream_response(
|
|
response, client.id, provider.id, request.model, start_time
|
|
),
|
|
media_type="text/event-stream"
|
|
)
|
|
else:
|
|
# Log successful request
|
|
latency_ms = int((time.time() - start_time) * 1000)
|
|
await async_logger.log(LogEntry(
|
|
client_key_id=client.id,
|
|
provider_id=provider.id,
|
|
model=request.model,
|
|
prompt_tokens=prompt_tokens,
|
|
completion_tokens=completion_tokens,
|
|
latency_ms=latency_ms,
|
|
success=True
|
|
))
|
|
|
|
# Return response
|
|
return response
|
|
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
# Log failed request
|
|
latency_ms = int((time.time() - start_time) * 1000)
|
|
await async_logger.log(LogEntry(
|
|
client_key_id=client.id,
|
|
provider_id=provider.id,
|
|
model=request.model,
|
|
latency_ms=latency_ms,
|
|
success=False,
|
|
error_message=str(e)
|
|
))
|
|
|
|
# Handle specific errors
|
|
error_str = str(e).lower()
|
|
if "rate" in error_str or "limit" in error_str:
|
|
raise HTTPException(status_code=429, detail={"error": str(e)})
|
|
raise HTTPException(status_code=500, detail={"error": str(e)})
|
|
|
|
|
|
@router.get("/models")
|
|
async def list_models(client: ClientKey = Depends(get_current_client)):
|
|
"""List available models."""
|
|
return {
|
|
"object": "list",
|
|
"data": [
|
|
{"id": "gpt-4o", "object": "model", "owned_by": "openai"},
|
|
{"id": "gpt-4o-mini", "object": "model", "owned_by": "openai"},
|
|
{"id": "o1", "object": "model", "owned_by": "openai"},
|
|
{"id": "o1-mini", "object": "model", "owned_by": "openai"},
|
|
{"id": "claude-3-opus", "object": "model", "owned_by": "anthropic"},
|
|
{"id": "claude-3-sonnet", "object": "model", "owned_by": "anthropic"},
|
|
{"id": "claude-3-haiku", "object": "model", "owned_by": "anthropic"},
|
|
]
|
|
}
|