From 464d812fa9461e5eb90e0dd4e8abf1e6da52a22e Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Mon, 15 Apr 2024 13:13:51 -0700
Subject: [PATCH] test - team based logging

---
 tests/test_team_logging.py | 172 +++++++++++++++++++++++++++++++++++++
 1 file changed, 172 insertions(+)
 create mode 100644 tests/test_team_logging.py

diff --git a/tests/test_team_logging.py b/tests/test_team_logging.py
new file mode 100644
index 0000000000..a6c8fa015e
--- /dev/null
+++ b/tests/test_team_logging.py
@@ -0,0 +1,172 @@
+# What this tests ?
+## Tests /models and /model/* endpoints
+
+import pytest
+import asyncio
+import aiohttp
+import os
+import dotenv
+from dotenv import load_dotenv
+
+load_dotenv()
+
+
+async def generate_key(session, models=[], team_id=None):
+    url = "http://0.0.0.0:4000/key/generate"
+    headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
+    data = {
+        "models": models,
+        "duration": None,
+        "team_id": team_id,
+    }
+
+    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():
+    """
+    -> Team 1 logs to project 1
+    -> Create Key
+    -> Make chat/completions call
+    -> Fetch logs from langfuse
+    """
+    try:
+        async with aiohttp.ClientSession() as session:
+
+            key = await generate_key(
+                session, models=["fake-openai-endpoint"], team_id="team-1"
+            )  # team-1 logs to project 1
+
+            import uuid
+
+            _trace_id = f"trace-{uuid.uuid4()}"
+            _request_metadata = {
+                "trace_id": _trace_id,
+            }
+
+            await chat_completion(
+                session,
+                key["key"],
+                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_PROJECT1_PUBLIC"),
+                secret_key=os.getenv("LANGFUSE_PROJECT1_SECRET"),
+            )
+
+            await asyncio.sleep(10)
+
+            print(f"searching for trace_id={_trace_id} on langfuse")
+
+            generations = langfuse_client.get_generations(trace_id=_trace_id).data
+            print(generations)
+            assert len(generations) == 2
+    except Exception as e:
+        print(e)
+
+
+@pytest.mark.asyncio
+async def test_team_2logging():
+    """
+    -> Team 1 logs to project 2
+    -> Create Key
+    -> Make chat/completions call
+    -> Fetch logs from langfuse
+    """
+    try:
+        async with aiohttp.ClientSession() as session:
+
+            key = await generate_key(
+                session, models=["fake-openai-endpoint"], team_id="team-2"
+            )  # team-1 logs to project 1
+
+            import uuid
+
+            _trace_id = f"trace-{uuid.uuid4()}"
+            _request_metadata = {
+                "trace_id": _trace_id,
+            }
+
+            await chat_completion(
+                session,
+                key["key"],
+                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_PROJECT2_PUBLIC"),
+                secret_key=os.getenv("LANGFUSE_PROJECT2_SECRET"),
+            )
+
+            await asyncio.sleep(10)
+
+            print(f"searching for trace_id={_trace_id} on langfuse")
+
+            generations = langfuse_client.get_generations(trace_id=_trace_id).data
+            print("Team 2 generations", generations)
+
+            # team-2 should have 1 generation with this trace id
+            assert len(generations) == 1
+
+            # team-1 should have 0 generations with this trace id
+            langfuse_client_1 = langfuse.Langfuse(
+                public_key=os.getenv("LANGFUSE_PROJECT1_PUBLIC"),
+                secret_key=os.getenv("LANGFUSE_PROJECT1_SECRET"),
+            )
+
+            generations_team_1 = langfuse_client_1.get_generations(
+                trace_id=_trace_id
+            ).data
+            print("Team 1 generations", generations_team_1)
+
+            assert len(generations_team_1) == 0
+
+    except Exception as e:
+        print(e)