Merge pull request #1376 from BerriAI/litellm_deployed_proxy_prisma

[Test+Fix] Use deployed proxy with Prisma

.circleci/config.yml

@@ -36,7 +36,7 @@ jobs:
             pip install numpydoc
             pip install traceloop-sdk==0.0.69
             pip install openai
-            pip install prisma
+            pip install prisma            
             pip install "httpx==0.24.1"
             pip install "anyio==3.7.1"
             pip install "asyncio==3.4.3"
@@ -44,6 +44,13 @@ jobs:
           paths:
             - ./venv
           key: v1-dependencies-{{ checksum ".circleci/requirements.txt" }}
+      - run:
+          name: Run prisma ./entrypoint.sh
+          command: |
+            set +e
+            chmod +x entrypoint.sh
+            ./entrypoint.sh
+            set -e
       - run:
           name: Black Formatting
           command: |

Dockerfile

@@ -34,7 +34,6 @@ RUN pip wheel --no-cache-dir --wheel-dir=/wheels/ -r requirements.txt
 
 # Runtime stage
 FROM $LITELLM_RUNTIME_IMAGE as runtime
-ARG with_database
 
 WORKDIR /app
 # Copy the current directory contents into the container at /app
@@ -46,12 +45,14 @@ COPY --from=builder /app/dist/*.whl .
 COPY --from=builder /wheels/ /wheels/
 
 # Install the built wheel using pip; again using a wildcard if it's the only file
-RUN pip install --no-cache-dir --find-links=/wheels/ -r requirements.txt \
-    && pip install *.whl \
-    && rm -f *.whl
+RUN pip install *.whl /wheels/* --no-index --find-links=/wheels/ && rm -f *.whl && rm -rf /wheels
+
+RUN chmod +x entrypoint.sh
 
 EXPOSE 4000/tcp
 
 # Set your entrypoint and command
-ENTRYPOINT ["litellm"]
-CMD ["--port", "4000"]
+CMD [ \
+    "sh", "-c", \
+    "if [ -n \"$DATABASE_URL\" ]; then ./entrypoint.sh; else litellm --port 4000 --num_workers 8; fi" \
+]

litellm/proxy/proxy_server.py

@@ -311,6 +311,7 @@ def prisma_setup(database_url: Optional[str]):
                 database_url=database_url, proxy_logging_obj=proxy_logging_obj
             )
         except Exception as e:
+            raise e
             verbose_proxy_logger.debug(
                 f"Error when initializing prisma, Ensure you run pip install prisma {str(e)}"
             )
@@ -510,10 +511,9 @@ class ProxyConfig:
     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'
-            
+        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
@@ -776,7 +776,6 @@ class ProxyConfig:
                 verbose_proxy_logger.debug(f"GOING INTO LITELLM.GET_SECRET!")
                 database_url = litellm.get_secret(database_url)
                 verbose_proxy_logger.debug(f"RETRIEVED DB URL: {database_url}")
-            prisma_setup(database_url=database_url)
             ## COST TRACKING ##
             cost_tracking()
             ### MASTER KEY ###
@@ -1169,7 +1168,9 @@ async def startup_event():
                 llm_router,
                 llm_model_list,
                 general_settings,
-            ) = await proxy_config.load_config(router=llm_router, config_file_path=worker_config)
+            ) = await proxy_config.load_config(
+                router=llm_router, config_file_path=worker_config
+            )
         else:
             await initialize(**worker_config)
     else:

litellm/tests/test_configs/test_config.yaml

@@ -1,5 +1,5 @@
 general_settings:
-  database_url: os.environ/PROXY_DATABASE_URL
+  database_url: os.environ/DATABASE_URL
   master_key: os.environ/PROXY_MASTER_KEY
 litellm_settings:
   drop_params: true

litellm/tests/test_deployed_proxy_keygen.py

@@ -0,0 +1,53 @@
+import sys, os, time
+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, requests
+import litellm
+from litellm import embedding, completion, completion_cost, Timeout
+from litellm import RateLimitError
+
+
+def test_add_new_key():
+    try:
+        # Your test data
+        test_data = {
+            "models": ["gpt-3.5-turbo", "gpt-4", "claude-2", "azure-model"],
+            "aliases": {"mistral-7b": "gpt-3.5-turbo"},
+            "duration": "20m",
+        }
+        print("testing proxy server")
+        # Your bearer token
+        token = os.getenv("PROXY_MASTER_KEY")
+
+        headers = {"Authorization": f"Bearer {token}"}
+
+        staging_endpoint = "https://litellm-litellm-pr-1376.up.railway.app"
+
+        # Your bearer token
+        token = os.getenv("PROXY_MASTER_KEY")
+
+        headers = {"Authorization": f"Bearer {token}"}
+
+        # Make a request to the staging endpoint
+        response = requests.post(
+            staging_endpoint + "/key/generate", json=test_data, headers=headers
+        )
+
+        print(f"response: {response.text}")
+        assert response.status_code == 200
+        result = response.json()
+    except Exception as e:
+        print(traceback.format_exc())
+        pytest.fail(f"An error occurred {e}")
+
+
+# test_add_new_key()