From d011a4b5eac967be4cc4e36df1167d48c7a6d511 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Fri, 27 Oct 2023 11:42:51 -0700
Subject: [PATCH] build(litellm_server/main.py): fix config loading

---
 litellm_server/main.py  | 15 ++++++++++++---
 litellm_server/utils.py |  7 +++----
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/litellm_server/main.py b/litellm_server/main.py
index cc29d96ce..8593631c3 100644
--- a/litellm_server/main.py
+++ b/litellm_server/main.py
@@ -29,7 +29,11 @@ llm_router: Optional[litellm.Router] = None
 llm_model_list: Optional[list] = None
 
 set_callbacks() # sets litellm callbacks for logging if they exist in the environment 
-llm_router, llm_model_list = load_router_config(router=llm_router)
+
+if "CONFIG_FILE_PATH" in os.environ:
+    llm_router, llm_model_list = load_router_config(router=llm_router, config_file_path=os.getenv("CONFIG_FILE_PATH"))
+else:
+    llm_router, llm_model_list = load_router_config(router=llm_router)
 #### API ENDPOINTS ####
 @router.get("/v1/models")
 @router.get("/models")  # if project requires model list
@@ -103,7 +107,12 @@ async def chat_completion(request: Request):
                 data["api_key"] = api_key
         ## CHECK CONFIG ## 
         if llm_model_list and data["model"] in [m["model_name"] for m in llm_model_list]:
-             return await router_completion(request=request)
+            for m in llm_model_list: 
+                if data["model"] == m["model_name"]: 
+                    for key, value in m["litellm_params"].items(): 
+                        data[key] = value
+                    break
+        print(f"data going into litellm completion: {data}")
         response = litellm.completion(
             **data
         )
@@ -164,4 +173,4 @@ async def home(request: Request):
     return "LiteLLM: RUNNING"
 
 
-app.include_router(router)
+app.include_router(router)
\ No newline at end of file
diff --git a/litellm_server/utils.py b/litellm_server/utils.py
index 8dee3df03..359f4ab5f 100644
--- a/litellm_server/utils.py
+++ b/litellm_server/utils.py
@@ -43,13 +43,12 @@ def set_callbacks():
 
 
 
-def load_router_config(router: Optional[litellm.Router]):
+def load_router_config(router: Optional[litellm.Router], config_file_path: Optional[str]='/app/config.yaml'):
     config = {}
-    config_file = '/app/config.yaml'
 
     try: 
-        if os.path.exists(config_file):
-            with open(config_file, 'r') as file:
+        if os.path.exists(config_file_path):
+            with open(config_file_path, 'r') as file:
                 config = yaml.safe_load(file)
         else:
             pass