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

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

27

28
29
30
31
32
from config import (
    MODEL_FILTER_ENABLED,
    MODEL_FILTER_LIST,
    DATA_DIR,
    LITELLM_PROXY_PORT,
33
    LITELLM_PROXY_HOST,
34
)
35

36
from litellm.utils import get_llm_provider
37

38
39
import asyncio
import subprocess
Timothy J. Baek's avatar
Timothy J. Baek committed
40
import yaml
Timothy J. Baek's avatar
Timothy J. Baek committed
41

42
app = FastAPI()
Timothy J. Baek's avatar
Timothy J. Baek committed
43

44
origins = ["*"]
Timothy J. Baek's avatar
Timothy J. Baek committed
45

46
47
48
49
50
51
52
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
Timothy J. Baek's avatar
Timothy J. Baek committed
53

54

Timothy J. Baek's avatar
Timothy J. Baek committed
55
56
57
58
59
60
61
LITELLM_CONFIG_DIR = f"{DATA_DIR}/litellm/config.yaml"

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

app.state.CONFIG = litellm_config

62
63
64
# Global variable to store the subprocess reference
background_process = None

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

66
67
async def run_background_process(command):
    global background_process
Timothy J. Baek's avatar
Timothy J. Baek committed
68
    log.info("run_background_process")
69
70
71

    try:
        # Log the command to be executed
Timothy J. Baek's avatar
Timothy J. Baek committed
72
        log.info(f"Executing command: {command}")
73
74
        # Execute the command and create a subprocess
        process = await asyncio.create_subprocess_exec(
75
            *command, stdout=subprocess.PIPE, stderr=subprocess.PIPE
76
77
        )
        background_process = process
Timothy J. Baek's avatar
Timothy J. Baek committed
78
        log.info("Subprocess started successfully.")
79
80
81
82
83

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

Timothy J. Baek's avatar
Timothy J. Baek committed
86
        # log.info output line by line
87
        async for line in process.stdout:
Timothy J. Baek's avatar
Timothy J. Baek committed
88
            log.info(line.decode().strip())
89
90
91

        # Wait for the process to finish
        returncode = await process.wait()
Timothy J. Baek's avatar
Timothy J. Baek committed
92
        log.info(f"Subprocess exited with return code {returncode}")
93
94
95
    except Exception as e:
        log.error(f"Failed to start subprocess: {e}")
        raise  # Optionally re-raise the exception if you want it to propagate
96
97
98


async def start_litellm_background():
Timothy J. Baek's avatar
Timothy J. Baek committed
99
    log.info("start_litellm_background")
100
    # Command to run in the background
101
102
103
104
105
106
107
108
109
110
111
    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
112

113
    await run_background_process(command)
Timothy J. Baek's avatar
Timothy J. Baek committed
114
115


116
async def shutdown_litellm_background():
Timothy J. Baek's avatar
Timothy J. Baek committed
117
    log.info("shutdown_litellm_background")
118
119
120
121
    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
122
        log.info("Subprocess terminated")
123
        background_process = None
124
125


Timothy J. Baek's avatar
Timothy J. Baek committed
126
@app.on_event("startup")
127
async def startup_event():
Timothy J. Baek's avatar
Timothy J. Baek committed
128
    log.info("startup_event")
Timothy J. Baek's avatar
Timothy J. Baek committed
129
    # TODO: Check config.yaml file and create one
130
    asyncio.create_task(start_litellm_background())
Timothy J. Baek's avatar
Timothy J. Baek committed
131
132


133
134
135
136
app.state.MODEL_FILTER_ENABLED = MODEL_FILTER_ENABLED
app.state.MODEL_FILTER_LIST = MODEL_FILTER_LIST


137
138
139
140
141
@app.get("/")
async def get_status():
    return {"status": True}


Timothy J. Baek's avatar
Timothy J. Baek committed
142
async def restart_litellm():
143
144
145
146
147
148
149
150
151
152
    """
    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
153
154

        asyncio.create_task(start_litellm_background())
155
156
157
158
159
160
161
        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
162
        log.info(f"Error restarting litellm service: {e}")
163
164
165
166
167
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e)
        )


Timothy J. Baek's avatar
Timothy J. Baek committed
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
@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

184
185
    model_config = ConfigDict(protected_namespaces=())

Timothy J. Baek's avatar
Timothy J. Baek committed
186
187
188
189
190
191
192
193
194
195
196
197

@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
198
199
200
@app.get("/models")
@app.get("/v1/models")
async def get_models(user=Depends(get_current_user)):
201
202
203
    while not background_process:
        await asyncio.sleep(0.1)

204
    url = f"http://localhost:{LITELLM_PROXY_PORT}/v1"
Timothy J. Baek's avatar
Timothy J. Baek committed
205
206
207
208
    r = None
    try:
        r = requests.request(method="GET", url=f"{url}/models")
        r.raise_for_status()
209

Timothy J. Baek's avatar
Timothy J. Baek committed
210
        data = r.json()
211

Timothy J. Baek's avatar
Timothy J. Baek committed
212
213
214
215
216
217
218
219
        if app.state.MODEL_FILTER_ENABLED:
            if user and user.role == "user":
                data["data"] = list(
                    filter(
                        lambda model: model["id"] in app.state.MODEL_FILTER_LIST,
                        data["data"],
                    )
                )
220

Timothy J. Baek's avatar
Timothy J. Baek committed
221
222
        return data
    except Exception as e:
223

Timothy J. Baek's avatar
Timothy J. Baek committed
224
225
226
227
228
229
230
231
232
        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}"
233

234
235
236
237
238
239
240
241
242
243
244
245
        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",
        }
246
247


248
249
250
251
252
253
254
255
256
@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

257
258
    model_config = ConfigDict(protected_namespaces=())

259
260
261
262
263

@app.post("/model/new")
async def add_model_to_config(
    form_data: AddLiteLLMModelForm, user=Depends(get_admin_user)
):
264
265
266
    try:
        get_llm_provider(model=form_data.model_name)
        app.state.CONFIG["model_list"].append(form_data.model_dump())
267

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

271
        await restart_litellm()
272

273
274
275
276
277
278
        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)
        )
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299


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

300
    return {"message": MESSAGES.MODEL_DELETED(form_data.id)}
301
302


Timothy J. Baek's avatar
Timothy J. Baek committed
303
304
305
@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()
306

307
    url = f"http://localhost:{LITELLM_PROXY_PORT}"
308

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

Timothy J. Baek's avatar
Timothy J. Baek committed
311
312
313
    headers = {}
    # headers["Authorization"] = f"Bearer {key}"
    headers["Content-Type"] = "application/json"
314

Timothy J. Baek's avatar
Timothy J. Baek committed
315
    r = None
316

Timothy J. Baek's avatar
Timothy J. Baek committed
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
    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
        )