From d9539e518e2d4d82ea2b6ac737de19147790e5ea Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Wed, 24 Jul 2024 10:08:25 -0700
Subject: [PATCH] build(docker-compose.yml): add prometheus scraper to docker
 compose

persists prometheus data across restarts
---
 docker-compose.yml               | 20 ++++++++++++--
 litellm/tests/test_completion.py | 47 ++++++++++++++++++++++++++++++++
 prometheus.yml                   |  7 +++++
 3 files changed, 72 insertions(+), 2 deletions(-)
 create mode 100644 prometheus.yml

diff --git a/docker-compose.yml b/docker-compose.yml
index be84462ef..6991bf7eb 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -9,8 +9,6 @@ services:
     #########################################
     ## Uncomment these lines to start proxy with a config.yaml file ##
     # volumes:
-    #   - ./proxy_server_config.yaml:/app/config.yaml
-    # command: [ "--config", "./config.yaml", "--port", "4000"]
     ###############################################
     ports:
       - "4000:4000" # Map the container port to the host, change the host port if necessary
@@ -33,5 +31,23 @@ services:
       interval: 1s
       timeout: 5s
       retries: 10
+  
+  prometheus:
+    image: prom/prometheus
+    volumes:
+      - prometheus_data:/prometheus
+      - ./prometheus.yml:/etc/prometheus/prometheus.yml
+    ports:
+      - "9090:9090"
+    command:
+      - '--config.file=/etc/prometheus/prometheus.yml'
+      - '--storage.tsdb.path=/prometheus'
+      - '--storage.tsdb.retention.time=15d'
+    restart: always
+
+volumes:
+  prometheus_data:
+    driver: local
+
 
 # ...rest of your docker-compose config if any
diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py
index 31b7b8355..dae5e7f80 100644
--- a/litellm/tests/test_completion.py
+++ b/litellm/tests/test_completion.py
@@ -2560,6 +2560,53 @@ def test_completion_anyscale_with_functions():
 # test_completion_anyscale_with_functions()
 
 
+def test_completion_azure_extra_headers():
+    # this tests if we can pass api_key to completion, when it's not in the env.
+    # DO NOT REMOVE THIS TEST. No MATTER WHAT Happens!
+    # If you want to remove it, speak to Ishaan!
+    # Ishaan will be very disappointed if this test is removed -> this is a standard way to pass api_key + the router + proxy use this
+    from httpx import Client
+    from openai import AzureOpenAI
+
+    from litellm.llms.custom_httpx.httpx_handler import HTTPHandler
+
+    http_client = Client()
+
+    with patch.object(http_client, "send", new=MagicMock()) as mock_client:
+        client = AzureOpenAI(
+            azure_endpoint=os.getenv("AZURE_API_BASE"),
+            api_version=litellm.AZURE_DEFAULT_API_VERSION,
+            api_key=os.getenv("AZURE_API_KEY"),
+            http_client=http_client,
+        )
+        try:
+            response = completion(
+                model="azure/chatgpt-v-2",
+                messages=messages,
+                client=client,
+                extra_headers={
+                    "Authorization": "my-bad-key",
+                    "Ocp-Apim-Subscription-Key": "hello-world-testing",
+                    "api-key": "my-bad-key",
+                },
+            )
+            print(response)
+            pytest.fail("Expected this to fail")
+        except Exception as e:
+            pass
+
+        mock_client.assert_called()
+
+        print(f"mock_client.call_args: {mock_client.call_args}")
+        request = mock_client.call_args[0][0]
+        print(request.method)  # This will print 'POST'
+        print(request.url)  # This will print the full URL
+        print(request.headers)  # This will print the full URL
+        auth_header = request.headers.get("Authorization")
+        print(auth_header)
+        assert auth_header == "my-bad-key"
+
+
 def test_completion_azure_key_completion_arg():
     # this tests if we can pass api_key to completion, when it's not in the env.
     # DO NOT REMOVE THIS TEST. No MATTER WHAT Happens!
diff --git a/prometheus.yml b/prometheus.yml
new file mode 100644
index 000000000..5cb4f90d7
--- /dev/null
+++ b/prometheus.yml
@@ -0,0 +1,7 @@
+global:
+  scrape_interval: 15s
+
+scrape_configs:
+  - job_name: 'litellm'
+    static_configs:
+      - targets: ['litellm:4000']  # Assuming Litellm exposes metrics at port 4000