Fix: Populate the node assistant's trace service with the existing callback

apps/pipelines/nodes/base.py

@@ -101,7 +101,9 @@ def _process(self, state: PipelineState) -> PipelineState:
     _config: RunnableConfig | None = None
     name: str = Field(title="Node Name", json_schema_extra={"ui:widget": "node_name"})
 
-    def process(self, node_id: str, incoming_edges: list, state: PipelineState, config) -> PipelineState:
+    def process(
+        self, node_id: str, incoming_edges: list, state: PipelineState, config: RunnableConfig
+    ) -> PipelineState:
         from apps.channels.datamodels import Attachment
 
         self._config = config

apps/pipelines/nodes/nodes.py

@@ -682,7 +682,7 @@ def _process(self, input, state: PipelineState, node_id: str, **kwargs) -> Pipel
         session: ExperimentSession | None = state.get("experiment_session")
         runnable = self._get_assistant_runnable(assistant, session=session, node_id=node_id)
         attachments = self._get_attachments(state)
-        chain_output: ChainOutput = runnable.invoke(input, config={}, attachments=attachments)
+        chain_output: ChainOutput = runnable.invoke(input, config=self._config, attachments=attachments)
         output = chain_output.output
 
         return PipelineState.from_node_output(
@@ -699,8 +699,11 @@ def _get_attachments(self, state) -> list:
         return [att for att in state.get("shared_state", {}).get("attachments", []) if att.upload_to_assistant]
 
     def _get_assistant_runnable(self, assistant: OpenAiAssistant, session: ExperimentSession, node_id: str):
+        trace_service = session.experiment.trace_service
+        trace_service.from_callback_manager(self._config.get("callbacks"))
+
         history_manager = PipelineHistoryManager.for_assistant()
-        adapter = AssistantAdapter.for_pipeline(session=session, node=self)
+        adapter = AssistantAdapter.for_pipeline(session=session, node=self, trace_service=trace_service)
         if assistant.tools_enabled:
             return AgentAssistantChat(adapter=adapter, history_manager=history_manager)
         else:

apps/service_providers/llm_service/adapters.py

@@ -180,15 +180,14 @@ def for_experiment(experiment: Experiment, session: ExperimentSession, trace_ser
         )
 
     @staticmethod
-    def for_pipeline(session: ExperimentSession, node: "AssistantNode") -> Self:
+    def for_pipeline(session: ExperimentSession, node: "AssistantNode", trace_service=None) -> Self:
         assistant = OpenAiAssistant.objects.get(id=node.assistant_id)
-        experiment = session.experiment
         return AssistantAdapter(
             session=session,
             assistant=assistant,
             citations_enabled=node.citations_enabled,
             input_formatter=node.input_formatter,
-            trace_service=experiment.trace_service,
+            trace_service=trace_service,
             save_message_metadata_only=True,
         )
 

apps/service_providers/tracing/service.py

@@ -1,3 +1,4 @@
+from langchain_core.callbacks.manager import CallbackManager
 from langchain_core.tracers import LangChainTracer
 from pydantic import BaseModel
 
@@ -29,6 +30,9 @@ def update_trace(self, metadata: dict):
     def get_current_trace_info(self) -> TraceInfo | None:
         return None
 
+    def from_callback_manager(self, callback_manager: CallbackManager):
+        pass
+
 
 class LangFuseTraceService(TraceService):
     """
@@ -51,6 +55,17 @@ def get_callback(self, participant_id: str, session_id: str):
         self._callback = CallbackHandler(user_id=participant_id, session_id=session_id, **self.config)
         return self._callback
 
+    def from_callback_manager(self, callback_manager: CallbackManager):
+        """
+        Populates the callback from the callback handler already configured in `callback_manager`. This allows the trace
+        service to reuse existing callbacks.
+        """
+        from langfuse.callback import CallbackHandler
+
+        for handler in callback_manager.handlers:
+            if isinstance(handler, CallbackHandler):
+                self._callback = handler
+
     def update_trace(self, metadata: dict):
         if not metadata:
             return