mark-xl1tc0/endpoints/llm.post.py

56 lines
1.8 KiB
Python

from fastapi import APIRouter, HTTPException, status
from typing import Dict
from pydantic import BaseModel
from helpers.generic_helpers import process_llm_request, validate_data, log_error
router = APIRouter()
class LLMRequest(BaseModel):
model: str
prompt: str
temperature: float = 0.7
max_tokens: int = 1000
class LLMResponse(BaseModel):
id: str
content: str
model: str
created: int
usage: Dict[str, int]
@router.post("/llm", status_code=status.HTTP_200_OK, response_model=LLMResponse)
async def generate_llm_response(request: LLMRequest):
"""
Generate a response from an LLM model using the litellm library.
This endpoint accepts a model name, prompt, and optional parameters,
then calls the actual LLM service to generate a response.
"""
try:
# Validate required fields
required_fields = ["model", "prompt"]
if not validate_data(request.dict(), required_fields):
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail="Missing required fields: model and prompt are required"
)
# Process the LLM request using the helper function
result = process_llm_request(
model=request.model,
prompt=request.prompt,
temperature=request.temperature,
max_tokens=request.max_tokens
)
return result
except HTTPException:
# Re-raise HTTP exceptions as they already have status codes
raise
except Exception as e:
log_error("Unexpected error in LLM endpoint", e)
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"Failed to process LLM request: {str(e)}"
)