From 2b7d48f7b409b78df95e41feef16738a0518eaad Mon Sep 17 00:00:00 2001
From: afel <andreas.feldl@netlight.com>
Date: Mon, 3 Jun 2024 22:11:57 +0200
Subject: [PATCH] add metadata from header changes

---
 litellm/integrations/langfuse.py | 25 ++++++++++++++++++++++---
 1 file changed, 22 insertions(+), 3 deletions(-)

diff --git a/litellm/integrations/langfuse.py b/litellm/integrations/langfuse.py
index 4d580f6666..ed03a784e2 100644
--- a/litellm/integrations/langfuse.py
+++ b/litellm/integrations/langfuse.py
@@ -69,6 +69,27 @@ class LangFuseLogger:
         else:
             self.upstream_langfuse = None
 
+    @staticmethod
+    def add_metadata_from_header(litellm_params) -> dict:
+        """
+        Adds metadata from proxy request headers to Langfuse logging if keys start with "trace_" 
+        and overwrites if already exists in litellm_params.metadata
+        """
+        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 trace_param_key in metadata:
+                    verbose_logger.warning(f"Overwriting Langfuse `{trace_param_key}` from request header")
+                else:
+                    verbose_logger.debug(f"Found Langfuse `{trace_param_key}` in request header")
+                metadata[trace_param_key] = proxy_headers.get(metadata_param_key)
+
+        return metadata
+
     # def log_error(kwargs, response_obj, start_time, end_time):
     #     generation = trace.generation(
     #         level ="ERROR" # can be any of DEBUG, DEFAULT, WARNING or ERROR
@@ -94,9 +115,7 @@ class LangFuseLogger:
 
             litellm_params = kwargs.get("litellm_params", {})
             litellm_call_id = kwargs.get("litellm_call_id", None)
-            metadata = (
-                litellm_params.get("metadata", {}) or {}
-            )  # if litellm_params['metadata'] == None
+            metadata = self.add_metadata_from_header(litellm_params)
             optional_params = copy.deepcopy(kwargs.get("optional_params", {}))
 
             prompt = {"messages": kwargs.get("messages")}