From 2174801cb05afd44e9cd9895e4d1158dae3ded64 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Sat, 10 Aug 2024 09:45:12 -0700
Subject: [PATCH] prometheus add basic testing for success

---
 litellm/tests/test_prometheus.py | 79 ++++++++++++++++++++++++++++++++
 1 file changed, 79 insertions(+)
 create mode 100644 litellm/tests/test_prometheus.py

diff --git a/litellm/tests/test_prometheus.py b/litellm/tests/test_prometheus.py
new file mode 100644
index 0000000000..73281da728
--- /dev/null
+++ b/litellm/tests/test_prometheus.py
@@ -0,0 +1,79 @@
+import io
+import os
+import sys
+
+sys.path.insert(0, os.path.abspath("../.."))
+
+import asyncio
+import logging
+import uuid
+
+import pytest
+from prometheus_client import REGISTRY
+
+import litellm
+from litellm import completion
+from litellm._logging import verbose_logger
+from litellm.integrations.prometheus import PrometheusLogger
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
+
+verbose_logger.setLevel(logging.DEBUG)
+
+litellm.set_verbose = True
+import time
+
+
+@pytest.mark.asyncio()
+async def test_async_prometheus_success_logging():
+    run_id = str(uuid.uuid4())
+    litellm.set_verbose = True
+    litellm.success_callback = ["prometheus"]
+    litellm.failure_callback = ["prometheus"]
+
+    response = await litellm.acompletion(
+        model="claude-instant-1.2",
+        messages=[{"role": "user", "content": "what llm are u"}],
+        max_tokens=10,
+        mock_response="hi",
+        temperature=0.2,
+        metadata={
+            "id": run_id,
+            "tags": ["tag1", "tag2"],
+            "user_api_key": "6eb81e014497d89f3cc1aa9da7c2b37bda6b7fea68e4b710d33d94201e68970c",
+            "user_api_key_alias": "ishaans-prometheus-key",
+            "user_api_end_user_max_budget": None,
+            "litellm_api_version": "1.40.19",
+            "global_max_parallel_requests": None,
+            "user_api_key_user_id": "admin",
+            "user_api_key_org_id": None,
+            "user_api_key_team_id": "dbe2f686-a686-4896-864a-4c3924458709",
+            "user_api_key_team_alias": "testing-team",
+        },
+    )
+    print(response)
+    await asyncio.sleep(3)
+
+    # get prometheus logger
+    from litellm.litellm_core_utils.litellm_logging import _in_memory_loggers
+
+    for callback in _in_memory_loggers:
+        if isinstance(callback, PrometheusLogger):
+            test_prometheus_logger = callback
+
+    print("done with success request")
+
+    print(
+        "vars of test_prometheus_logger",
+        vars(test_prometheus_logger.litellm_requests_metric),
+    )
+
+    # Get the metrics
+    metrics = {}
+    for metric in REGISTRY.collect():
+        for sample in metric.samples:
+            metrics[sample.name] = sample.value
+
+    print("metrics from prometheus", metrics)
+
+    assert metrics["litellm_requests_metric_total"] == 1.0
+    assert metrics["litellm_total_tokens_total"] == 30.0