Include optional metadata on tool and thinking-step SSE payloads.

This commit is contained in:
CREDO23 2026-05-08 22:47:58 +02:00
parent f1d80ffe5d
commit e802de2333

View file

@ -456,6 +456,8 @@ class VercelStreamingService:
title: str, title: str,
status: str = "in_progress", status: str = "in_progress",
items: list[str] | None = None, items: list[str] | None = None,
*,
metadata: dict[str, Any] | None = None,
) -> str: ) -> str:
""" """
Format a thinking step for chain-of-thought display (SurfSense specific). Format a thinking step for chain-of-thought display (SurfSense specific).
@ -469,15 +471,15 @@ class VercelStreamingService:
Returns: Returns:
str: SSE formatted thinking step data part str: SSE formatted thinking step data part
""" """
return self.format_data( payload: dict[str, Any] = {
"thinking-step", "id": step_id,
{ "title": title,
"id": step_id, "status": status,
"title": title, "items": items or [],
"status": status, }
"items": items or [], if metadata:
}, payload["metadata"] = metadata
) return self.format_data("thinking-step", payload)
def format_thread_title_update(self, thread_id: int, title: str) -> str: def format_thread_title_update(self, thread_id: int, title: str) -> str:
""" """
@ -601,6 +603,7 @@ class VercelStreamingService:
tool_name: str, tool_name: str,
*, *,
langchain_tool_call_id: str | None = None, langchain_tool_call_id: str | None = None,
metadata: dict[str, Any] | None = None,
) -> str: ) -> str:
""" """
Format the start of tool input streaming. Format the start of tool input streaming.
@ -635,6 +638,8 @@ class VercelStreamingService:
} }
if langchain_tool_call_id: if langchain_tool_call_id:
payload["langchainToolCallId"] = langchain_tool_call_id payload["langchainToolCallId"] = langchain_tool_call_id
if metadata:
payload["metadata"] = metadata
return self._format_sse(payload) return self._format_sse(payload)
def format_tool_input_delta(self, tool_call_id: str, input_text_delta: str) -> str: def format_tool_input_delta(self, tool_call_id: str, input_text_delta: str) -> str:
@ -666,6 +671,7 @@ class VercelStreamingService:
input_data: dict[str, Any], input_data: dict[str, Any],
*, *,
langchain_tool_call_id: str | None = None, langchain_tool_call_id: str | None = None,
metadata: dict[str, Any] | None = None,
) -> str: ) -> str:
""" """
Format the completion of tool input. Format the completion of tool input.
@ -691,6 +697,8 @@ class VercelStreamingService:
} }
if langchain_tool_call_id: if langchain_tool_call_id:
payload["langchainToolCallId"] = langchain_tool_call_id payload["langchainToolCallId"] = langchain_tool_call_id
if metadata:
payload["metadata"] = metadata
return self._format_sse(payload) return self._format_sse(payload)
def format_tool_output_available( def format_tool_output_available(
@ -699,6 +707,7 @@ class VercelStreamingService:
output: Any, output: Any,
*, *,
langchain_tool_call_id: str | None = None, langchain_tool_call_id: str | None = None,
metadata: dict[str, Any] | None = None,
) -> str: ) -> str:
""" """
Format tool execution output. Format tool execution output.
@ -725,6 +734,8 @@ class VercelStreamingService:
} }
if langchain_tool_call_id: if langchain_tool_call_id:
payload["langchainToolCallId"] = langchain_tool_call_id payload["langchainToolCallId"] = langchain_tool_call_id
if metadata:
payload["metadata"] = metadata
return self._format_sse(payload) return self._format_sse(payload)
# ========================================================================= # =========================================================================