diff --git a/litellm/tests/test_alangfuse.py b/litellm/tests/test_alangfuse.py index 31f1f7bf867..97d6baaaeee 100644 --- a/litellm/tests/test_alangfuse.py +++ b/litellm/tests/test_alangfuse.py @@ -242,12 +242,24 @@ async def test_langfuse_masked_input_output(langfuse_client): response = await create_async_task( model="gpt-3.5-turbo", messages=[{"role": "user", "content": "This is a test"}], - metadata={"trace_id": _unique_trace_name, "mask_input": mask_value, "mask_output": mask_value}, - mock_response="This is a test response" + metadata={ + "trace_id": _unique_trace_name, + "mask_input": mask_value, + "mask_output": mask_value, + }, + mock_response="This is a test response", ) print(response) - expected_input = "redacted-by-litellm" if mask_value else {'messages': [{'content': 'This is a test', 'role': 'user'}]} - expected_output = "redacted-by-litellm" if mask_value else {'content': 'This is a test response', 'role': 'assistant'} + expected_input = ( + "redacted-by-litellm" + if mask_value + else {"messages": [{"content": "This is a test", "role": "user"}]} + ) + expected_output = ( + "redacted-by-litellm" + if mask_value + else {"content": "This is a test response", "role": "assistant"} + ) langfuse_client.flush() await asyncio.sleep(2) @@ -262,6 +274,7 @@ async def test_langfuse_masked_input_output(langfuse_client): assert generations[0].input == expected_input assert generations[0].output == expected_output + @pytest.mark.asyncio async def test_langfuse_logging_metadata(langfuse_client): """ @@ -523,7 +536,7 @@ def test_langfuse_logging_function_calling(): # test_langfuse_logging_function_calling() -def test_langfuse_existing_trace_id(): +def test_aaalangfuse_existing_trace_id(): """ When existing trace id is passed, don't set trace params -> prevents overwriting the trace @@ -577,7 +590,7 @@ def test_langfuse_existing_trace_id(): "verbose": False, "custom_llm_provider": "openai", "api_base": "https://api.openai.com/v1/", - "litellm_call_id": "508113a1-c6f1-48ce-a3e1-01c6cce9330e", + "litellm_call_id": None, "model_alias_map": {}, "completion_call_id": None, "metadata": None, @@ -593,7 +606,7 @@ def test_langfuse_existing_trace_id(): "stream": False, "user": None, "call_type": "completion", - "litellm_call_id": "508113a1-c6f1-48ce-a3e1-01c6cce9330e", + "litellm_call_id": None, "completion_start_time": "2024-05-01 07:31:29.903685", "temperature": 0.1, "extra_body": {}, @@ -633,6 +646,8 @@ def test_langfuse_existing_trace_id(): trace_id = langfuse_response_object["trace_id"] + assert trace_id is not None + langfuse_client.flush() time.sleep(2)