main.py 10.1 KB
Newer Older
1
2
import sys

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

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

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

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

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

24
25
import os

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

29

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

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

warnings.simplefilter("ignore")

43
from litellm.utils import get_llm_provider
44

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

49
app = FastAPI()
Timothy J. Baek's avatar
Timothy J. Baek committed
50

51
origins = ["*"]
Timothy J. Baek's avatar
Timothy J. Baek committed
52

53
54
55
56
57
58
59
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
Timothy J. Baek's avatar
Timothy J. Baek committed
60

61

Timothy J. Baek's avatar
Timothy J. Baek committed
62
63
64
65
66
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
67
68

app.state.ENABLE = ENABLE_LITELLM
Timothy J. Baek's avatar
Timothy J. Baek committed
69
70
app.state.CONFIG = litellm_config

71
72
73
# Global variable to store the subprocess reference
background_process = None

74
75
76
77
78
79
80
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
81

82
83
async def run_background_process(command):
    global background_process
Timothy J. Baek's avatar
Timothy J. Baek committed
84
    log.info("run_background_process")
85
86
87

    try:
        # Log the command to be executed
Timothy J. Baek's avatar
Timothy J. Baek committed
88
        log.info(f"Executing command: {command}")
89
90
        # 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}
91
92
        # Execute the command and create a subprocess
        process = await asyncio.create_subprocess_exec(
93
            *command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env
94
95
        )
        background_process = process
Timothy J. Baek's avatar
Timothy J. Baek committed
96
        log.info("Subprocess started successfully.")
97
98
99
100
101

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

Timothy J. Baek's avatar
Timothy J. Baek committed
104
        # log.info output line by line
105
        async for line in process.stdout:
Timothy J. Baek's avatar
Timothy J. Baek committed
106
            log.info(line.decode().strip())
107
108
109

        # Wait for the process to finish
        returncode = await process.wait()
Timothy J. Baek's avatar
Timothy J. Baek committed
110
        log.info(f"Subprocess exited with return code {returncode}")
111
112
113
    except Exception as e:
        log.error(f"Failed to start subprocess: {e}")
        raise  # Optionally re-raise the exception if you want it to propagate
114
115
116


async def start_litellm_background():
Timothy J. Baek's avatar
Timothy J. Baek committed
117
    log.info("start_litellm_background")
118
    # Command to run in the background
119
120
121
122
123
124
125
126
127
128
129
    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
130

131
    await run_background_process(command)
Timothy J. Baek's avatar
Timothy J. Baek committed
132
133


134
async def shutdown_litellm_background():
Timothy J. Baek's avatar
Timothy J. Baek committed
135
    log.info("shutdown_litellm_background")
136
137
138
139
    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
140
        log.info("Subprocess terminated")
141
        background_process = None
142
143


Timothy J. Baek's avatar
Timothy J. Baek committed
144
@app.on_event("startup")
145
async def startup_event():
Timothy J. Baek's avatar
Timothy J. Baek committed
146
    log.info("startup_event")
Timothy J. Baek's avatar
Timothy J. Baek committed
147
    # TODO: Check config.yaml file and create one
148
    asyncio.create_task(start_litellm_background())
Timothy J. Baek's avatar
Timothy J. Baek committed
149
150


Timothy J. Baek's avatar
Timothy J. Baek committed
151
app.state.ENABLE_MODEL_FILTER = ENABLE_MODEL_FILTER
152
153
154
app.state.MODEL_FILTER_LIST = MODEL_FILTER_LIST


155
156
157
158
159
@app.get("/")
async def get_status():
    return {"status": True}


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

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


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

202
203
    model_config = ConfigDict(protected_namespaces=())

Timothy J. Baek's avatar
Timothy J. Baek committed
204
205
206
207
208
209
210
211
212
213
214
215

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

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

Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
241
242
243
244
245
246
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",
                    }
                    for model in app.state.CONFIG["model_list"]
                ],
                "object": "list",
            }
    else:
267
        return {
Timothy J. Baek's avatar
fix  
Timothy J. Baek committed
268
            "data": [],
269
270
            "object": "list",
        }
271
272


273
274
275
276
277
278
279
280
281
@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

282
283
    model_config = ConfigDict(protected_namespaces=())

284
285
286
287
288

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

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

296
        await restart_litellm()
297

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


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

325
    return {"message": MESSAGES.MODEL_DELETED(form_data.id)}
326
327


Timothy J. Baek's avatar
Timothy J. Baek committed
328
329
330
@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()
331

332
    url = f"http://localhost:{LITELLM_PROXY_PORT}"
333

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

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

Timothy J. Baek's avatar
Timothy J. Baek committed
340
    r = None
341

Timothy J. Baek's avatar
Timothy J. Baek committed
342
343
344
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
    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
        )