test - setting up langfuse callback on proxy

tests/test_adding_callbacks.py

@@ -0,0 +1,114 @@
+# What this tests ?
+## Tests /config/update + Test /chat/completions -> assert logs are sent to Langfuse
+
+import pytest
+import asyncio
+import aiohttp
+import os
+import dotenv
+from dotenv import load_dotenv
+import pytest
+
+load_dotenv()
+
+
+async def config_update(session):
+    url = "http://0.0.0.0:4000/config/update"
+    headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
+    data = {
+        "litellm_settings": {
+            "success_callback": ["langfuse"],
+        },
+        "environment_variables": {
+            "LANGFUSE_PUBLIC_KEY": os.environ["LANGFUSE_PUBLIC_KEY"],
+            "LANGFUSE_SECRET_KEY": os.environ["LANGFUSE_SECRET_KEY"],
+        },
+    }
+
+    async with session.post(url, headers=headers, json=data) as response:
+        status = response.status
+        response_text = await response.text()
+
+        print(response_text)
+        print()
+
+        if status != 200:
+            raise Exception(f"Request did not return a 200 status code: {status}")
+        return await response.json()
+
+
+async def chat_completion(session, key, model="azure-gpt-3.5", request_metadata=None):
+    url = "http://0.0.0.0:4000/chat/completions"
+    headers = {
+        "Authorization": f"Bearer {key}",
+        "Content-Type": "application/json",
+    }
+    data = {
+        "model": model,
+        "messages": [
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": "Hello!"},
+        ],
+        "metadata": request_metadata,
+    }
+
+    print("data sent in test=", data)
+
+    async with session.post(url, headers=headers, json=data) as response:
+        status = response.status
+        response_text = await response.text()
+
+        print(response_text)
+        print()
+
+        if status != 200:
+            raise Exception(f"Request did not return a 200 status code: {status}")
+
+
+@pytest.mark.asyncio
+async def test_team_logging():
+    """
+    1. Add Langfuse as a callback with /config/update
+    2. Call /chat/completions
+    3. Assert the logs are sent to Langfuse
+    """
+    try:
+        async with aiohttp.ClientSession() as session:
+
+            # Add Langfuse as a callback with /config/update
+            await config_update(session)
+
+            # 2. Call /chat/completions with a specific trace id
+            import uuid
+
+            _trace_id = f"trace-{uuid.uuid4()}"
+            _request_metadata = {
+                "trace_id": _trace_id,
+            }
+
+            await chat_completion(
+                session,
+                key="sk-1234",
+                model="fake-openai-endpoint",
+                request_metadata=_request_metadata,
+            )
+
+            # Test - if the logs were sent to the correct team on langfuse
+            import langfuse
+
+            langfuse_client = langfuse.Langfuse(
+                public_key=os.getenv("LANGFUSE_PUBLIC_KEY"),
+                secret_key=os.getenv("LANGFUSE_SECRET_KEY"),
+            )
+
+            await asyncio.sleep(10)
+
+            print(f"searching for trace_id={_trace_id} on langfuse")
+
+            generations = langfuse_client.get_generations(trace_id=_trace_id).data
+
+            # 1 generation with this trace id
+            assert len(generations) == 1
+
+    except Exception as e:
+        pytest.fail("Team 2 logging failed: " + str(e))