main.py 10.9 KB
Newer Older
1
import sys
2
from contextlib import asynccontextmanager
3

Timothy J. Baek's avatar
Timothy J. Baek committed
4
from fastapi import FastAPI, Depends, HTTPException
5
6
from fastapi.routing import APIRoute
from fastapi.middleware.cors import CORSMiddleware
Timothy J. Baek's avatar
Timothy J. Baek committed
7

8
import logging
9
from fastapi import FastAPI, Request, Depends, status, Response
Timothy J. Baek's avatar
Timothy J. Baek committed
10
from fastapi.responses import JSONResponse
11
12
13
14

from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.responses import StreamingResponse
import json
15
import time
Timothy J. Baek's avatar
Timothy J. Baek committed
16
import requests
17

18
from pydantic import BaseModel, ConfigDict
Timothy J. Baek's avatar
Timothy J. Baek committed
19
20
from typing import Optional, List

21
from utils.utils import get_verified_user, get_current_user, get_admin_user
22
from config import SRC_LOG_LEVELS, ENV, MODEL_CONFIG
23
from constants import MESSAGES
24

25
26
import os

27
28
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["LITELLM"])
Timothy J. Baek's avatar
Timothy J. Baek committed
29

30

31
from config import (
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
32
    ENABLE_LITELLM,
Timothy J. Baek's avatar
Timothy J. Baek committed
33
    ENABLE_MODEL_FILTER,
34
35
36
    MODEL_FILTER_LIST,
    DATA_DIR,
    LITELLM_PROXY_PORT,
37
    LITELLM_PROXY_HOST,
38
)
39

Timothy J. Baek's avatar
Timothy J. Baek committed
40
41
42
43
import warnings

warnings.simplefilter("ignore")

44
from litellm.utils import get_llm_provider
45

46
47
import asyncio
import subprocess
Timothy J. Baek's avatar
Timothy J. Baek committed
48
import yaml
Timothy J. Baek's avatar
Timothy J. Baek committed
49

50
51
52
53
54
55
56
57
58
59

@asynccontextmanager
async def lifespan(app: FastAPI):
    log.info("startup_event")
    # TODO: Check config.yaml file and create one
    asyncio.create_task(start_litellm_background())
    yield


app = FastAPI(lifespan=lifespan)
Timothy J. Baek's avatar
Timothy J. Baek committed
60

61
origins = ["*"]
Timothy J. Baek's avatar
Timothy J. Baek committed
62

63
64
65
66
67
68
69
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
Timothy J. Baek's avatar
Timothy J. Baek committed
70

71

Timothy J. Baek's avatar
Timothy J. Baek committed
72
73
74
75
76
LITELLM_CONFIG_DIR = f"{DATA_DIR}/litellm/config.yaml"

with open(LITELLM_CONFIG_DIR, "r") as file:
    litellm_config = yaml.safe_load(file)

Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
77
78

app.state.ENABLE = ENABLE_LITELLM
Timothy J. Baek's avatar
Timothy J. Baek committed
79
app.state.CONFIG = litellm_config
80
app.state.MODEL_CONFIG = MODEL_CONFIG.value.get("litellm", [])
Timothy J. Baek's avatar
Timothy J. Baek committed
81

82
83
84
# Global variable to store the subprocess reference
background_process = None

85
86
87
88
89
90
91
CONFLICT_ENV_VARS = [
    # Uvicorn uses PORT, so LiteLLM might use it as well
    "PORT",
    # LiteLLM uses DATABASE_URL for Prisma connections
    "DATABASE_URL",
]

Timothy J. Baek's avatar
Timothy J. Baek committed
92

93
94
async def run_background_process(command):
    global background_process
Timothy J. Baek's avatar
Timothy J. Baek committed
95
    log.info("run_background_process")
96
97
98

    try:
        # Log the command to be executed
Timothy J. Baek's avatar
Timothy J. Baek committed
99
        log.info(f"Executing command: {command}")
100
101
        # Filter environment variables known to conflict with litellm
        env = {k: v for k, v in os.environ.items() if k not in CONFLICT_ENV_VARS}
102
103
        # Execute the command and create a subprocess
        process = await asyncio.create_subprocess_exec(
104
            *command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env
105
106
        )
        background_process = process
Timothy J. Baek's avatar
Timothy J. Baek committed
107
        log.info("Subprocess started successfully.")
108
109
110
111
112

        # Capture STDERR for debugging purposes
        stderr_output = await process.stderr.read()
        stderr_text = stderr_output.decode().strip()
        if stderr_text:
Timothy J. Baek's avatar
Timothy J. Baek committed
113
            log.info(f"Subprocess STDERR: {stderr_text}")
114

Timothy J. Baek's avatar
Timothy J. Baek committed
115
        # log.info output line by line
116
        async for line in process.stdout:
Timothy J. Baek's avatar
Timothy J. Baek committed
117
            log.info(line.decode().strip())
118
119
120

        # Wait for the process to finish
        returncode = await process.wait()
Timothy J. Baek's avatar
Timothy J. Baek committed
121
        log.info(f"Subprocess exited with return code {returncode}")
122
123
124
    except Exception as e:
        log.error(f"Failed to start subprocess: {e}")
        raise  # Optionally re-raise the exception if you want it to propagate
125
126
127


async def start_litellm_background():
Timothy J. Baek's avatar
Timothy J. Baek committed
128
    log.info("start_litellm_background")
129
    # Command to run in the background
130
131
132
133
134
135
136
137
138
139
140
    command = [
        "litellm",
        "--port",
        str(LITELLM_PROXY_PORT),
        "--host",
        LITELLM_PROXY_HOST,
        "--telemetry",
        "False",
        "--config",
        LITELLM_CONFIG_DIR,
    ]
Timothy J. Baek's avatar
Timothy J. Baek committed
141

142
    await run_background_process(command)
Timothy J. Baek's avatar
Timothy J. Baek committed
143
144


145
async def shutdown_litellm_background():
Timothy J. Baek's avatar
Timothy J. Baek committed
146
    log.info("shutdown_litellm_background")
147
148
149
150
    global background_process
    if background_process:
        background_process.terminate()
        await background_process.wait()  # Ensure the process has terminated
Timothy J. Baek's avatar
Timothy J. Baek committed
151
        log.info("Subprocess terminated")
152
        background_process = None
153
154


Timothy J. Baek's avatar
Timothy J. Baek committed
155
app.state.ENABLE_MODEL_FILTER = ENABLE_MODEL_FILTER
156
157
158
app.state.MODEL_FILTER_LIST = MODEL_FILTER_LIST


159
160
161
162
163
@app.get("/")
async def get_status():
    return {"status": True}


Timothy J. Baek's avatar
Timothy J. Baek committed
164
async def restart_litellm():
165
166
167
168
169
170
171
172
173
174
    """
    Endpoint to restart the litellm background service.
    """
    log.info("Requested restart of litellm service.")
    try:
        # Shut down the existing process if it is running
        await shutdown_litellm_background()
        log.info("litellm service shutdown complete.")

        # Restart the background service
Timothy J. Baek's avatar
Timothy J. Baek committed
175
176

        asyncio.create_task(start_litellm_background())
177
178
179
180
181
182
183
        log.info("litellm service restart complete.")

        return {
            "status": "success",
            "message": "litellm service restarted successfully.",
        }
    except Exception as e:
Timothy J. Baek's avatar
Timothy J. Baek committed
184
        log.info(f"Error restarting litellm service: {e}")
185
186
187
188
189
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e)
        )


Timothy J. Baek's avatar
Timothy J. Baek committed
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
@app.get("/restart")
async def restart_litellm_handler(user=Depends(get_admin_user)):
    return await restart_litellm()


@app.get("/config")
async def get_config(user=Depends(get_admin_user)):
    return app.state.CONFIG


class LiteLLMConfigForm(BaseModel):
    general_settings: Optional[dict] = None
    litellm_settings: Optional[dict] = None
    model_list: Optional[List[dict]] = None
    router_settings: Optional[dict] = None

206
207
    model_config = ConfigDict(protected_namespaces=())

Timothy J. Baek's avatar
Timothy J. Baek committed
208
209
210
211
212
213
214
215
216
217
218
219

@app.post("/config/update")
async def update_config(form_data: LiteLLMConfigForm, user=Depends(get_admin_user)):
    app.state.CONFIG = form_data.model_dump(exclude_none=True)

    with open(LITELLM_CONFIG_DIR, "w") as file:
        yaml.dump(app.state.CONFIG, file)

    await restart_litellm()
    return app.state.CONFIG


Timothy J. Baek's avatar
Timothy J. Baek committed
220
221
222
@app.get("/models")
@app.get("/v1/models")
async def get_models(user=Depends(get_current_user)):
223

Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
    if app.state.ENABLE:
        while not background_process:
            await asyncio.sleep(0.1)

        url = f"http://localhost:{LITELLM_PROXY_PORT}/v1"
        r = None
        try:
            r = requests.request(method="GET", url=f"{url}/models")
            r.raise_for_status()

            data = r.json()

            if app.state.ENABLE_MODEL_FILTER:
                if user and user.role == "user":
                    data["data"] = list(
                        filter(
                            lambda model: model["id"] in app.state.MODEL_FILTER_LIST,
                            data["data"],
                        )
Timothy J. Baek's avatar
Timothy J. Baek committed
243
                    )
244

245
246
            for model in data["data"]:
                add_custom_info_to_model(model)
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
            return data
        except Exception as e:

            log.exception(e)
            error_detail = "Open WebUI: Server Connection Error"
            if r is not None:
                try:
                    res = r.json()
                    if "error" in res:
                        error_detail = f"External: {res['error']}"
                except:
                    error_detail = f"External: {e}"

            return {
                "data": [
                    {
                        "id": model["model_name"],
                        "object": "model",
                        "created": int(time.time()),
                        "owned_by": "openai",
267
268
269
270
271
272
273
274
                        "custom_info": next(
                            (
                                item
                                for item in app.state.MODEL_CONFIG
                                if item["name"] == model["model_name"]
                            ),
                            {},
                        ),
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
275
276
277
278
279
280
                    }
                    for model in app.state.CONFIG["model_list"]
                ],
                "object": "list",
            }
    else:
281
        return {
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
282
            "data": [],
283
284
            "object": "list",
        }
285
286


287
288
def add_custom_info_to_model(model: dict):
    model["custom_info"] = next(
289
        (item for item in app.state.MODEL_CONFIG if item["id"] == model["id"]), {}
290
291
292
    )


293
294
295
296
297
298
299
300
301
@app.get("/model/info")
async def get_model_list(user=Depends(get_admin_user)):
    return {"data": app.state.CONFIG["model_list"]}


class AddLiteLLMModelForm(BaseModel):
    model_name: str
    litellm_params: dict

302
303
    model_config = ConfigDict(protected_namespaces=())

304
305
306
307
308

@app.post("/model/new")
async def add_model_to_config(
    form_data: AddLiteLLMModelForm, user=Depends(get_admin_user)
):
309
310
311
    try:
        get_llm_provider(model=form_data.model_name)
        app.state.CONFIG["model_list"].append(form_data.model_dump())
312

313
314
        with open(LITELLM_CONFIG_DIR, "w") as file:
            yaml.dump(app.state.CONFIG, file)
315

316
        await restart_litellm()
317

318
319
320
321
322
323
        return {"message": MESSAGES.MODEL_ADDED(form_data.model_name)}
    except Exception as e:
        print(e)
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e)
        )
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344


class DeleteLiteLLMModelForm(BaseModel):
    id: str


@app.post("/model/delete")
async def delete_model_from_config(
    form_data: DeleteLiteLLMModelForm, user=Depends(get_admin_user)
):
    app.state.CONFIG["model_list"] = [
        model
        for model in app.state.CONFIG["model_list"]
        if model["model_name"] != form_data.id
    ]

    with open(LITELLM_CONFIG_DIR, "w") as file:
        yaml.dump(app.state.CONFIG, file)

    await restart_litellm()

345
    return {"message": MESSAGES.MODEL_DELETED(form_data.id)}
346
347


Timothy J. Baek's avatar
Timothy J. Baek committed
348
349
350
@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
async def proxy(path: str, request: Request, user=Depends(get_verified_user)):
    body = await request.body()
351

352
    url = f"http://localhost:{LITELLM_PROXY_PORT}"
353

Timothy J. Baek's avatar
Timothy J. Baek committed
354
    target_url = f"{url}/{path}"
355

Timothy J. Baek's avatar
Timothy J. Baek committed
356
357
358
    headers = {}
    # headers["Authorization"] = f"Bearer {key}"
    headers["Content-Type"] = "application/json"
359

Timothy J. Baek's avatar
Timothy J. Baek committed
360
    r = None
361

Timothy J. Baek's avatar
Timothy J. Baek committed
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
    try:
        r = requests.request(
            method=request.method,
            url=target_url,
            data=body,
            headers=headers,
            stream=True,
        )

        r.raise_for_status()

        # Check if response is SSE
        if "text/event-stream" in r.headers.get("Content-Type", ""):
            return StreamingResponse(
                r.iter_content(chunk_size=8192),
                status_code=r.status_code,
                headers=dict(r.headers),
            )
        else:
            response_data = r.json()
            return response_data
    except Exception as e:
        log.exception(e)
        error_detail = "Open WebUI: Server Connection Error"
        if r is not None:
            try:
                res = r.json()
                if "error" in res:
                    error_detail = f"External: {res['error']['message'] if 'message' in res['error'] else res['error']}"
            except:
                error_detail = f"External: {e}"

        raise HTTPException(
            status_code=r.status_code if r else 500, detail=error_detail
        )