From 27f7af2160383edf570be94ef2f2fb482622b668 Mon Sep 17 00:00:00 2001 From: Eduard van Valkenburg Date: Tue, 22 Jul 2025 23:37:18 +0200 Subject: [PATCH] Python: moved prepare tools into class (#215) * moved prepare tools into class * moved test * changed tool handling * fix test * second fix --- .../agent_framework_foundry/_chat_client.py | 11 ++-- .../packages/main/agent_framework/_clients.py | 55 ++++++++++--------- python/packages/main/tests/test_clients.py | 45 +++++++++++++++ python/packages/main/tests/test_types.py | 48 ---------------- 4 files changed, 80 insertions(+), 79 deletions(-) diff --git a/python/packages/foundry/agent_framework_foundry/_chat_client.py b/python/packages/foundry/agent_framework_foundry/_chat_client.py index fc6f17e218..2869471c62 100644 --- a/python/packages/foundry/agent_framework_foundry/_chat_client.py +++ b/python/packages/foundry/agent_framework_foundry/_chat_client.py @@ -8,7 +8,7 @@ from typing import Any, ClassVar from agent_framework import ( AFBaseSettings, AIContents, - AITool, + AIFunction, ChatClientBase, ChatMessage, ChatOptions, @@ -25,7 +25,7 @@ from agent_framework import ( UsageDetails, use_tool_calling, ) -from agent_framework._clients import tool_to_json_schema_spec +from agent_framework._clients import ai_function_to_json_schema_spec from agent_framework.exceptions import ServiceInitializationError from azure.ai.agents.models import ( AgentsNamedToolChoice, @@ -455,13 +455,14 @@ class FoundryChatClient(ChatClientBase): run_options["parallel_tool_calls"] = chat_options.allow_multiple_tool_calls if chat_options.tools is not None: + # TODO (eavanvalkenburg): replace with _prepare_tools_and_tool_choice overload tool_definitions: list[MutableMapping[str, Any]] = [] for tool in chat_options.tools: - if isinstance(tool, AITool): - tool_definitions.append(tool_to_json_schema_spec(tool)) + if isinstance(tool, AIFunction): + tool_definitions.append(ai_function_to_json_schema_spec(tool)) else: - tool_definitions.append(tool) + tool_definitions.append(tool) # type: ignore if len(tool_definitions) > 0: run_options["tools"] = tool_definitions diff --git a/python/packages/main/agent_framework/_clients.py b/python/packages/main/agent_framework/_clients.py index 9d1ee4a5d3..a73154b4b8 100644 --- a/python/packages/main/agent_framework/_clients.py +++ b/python/packages/main/agent_framework/_clients.py @@ -75,35 +75,18 @@ async def _auto_invoke_function( ) -def tool_to_json_schema_spec(tool: AITool) -> dict[str, Any]: - """Convert a AITool to the JSON Schema function specification format.""" +def ai_function_to_json_schema_spec(function: AIFunction[BaseModel, Any]) -> dict[str, Any]: + """Convert a AIFunction to the JSON Schema function specification format.""" return { "type": "function", "function": { - "name": tool.name, - "description": tool.description, - "parameters": tool.parameters(), + "name": function.name, + "description": function.description, + "parameters": function.parameters(), }, } -def _prepare_tools_and_tool_choice(chat_options: ChatOptions) -> None: - """Prepare the tools and tool choice for the chat options.""" - chat_tool_mode: ChatToolMode | None = chat_options.tool_choice # type: ignore - if chat_tool_mode is None or chat_tool_mode == ChatToolMode.NONE: - chat_options.tools = None - chat_options.tool_choice = ChatToolMode.NONE.mode - return - chat_options.tools = [ - (tool_to_json_schema_spec(t) if isinstance(t, AITool) else t) - for t in chat_options._ai_tools or [] # type: ignore[reportPrivateUsage] - ] - if not chat_options.tools: - chat_options.tool_choice = ChatToolMode.NONE.mode - else: - chat_options.tool_choice = chat_tool_mode.mode - - def _tool_call_non_streaming(func: TInnerGetResponse) -> TInnerGetResponse: """Decorate the internal _inner_get_response method to enable tool calls.""" @@ -163,7 +146,7 @@ def _tool_call_non_streaming(func: TInnerGetResponse) -> TInnerGetResponse: # Failsafe: give up on tools, ask model for plain answer chat_options.tool_choice = "none" - _prepare_tools_and_tool_choice(chat_options=chat_options) + self._prepare_tools_and_tool_choice(chat_options=chat_options) # type: ignore[reportPrivateUsage] response = await func(self, messages=messages, chat_options=chat_options) if fcc_messages: for msg in reversed(fcc_messages): @@ -231,7 +214,7 @@ def _tool_call_streaming(func: TInnerGetStreamingResponse) -> TInnerGetStreaming # Failsafe: give up on tools, ask model for plain answer chat_options.tool_choice = "none" - _prepare_tools_and_tool_choice(chat_options=chat_options) + self._prepare_tools_and_tool_choice(chat_options=chat_options) # type: ignore[reportPrivateUsage] async for update in func(self, messages=messages, chat_options=chat_options, **kwargs): yield update @@ -542,7 +525,7 @@ class ChatClientBase(AFBaseModel, ABC): additional_properties=additional_properties or {}, ) prepped_messages = self._prepare_messages(messages) - _prepare_tools_and_tool_choice(chat_options=chat_options) + self._prepare_tools_and_tool_choice(chat_options=chat_options) return await self._inner_get_response(messages=prepped_messages, chat_options=chat_options, **kwargs) async def get_streaming_response( @@ -623,12 +606,32 @@ class ChatClientBase(AFBaseModel, ABC): **kwargs, ) prepped_messages = self._prepare_messages(messages) - _prepare_tools_and_tool_choice(chat_options=chat_options) + self._prepare_tools_and_tool_choice(chat_options=chat_options) async for update in self._inner_get_streaming_response( messages=prepped_messages, chat_options=chat_options, **kwargs ): yield update + def _prepare_tools_and_tool_choice(self, chat_options: ChatOptions) -> None: + """Prepare the tools and tool choice for the chat options. + + This function should be overridden by subclasses to customize tool handling. + Because it currently parses only AIFunctions. + """ + chat_tool_mode: ChatToolMode | None = chat_options.tool_choice # type: ignore + if chat_tool_mode is None or chat_tool_mode == ChatToolMode.NONE: + chat_options.tools = None + chat_options.tool_choice = ChatToolMode.NONE.mode + return + chat_options.tools = [ + (ai_function_to_json_schema_spec(t) if isinstance(t, AIFunction) else t) # type: ignore[reportUnknownArgumentType] + for t in chat_options._ai_tools or [] # type: ignore[reportPrivateUsage] + ] + if not chat_options.tools: + chat_options.tool_choice = ChatToolMode.NONE.mode + else: + chat_options.tool_choice = chat_tool_mode.mode + # region: Embedding Client diff --git a/python/packages/main/tests/test_clients.py b/python/packages/main/tests/test_clients.py index 36238a44b5..cea5b8b2a1 100644 --- a/python/packages/main/tests/test_clients.py +++ b/python/packages/main/tests/test_clients.py @@ -305,3 +305,48 @@ async def test_base_client_with_streaming_function_calling_disabled(chat_client_ updates.append(update) assert len(updates) == 1 assert exec_counter == 0 + + +def test_chat_options_parsing_tools(chat_client_base, ai_function_tool) -> None: + """Test that chat options can parse tools correctly.""" + + def echo() -> str: + """Echo the input.""" + return "Echo" + + dict_function = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Retrieves current weather for the given location.", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "City and country e.g. Bogotá, Colombia"}, + "units": { + "type": "string", + "enum": ["celsius", "fahrenheit"], + "description": "Units the temperature will be returned in.", + }, + }, + "required": ["location", "units"], + "additionalProperties": False, + }, + "strict": True, + }, + } + + options = ChatOptions(tools=[ai_function_tool, echo, dict_function], tool_choice="auto") + assert len(options.tools) == 3 + assert options.tools[0] == ai_function_tool + assert options.tools[1] != echo + assert options.tools[2] == dict_function + # after prepare, the tools should be represented as dicts + # while ai_tools is still the same. + chat_client_base._prepare_tools_and_tool_choice(chat_options=options) + assert options._ai_tools[0] == ai_function_tool + assert options._ai_tools[2] == dict_function + assert len(options.tools) == 3 + assert options.tools[0]["function"]["name"] == "simple_function" + assert options.tools[1]["function"]["name"] == "echo" + assert options.tools[2]["function"]["name"] == "get_weather" diff --git a/python/packages/main/tests/test_types.py b/python/packages/main/tests/test_types.py index 9ef498e2ab..b6e1bdba93 100644 --- a/python/packages/main/tests/test_types.py +++ b/python/packages/main/tests/test_types.py @@ -539,54 +539,6 @@ def test_chat_options_and(ai_function_tool, ai_tool) -> None: assert options3.tools == [ai_function_tool, ai_tool] -def test_chat_options_parsing_tools(ai_function_tool, ai_tool) -> None: - from agent_framework._clients import _prepare_tools_and_tool_choice - - def echo() -> str: - """Echo the input.""" - return "Echo" - - dict_function = { - "type": "function", - "function": { - "name": "get_weather", - "description": "Retrieves current weather for the given location.", - "parameters": { - "type": "object", - "properties": { - "location": {"type": "string", "description": "City and country e.g. Bogotá, Colombia"}, - "units": { - "type": "string", - "enum": ["celsius", "fahrenheit"], - "description": "Units the temperature will be returned in.", - }, - }, - "required": ["location", "units"], - "additionalProperties": False, - }, - "strict": True, - }, - } - - options = ChatOptions(tools=[ai_function_tool, ai_tool, echo, dict_function], tool_choice="auto") - assert len(options.tools) == 4 - assert options.tools[0] == ai_function_tool - assert options.tools[1] == ai_tool - assert options.tools[2] != echo - assert options.tools[3] == dict_function - # after prepare, the tools should be represented as dicts - # while ai_tools is still the same. - _prepare_tools_and_tool_choice(options) - assert options._ai_tools[0] == ai_function_tool - assert options._ai_tools[1] == ai_tool - assert options._ai_tools[3] == dict_function - assert len(options.tools) == 4 - assert options.tools[0]["function"]["name"] == "simple_function" - assert options.tools[1]["function"]["name"] == "generic_tool" - assert options.tools[2]["function"]["name"] == "echo" - assert options.tools[3]["function"]["name"] == "get_weather" - - # region Agent Response Fixtures