Added in_progress, searching, and completed events for knowledge_search tool execution to match OpenAI Responses API spec. Mirrors existing web_search event pattern.

This commit is contained in:
Ashwin Bharambe 2025-10-11 14:32:44 -07:00
parent 3f1f7c3f7f
commit 9dcee00ef7

View file

@ -11,6 +11,9 @@ from collections.abc import AsyncIterator
from llama_stack.apis.agents.openai_responses import ( from llama_stack.apis.agents.openai_responses import (
OpenAIResponseInputToolFileSearch, OpenAIResponseInputToolFileSearch,
OpenAIResponseInputToolMCP, OpenAIResponseInputToolMCP,
OpenAIResponseObjectStreamResponseFileSearchCallCompleted,
OpenAIResponseObjectStreamResponseFileSearchCallInProgress,
OpenAIResponseObjectStreamResponseFileSearchCallSearching,
OpenAIResponseObjectStreamResponseMcpCallCompleted, OpenAIResponseObjectStreamResponseMcpCallCompleted,
OpenAIResponseObjectStreamResponseMcpCallFailed, OpenAIResponseObjectStreamResponseMcpCallFailed,
OpenAIResponseObjectStreamResponseMcpCallInProgress, OpenAIResponseObjectStreamResponseMcpCallInProgress,
@ -221,7 +224,13 @@ class ToolExecutor:
output_index=output_index, output_index=output_index,
sequence_number=sequence_number, sequence_number=sequence_number,
) )
# Note: knowledge_search and other custom tools don't have specific streaming events in OpenAI spec elif function_name == "knowledge_search":
sequence_number += 1
progress_event = OpenAIResponseObjectStreamResponseFileSearchCallInProgress(
item_id=item_id,
output_index=output_index,
sequence_number=sequence_number,
)
if progress_event: if progress_event:
yield ToolExecutionResult(stream_event=progress_event, sequence_number=sequence_number) yield ToolExecutionResult(stream_event=progress_event, sequence_number=sequence_number)
@ -236,6 +245,16 @@ class ToolExecutor:
) )
yield ToolExecutionResult(stream_event=searching_event, sequence_number=sequence_number) yield ToolExecutionResult(stream_event=searching_event, sequence_number=sequence_number)
# For file search, emit searching event
if function_name == "knowledge_search":
sequence_number += 1
searching_event = OpenAIResponseObjectStreamResponseFileSearchCallSearching(
item_id=item_id,
output_index=output_index,
sequence_number=sequence_number,
)
yield ToolExecutionResult(stream_event=searching_event, sequence_number=sequence_number)
async def _execute_tool( async def _execute_tool(
self, self,
function_name: str, function_name: str,
@ -322,7 +341,13 @@ class ToolExecutor:
output_index=output_index, output_index=output_index,
sequence_number=sequence_number, sequence_number=sequence_number,
) )
# Note: knowledge_search and other custom tools don't have specific completion events in OpenAI spec elif function_name == "knowledge_search":
sequence_number += 1
completion_event = OpenAIResponseObjectStreamResponseFileSearchCallCompleted(
item_id=item_id,
output_index=output_index,
sequence_number=sequence_number,
)
if completion_event: if completion_event:
yield ToolExecutionResult(stream_event=completion_event, sequence_number=sequence_number) yield ToolExecutionResult(stream_event=completion_event, sequence_number=sequence_number)