diff --git a/surfsense_backend/app/tasks/chat/streaming/relay/event_relay.py b/surfsense_backend/app/tasks/chat/streaming/relay/event_relay.py index 872998926..03d6a66e6 100644 --- a/surfsense_backend/app/tasks/chat/streaming/relay/event_relay.py +++ b/surfsense_backend/app/tasks/chat/streaming/relay/event_relay.py @@ -115,6 +115,7 @@ class EventRelay: state.current_text_id = None completion_event, new_active = complete_active_thinking_step( + state=state, streaming_service=self.streaming_service, content_builder=content_builder, last_active_step_id=state.last_active_step_id, diff --git a/surfsense_backend/app/tasks/chat/streaming/relay/thinking_step_completion.py b/surfsense_backend/app/tasks/chat/streaming/relay/thinking_step_completion.py index a0be71281..ad0930341 100644 --- a/surfsense_backend/app/tasks/chat/streaming/relay/thinking_step_completion.py +++ b/surfsense_backend/app/tasks/chat/streaming/relay/thinking_step_completion.py @@ -4,11 +4,13 @@ from __future__ import annotations from typing import Any +from .state import AgentEventRelayState from .thinking_step_sse import emit_thinking_step_frame def complete_active_thinking_step( *, + state: AgentEventRelayState, streaming_service: Any, content_builder: Any | None, last_active_step_id: str | None, @@ -26,6 +28,7 @@ def complete_active_thinking_step( title=last_active_step_title, status="completed", items=last_active_step_items if last_active_step_items else None, + metadata=state.span_metadata_if_active(), ) return event, None return None, last_active_step_id diff --git a/surfsense_backend/app/tasks/chat/streaming/relay/thinking_step_sse.py b/surfsense_backend/app/tasks/chat/streaming/relay/thinking_step_sse.py index 9e8c08dd5..6737f536b 100644 --- a/surfsense_backend/app/tasks/chat/streaming/relay/thinking_step_sse.py +++ b/surfsense_backend/app/tasks/chat/streaming/relay/thinking_step_sse.py @@ -13,12 +13,16 @@ def emit_thinking_step_frame( title: str, status: str = "in_progress", items: list[str] | None = None, + metadata: dict[str, Any] | None = None, ) -> str: if content_builder is not None: - content_builder.on_thinking_step(step_id, title, status, items) + content_builder.on_thinking_step( + step_id, title, status, items, metadata=metadata + ) return streaming_service.format_thinking_step( step_id=step_id, title=title, status=status, items=items, + metadata=metadata, )