test: fix test

This commit is contained in:
Krrish Dholakia 2024-05-17 22:35:34 -07:00
parent 1cecdc4690
commit 5d3fe52a08

View file

@ -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)