push

Author: hassiebp

langfuse/callback/langchain.py

@@ -1,5 +1,4 @@
 import typing
-from collections import defaultdict
 
 import pydantic
 
@@ -49,12 +48,11 @@
 
 
 class LangchainCallbackHandler(LangchainBaseCallbackHandler):
-    def __init__(self) -> None:
-        self.langfuse = get_client()
+    def __init__(self, *, public_key: Optional[str] = None) -> None:
+        self.langfuse_client = get_client(public_key=public_key)
 
         self.runs: Dict[UUID, Union[LangfuseSpan, LangfuseGeneration]] = {}
         self.prompt_to_parent_run_map = {}
-        self.trace_updates = defaultdict(dict)
         self.updated_completion_start_time_memo = set()
 
     def on_llm_new_token(
@@ -168,7 +166,7 @@ def on_chain_start(
             }
 
             if parent_run_id is None:
-                self.runs[run_id] = self.langfuse.start_span(**content)
+                self.runs[run_id] = self.langfuse_client.start_span(**content)
             else:
                 self.runs[run_id] = cast(
                     LangfuseSpan, self.runs[parent_run_id]
@@ -428,7 +426,7 @@ def on_retriever_start(
                     "level": "DEBUG" if tags and LANGSMITH_TAG_HIDDEN in tags else None,
                 }
 
-                self.runs[run_id] = self.langfuse.start_span(**content)
+                self.runs[run_id] = self.langfuse_client.start_span(**content)
             else:
                 self.runs[run_id] = cast(
                     LangfuseSpan, self.runs[parent_run_id]
@@ -551,7 +549,7 @@ def __on_llm_action(
                     LangfuseSpan, self.runs[parent_run_id]
                 ).start_generation(**content)
             else:
-                self.runs[run_id] = self.langfuse.start_generation(**content)
+                self.runs[run_id] = self.langfuse_client.start_generation(**content)
 
         except Exception as e:
             langfuse_logger.exception(e)

langfuse/openai.py

@@ -147,19 +147,12 @@ class OpenAiDefinition:
 class OpenAiArgsExtractor:
     def __init__(
         self,
-        name=None,
         metadata=None,
-        trace_id=None,
-        session_id=None,
-        user_id=None,
-        tags=None,
-        parent_observation_id=None,
         langfuse_prompt=None,  # we cannot use prompt because it's an argument of the old OpenAI completions API
         langfuse_public_key=None,
         **kwargs,
     ):
         self.args = {}
-        self.args["name"] = name
         self.args["metadata"] = (
             metadata
             if "response_format" not in kwargs
@@ -172,12 +165,8 @@ def __init__(
             }
         )
         self.args["langfuse_public_key"] = langfuse_public_key
-        self.args["trace_id"] = trace_id
-        self.args["parent_observation_id"] = parent_observation_id
-        self.args["session_id"] = session_id
-        self.args["user_id"] = user_id
-        self.args["tags"] = tags
         self.args["langfuse_prompt"] = langfuse_prompt
+
         self.kwargs = kwargs
 
     def get_langfuse_args(self):
@@ -689,12 +678,6 @@ def _wrap(open_ai_resource: OpenAiDefinition, wrapped, args, kwargs):
         name=langfuse_data["name"],
         input=langfuse_data.get("input", None),
         metadata=langfuse_data.get("metadata", None),
-        trace_context={
-            "trace_id": cast(str, langfuse_data.get("trace_id", None)),
-            "parent_span_id": cast(
-                str, langfuse_data.get("parent_observation_id", None)
-            ),
-        },
         model_parameters=langfuse_data.get("model_parameters", None),
         model=langfuse_data.get("model", None),
         prompt=langfuse_data.get("langfuse_prompt", None),
@@ -722,8 +705,7 @@ def _wrap(open_ai_resource: OpenAiDefinition, wrapped, args, kwargs):
                 model=model,
                 output=completion,
                 usage_details=usage,
-            )
-            generation.end()
+            ).end()
 
         return openai_response
     except Exception as ex:
@@ -734,9 +716,7 @@ def _wrap(open_ai_resource: OpenAiDefinition, wrapped, args, kwargs):
             level="ERROR",
             model=model,
             cost_details={"input": 0, "output": 0, "total": 0},
-        )
-
-        generation.end()
+        ).end()
 
         raise ex
 
@@ -786,8 +766,7 @@ async def _wrap_async(open_ai_resource: OpenAiDefinition, wrapped, args, kwargs)
                 output=completion,
                 usage=usage,  # backward compat for all V2 self hosters
                 usage_details=usage,
-            )
-            generation.end()
+            ).end()
 
         return openai_response
     except Exception as ex:
@@ -798,9 +777,7 @@ async def _wrap_async(open_ai_resource: OpenAiDefinition, wrapped, args, kwargs)
             level="ERROR",
             model=model,
             cost_details={"input": 0, "output": 0, "total": 0},
-        )
-
-        generation.end()
+        ).end()
 
         raise ex