diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 5d783cb0368..ebf7a5fe984 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -1923,7 +1923,14 @@ def _response_has_visible_content(response: ChatResponse[Any]) -> bool: return False +def _drop_function_call_contents_from_response(response: ChatResponse[Any]) -> None: + for message in response.messages: + if any(content.type == "function_call" for content in message.contents): + message.contents = [content for content in message.contents if content.type != "function_call"] + + def _ensure_function_invocation_limit_fallback_response(response: ChatResponse[Any]) -> ChatResponse[Any]: + _drop_function_call_contents_from_response(response) if _response_has_visible_content(response): return response @@ -1948,6 +1955,31 @@ def _function_invocation_limit_fallback_update() -> ChatResponseUpdate: ) +def _drop_function_call_contents_from_update(update: ChatResponseUpdate) -> ChatResponseUpdate | None: + if not any(content.type == "function_call" for content in update.contents): + return update + + update.contents = [content for content in update.contents if content.type != "function_call"] + if update.contents or _update_has_meaningful_metadata(update): + return update + return None + + +def _update_has_meaningful_metadata(update: ChatResponseUpdate) -> bool: + return any(( + update.author_name is not None, + update.response_id is not None, + update.message_id is not None, + update.conversation_id is not None, + update.model is not None, + update.created_at is not None, + update.finish_reason is not None, + update.continuation_token is not None, + bool(update.additional_properties), + update.raw_representation is not None, + )) + + def _extract_tools( options: dict[str, Any] | None, ) -> ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None: @@ -2823,20 +2855,31 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]: await inner_stream # Collect result hooks from the inner stream to run later stream_result_hooks[:] = _get_result_hooks_from_stream(inner_stream) + dropping_post_limit_function_calls = ( + mutable_options.get("tool_choice") == "none" + and max_function_calls is not None + and total_function_calls >= max_function_calls + ) # Yield updates from the inner stream, letting it collect them async for update in inner_stream: + if dropping_post_limit_function_calls: + update = _drop_function_call_contents_from_update(update) + if update is None: + continue yield update # Get the finalized response from the inner stream # This triggers the inner stream's finalizer and result hooks response = await inner_stream.get_final_response() - response_had_visible_content = _response_has_visible_content(response) function_call_limit_reached = ( mutable_options.get("tool_choice") == "none" and max_function_calls is not None and total_function_calls >= max_function_calls ) + if function_call_limit_reached: + _drop_function_call_contents_from_response(response) + response_had_visible_content = _response_has_visible_content(response) if function_call_limit_reached: response = _ensure_function_invocation_limit_fallback_response(response) _update_continuation_state( diff --git a/python/packages/core/tests/core/test_function_invocation_logic.py b/python/packages/core/tests/core/test_function_invocation_logic.py index ae01d82a226..eaffb4095e5 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -2,7 +2,7 @@ import asyncio from collections.abc import AsyncIterable, Awaitable, Callable, Sequence -from typing import Any +from typing import Any, Literal import pytest @@ -96,6 +96,47 @@ def _finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse: chat_client_base._get_streaming_response = _get_streaming_response +def _force_function_call_tool_choice_none_stream( + chat_client_base: Any, + *, + call_id: str, + name: str, + arguments: str, + additional_properties: dict[str, Any] | None = None, + finish_reason: Literal["stop", "length", "tool_calls", "content_filter"] | None = None, +) -> None: + original_get_streaming_response = chat_client_base._get_streaming_response + + def _get_streaming_response( + *, + messages: Sequence[Message], + options: dict[str, Any], + **kwargs: Any, + ) -> ResponseStream[ChatResponseUpdate, ChatResponse]: + if options.get("tool_choice") != "none": + return original_get_streaming_response(messages=messages, options=options, **kwargs) + + updates = ( + ChatResponseUpdate( + contents=[Content.from_function_call(call_id=call_id, name=name, arguments=arguments)], + role="assistant", + additional_properties=additional_properties, + finish_reason=finish_reason, + ), + ) + + async def _stream() -> AsyncIterable[ChatResponseUpdate]: + for update in updates: + yield update + + def _finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse: + return ChatResponse.from_updates(updates, output_format_type=options.get("response_format")) + + return ResponseStream(_stream(), finalizer=_finalize) + + chat_client_base._get_streaming_response = _get_streaming_response + + async def test_base_client_with_function_calling(chat_client_base: SupportsChatGetResponse): exec_counter = 0 @@ -2374,6 +2415,7 @@ def test_replace_approval_contents_with_results_allows_reused_call_id_after_comp ("call_reused", "second output"), ] + def test_replace_approval_contents_with_results_uses_result_call_ids_for_placeholders() -> None: from agent_framework._tools import _collect_approval_responses, _replace_approval_contents_with_results @@ -3284,6 +3326,109 @@ def lookup_func(key: str) -> str: assert updates[-1].text == _EXPECTED_FUNCTION_INVOCATION_LIMIT_FALLBACK_TEXT +@pytest.mark.parametrize("max_iterations", [10]) +async def test_streaming_max_function_calls_drops_post_limit_function_call_update( + chat_client_base: SupportsChatGetResponse, +): + """Function calls streamed after max_function_calls is reached should not be forwarded.""" + _force_function_call_tool_choice_none_stream( + chat_client_base, + call_id="call_2", + name="lookup", + arguments='{"key": "b"}', + ) + exec_counter = 0 + + @tool(name="lookup", approval_mode="never_require") + def lookup_func(key: str) -> str: + nonlocal exec_counter + exec_counter += 1 + return f"Value for {key}" + + chat_client_base.streaming_responses = [ # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + [ + ChatResponseUpdate( + contents=[Content.from_function_call(call_id="call_1", name="lookup", arguments='{"key": "a"}')], + role="assistant", + ), + ], + ] + chat_client_base.function_invocation_configuration["max_function_calls"] = 1 # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + + updates = [] + async for update in chat_client_base.get_response( # type: ignore[call-overload] # pyrefly: ignore[no-matching-overload] + [Message(role="user", contents=["look up key"])], + options={"tool_choice": "auto", "tools": [lookup_func]}, + stream=True, + ): + updates.append(update) + + function_call_ids = { + content.call_id for update in updates for content in update.contents if content.type == "function_call" + } + function_result_ids = { + content.call_id for update in updates for content in update.contents if content.type == "function_result" + } + + assert exec_counter == 1 + assert "call_1" in function_call_ids + assert "call_1" in function_result_ids + assert "call_2" not in function_call_ids + assert "call_2" not in function_result_ids + assert updates[-1].role == "assistant" + assert updates[-1].text == _EXPECTED_FUNCTION_INVOCATION_LIMIT_FALLBACK_TEXT + + +@pytest.mark.parametrize("max_iterations", [10]) +async def test_streaming_max_function_calls_preserves_post_limit_update_metadata( + chat_client_base: SupportsChatGetResponse, +): + """Post-limit function call chunks should retain metadata after call content is stripped.""" + _force_function_call_tool_choice_none_stream( + chat_client_base, + call_id="call_2", + name="lookup", + arguments='{"key": "b"}', + additional_properties={"provider_metadata": "keep"}, + finish_reason="stop", + ) + exec_counter = 0 + + @tool(name="lookup", approval_mode="never_require") + def lookup_func(key: str) -> str: + nonlocal exec_counter + exec_counter += 1 + return f"Value for {key}" + + chat_client_base.streaming_responses = [ # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + [ + ChatResponseUpdate( + contents=[Content.from_function_call(call_id="call_1", name="lookup", arguments='{"key": "a"}')], + role="assistant", + ), + ], + ] + chat_client_base.function_invocation_configuration["max_function_calls"] = 1 # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + + updates = [] + async for update in chat_client_base.get_response( # type: ignore[call-overload] # pyrefly: ignore[no-matching-overload] + [Message(role="user", contents=["look up key"])], + options={"tool_choice": "auto", "tools": [lookup_func]}, + stream=True, + ): + updates.append(update) + + metadata_updates = [ + update + for update in updates + if update.additional_properties == {"provider_metadata": "keep"} and update.finish_reason == "stop" + ] + + assert exec_counter == 1 + assert len(metadata_updates) == 1 + assert metadata_updates[0].contents == [] + + async def test_streaming_function_invocation_config_enabled_false(chat_client_base: SupportsChatGetResponse): """Test that setting enabled=False disables function invocation in streaming mode.""" exec_counter = 0