address review comments

This commit is contained in:
afel 2024-06-06 08:01:42 +02:00
parent 2b7d48f7b4
commit aad0ea80f6

View file

@ -70,18 +70,19 @@ class LangFuseLogger:
self.upstream_langfuse = None
@staticmethod
def add_metadata_from_header(litellm_params) -> dict:
def add_metadata_from_header(litellm_params: dict, metadata: dict) -> dict:
"""
Adds metadata from proxy request headers to Langfuse logging if keys start with "trace_"
and overwrites if already exists in litellm_params.metadata
Adds metadata from proxy request headers to Langfuse logging if keys start with "langfuse_"
and overwrites litellm_params.metadata if already included.
For example if you want to append your trace to an existing `trace_id` via header, send
`headers: { ..., langfuse_existing_trace_id: your-existing-trace-id }` via proxy request.
"""
metadata = litellm_params.get("metadata", {})
proxy_headers = litellm_params.get("proxy_server_request", {}).get("headers", {})
# Update the following keys for this trace
for metadata_param_key in proxy_headers:
if metadata_param_key.startswith("trace_"):
trace_param_key = metadata_param_key.replace("trace_", "", 1)
if metadata_param_key.startswith("langfuse_"):
trace_param_key = metadata_param_key.replace("langfuse_", "", 1)
if trace_param_key in metadata:
verbose_logger.warning(f"Overwriting Langfuse `{trace_param_key}` from request header")
else:
@ -115,7 +116,10 @@ class LangFuseLogger:
litellm_params = kwargs.get("litellm_params", {})
litellm_call_id = kwargs.get("litellm_call_id", None)
metadata = self.add_metadata_from_header(litellm_params)
metadata = (
litellm_params.get("metadata", {}) or {}
) # if litellm_params['metadata'] == None
metadata = self.add_metadata_from_header(litellm_params, metadata)
optional_params = copy.deepcopy(kwargs.get("optional_params", {}))
prompt = {"messages": kwargs.get("messages")}