From 8f8f961941f9e3ed0ae34e0865b4bc1f0f65a59a Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Mon, 8 Jan 2024 16:47:07 +0530
Subject: [PATCH 1/2] fix(proxy_server.py): add support for passing in config
 file via worker_config directly + testing

---
 litellm/proxy/proxy_server.py       | 17 ++++++++++-
 litellm/tests/test_proxy_startup.py | 45 +++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+), 1 deletion(-)
 create mode 100644 litellm/tests/test_proxy_startup.py

diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 6feba0bb526..1cbe4b887d7 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -507,6 +507,14 @@ class ProxyConfig:
     def __init__(self) -> None:
         pass
 
+    def is_yaml(self, config_file_path: str) -> bool:
+        if not os.path.isfile(config_file_path):
+            return False
+        
+        _, file_extension = os.path.splitext(config_file_path)
+        return file_extension.lower() == '.yaml' or file_extension.lower() == '.yml'
+            
+
     async def get_config(self, config_file_path: Optional[str] = None) -> dict:
         global prisma_client, user_config_file_path
 
@@ -1156,7 +1164,14 @@ async def startup_event():
     verbose_proxy_logger.debug(f"worker_config: {worker_config}")
     # check if it's a valid file path
     if os.path.isfile(worker_config):
-        await initialize(**worker_config)
+        if proxy_config.is_yaml(config_file_path=worker_config):
+            (
+                llm_router,
+                llm_model_list,
+                general_settings,
+            ) = await proxy_config.load_config(router=llm_router, config_file_path=worker_config)
+        else:
+            await initialize(**worker_config)
     else:
         # if not, assume it's a json string
         worker_config = json.loads(os.getenv("WORKER_CONFIG"))
diff --git a/litellm/tests/test_proxy_startup.py b/litellm/tests/test_proxy_startup.py
new file mode 100644
index 00000000000..12e888e7a24
--- /dev/null
+++ b/litellm/tests/test_proxy_startup.py
@@ -0,0 +1,45 @@
+# What this tests 
+## This tests the proxy server startup
+import sys, os, json
+import traceback
+from dotenv import load_dotenv
+
+load_dotenv()
+import os, io
+
+# this file is to test litellm/proxy
+
+sys.path.insert(
+    0, os.path.abspath("../..")
+)  # Adds the parent directory to the system path
+import pytest, logging, asyncio
+import litellm
+from litellm.proxy.proxy_server import (
+    router,
+    save_worker_config,
+    initialize,
+    startup_event,
+    llm_model_list
+)
+
+def test_proxy_gunicorn_startup():
+    """
+    gunicorn startup requires the config to be passed in via environment variables
+
+    We support saving either the config or the dict as an environment variable. 
+
+    Test both approaches
+    """
+
+    filepath = os.path.dirname(os.path.abspath(__file__))
+    # test with worker_config = config yaml 
+    config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml"
+    os.environ["WORKER_CONFIG"] = config_fp
+    asyncio.run(startup_event())
+    # test with worker_config = dict 
+    worker_config = {"config": config_fp}
+    os.environ["WORKER_CONFIG"] = json.dumps(worker_config)
+    asyncio.run(startup_event())
+
+
+# test_proxy_gunicorn_startup()
\ No newline at end of file

From e305dcf0a63428adecfc34067217af6b2a905a3f Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Mon, 8 Jan 2024 17:58:37 +0530
Subject: [PATCH 2/2] test(test_proxy_startup.py): separate tests

---
 litellm/tests/test_proxy_startup.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/litellm/tests/test_proxy_startup.py b/litellm/tests/test_proxy_startup.py
index 12e888e7a24..a24e68b39e2 100644
--- a/litellm/tests/test_proxy_startup.py
+++ b/litellm/tests/test_proxy_startup.py
@@ -22,7 +22,7 @@ from litellm.proxy.proxy_server import (
     llm_model_list
 )
 
-def test_proxy_gunicorn_startup():
+def test_proxy_gunicorn_startup_direct_config():
     """
     gunicorn startup requires the config to be passed in via environment variables
 
@@ -36,6 +36,11 @@ def test_proxy_gunicorn_startup():
     config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml"
     os.environ["WORKER_CONFIG"] = config_fp
     asyncio.run(startup_event())
+
+def test_proxy_gunicorn_startup_config_dict():
+    filepath = os.path.dirname(os.path.abspath(__file__))
+    # test with worker_config = config yaml 
+    config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml"
     # test with worker_config = dict 
     worker_config = {"config": config_fp}
     os.environ["WORKER_CONFIG"] = json.dumps(worker_config)