main.py 1.23 KB
Newer Older
1
2
import logging

Timothy J. Baek's avatar
Timothy J. Baek committed
3
4
5
6
7
8
from litellm.proxy.proxy_server import ProxyConfig, initialize
from litellm.proxy.proxy_server import app

from fastapi import FastAPI, Request, Depends, status
from fastapi.responses import JSONResponse
from utils.utils import get_http_authorization_cred, get_current_user
9
10
11
12
from config import SRC_LOG_LEVELS, ENV

log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["LITELLM"])
Timothy J. Baek's avatar
Timothy J. Baek committed
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

proxy_config = ProxyConfig()


async def config():
    router, model_list, general_settings = await proxy_config.load_config(
        router=None, config_file_path="./data/litellm/config.yaml"
    )

    await initialize(config="./data/litellm/config.yaml", telemetry=False)


async def startup():
    await config()


@app.on_event("startup")
async def on_startup():
    await startup()


@app.middleware("http")
async def auth_middleware(request: Request, call_next):
    auth_header = request.headers.get("Authorization", "")

    if ENV != "dev":
        try:
            user = get_current_user(get_http_authorization_cred(auth_header))
41
            log.debug(f"user: {user}")
Timothy J. Baek's avatar
Timothy J. Baek committed
42
43
44
45
46
        except Exception as e:
            return JSONResponse(status_code=400, content={"detail": str(e)})

    response = await call_next(request)
    return response