diff --git a/src/quickapp/common/chat_completion_stream/chat_stream_sink_factory.py b/src/quickapp/common/chat_completion_stream/chat_stream_sink_factory.py index f03f9cb0..46ca8fb1 100644 --- a/src/quickapp/common/chat_completion_stream/chat_stream_sink_factory.py +++ b/src/quickapp/common/chat_completion_stream/chat_stream_sink_factory.py @@ -88,6 +88,7 @@ def create(self, config: ChatStreamConfig) -> ChatStreamPipeline: stream_content=config.stream_content, propagate_stages=config.propagate_stages, tools_by_name=_tools_by_name(self._tools()), + excluded_attachment_urls=config.excluded_attachment_urls, ), StageWrapperUiSink( stage_wrapper=config.stage_wrapper, diff --git a/src/quickapp/common/chat_completion_stream/choice_ui_stream_sink.py b/src/quickapp/common/chat_completion_stream/choice_ui_stream_sink.py index 0f0da2d7..a34ff055 100644 --- a/src/quickapp/common/chat_completion_stream/choice_ui_stream_sink.py +++ b/src/quickapp/common/chat_completion_stream/choice_ui_stream_sink.py @@ -59,12 +59,14 @@ def __init__( stream_content: bool = True, propagate_stages: bool = False, tools_by_name: dict[str, StagedBaseTool] | None = None, + excluded_attachment_urls: set[str] | None = None, ) -> None: self._accumulator = accumulator self._destination = destination self._stream_content = stream_content self._propagate_stages = propagate_stages self._tools_by_name = tools_by_name or {} + self._excluded_attachment_urls: set[str] = excluded_attachment_urls or set() self._stages_by_index: dict[int, Stage] = {} self._tool_stages_by_index: dict[int, _StreamingToolStageState] = {} self._suppressed_tool_indexes: set[int] = set() @@ -115,7 +117,9 @@ def _apply_custom(self, norm: NormalizedCustomContent) -> None: destination = self._destination assert destination is not None if norm.attachments: - self._add_attachments(destination, norm.attachments) + to_add = [a for a in norm.attachments if a.url not in self._excluded_attachment_urls] + if to_add: + self._add_attachments(destination, to_add) for position, raw in norm.stage_entries: stage_delta = as_stage_delta(raw) if self._propagate_stages: diff --git a/src/quickapp/common/chat_completion_stream/handler.py b/src/quickapp/common/chat_completion_stream/handler.py index 1847e4f6..d248af16 100644 --- a/src/quickapp/common/chat_completion_stream/handler.py +++ b/src/quickapp/common/chat_completion_stream/handler.py @@ -5,7 +5,7 @@ from injector import inject from openai import APIError, BadRequestError from openai.types.chat import ChatCompletionChunk -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, Field from quickapp.common.base_stage_wrapper import BaseStageWrapper from quickapp.common.chat_completion_stream.chat_stream_sink_factory import ChatStreamSinkFactory @@ -36,6 +36,7 @@ class ChatStreamConfig(BaseModel): stage_wrapper: BaseStageWrapper | None = None stream_content: bool = True propagate_stages: bool = False + excluded_attachment_urls: set[str] = Field(default_factory=set) class ChatCompletionStreamHandler: diff --git a/src/quickapp/core/agent/orchestrator.py b/src/quickapp/core/agent/orchestrator.py index 8d40abe0..2b19d4aa 100644 --- a/src/quickapp/core/agent/orchestrator.py +++ b/src/quickapp/core/agent/orchestrator.py @@ -105,6 +105,7 @@ def __init__( request_async_close_registry ) self.__propagated_attachment_urls: set[str] = set() + self.__all_tool_attachment_urls: set[str] = set() @property def iteration_count(self) -> int: @@ -267,6 +268,9 @@ async def _execute_internal_tool_calls( for tool_call_result in tool_call_results: tool_call_result_message = tool_call_result.to_tool_message() self.__messages_context.append_message(tool_call_result_message) + for attachment in tool_call_result.attachments or []: + if attachment.url: + self.__all_tool_attachment_urls.add(attachment.url) for attachment in tool_call_result.propagate_to_choice: url = attachment.url if url is not None: @@ -313,6 +317,7 @@ async def accumulate_stream( destination=self.__choice, stream_content=True, propagate_stages=self.__propagate_orchestrator_stages, + excluded_attachment_urls=self.__all_tool_attachment_urls, ), ) except ChatStreamHandlerError: diff --git a/src/tests/unit_tests/agent_tests/test_orchestrator.py b/src/tests/unit_tests/agent_tests/test_orchestrator.py index 0d72d4c4..421bc09c 100644 --- a/src/tests/unit_tests/agent_tests/test_orchestrator.py +++ b/src/tests/unit_tests/agent_tests/test_orchestrator.py @@ -317,6 +317,7 @@ async def test_invoke_with_tool_calls_executes_tools_and_updates_state_and_messa tool_message = Message(role=Role.TOOL, content="tool output", tool_call_id="tc-1") tool_result = Mock() tool_result.to_tool_message = Mock(return_value=tool_message) + tool_result.attachments = None # propagate_to_choice contains attachments with model_dump() attach = Mock() @@ -885,6 +886,7 @@ async def test_invoke_terminal_flow_strips_get_content_attachments_in_saved_hist ) tool_result = Mock() tool_result.to_tool_message = Mock(return_value=tool_message) + tool_result.attachments = None tool_result.propagate_to_choice = [] tool_result.usage = [] @@ -995,6 +997,7 @@ async def test_invoke_interrupted_flow_keeps_get_content_attachments_in_saved_hi ) tool_result = Mock() tool_result.to_tool_message = Mock(return_value=tool_message) + tool_result.attachments = None tool_result.propagate_to_choice = [] tool_result.usage = [] @@ -1088,6 +1091,7 @@ async def test_propagation_deduplicates_repeated_urls(): tool_result.to_tool_message = Mock( return_value=Message(role=Role.TOOL, content="out", tool_call_id="tc-1") ) + tool_result.attachments = None tool_result.usage = None tool_result.propagate_to_choice = [ Attachment(url=same_url, type="text/csv"), @@ -1109,6 +1113,7 @@ async def test_propagation_keeps_urlless_attachments(): tool_result.to_tool_message = Mock( return_value=Message(role=Role.TOOL, content="out", tool_call_id="tc-1") ) + tool_result.attachments = None tool_result.usage = None tool_result.propagate_to_choice = [ Attachment(data="abc", type="image/png"), diff --git a/src/tests/unit_tests/agent_tests/test_orchestrator_external_tools.py b/src/tests/unit_tests/agent_tests/test_orchestrator_external_tools.py index 9eeb554d..d9dba171 100644 --- a/src/tests/unit_tests/agent_tests/test_orchestrator_external_tools.py +++ b/src/tests/unit_tests/agent_tests/test_orchestrator_external_tools.py @@ -183,6 +183,7 @@ def create_function_tool_call(self, id, name, arguments=None): tool_msg = Message(role=Role.TOOL, content="server result", tool_call_id="id-i") tool_result = Mock() tool_result.to_tool_message = Mock(return_value=tool_msg) + tool_result.attachments = None tool_result.propagate_to_choice = [] tool_result.usage = None @@ -243,6 +244,7 @@ async def test_all_internal_tools_loop_continues(): tool_msg = Message(role=Role.TOOL, content="ok", tool_call_id="id-1") tool_result = Mock() tool_result.to_tool_message = Mock(return_value=tool_msg) + tool_result.attachments = None tool_result.propagate_to_choice = [] tool_result.usage = None tool_executor = Mock( @@ -304,6 +306,7 @@ async def test_no_external_tools_configured_existing_behavior_unchanged(): tool_msg = Message(role=Role.TOOL, content="ok", tool_call_id="id-1") tool_result = Mock() tool_result.to_tool_message = Mock(return_value=tool_msg) + tool_result.attachments = None tool_result.propagate_to_choice = [] tool_result.usage = None tool_executor = Mock( @@ -363,6 +366,7 @@ async def test_mixed_batch_persists_history_without_external_tool_calls(): tool_msg = Message(role=Role.TOOL, content="server result", tool_call_id="id-i") tool_result = Mock() tool_result.to_tool_message = Mock(return_value=tool_msg) + tool_result.attachments = None tool_result.propagate_to_choice = [] tool_result.usage = None tool_executor = Mock( diff --git a/src/tests/unit_tests/chat_completion_stream_tests/test_stream_sinks.py b/src/tests/unit_tests/chat_completion_stream_tests/test_stream_sinks.py index f74fc59d..f47f6ab0 100644 --- a/src/tests/unit_tests/chat_completion_stream_tests/test_stream_sinks.py +++ b/src/tests/unit_tests/chat_completion_stream_tests/test_stream_sinks.py @@ -1,10 +1,14 @@ """Unit tests for chat-stream DI sinks.""" +from aidial_sdk.chat_completion import Attachment from openai.types.chat.chat_completion_chunk import ChoiceDeltaToolCall, ChoiceDeltaToolCallFunction from quickapp.common.chat_completion_stream.accumulation_stream_sink import AccumulationSink from quickapp.common.chat_completion_stream.choice_ui_stream_sink import ChoiceUiSink -from quickapp.common.chat_completion_stream.models import NormalizedChoiceDelta +from quickapp.common.chat_completion_stream.models import ( + NormalizedChoiceDelta, + NormalizedCustomContent, +) from quickapp.common.chat_completion_stream.stage_wrapper_ui_stream_sink import StageWrapperUiSink from quickapp.common.chat_completion_stream.stream_result import ChatStreamAccumulator from tests.unit_tests.stream_test_doubles import DummyStageWrapper, SpyChoice @@ -70,3 +74,67 @@ def test_choice_ui_opens_tool_stage_stage_wrapper_does_not(): StageWrapperUiSink(stage_wrapper=wrap).on_delta(tool_delta) wrap.stage_mock.create_stage.assert_not_called() + + +def _custom_delta(url: str, mime_type: str = "image/png") -> NormalizedChoiceDelta: + return NormalizedChoiceDelta( + custom=NormalizedCustomContent( + attachments=[Attachment(url=url, type=mime_type)], + stage_entries=[], + state=None, + ) + ) + + +def test_choice_ui_sink_filters_excluded_attachment_url_from_llm_echo(): + choice = SpyChoice() + excluded_url = "https://dial-core/uploads/image.png" + sink = ChoiceUiSink( + ChatStreamAccumulator(), + destination=choice, + excluded_attachment_urls={excluded_url}, + ) + sink.on_delta(_custom_delta(excluded_url)) + assert choice.add_attachment_kwargs == [] + + +def test_choice_ui_sink_allows_attachment_not_in_excluded_urls(): + choice = SpyChoice() + excluded_url = "https://dial-core/uploads/image.png" + other_url = "https://dial-core/uploads/other.png" + sink = ChoiceUiSink( + ChatStreamAccumulator(), + destination=choice, + excluded_attachment_urls={excluded_url}, + ) + sink.on_delta(_custom_delta(other_url)) + assert len(choice.add_attachment_kwargs) == 1 + recorded = choice.add_attachment_kwargs[0] + actual_url = recorded.get("url") or recorded["args"][0].url + assert actual_url == other_url + + +def test_choice_ui_sink_partial_exclusion_filters_only_matching_urls(): + choice = SpyChoice() + excluded_url = "https://dial-core/uploads/excluded.png" + kept_url = "https://dial-core/uploads/kept.png" + sink = ChoiceUiSink( + ChatStreamAccumulator(), + destination=choice, + excluded_attachment_urls={excluded_url}, + ) + delta = NormalizedChoiceDelta( + custom=NormalizedCustomContent( + attachments=[ + Attachment(url=excluded_url, type="image/png"), + Attachment(url=kept_url, type="image/png"), + ], + stage_entries=[], + state=None, + ) + ) + sink.on_delta(delta) + assert len(choice.add_attachment_kwargs) == 1 + recorded = choice.add_attachment_kwargs[0] + actual_url = recorded.get("url") or recorded["args"][0].url + assert actual_url == kept_url