From 0b4737b13772aa6f40375dd35a38af51d7d2f59c Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Wed, 22 Jan 2025 09:20:46 -0800
Subject: [PATCH] test_langfuse_logging_completion_with_langfuse_metadata

---
 .../completion_with_langfuse_metadata.json    | 104 ++++++++++++++++++
 .../test_langfuse_e2e_test.py                 |  29 +++++
 2 files changed, 133 insertions(+)
 create mode 100644 tests/logging_callback_tests/langfuse_expected_request_body/completion_with_langfuse_metadata.json

diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_langfuse_metadata.json b/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_langfuse_metadata.json
new file mode 100644
index 00000000000..d4882c962d8
--- /dev/null
+++ b/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_langfuse_metadata.json
@@ -0,0 +1,104 @@
+{
+    "batch": [
+        {
+            "id": "872a0a1c-4328-431b-80b6-fd55a8a44477",
+            "type": "trace-create",
+            "body": {
+                "id": "litellm-test-533ffb2d-a0a3-45b5-911c-7940466cdc8e",
+                "timestamp": "2025-01-22T17:19:11.234960Z",
+                "name": "test_trace_name",
+                "userId": "test_user_id",
+                "input": {
+                    "messages": [
+                        {
+                            "role": "user",
+                            "content": "Hello!"
+                        }
+                    ]
+                },
+                "output": {
+                    "content": "Hello! How can I assist you today?",
+                    "role": "assistant",
+                    "tool_calls": null,
+                    "function_call": null
+                },
+                "sessionId": "test_session_id",
+                "version": "test_trace_version",
+                "metadata": {
+                    "test_key": "test_value"
+                },
+                "tags": [
+                    "test_tag",
+                    "test_tag_2"
+                ]
+            },
+            "timestamp": "2025-01-22T17:19:11.235169Z"
+        },
+        {
+            "id": "18d6f044-e522-4376-96e0-7eec765677ed",
+            "type": "generation-create",
+            "body": {
+                "traceId": "litellm-test-533ffb2d-a0a3-45b5-911c-7940466cdc8e",
+                "name": "test_generation_name",
+                "startTime": "2025-01-22T09:19:10.957072-08:00",
+                "metadata": {
+                    "tags": [
+                        "test_tag",
+                        "test_tag_2"
+                    ],
+                    "parent_observation_id": "test_parent_observation_id",
+                    "version": "test_version",
+                    "hidden_params": {
+                        "model_id": null,
+                        "cache_key": null,
+                        "api_base": "https://api.openai.com",
+                        "response_cost": 5.4999999999999995e-05,
+                        "additional_headers": {},
+                        "litellm_overhead_time_ms": null
+                    },
+                    "litellm_response_cost": 5.4999999999999995e-05,
+                    "cache_hit": false,
+                    "requester_metadata": {}
+                },
+                "input": {
+                    "messages": [
+                        {
+                            "role": "user",
+                            "content": "Hello!"
+                        }
+                    ]
+                },
+                "output": {
+                    "content": "Hello! How can I assist you today?",
+                    "role": "assistant",
+                    "tool_calls": null,
+                    "function_call": null
+                },
+                "level": "DEFAULT",
+                "parentObservationId": "test_parent_observation_id",
+                "version": "test_version",
+                "id": "time-09-19-10-957072_chatcmpl-4da65aba-32e4-400d-aaa2-6bfe096d8141",
+                "endTime": "2025-01-22T09:19:11.234200-08:00",
+                "completionStartTime": "2025-01-22T09:19:11.234200-08:00",
+                "model": "gpt-3.5-turbo",
+                "modelParameters": {
+                    "extra_body": "{}"
+                },
+                "usage": {
+                    "input": 10,
+                    "output": 20,
+                    "unit": "TOKENS",
+                    "totalCost": 5.4999999999999995e-05
+                }
+            },
+            "timestamp": "2025-01-22T17:19:11.235541Z"
+        }
+    ],
+    "metadata": {
+        "batch_size": 2,
+        "sdk_integration": "litellm",
+        "sdk_name": "python",
+        "sdk_version": "2.44.1",
+        "public_key": "pk-lf-e02aaea3-8668-4c9f-8c69-771a4ea1f5c9"
+    }
+}
\ No newline at end of file
diff --git a/tests/logging_callback_tests/test_langfuse_e2e_test.py b/tests/logging_callback_tests/test_langfuse_e2e_test.py
index 4fa6f7d34a1..3ba5a776bc1 100644
--- a/tests/logging_callback_tests/test_langfuse_e2e_test.py
+++ b/tests/logging_callback_tests/test_langfuse_e2e_test.py
@@ -216,3 +216,32 @@ class TestLangfuseLogging:
                 "completion_with_tags_stream.json",
                 setup["trace_id"],
             )
+
+    @pytest.mark.asyncio
+    async def test_langfuse_logging_completion_with_langfuse_metadata(self, mock_setup):
+        """Test Langfuse logging for chat completion with metadata for langfuse"""
+        setup = await mock_setup  # Await the fixture
+        with patch("httpx.Client.post", setup["mock_post"]):
+            await litellm.acompletion(
+                model="gpt-3.5-turbo",
+                messages=[{"role": "user", "content": "Hello!"}],
+                mock_response="Hello! How can I assist you today?",
+                metadata={
+                    "trace_id": setup["trace_id"],
+                    "tags": ["test_tag", "test_tag_2"],
+                    "generation_name": "test_generation_name",
+                    "parent_observation_id": "test_parent_observation_id",
+                    "version": "test_version",
+                    "trace_user_id": "test_user_id",
+                    "session_id": "test_session_id",
+                    "trace_name": "test_trace_name",
+                    "trace_metadata": {"test_key": "test_value"},
+                    "trace_version": "test_trace_version",
+                    "trace_release": "test_trace_release",
+                },
+            )
+            await self._verify_langfuse_call(
+                setup["mock_post"],
+                "completion_with_langfuse_metadata.json",
+                setup["trace_id"],
+            )