main.py 10.2 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
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

Timothy J. Baek's avatar
Timothy J. Baek committed
78
79
80
81
app.state.ENABLE_MODEL_FILTER = ENABLE_MODEL_FILTER
app.state.MODEL_FILTER_LIST = MODEL_FILTER_LIST


Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
82
app.state.ENABLE = ENABLE_LITELLM
Timothy J. Baek's avatar
Timothy J. Baek committed
83
84
app.state.CONFIG = litellm_config

85
86
87
# Global variable to store the subprocess reference
background_process = None

88
89
90
91
92
93
94
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
95

96
97
async def run_background_process(command):
    global background_process
Timothy J. Baek's avatar
Timothy J. Baek committed
98
    log.info("run_background_process")
99
100
101

    try:
        # Log the command to be executed
Timothy J. Baek's avatar
Timothy J. Baek committed
102
        log.info(f"Executing command: {command}")
103
104
        # 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}
105
106
        # Execute the command and create a subprocess
        process = await asyncio.create_subprocess_exec(
107
            *command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env
108
109
        )
        background_process = process
Timothy J. Baek's avatar
Timothy J. Baek committed
110
        log.info("Subprocess started successfully.")
111
112
113
114
115

        # 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
116
            log.info(f"Subprocess STDERR: {stderr_text}")
117

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

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


async def start_litellm_background():
Timothy J. Baek's avatar
Timothy J. Baek committed
131
    log.info("start_litellm_background")
132
    # Command to run in the background
133
134
135
136
137
138
139
140
141
142
143
    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
144

145
    await run_background_process(command)
Timothy J. Baek's avatar
Timothy J. Baek committed
146
147


148
async def shutdown_litellm_background():
Timothy J. Baek's avatar
Timothy J. Baek committed
149
    log.info("shutdown_litellm_background")
150
151
152
153
    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
154
        log.info("Subprocess terminated")
155
        background_process = None
156
157


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


Timothy J. Baek's avatar
Timothy J. Baek committed
163
async def restart_litellm():
164
165
166
167
168
169
170
171
172
173
    """
    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
174
175

        asyncio.create_task(start_litellm_background())
176
177
178
179
180
181
182
        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
183
        log.info(f"Error restarting litellm service: {e}")
184
185
186
187
188
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e)
        )


Timothy J. Baek's avatar
Timothy J. Baek committed
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
@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

205
206
    model_config = ConfigDict(protected_namespaces=())

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

@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
219
220
221
@app.get("/models")
@app.get("/v1/models")
async def get_models(user=Depends(get_current_user)):
222

Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
    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
242
                    )
243

Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
            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",
                    }
                    for model in app.state.CONFIG["model_list"]
                ],
                "object": "list",
            }
    else:
270
        return {
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
271
            "data": [],
272
273
            "object": "list",
        }
274
275


276
277
278
279
280
281
282
283
284
@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

285
286
    model_config = ConfigDict(protected_namespaces=())

287
288
289
290
291

@app.post("/model/new")
async def add_model_to_config(
    form_data: AddLiteLLMModelForm, user=Depends(get_admin_user)
):
292
293
294
    try:
        get_llm_provider(model=form_data.model_name)
        app.state.CONFIG["model_list"].append(form_data.model_dump())
295

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

299
        await restart_litellm()
300

301
302
303
304
305
306
        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)
        )
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327


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()

328
    return {"message": MESSAGES.MODEL_DELETED(form_data.id)}
329
330


Timothy J. Baek's avatar
Timothy J. Baek committed
331
332
333
@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()
334

335
    url = f"http://localhost:{LITELLM_PROXY_PORT}"
336

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

Timothy J. Baek's avatar
Timothy J. Baek committed
339
340
341
    headers = {}
    # headers["Authorization"] = f"Bearer {key}"
    headers["Content-Type"] = "application/json"
342

Timothy J. Baek's avatar
Timothy J. Baek committed
343
    r = None
344

Timothy J. Baek's avatar
Timothy J. Baek committed
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
    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
        )