File size: 2,074 Bytes
ea89131
 
 
9ea68b1
ea89131
9b0f62b
9ea68b1
 
9b0f62b
 
 
ea89131
 
 
 
 
 
 
 
 
 
 
9b0f62b
 
 
 
ea89131
 
9b0f62b
 
ea89131
9b0f62b
ea89131
 
 
 
 
9b0f62b
 
ea89131
d866e5b
ea89131
9b0f62b
 
 
d866e5b
9b0f62b
 
ea89131
9b0f62b
d866e5b
ea89131
 
 
 
 
 
 
 
 
 
9b0f62b
ea89131
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
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