From aa57c63091a26f4be8b2d30e2ce4aba11cecb5a6 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Wed, 25 Oct 2023 13:33:24 -0700
Subject: [PATCH] build(openai_proxy): docker build fixes

---
 .gitignore                  |  3 ++-
 Dockerfile                  |  7 ++++++-
 openai_proxy/main.py        |  5 ++++-
 openai_proxy/utils.py       | 18 +++++++++++++++++-
 requirements.txt            |  1 +
 router_config_template.yaml |  8 ++++----
 6 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/.gitignore b/.gitignore
index 47fccd9c5..313241e4c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -11,4 +11,5 @@ secrets.toml
 .gitignore
 litellm/proxy/litellm_secrets.toml
 litellm/proxy/api_log.json
-.idea/
\ No newline at end of file
+.idea/
+router_config.yaml
diff --git a/Dockerfile b/Dockerfile
index 5b634c16b..de6c4a59d 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,6 +1,11 @@
 FROM python:3.10
 
-ENV LITELLM_CONFIG_PATH="/litellm.secrets.toml"
+# Define a build argument for the config file path
+ARG CONFIG_FILE
+
+# Copy the custom config file (if provided) into the Docker image
+COPY $CONFIG_FILE /app/config.yaml
+
 COPY . /app
 WORKDIR /app
 RUN pip install -r requirements.txt
diff --git a/openai_proxy/main.py b/openai_proxy/main.py
index 2ee191baa..1782e28c1 100644
--- a/openai_proxy/main.py
+++ b/openai_proxy/main.py
@@ -6,7 +6,10 @@ from fastapi.middleware.cors import CORSMiddleware
 import json
 import os
 from typing import Optional
-from utils import set_callbacks, load_router_config
+try:
+    from utils import set_callbacks, load_router_config
+except ImportError:
+    from openai_proxy.utils import set_callbacks, load_router_config
 import dotenv
 dotenv.load_dotenv() # load env variables
 
diff --git a/openai_proxy/utils.py b/openai_proxy/utils.py
index 2f31c10c3..be2ac3275 100644
--- a/openai_proxy/utils.py
+++ b/openai_proxy/utils.py
@@ -1,4 +1,20 @@
 import os, litellm
+import pkg_resources
+
+def get_package_version(package_name):
+    try:
+        package = pkg_resources.get_distribution(package_name)
+        return package.version
+    except pkg_resources.DistributionNotFound:
+        return None
+
+# Usage example
+package_name = "litellm"
+version = get_package_version(package_name)
+if version:
+    print(f"The version of {package_name} is {version}")
+else:
+    print(f"{package_name} is not installed")
 import yaml
 import dotenv
 from typing import Optional
@@ -29,7 +45,7 @@ def set_callbacks():
 
 def load_router_config(router: Optional[litellm.Router]):
     config = {}
-    config_file = 'config.yaml'
+    config_file = '/app/config.yaml'
 
     if os.path.exists(config_file):
         with open(config_file, 'r') as file:
diff --git a/requirements.txt b/requirements.txt
index 08d063f15..6591bb852 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,4 @@
+litellm
 openai
 fastapi
 uvicorn
diff --git a/router_config_template.yaml b/router_config_template.yaml
index 0fb951ba7..ed6e7adeb 100644
--- a/router_config_template.yaml
+++ b/router_config_template.yaml
@@ -1,7 +1,7 @@
 model_list:
   - model_name: gpt-3.5-turbo
     litellm_params:
-      model: azure/chatgpt-v-2
+      model: azure/chatgpt-v-2 # azure/<your-deployment-name>
       api_key: your_azure_api_key
       api_version: your_azure_api_version
       api_base: your_azure_api_base
@@ -23,6 +23,6 @@ model_list:
     rpm: 9000 # REPLACE with your openai rpm
 
 environment_variables:
-	REDIS_HOST: your_redis_host
-	REDIS_PASSWORD: your_redis_password
-	REDIS_PORT: your_redis_port 
\ No newline at end of file
+  REDIS_HOST: your_redis_host
+  REDIS_PASSWORD: your_redis_password
+  REDIS_PORT: your_redis_port 
\ No newline at end of file