fix(langfuse.py): fix trace param overwriting when existing trace id is given

This commit is contained in:
Krrish Dholakia 2024-05-01 08:44:46 -07:00
parent d0f9f8c0ed
commit 0ab6b4bb22
2 changed files with 31 additions and 7 deletions

View file

@ -279,7 +279,7 @@ class LangFuseLogger:
trace_name = f"litellm-{kwargs.get('call_type', 'completion')}" trace_name = f"litellm-{kwargs.get('call_type', 'completion')}"
if existing_trace_id is not None: if existing_trace_id is not None:
trace_params = {"trace_id": existing_trace_id} trace_params = {"id": existing_trace_id}
else: # don't overwrite an existing trace else: # don't overwrite an existing trace
trace_params = { trace_params = {
"name": trace_name, "name": trace_name,
@ -289,10 +289,10 @@ class LangFuseLogger:
"session_id": metadata.get("session_id", None), "session_id": metadata.get("session_id", None),
} }
if level == "ERROR": if level == "ERROR":
trace_params["status_message"] = output trace_params["status_message"] = output
else: else:
trace_params["output"] = output trace_params["output"] = output
cost = kwargs.get("response_cost", None) cost = kwargs.get("response_cost", None)
print_verbose(f"trace: {cost}") print_verbose(f"trace: {cost}")
@ -350,7 +350,8 @@ class LangFuseLogger:
kwargs["cache_hit"] = False kwargs["cache_hit"] = False
tags.append(f"cache_hit:{kwargs['cache_hit']}") tags.append(f"cache_hit:{kwargs['cache_hit']}")
clean_metadata["cache_hit"] = kwargs["cache_hit"] clean_metadata["cache_hit"] = kwargs["cache_hit"]
trace_params.update({"tags": tags}) if existing_trace_id is None:
trace_params.update({"tags": tags})
proxy_server_request = litellm_params.get("proxy_server_request", None) proxy_server_request = litellm_params.get("proxy_server_request", None)
if proxy_server_request: if proxy_server_request:
@ -372,6 +373,7 @@ class LangFuseLogger:
print_verbose(f"trace_params: {trace_params}") print_verbose(f"trace_params: {trace_params}")
print(f"trace_params: {trace_params}")
trace = self.Langfuse.trace(**trace_params) trace = self.Langfuse.trace(**trace_params)
generation_id = None generation_id = None

View file

@ -470,8 +470,26 @@ def test_langfuse_existing_trace_id():
new_metadata = {"existing_trace_id": trace_id} new_metadata = {"existing_trace_id": trace_id}
new_messages = [{"role": "user", "content": "What do you know?"}] new_messages = [{"role": "user", "content": "What do you know?"}]
new_response_obj = litellm.ModelResponse(
id="chatcmpl-9K5HUAbVRqFrMZKXL0WoC295xhguY",
choices=[
litellm.Choices(
finish_reason="stop",
index=0,
message=litellm.Message(
content="What do I know?",
role="assistant",
),
)
],
created=1714573888,
model="gpt-3.5-turbo-0125",
object="chat.completion",
system_fingerprint="fp_3b956da36b",
usage=litellm.Usage(completion_tokens=37, prompt_tokens=14, total_tokens=51),
)
langfuse_args = { langfuse_args = {
"response_obj": response_obj, "response_obj": new_response_obj,
"kwargs": { "kwargs": {
"model": "gpt-3.5-turbo", "model": "gpt-3.5-turbo",
"litellm_params": { "litellm_params": {
@ -529,6 +547,10 @@ def test_langfuse_existing_trace_id():
langfuse_response_object = langfuse_Logger.log_event(**langfuse_args) langfuse_response_object = langfuse_Logger.log_event(**langfuse_args)
new_trace_id = langfuse_response_object["trace_id"]
assert new_trace_id == trace_id
langfuse_client.flush() langfuse_client.flush()
time.sleep(2) time.sleep(2)