llm-proxy / app /main.py
knighter75's picture
Fix actual models
ea89131
Raw
History Blame Contribute Delete
2.07 kB
from contextlib import asynccontextmanager
from uuid import uuid4
from dotenv import load_dotenv
from fastapi import Depends, FastAPI, HTTPException
load_dotenv()
from .auth import verify_api_key
from .factory import ProviderFactory
from .models import ChatRequest, ChatResponse
from .providers.hf_openai import configured_models
@asynccontextmanager
async def lifespan(_: FastAPI):
yield
if ProviderFactory._instance is not None:
await ProviderFactory._instance.client.close()
app = FastAPI(title="LLM API Proxy", version="2.0.0", lifespan=lifespan)
@app.get("/")
async def root():
return {"message": "LLM API Proxy is running", "version": "2.0.0"}
@app.get("/v1/models")
async def list_models(_: str = Depends(verify_api_key)):
return {
"object": "list",
"data": [
{"id": alias, "object": "model", "owned_by": "huggingface", "hf_model": model_id}
for alias, model_id in configured_models().items()
],
}
@app.post("/v1/chat/completions")
async def chat_completion(request: ChatRequest, _: str = Depends(verify_api_key)):
try:
provider = ProviderFactory.get_provider(request.model)
result = await provider.generate(
messages=[{"role": m.role, "content": m.content} for m in request.messages],
max_tokens=request.max_tokens,
temperature=request.temperature,
model=request.model,
)
return ChatResponse(
id=f"chatcmpl-{uuid4().hex}",
choices=[
{
"index": 0,
"message": {"role": "assistant", "content": result["content"]},
"finish_reason": "stop",
}
],
usage={"total_tokens": result["total_tokens"]},
model=request.model,
)
except ValueError as exc:
raise HTTPException(status_code=400, detail=str(exc)) from exc
except Exception as exc:
raise HTTPException(status_code=502, detail=f"Hugging Face error: {exc}") from exc