main.py 9.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362
  1. import sys
  2. from fastapi import FastAPI, Depends, HTTPException
  3. from fastapi.routing import APIRoute
  4. from fastapi.middleware.cors import CORSMiddleware
  5. import logging
  6. from fastapi import FastAPI, Request, Depends, status, Response
  7. from fastapi.responses import JSONResponse
  8. from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
  9. from starlette.responses import StreamingResponse
  10. import json
  11. import time
  12. import requests
  13. from pydantic import BaseModel, ConfigDict
  14. from typing import Optional, List
  15. from utils.utils import get_verified_user, get_current_user, get_admin_user
  16. from config import SRC_LOG_LEVELS, ENV
  17. from constants import MESSAGES
  18. import os
  19. log = logging.getLogger(__name__)
  20. log.setLevel(SRC_LOG_LEVELS["LITELLM"])
  21. from config import (
  22. ENABLE_MODEL_FILTER,
  23. MODEL_FILTER_LIST,
  24. DATA_DIR,
  25. LITELLM_PROXY_PORT,
  26. LITELLM_PROXY_HOST,
  27. )
  28. from litellm.utils import get_llm_provider
  29. import asyncio
  30. import subprocess
  31. import yaml
  32. app = FastAPI()
  33. origins = ["*"]
  34. app.add_middleware(
  35. CORSMiddleware,
  36. allow_origins=origins,
  37. allow_credentials=True,
  38. allow_methods=["*"],
  39. allow_headers=["*"],
  40. )
  41. LITELLM_CONFIG_DIR = f"{DATA_DIR}/litellm/config.yaml"
  42. with open(LITELLM_CONFIG_DIR, "r") as file:
  43. litellm_config = yaml.safe_load(file)
  44. app.state.CONFIG = litellm_config
  45. # Global variable to store the subprocess reference
  46. background_process = None
  47. CONFLICT_ENV_VARS = [
  48. # Uvicorn uses PORT, so LiteLLM might use it as well
  49. "PORT",
  50. # LiteLLM uses DATABASE_URL for Prisma connections
  51. "DATABASE_URL",
  52. ]
  53. async def run_background_process(command):
  54. global background_process
  55. log.info("run_background_process")
  56. try:
  57. # Log the command to be executed
  58. log.info(f"Executing command: {command}")
  59. # Filter environment variables known to conflict with litellm
  60. env = {k: v for k, v in os.environ.items() if k not in CONFLICT_ENV_VARS}
  61. # Execute the command and create a subprocess
  62. process = await asyncio.create_subprocess_exec(
  63. *command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env
  64. )
  65. background_process = process
  66. log.info("Subprocess started successfully.")
  67. # Capture STDERR for debugging purposes
  68. stderr_output = await process.stderr.read()
  69. stderr_text = stderr_output.decode().strip()
  70. if stderr_text:
  71. log.info(f"Subprocess STDERR: {stderr_text}")
  72. # log.info output line by line
  73. async for line in process.stdout:
  74. log.info(line.decode().strip())
  75. # Wait for the process to finish
  76. returncode = await process.wait()
  77. log.info(f"Subprocess exited with return code {returncode}")
  78. except Exception as e:
  79. log.error(f"Failed to start subprocess: {e}")
  80. raise # Optionally re-raise the exception if you want it to propagate
  81. async def start_litellm_background():
  82. log.info("start_litellm_background")
  83. # Command to run in the background
  84. command = [
  85. "litellm",
  86. "--port",
  87. str(LITELLM_PROXY_PORT),
  88. "--host",
  89. LITELLM_PROXY_HOST,
  90. "--telemetry",
  91. "False",
  92. "--config",
  93. LITELLM_CONFIG_DIR,
  94. ]
  95. await run_background_process(command)
  96. async def shutdown_litellm_background():
  97. log.info("shutdown_litellm_background")
  98. global background_process
  99. if background_process:
  100. background_process.terminate()
  101. await background_process.wait() # Ensure the process has terminated
  102. log.info("Subprocess terminated")
  103. background_process = None
  104. @app.on_event("startup")
  105. async def startup_event():
  106. log.info("startup_event")
  107. # TODO: Check config.yaml file and create one
  108. asyncio.create_task(start_litellm_background())
  109. app.state.ENABLE_MODEL_FILTER = ENABLE_MODEL_FILTER
  110. app.state.MODEL_FILTER_LIST = MODEL_FILTER_LIST
  111. @app.get("/")
  112. async def get_status():
  113. return {"status": True}
  114. async def restart_litellm():
  115. """
  116. Endpoint to restart the litellm background service.
  117. """
  118. log.info("Requested restart of litellm service.")
  119. try:
  120. # Shut down the existing process if it is running
  121. await shutdown_litellm_background()
  122. log.info("litellm service shutdown complete.")
  123. # Restart the background service
  124. asyncio.create_task(start_litellm_background())
  125. log.info("litellm service restart complete.")
  126. return {
  127. "status": "success",
  128. "message": "litellm service restarted successfully.",
  129. }
  130. except Exception as e:
  131. log.info(f"Error restarting litellm service: {e}")
  132. raise HTTPException(
  133. status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e)
  134. )
  135. @app.get("/restart")
  136. async def restart_litellm_handler(user=Depends(get_admin_user)):
  137. return await restart_litellm()
  138. @app.get("/config")
  139. async def get_config(user=Depends(get_admin_user)):
  140. return app.state.CONFIG
  141. class LiteLLMConfigForm(BaseModel):
  142. general_settings: Optional[dict] = None
  143. litellm_settings: Optional[dict] = None
  144. model_list: Optional[List[dict]] = None
  145. router_settings: Optional[dict] = None
  146. model_config = ConfigDict(protected_namespaces=())
  147. @app.post("/config/update")
  148. async def update_config(form_data: LiteLLMConfigForm, user=Depends(get_admin_user)):
  149. app.state.CONFIG = form_data.model_dump(exclude_none=True)
  150. with open(LITELLM_CONFIG_DIR, "w") as file:
  151. yaml.dump(app.state.CONFIG, file)
  152. await restart_litellm()
  153. return app.state.CONFIG
  154. @app.get("/models")
  155. @app.get("/v1/models")
  156. async def get_models(user=Depends(get_current_user)):
  157. while not background_process:
  158. await asyncio.sleep(0.1)
  159. url = f"http://localhost:{LITELLM_PROXY_PORT}/v1"
  160. r = None
  161. try:
  162. r = requests.request(method="GET", url=f"{url}/models")
  163. r.raise_for_status()
  164. data = r.json()
  165. if app.state.ENABLE_MODEL_FILTER:
  166. if user and user.role == "user":
  167. data["data"] = list(
  168. filter(
  169. lambda model: model["id"] in app.state.MODEL_FILTER_LIST,
  170. data["data"],
  171. )
  172. )
  173. return data
  174. except Exception as e:
  175. log.exception(e)
  176. error_detail = "Open WebUI: Server Connection Error"
  177. if r is not None:
  178. try:
  179. res = r.json()
  180. if "error" in res:
  181. error_detail = f"External: {res['error']}"
  182. except:
  183. error_detail = f"External: {e}"
  184. return {
  185. "data": [
  186. {
  187. "id": model["model_name"],
  188. "object": "model",
  189. "created": int(time.time()),
  190. "owned_by": "openai",
  191. }
  192. for model in app.state.CONFIG["model_list"]
  193. ],
  194. "object": "list",
  195. }
  196. @app.get("/model/info")
  197. async def get_model_list(user=Depends(get_admin_user)):
  198. return {"data": app.state.CONFIG["model_list"]}
  199. class AddLiteLLMModelForm(BaseModel):
  200. model_name: str
  201. litellm_params: dict
  202. model_config = ConfigDict(protected_namespaces=())
  203. @app.post("/model/new")
  204. async def add_model_to_config(
  205. form_data: AddLiteLLMModelForm, user=Depends(get_admin_user)
  206. ):
  207. try:
  208. get_llm_provider(model=form_data.model_name)
  209. app.state.CONFIG["model_list"].append(form_data.model_dump())
  210. with open(LITELLM_CONFIG_DIR, "w") as file:
  211. yaml.dump(app.state.CONFIG, file)
  212. await restart_litellm()
  213. return {"message": MESSAGES.MODEL_ADDED(form_data.model_name)}
  214. except Exception as e:
  215. print(e)
  216. raise HTTPException(
  217. status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e)
  218. )
  219. class DeleteLiteLLMModelForm(BaseModel):
  220. id: str
  221. @app.post("/model/delete")
  222. async def delete_model_from_config(
  223. form_data: DeleteLiteLLMModelForm, user=Depends(get_admin_user)
  224. ):
  225. app.state.CONFIG["model_list"] = [
  226. model
  227. for model in app.state.CONFIG["model_list"]
  228. if model["model_name"] != form_data.id
  229. ]
  230. with open(LITELLM_CONFIG_DIR, "w") as file:
  231. yaml.dump(app.state.CONFIG, file)
  232. await restart_litellm()
  233. return {"message": MESSAGES.MODEL_DELETED(form_data.id)}
  234. @app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
  235. async def proxy(path: str, request: Request, user=Depends(get_verified_user)):
  236. body = await request.body()
  237. url = f"http://localhost:{LITELLM_PROXY_PORT}"
  238. target_url = f"{url}/{path}"
  239. headers = {}
  240. # headers["Authorization"] = f"Bearer {key}"
  241. headers["Content-Type"] = "application/json"
  242. r = None
  243. try:
  244. r = requests.request(
  245. method=request.method,
  246. url=target_url,
  247. data=body,
  248. headers=headers,
  249. stream=True,
  250. )
  251. r.raise_for_status()
  252. # Check if response is SSE
  253. if "text/event-stream" in r.headers.get("Content-Type", ""):
  254. return StreamingResponse(
  255. r.iter_content(chunk_size=8192),
  256. status_code=r.status_code,
  257. headers=dict(r.headers),
  258. )
  259. else:
  260. response_data = r.json()
  261. return response_data
  262. except Exception as e:
  263. log.exception(e)
  264. error_detail = "Open WebUI: Server Connection Error"
  265. if r is not None:
  266. try:
  267. res = r.json()
  268. if "error" in res:
  269. error_detail = f"External: {res['error']['message'] if 'message' in res['error'] else res['error']}"
  270. except:
  271. error_detail = f"External: {e}"
  272. raise HTTPException(
  273. status_code=r.status_code if r else 500, detail=error_detail
  274. )