From 14dd0d11fb6d670c0e6d3c430541810c484dcdb1 Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Sat, 24 Feb 2024 00:21:53 -0800
Subject: [PATCH] feat: litellm yaml

---
 backend/.gitignore |  2 ++
 backend/main.py    | 25 +++++++++++++++++++++++--
 test.json          |  6 ++++++
 3 files changed, 31 insertions(+), 2 deletions(-)
 create mode 100644 test.json

diff --git a/backend/.gitignore b/backend/.gitignore
index 16180123c..1a2bc493f 100644
--- a/backend/.gitignore
+++ b/backend/.gitignore
@@ -8,4 +8,6 @@ _test
 Pipfile
 data/*
 !data/config.json
+!data/litellm/config.yaml
+
 .webui_secret_key
\ No newline at end of file
diff --git a/backend/main.py b/backend/main.py
index 94145a975..a432e9ed6 100644
--- a/backend/main.py
+++ b/backend/main.py
@@ -2,7 +2,8 @@ from bs4 import BeautifulSoup
 import json
 import markdown
 import time
-
+import os
+import sys
 
 from fastapi import FastAPI, Request, Depends
 from fastapi.staticfiles import StaticFiles
@@ -13,6 +14,7 @@ from fastapi.middleware.cors import CORSMiddleware
 from starlette.exceptions import HTTPException as StarletteHTTPException
 
 
+from litellm.proxy.proxy_server import ProxyConfig, initialize
 from litellm.proxy.proxy_server import app as litellm_app
 
 from apps.ollama.main import app as ollama_app
@@ -38,6 +40,21 @@ class SPAStaticFiles(StaticFiles):
                 raise ex
 
 
+proxy_config = ProxyConfig()
+
+
+async def config():
+    router, model_list, general_settings = await proxy_config.load_config(
+        router=None, config_file_path="./data/litellm/config.yaml"
+    )
+
+    await initialize(config="./data/litellm/config.yaml", telemetry=False)
+
+
+async def startup():
+    await config()
+
+
 app = FastAPI(docs_url="/docs" if ENV == "dev" else None, redoc_url=None)
 
 origins = ["*"]
@@ -51,6 +68,11 @@ app.add_middleware(
 )
 
 
+@app.on_event("startup")
+async def on_startup():
+    await startup()
+
+
 @app.middleware("http")
 async def check_url(request: Request, call_next):
     start_time = int(time.time())
@@ -79,7 +101,6 @@ async def auth_middleware(request: Request, call_next):
 app.mount("/api/v1", webui_app)
 app.mount("/litellm/api", litellm_app)
 
-
 app.mount("/ollama/api", ollama_app)
 app.mount("/openai/api", openai_app)
 
diff --git a/test.json b/test.json
new file mode 100644
index 000000000..dccb6b269
--- /dev/null
+++ b/test.json
@@ -0,0 +1,6 @@
+{
+    "model_name": "string",
+    "litellm_params": {
+        "model": "ollama/mistral"
+    }
+}
\ No newline at end of file