push

Author: hassiebp

langfuse/_client/client.py

@@ -84,10 +84,13 @@
     CreateTextPromptRequest,
     Dataset,
     DatasetItem,
+    DatasetRunWithItems,
     DatasetStatus,
+    DeleteDatasetRunResponse,
     Error,
     MapValue,
     NotFoundError,
+    PaginatedDatasetRuns,
     Prompt_Chat,
     Prompt_Text,
     ScoreBody,
@@ -2204,10 +2207,13 @@ def get_dataset_run(
             DatasetRunWithItems: The dataset run with its items.
         """
         try:
-            return self.api.datasets.get_run(
-                dataset_name=self._url_encode(dataset_name),
-                run_name=self._url_encode(run_name),
-                request_options=None,
+            return cast(
+                DatasetRunWithItems,
+                self.api.datasets.get_run(
+                    dataset_name=self._url_encode(dataset_name),
+                    run_name=self._url_encode(run_name),
+                    request_options=None,
+                ),
             )
         except Error as e:
             handle_fern_exception(e)
@@ -2231,11 +2237,14 @@ def get_dataset_runs(
             PaginatedDatasetRuns: Paginated list of dataset runs.
         """
         try:
-            return self.api.datasets.get_runs(
-                dataset_name=self._url_encode(dataset_name),
-                page=page,
-                limit=limit,
-                request_options=None,
+            return cast(
+                PaginatedDatasetRuns,
+                self.api.datasets.get_runs(
+                    dataset_name=self._url_encode(dataset_name),
+                    page=page,
+                    limit=limit,
+                    request_options=None,
+                ),
             )
         except Error as e:
             handle_fern_exception(e)
@@ -2254,10 +2263,13 @@ def delete_dataset_run(
             DeleteDatasetRunResponse: Confirmation of deletion.
         """
         try:
-            return self.api.datasets.delete_run(
-                dataset_name=self._url_encode(dataset_name),
-                run_name=self._url_encode(run_name),
-                request_options=None,
+            return cast(
+                DeleteDatasetRunResponse,
+                self.api.datasets.delete_run(
+                    dataset_name=self._url_encode(dataset_name),
+                    run_name=self._url_encode(run_name),
+                    request_options=None,
+                ),
             )
         except Error as e:
             handle_fern_exception(e)

langfuse/langchain/CallbackHandler.py

@@ -34,7 +34,6 @@
 from langfuse.langchain.utils import _extract_model_name
 from langfuse.logger import langfuse_logger
 from langfuse.types import TraceContext
-from langfuse.types import TraceContext
 
 try:
     import langchain
@@ -843,19 +842,15 @@ def __on_llm_action(
             model_name = self._parse_model_and_log_errors(
                 serialized=serialized, metadata=metadata, kwargs=kwargs
             )
-            registered_prompt = (
-                self._prompt_to_parent_run_map.get(parent_run_id)
-                if parent_run_id is not None
-                else None
-            )
 
-                if registered_prompt:
-                    self._deregister_langfuse_prompt(current_parent_run_id)
-                    break
-                else:
-                    current_parent_run_id = self._child_to_parent_run_id_map.get(
-                        current_parent_run_id, None
-                    )
+            registered_prompt = None
+            current_parent_run_id = parent_run_id
+
+            # Check all parents for registered prompt
+            while current_parent_run_id is not None:
+                registered_prompt = self._prompt_to_parent_run_map.get(
+                    current_parent_run_id
+                )
 
             content = {
                 "name": self.get_langchain_run_name(serialized, **kwargs),
@@ -1012,7 +1007,7 @@ def on_llm_error(
             langfuse_logger.exception(e)
 
     def _reset(self) -> None:
-        self._child_to_parent_run_id_map = {}
+        self._child_to_parent_run_id_map: Dict[UUID, Optional[UUID]] = {}
 
     def __join_tags_and_metadata(
         self,