agent/agent-framework
1
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Python: [BREAKING] changed AIFunction to FunctionTool and @ai_function to @tool (#3413)

Browse Source 
* changed AIFunction to FunctionTool and @ai_function to @tool

* test and mypy fixes

* mypy fix

* switch function tool to always_require

* fix noop

* fix github copilot imports

* test fixes

* fix ollama test

* fixes for tests

* fix tests

* reverted change to always_require and extended timeout

* fix test
This commit is contained in:
Eduard van Valkenburg
2026-01-28 15:53:53 +01:00
committed by GitHub
Unverified
parent 15b43f2abe
commit a7d924a7d2
255 changed files with 1202 additions and 1290 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/core/tests/azure/test_azure_assistants_client.py
+2
View File
@@ -18,6 +18,7 @@ from agent_framework import (
ChatResponse,
ChatResponseUpdate,
HostedCodeInterpreterTool,
tool,
)
from agent_framework.azure import AzureOpenAIAssistantsClient
from agent_framework.exceptions import ServiceInitializationError
@@ -253,6 +254,7 @@ def test_azure_assistants_client_serialize(azure_openai_unit_test_env: dict[str,
assert "User-Agent" not in dumped_settings["default_headers"]
@tool(approval_mode="never_require")
def get_weather(
location: Annotated[str, Field(description="The location to get the weather for.")],
) -> str:
python/packages/core/tests/azure/test_azure_chat_client.py
+3 -3
View File
@@ -25,7 +25,7 @@ from agent_framework import (
ChatMessage,
ChatResponse,
ChatResponseUpdate,
ai_function,
tool,
)
from agent_framework._telemetry import USER_AGENT_KEY
from agent_framework.azure import AzureOpenAIChatClient
@@ -631,7 +631,7 @@ async def test_streaming_with_none_delta(
assert any(msg.contents for msg in results)
@ai_function
@tool(approval_mode="never_require")
def get_story_text() -> str:
"""Returns a story about Emily and David."""
return (
@@ -642,7 +642,7 @@ def get_story_text() -> str:
)
@ai_function
@tool(approval_mode="never_require")
def get_weather(location: str) -> str:
"""Get the current weather for a location."""
return f"The weather in {location} is sunny and 72°F."
python/packages/core/tests/azure/test_azure_responses_client.py
+2 -2
View File
@@ -20,7 +20,7 @@ from agent_framework import (
HostedFileSearchTool,
HostedMCPTool,
HostedWebSearchTool,
ai_function,
tool,
)
from agent_framework.azure import AzureOpenAIResponsesClient
from agent_framework.exceptions import ServiceInitializationError
@@ -41,7 +41,7 @@ class OutputStruct(BaseModel):
weather: str
@ai_function
@tool(approval_mode="never_require")
async def get_weather(location: Annotated[str, "The location as a city name"]) -> str:
"""Get the current weather in a given location."""
# Implementation of the tool to get weather
python/packages/core/tests/core/conftest.py
+3 -3
View File
@@ -23,7 +23,7 @@ from agent_framework import (
Content,
Role,
ToolProtocol,
ai_function,
tool,
use_chat_middleware,
use_function_invocation,
)
@@ -65,10 +65,10 @@ def ai_tool() -> ToolProtocol:
@fixture
def ai_function_tool() -> ToolProtocol:
def tool_tool() -> ToolProtocol:
"""Returns a executable ToolProtocol."""
@ai_function
@tool(approval_mode="never_require")
def simple_function(x: int, y: int) -> int:
"""A simple function that adds two numbers."""
return x + y
python/packages/core/tests/core/test_agents.py
+9 -13
View File
@@ -23,7 +23,7 @@ from agent_framework import (
ContextProvider,
HostedCodeInterpreterTool,
Role,
ai_function,
tool,
)
from agent_framework._mcp import MCPTool
from agent_framework.exceptions import AgentExecutionException
@@ -423,7 +423,7 @@ async def test_chat_agent_as_tool_no_name(chat_client: ChatClientProtocol) -> No
async def test_chat_agent_as_tool_function_execution(chat_client: ChatClientProtocol) -> None:
"""Test that the generated AIFunction can be executed."""
"""Test that the generated FunctionTool can be executed."""
agent = ChatAgent(chat_client=chat_client, name="TestAgent", description="Test agent")
tool = agent.as_tool()
@@ -564,7 +564,7 @@ async def test_agent_tool_receives_thread_in_kwargs(chat_client_base: Any) -> No
captured: dict[str, Any] = {}
@ai_function(name="echo_thread_info")
@tool(name="echo_thread_info", approval_mode="never_require")
def echo_thread_info(text: str, **kwargs: Any) -> str: # type: ignore[reportUnknownParameterType]
thread = kwargs.get("thread")
captured["has_thread"] = thread is not None
@@ -596,9 +596,7 @@ async def test_agent_tool_receives_thread_in_kwargs(chat_client_base: Any) -> No
assert captured.get("has_message_store") is True
async def test_chat_agent_tool_choice_run_level_overrides_agent_level(
chat_client_base: Any, ai_function_tool: Any
) -> None:
async def test_chat_agent_tool_choice_run_level_overrides_agent_level(chat_client_base: Any, tool_tool: Any) -> None:
"""Verify that tool_choice passed to run() overrides agent-level tool_choice."""
captured_options: list[dict[str, Any]] = []
@@ -617,7 +615,7 @@ async def test_chat_agent_tool_choice_run_level_overrides_agent_level(
# Create agent with agent-level tool_choice="auto" and a tool (tools required for tool_choice to be meaningful)
agent = ChatAgent(
chat_client=chat_client_base,
tools=[ai_function_tool],
tools=[tool_tool],
options={"tool_choice": "auto"},
)
@@ -630,7 +628,7 @@ async def test_chat_agent_tool_choice_run_level_overrides_agent_level(
async def test_chat_agent_tool_choice_agent_level_used_when_run_level_not_specified(
chat_client_base: Any, ai_function_tool: Any
chat_client_base: Any, tool_tool: Any
) -> None:
"""Verify that agent-level tool_choice is used when run() doesn't specify one."""
from agent_framework import ChatOptions
@@ -650,7 +648,7 @@ async def test_chat_agent_tool_choice_agent_level_used_when_run_level_not_specif
# Create agent with agent-level tool_choice="required" and a tool
agent = ChatAgent(
chat_client=chat_client_base,
tools=[ai_function_tool],
tools=[tool_tool],
default_options={"tool_choice": "required"},
)
@@ -664,9 +662,7 @@ async def test_chat_agent_tool_choice_agent_level_used_when_run_level_not_specif
assert captured_options[0]["tool_choice"] == "required"
async def test_chat_agent_tool_choice_none_at_run_preserves_agent_level(
chat_client_base: Any, ai_function_tool: Any
) -> None:
async def test_chat_agent_tool_choice_none_at_run_preserves_agent_level(chat_client_base: Any, tool_tool: Any) -> None:
"""Verify that tool_choice=None at run() uses agent-level default."""
from agent_framework import ChatOptions
@@ -685,7 +681,7 @@ async def test_chat_agent_tool_choice_none_at_run_preserves_agent_level(
# Create agent with agent-level tool_choice="auto" and a tool
agent = ChatAgent(
chat_client=chat_client_base,
tools=[ai_function_tool],
tools=[tool_tool],
default_options={"tool_choice": "auto"},
)
python/packages/core/tests/core/test_chat_agent_integration.py
-433
View File
@@ -1,433 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
import json
import os
from typing import Annotated
import pytest
from pydantic import BaseModel
from agent_framework import (
AgentResponse,
AgentResponseUpdate,
AgentThread,
ChatAgent,
HostedCodeInterpreterTool,
HostedImageGenerationTool,
HostedMCPTool,
MCPStreamableHTTPTool,
ai_function,
)
from agent_framework.openai import OpenAIResponsesClient
skip_if_openai_integration_tests_disabled = pytest.mark.skipif(
os.getenv("RUN_INTEGRATION_TESTS", "false").lower() != "true"
or os.getenv("OPENAI_API_KEY", "") in ("", "test-dummy-key"),
reason="No real OPENAI_API_KEY provided; skipping integration tests."
if os.getenv("RUN_INTEGRATION_TESTS", "false").lower() == "true"
else "Integration tests are disabled.",
)
@ai_function
async def get_weather(location: Annotated[str, "The location as a city name"]) -> str:
"""Get the current weather in a given location."""
# Implementation of the tool to get weather
return f"The current weather in {location} is sunny."
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_agent_basic_run_streaming():
"""Test OpenAI Responses Client agent basic streaming functionality with OpenAIResponsesClient."""
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
) as agent:
# Test streaming run
full_text = ""
async for chunk in agent.run_stream("Please respond with exactly: 'This is a streaming response test.'"):
assert isinstance(chunk, AgentResponseUpdate)
if chunk.text:
full_text += chunk.text
assert len(full_text) > 0
assert "streaming response test" in full_text.lower()
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_agent_thread_persistence():
"""Test OpenAI Responses Client agent thread persistence across runs with OpenAIResponsesClient."""
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant with good memory.",
) as agent:
# Create a new thread that will be reused
thread = agent.get_new_thread()
# First interaction
first_response = await agent.run("My favorite programming language is Python. Remember this.", thread=thread)
assert isinstance(first_response, AgentResponse)
assert first_response.text is not None
# Second interaction - test memory
second_response = await agent.run("What is my favorite programming language?", thread=thread)
assert isinstance(second_response, AgentResponse)
assert second_response.text is not None
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_agent_thread_storage_with_store_true():
"""Test OpenAI Responses Client agent with store=True to verify service_thread_id is returned."""
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant.",
) as agent:
# Create a new thread
thread = AgentThread()
# Initially, service_thread_id should be None
assert thread.service_thread_id is None
# Run with store=True to store messages on OpenAI side
response = await agent.run(
"Hello! Please remember that my name is Alex.",
thread=thread,
options={"store": True},
)
# Validate response
assert isinstance(response, AgentResponse)
assert response.text is not None
assert len(response.text) > 0
# After store=True, service_thread_id should be populated
assert thread.service_thread_id is not None
assert isinstance(thread.service_thread_id, str)
assert len(thread.service_thread_id) > 0
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_agent_existing_thread():
"""Test OpenAI Responses Client agent with existing thread to continue conversations across agent instances."""
# First conversation - capture the thread
preserved_thread = None
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant with good memory.",
) as first_agent:
# Start a conversation and capture the thread
thread = first_agent.get_new_thread()
first_response = await first_agent.run("My hobby is photography. Remember this.", thread=thread)
assert isinstance(first_response, AgentResponse)
assert first_response.text is not None
# Preserve the thread for reuse
preserved_thread = thread
# Second conversation - reuse the thread in a new agent instance
if preserved_thread:
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant with good memory.",
) as second_agent:
# Reuse the preserved thread
second_response = await second_agent.run("What is my hobby?", thread=preserved_thread)
assert isinstance(second_response, AgentResponse)
assert second_response.text is not None
assert "photography" in second_response.text.lower()
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_agent_hosted_code_interpreter_tool():
"""Test OpenAI Responses Client agent with HostedCodeInterpreterTool through OpenAIResponsesClient."""
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant that can execute Python code.",
tools=[HostedCodeInterpreterTool()],
) as agent:
# Test code interpreter functionality
response = await agent.run("Calculate the sum of numbers from 1 to 10 using Python code.")
assert isinstance(response, AgentResponse)
assert response.text is not None
assert len(response.text) > 0
# Should contain calculation result (sum of 1-10 = 55) or code execution content
contains_relevant_content = any(
term in response.text.lower() for term in ["55", "sum", "code", "python", "calculate", "10"]
)
assert contains_relevant_content or len(response.text.strip()) > 10
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_agent_image_generation_tool():
"""Test OpenAI Responses Client agent with raw image_generation tool through OpenAIResponsesClient."""
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant that can generate images.",
tools=HostedImageGenerationTool(options={"image_size": "1024x1024", "media_type": "png"}),
) as agent:
# Test image generation functionality
response = await agent.run("Generate an image of a cute red panda sitting on a tree branch in a forest.")
assert isinstance(response, AgentResponse)
assert response.messages
# Verify we got image content - look for ImageGenerationToolResultContent
image_content_found = False
for message in response.messages:
for content in message.contents:
if content.type == "image_generation_tool_result" and content.outputs:
image_content_found = True
break
if image_content_found:
break
# The test passes if we got image content
assert image_content_found, "Expected to find image content in response"
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_agent_level_tool_persistence():
"""Test that agent-level tools persist across multiple runs with OpenAI Responses Client."""
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant that uses available tools.",
tools=[get_weather], # Agent-level tool
) as agent:
# First run - agent-level tool should be available
first_response = await agent.run("What's the weather like in Chicago?")
assert isinstance(first_response, AgentResponse)
assert first_response.text is not None
# Should use the agent-level weather tool
assert any(term in first_response.text.lower() for term in ["chicago", "sunny", "72"])
# Second run - agent-level tool should still be available (persistence test)
second_response = await agent.run("What's the weather in Miami?")
assert isinstance(second_response, AgentResponse)
assert second_response.text is not None
# Should use the agent-level weather tool again
assert any(term in second_response.text.lower() for term in ["miami", "sunny", "72"])
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_run_level_tool_isolation():
"""Test that run-level tools are isolated to specific runs and don't persist with OpenAI Responses Client."""
# Counter to track how many times the weather tool is called
call_count = 0
@ai_function
async def get_weather_with_counter(
location: Annotated[str, "The location as a city name"],
) -> str:
"""Get the current weather in a given location."""
nonlocal call_count
call_count += 1
return f"The weather in {location} is sunny and 72°F."
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant.",
) as agent:
# First run - use run-level tool
first_response = await agent.run(
"What's the weather like in Chicago?",
tools=[get_weather_with_counter], # Run-level tool
)
assert isinstance(first_response, AgentResponse)
assert first_response.text is not None
# Should use the run-level weather tool (call count should be 1)
assert call_count == 1
assert any(term in first_response.text.lower() for term in ["chicago", "sunny", "72"])
# Second run - run-level tool should NOT persist (key isolation test)
second_response = await agent.run("What's the weather like in Miami?")
assert isinstance(second_response, AgentResponse)
assert second_response.text is not None
# Should NOT use the weather tool since it was only run-level in previous call
# Call count should still be 1 (no additional calls)
assert call_count == 1
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_agent_chat_options_agent_level() -> None:
"""Integration test for comprehensive ChatOptions parameter coverage with OpenAI Response Agent."""
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant.",
tools=[get_weather],
default_options={
"max_tokens": 100,
"temperature": 0.7,
"top_p": 0.9,
"seed": 123,
"user": "comprehensive-test-user",
"tool_choice": "auto",
},
) as agent:
response = await agent.run(
"Provide a brief, helpful response.",
)
assert isinstance(response, AgentResponse)
assert response.text is not None
assert len(response.text) > 0
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_agent_hosted_mcp_tool() -> None:
"""Integration test for HostedMCPTool with OpenAI Response Agent using Microsoft Learn MCP."""
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant that can help with microsoft documentation questions.",
tools=HostedMCPTool(
name="Microsoft Learn MCP",
url="https://learn.microsoft.com/api/mcp",
description="A Microsoft Learn MCP server for documentation questions",
approval_mode="never_require",
),
) as agent:
response = await agent.run(
"How to create an Azure storage account using az cli?",
# this needs to be high enough to handle the full MCP tool response.
options={"max_tokens": 5000},
)
assert isinstance(response, AgentResponse)
assert response.text
# Should contain Azure-related content since it's asking about Azure CLI
assert any(term in response.text.lower() for term in ["azure", "storage", "account", "cli"])
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_agent_local_mcp_tool() -> None:
"""Integration test for MCPStreamableHTTPTool with OpenAI Response Agent using Microsoft Learn MCP."""
mcp_tool = MCPStreamableHTTPTool(
name="Microsoft Learn MCP",
url="https://learn.microsoft.com/api/mcp",
)
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant that can help with microsoft documentation questions.",
tools=[mcp_tool],
) as agent:
response = await agent.run(
"How to create an Azure storage account using az cli?",
options={"max_tokens": 200},
)
assert isinstance(response, AgentResponse)
assert response.text is not None
assert len(response.text) > 0
# Should contain Azure-related content since it's asking about Azure CLI
assert any(term in response.text.lower() for term in ["azure", "storage", "account", "cli"])
class ReleaseBrief(BaseModel):
"""Structured output model for release brief testing."""
title: str
summary: str
highlights: list[str]
model_config = {"extra": "forbid"}
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_agent_with_response_format_pydantic() -> None:
"""Integration test for response_format with Pydantic model using OpenAI Responses Client."""
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="You are a helpful assistant that returns structured JSON responses.",
) as agent:
response = await agent.run(
"Summarize the following release notes into a ReleaseBrief:\n\n"
"Version 2.0 Release Notes:\n"
"- Added new streaming API for real-time responses\n"
"- Improved error handling with detailed messages\n"
"- Performance boost of 50% in batch processing\n"
"- Fixed memory leak in connection pooling",
options={
"response_format": ReleaseBrief,
},
)
# Validate response
assert isinstance(response, AgentResponse)
assert response.value is not None
assert isinstance(response.value, ReleaseBrief)
# Validate structured output fields
brief = response.value
assert len(brief.title) > 0
assert len(brief.summary) > 0
assert len(brief.highlights) > 0
@pytest.mark.flaky
@skip_if_openai_integration_tests_disabled
async def test_openai_responses_client_agent_with_runtime_json_schema() -> None:
"""Integration test for response_format with runtime JSON schema using OpenAI Responses Client."""
runtime_schema = {
"title": "WeatherDigest",
"type": "object",
"properties": {
"location": {"type": "string"},
"conditions": {"type": "string"},
"temperature_c": {"type": "number"},
"advisory": {"type": "string"},
},
"required": ["location", "conditions", "temperature_c", "advisory"],
"additionalProperties": False,
}
async with ChatAgent(
chat_client=OpenAIResponsesClient(),
instructions="Return only JSON that matches the provided schema. Do not add commentary.",
) as agent:
response = await agent.run(
"Give a brief weather digest for Seattle.",
options={
"response_format": {
"type": "json_schema",
"json_schema": {
"name": runtime_schema["title"],
"strict": True,
"schema": runtime_schema,
},
},
},
)
# Validate response
assert isinstance(response, AgentResponse)
assert response.text is not None
# Parse JSON and validate structure
parsed = json.loads(response.text)
assert "location" in parsed
assert "conditions" in parsed
assert "temperature_c" in parsed
assert "advisory" in parsed
python/packages/core/tests/core/test_function_invocation_logic.py
+60 -59
View File
@@ -1,5 +1,6 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
from collections.abc import Awaitable, Callable
from typing import Any
@@ -13,7 +14,7 @@ from agent_framework import (
ChatResponseUpdate,
Content,
Role,
ai_function,
tool,
)
from agent_framework._middleware import FunctionInvocationContext, FunctionMiddleware
@@ -21,7 +22,7 @@ from agent_framework._middleware import FunctionInvocationContext, FunctionMiddl
async def test_base_client_with_function_calling(chat_client_base: ChatClientProtocol):
exec_counter = 0
@ai_function(name="test_function")
@tool(name="test_function", approval_mode="never_require")
def ai_func(arg1: str) -> str:
nonlocal exec_counter
exec_counter += 1
@@ -57,7 +58,7 @@ async def test_base_client_with_function_calling(chat_client_base: ChatClientPro
async def test_base_client_with_function_calling_resets(chat_client_base: ChatClientProtocol):
exec_counter = 0
@ai_function(name="test_function")
@tool(name="test_function", approval_mode="never_require")
def ai_func(arg1: str) -> str:
nonlocal exec_counter
exec_counter += 1
@@ -99,7 +100,7 @@ async def test_base_client_with_function_calling_resets(chat_client_base: ChatCl
async def test_base_client_with_streaming_function_calling(chat_client_base: ChatClientProtocol):
exec_counter = 0
@ai_function(name="test_function")
@tool(name="test_function", approval_mode="never_require")
def ai_func(arg1: str) -> str:
nonlocal exec_counter
exec_counter += 1
@@ -142,7 +143,7 @@ async def test_function_invocation_inside_aiohttp_server(chat_client_base: ChatC
exec_counter = 0
@ai_function(name="start_todo_investigation")
@tool(name="start_todo_investigation", approval_mode="never_require")
def ai_func(user_query: str) -> str:
nonlocal exec_counter
exec_counter += 1
@@ -199,7 +200,7 @@ async def test_function_invocation_in_threaded_aiohttp_app(chat_client_base: Cha
exec_counter = 0
@ai_function(name="start_threaded_investigation")
@tool(name="start_threaded_investigation", approval_mode="never_require")
def ai_func(user_query: str) -> str:
nonlocal exec_counter
exec_counter += 1
@@ -319,13 +320,13 @@ async def test_function_invocation_scenarios(
# Simulate a service-side thread with conversation_id
conversation_id = "test-thread-123"
@ai_function(name="no_approval_func")
@tool(name="no_approval_func", approval_mode="never_require")
def func_no_approval(arg1: str) -> str:
nonlocal exec_counter
exec_counter += 1
return f"Processed {arg1}"
@ai_function(name="approval_func", approval_mode="always_require")
@tool(name="approval_func", approval_mode="always_require")
def func_with_approval(arg1: str) -> str:
nonlocal exec_counter
exec_counter += 1
@@ -473,13 +474,13 @@ async def test_rejected_approval(chat_client_base: ChatClientProtocol):
exec_counter_approved = 0
exec_counter_rejected = 0
@ai_function(name="approved_func", approval_mode="always_require")
@tool(name="approved_func", approval_mode="always_require")
def func_approved(arg1: str) -> str:
nonlocal exec_counter_approved
exec_counter_approved += 1
return f"Approved {arg1}"
@ai_function(name="rejected_func", approval_mode="always_require")
@tool(name="rejected_func", approval_mode="always_require")
def func_rejected(arg1: str) -> str:
nonlocal exec_counter_rejected
exec_counter_rejected += 1
@@ -569,7 +570,7 @@ async def test_approval_requests_in_assistant_message(chat_client_base: ChatClie
"""Approval requests should be added to the assistant message that contains the function call."""
exec_counter = 0
@ai_function(name="test_func", approval_mode="always_require")
@tool(name="test_func", approval_mode="always_require")
def func_with_approval(arg1: str) -> str:
nonlocal exec_counter
exec_counter += 1
@@ -604,7 +605,7 @@ async def test_persisted_approval_messages_replay_correctly(chat_client_base: Ch
exec_counter = 0
@ai_function(name="test_func", approval_mode="always_require")
@tool(name="test_func", approval_mode="always_require")
def func_with_approval(arg1: str) -> str:
nonlocal exec_counter
exec_counter += 1
@@ -656,7 +657,7 @@ async def test_persisted_approval_messages_replay_correctly(chat_client_base: Ch
async def test_no_duplicate_function_calls_after_approval_processing(chat_client_base: ChatClientProtocol):
"""Processing approval should not create duplicate function calls in messages."""
@ai_function(name="test_func", approval_mode="always_require")
@tool(name="test_func", approval_mode="always_require")
def func_with_approval(arg1: str) -> str:
return f"Result {arg1}"
@@ -700,7 +701,7 @@ async def test_no_duplicate_function_calls_after_approval_processing(chat_client
async def test_rejection_result_uses_function_call_id(chat_client_base: ChatClientProtocol):
"""Rejection error result should use the function call's call_id, not the approval's id."""
@ai_function(name="test_func", approval_mode="always_require")
@tool(name="test_func", approval_mode="always_require")
def func_with_approval(arg1: str) -> str:
return f"Result {arg1}"
@@ -745,7 +746,7 @@ async def test_max_iterations_limit(chat_client_base: ChatClientProtocol):
"""Test that MAX_ITERATIONS in additional_properties limits function call loops."""
exec_counter = 0
@ai_function(name="test_function")
@tool(name="test_function", approval_mode="never_require")
def ai_func(arg1: str) -> str:
nonlocal exec_counter
exec_counter += 1
@@ -790,7 +791,7 @@ async def test_function_invocation_config_enabled_false(chat_client_base: ChatCl
"""Test that setting enabled=False disables function invocation."""
exec_counter = 0
@ai_function(name="test_function")
@tool(name="test_function")
def ai_func(arg1: str) -> str:
nonlocal exec_counter
exec_counter += 1
@@ -814,7 +815,7 @@ async def test_function_invocation_config_enabled_false(chat_client_base: ChatCl
async def test_function_invocation_config_max_consecutive_errors(chat_client_base: ChatClientProtocol):
"""Test that max_consecutive_errors_per_request limits error retries."""
@ai_function(name="error_function")
@tool(name="error_function", approval_mode="never_require")
def error_func(arg1: str) -> str:
raise ValueError("Function error")
@@ -882,7 +883,7 @@ async def test_function_invocation_config_terminate_on_unknown_calls_false(chat_
"""Test that terminate_on_unknown_calls=False returns error message for unknown functions."""
exec_counter = 0
@ai_function(name="known_function")
@tool(name="known_function")
def known_func(arg1: str) -> str:
nonlocal exec_counter
exec_counter += 1
@@ -917,7 +918,7 @@ async def test_function_invocation_config_terminate_on_unknown_calls_true(chat_c
"""Test that terminate_on_unknown_calls=True stops execution on unknown functions."""
exec_counter = 0
@ai_function(name="known_function")
@tool(name="known_function")
def known_func(arg1: str) -> str:
nonlocal exec_counter
exec_counter += 1
@@ -949,13 +950,13 @@ async def test_function_invocation_config_additional_tools(chat_client_base: Cha
exec_counter_visible = 0
exec_counter_hidden = 0
@ai_function(name="visible_function")
@tool(name="visible_function")
def visible_func(arg1: str) -> str:
nonlocal exec_counter_visible
exec_counter_visible += 1
return f"Visible {arg1}"
@ai_function(name="hidden_function")
@tool(name="hidden_function")
def hidden_func(arg1: str) -> str:
nonlocal exec_counter_hidden
exec_counter_hidden += 1
@@ -996,7 +997,7 @@ async def test_function_invocation_config_additional_tools(chat_client_base: Cha
async def test_function_invocation_config_include_detailed_errors_false(chat_client_base: ChatClientProtocol):
"""Test that include_detailed_errors=False returns generic error messages."""
@ai_function(name="error_function")
@tool(name="error_function", approval_mode="never_require")
def error_func(arg1: str) -> str:
raise ValueError("Specific error message that should not appear")
@@ -1030,7 +1031,7 @@ async def test_function_invocation_config_include_detailed_errors_false(chat_cli
async def test_function_invocation_config_include_detailed_errors_true(chat_client_base: ChatClientProtocol):
"""Test that include_detailed_errors=True returns detailed error information."""
@ai_function(name="error_function")
@tool(name="error_function", approval_mode="never_require")
def error_func(arg1: str) -> str:
raise ValueError("Specific error message that should appear")
@@ -1100,7 +1101,7 @@ async def test_function_invocation_config_validation_max_consecutive_errors():
async def test_argument_validation_error_with_detailed_errors(chat_client_base: ChatClientProtocol):
"""Test that argument validation errors include details when include_detailed_errors=True."""
@ai_function(name="typed_function")
@tool(name="typed_function", approval_mode="never_require")
def typed_func(arg1: int) -> str: # Expects int, not str
return f"Got {arg1}"
@@ -1134,7 +1135,7 @@ async def test_argument_validation_error_with_detailed_errors(chat_client_base:
async def test_argument_validation_error_without_detailed_errors(chat_client_base: ChatClientProtocol):
"""Test that argument validation errors are generic when include_detailed_errors=False."""
@ai_function(name="typed_function")
@tool(name="typed_function", approval_mode="never_require")
def typed_func(arg1: int) -> str: # Expects int, not str
return f"Got {arg1}"
@@ -1168,7 +1169,7 @@ async def test_argument_validation_error_without_detailed_errors(chat_client_bas
async def test_hosted_tool_approval_response(chat_client_base: ChatClientProtocol):
"""Test handling of approval responses for hosted tools (tools not in tool_map)."""
@ai_function(name="local_function")
@tool(name="local_function")
def local_func(arg1: str) -> str:
return f"Local {arg1}"
@@ -1201,7 +1202,7 @@ async def test_hosted_tool_approval_response(chat_client_base: ChatClientProtoco
async def test_unapproved_tool_execution_raises_exception(chat_client_base: ChatClientProtocol):
"""Test that attempting to execute an unapproved tool raises ToolException."""
@ai_function(name="test_function", approval_mode="always_require")
@tool(name="test_function", approval_mode="always_require")
def test_func(arg1: str) -> str:
return f"Result {arg1}"
@@ -1256,7 +1257,7 @@ async def test_approved_function_call_with_error_without_detailed_errors(chat_cl
exec_counter = 0
@ai_function(name="error_func", approval_mode="always_require")
@tool(name="error_func", approval_mode="always_require")
def error_func(arg1: str) -> str:
nonlocal exec_counter
exec_counter += 1
@@ -1319,7 +1320,7 @@ async def test_approved_function_call_with_error_with_detailed_errors(chat_clien
exec_counter = 0
@ai_function(name="error_func", approval_mode="always_require")
@tool(name="error_func", approval_mode="always_require")
def error_func(arg1: str) -> str:
nonlocal exec_counter
exec_counter += 1
@@ -1380,7 +1381,7 @@ async def test_approved_function_call_with_validation_error(chat_client_base: Ch
exec_counter = 0
@ai_function(name="typed_func", approval_mode="always_require")
@tool(name="typed_func", approval_mode="always_require")
def typed_func(arg1: int) -> str: # Expects int, not str
nonlocal exec_counter
exec_counter += 1
@@ -1441,7 +1442,7 @@ async def test_approved_function_call_successful_execution(chat_client_base: Cha
exec_counter = 0
@ai_function(name="success_func", approval_mode="always_require")
@tool(name="success_func", approval_mode="always_require")
def success_func(arg1: str) -> str:
nonlocal exec_counter
exec_counter += 1
@@ -1493,10 +1494,10 @@ async def test_approved_function_call_successful_execution(chat_client_base: Cha
async def test_declaration_only_tool(chat_client_base: ChatClientProtocol):
"""Test that declaration_only tools without implementation (func=None) are not executed."""
from agent_framework import AIFunction
from agent_framework import FunctionTool
# Create a truly declaration-only function with no implementation
declaration_func = AIFunction(
declaration_func = FunctionTool(
name="declaration_func",
func=None,
description="A declaration-only function for testing",
@@ -1547,14 +1548,14 @@ async def test_multiple_function_calls_parallel_execution(chat_client_base: Chat
exec_order = []
@ai_function(name="func1")
@tool(name="func1", approval_mode="never_require")
async def func1(arg1: str) -> str:
exec_order.append("func1_start")
await asyncio.sleep(0.01) # Small delay
exec_order.append("func1_end")
return f"Result1 {arg1}"
@ai_function(name="func2")
@tool(name="func2", approval_mode="never_require")
async def func2(arg1: str) -> str:
exec_order.append("func2_start")
await asyncio.sleep(0.01) # Small delay
@@ -1587,10 +1588,11 @@ async def test_multiple_function_calls_parallel_execution(chat_client_base: Chat
assert len(results) == 2
async def test_callable_function_converted_to_ai_function(chat_client_base: ChatClientProtocol):
"""Test that plain callable functions are converted to AIFunction."""
async def test_callable_function_converted_to_tool(chat_client_base: ChatClientProtocol):
"""Test that plain callable functions are converted to FunctionTool."""
exec_counter = 0
@tool(approval_mode="never_require")
def plain_function(arg1: str) -> str:
"""A plain function without decorator."""
nonlocal exec_counter
@@ -1621,7 +1623,7 @@ async def test_callable_function_converted_to_ai_function(chat_client_base: Chat
async def test_conversation_id_handling(chat_client_base: ChatClientProtocol):
"""Test that conversation_id is properly handled and messages are cleared."""
@ai_function(name="test_function")
@tool(name="test_function", approval_mode="never_require")
def test_func(arg1: str) -> str:
return f"Result {arg1}"
@@ -1653,7 +1655,7 @@ async def test_conversation_id_handling(chat_client_base: ChatClientProtocol):
async def test_function_result_appended_to_existing_assistant_message(chat_client_base: ChatClientProtocol):
"""Test that function results are appended to existing assistant message when appropriate."""
@ai_function(name="test_function")
@tool(name="test_function", approval_mode="never_require")
def test_func(arg1: str) -> str:
return f"Result {arg1}"
@@ -1685,7 +1687,7 @@ async def test_error_recovery_resets_counter(chat_client_base: ChatClientProtoco
call_count = 0
@ai_function(name="sometimes_fails")
@tool(name="sometimes_fails", approval_mode="never_require")
def sometimes_fails(arg1: str) -> str:
nonlocal call_count
call_count += 1
@@ -1741,7 +1743,7 @@ async def test_streaming_approval_request_generated(chat_client_base: ChatClient
"""Test that approval requests are generated correctly in streaming mode."""
exec_counter = 0
@ai_function(name="test_func", approval_mode="always_require")
@tool(name="test_func", approval_mode="always_require")
def func_with_approval(arg1: str) -> str:
nonlocal exec_counter
exec_counter += 1
@@ -1777,7 +1779,7 @@ async def test_streaming_max_iterations_limit(chat_client_base: ChatClientProtoc
"""Test that MAX_ITERATIONS in streaming mode limits function call loops."""
exec_counter = 0
@ai_function(name="test_function")
@tool(name="test_function", approval_mode="never_require")
def ai_func(arg1: str) -> str:
nonlocal exec_counter
exec_counter += 1
@@ -1829,7 +1831,7 @@ async def test_streaming_function_invocation_config_enabled_false(chat_client_ba
"""Test that setting enabled=False disables function invocation in streaming mode."""
exec_counter = 0
@ai_function(name="test_function")
@tool(name="test_function", approval_mode="never_require")
def ai_func(arg1: str) -> str:
nonlocal exec_counter
exec_counter += 1
@@ -1857,7 +1859,7 @@ async def test_streaming_function_invocation_config_enabled_false(chat_client_ba
async def test_streaming_function_invocation_config_max_consecutive_errors(chat_client_base: ChatClientProtocol):
"""Test that max_consecutive_errors_per_request limits error retries in streaming mode."""
@ai_function(name="error_function")
@tool(name="error_function", approval_mode="never_require")
def error_func(arg1: str) -> str:
raise ValueError("Function error")
@@ -1920,7 +1922,7 @@ async def test_streaming_function_invocation_config_terminate_on_unknown_calls_f
"""Test that terminate_on_unknown_calls=False returns error message for unknown functions in streaming mode."""
exec_counter = 0
@ai_function(name="known_function")
@tool(name="known_function", approval_mode="never_require")
def known_func(arg1: str) -> str:
nonlocal exec_counter
exec_counter += 1
@@ -1963,7 +1965,7 @@ async def test_streaming_function_invocation_config_terminate_on_unknown_calls_t
"""Test that terminate_on_unknown_calls=True stops execution on unknown functions in streaming mode."""
exec_counter = 0
@ai_function(name="known_function")
@tool(name="known_function", approval_mode="never_require")
def known_func(arg1: str) -> str:
nonlocal exec_counter
exec_counter += 1
@@ -1996,7 +1998,7 @@ async def test_streaming_function_invocation_config_terminate_on_unknown_calls_t
async def test_streaming_function_invocation_config_include_detailed_errors_true(chat_client_base: ChatClientProtocol):
"""Test that include_detailed_errors=True returns detailed error information in streaming mode."""
@ai_function(name="error_function")
@tool(name="error_function", approval_mode="never_require")
def error_func(arg1: str) -> str:
raise ValueError("Specific error message that should appear")
@@ -2036,7 +2038,7 @@ async def test_streaming_function_invocation_config_include_detailed_errors_fals
):
"""Test that include_detailed_errors=False returns generic error messages in streaming mode."""
@ai_function(name="error_function")
@tool(name="error_function", approval_mode="never_require")
def error_func(arg1: str) -> str:
raise ValueError("Specific error message that should not appear")
@@ -2074,7 +2076,7 @@ async def test_streaming_function_invocation_config_include_detailed_errors_fals
async def test_streaming_argument_validation_error_with_detailed_errors(chat_client_base: ChatClientProtocol):
"""Test that argument validation errors include details when include_detailed_errors=True in streaming mode."""
@ai_function(name="typed_function")
@tool(name="typed_function", approval_mode="never_require")
def typed_func(arg1: int) -> str: # Expects int, not str
return f"Got {arg1}"
@@ -2112,7 +2114,7 @@ async def test_streaming_argument_validation_error_with_detailed_errors(chat_cli
async def test_streaming_argument_validation_error_without_detailed_errors(chat_client_base: ChatClientProtocol):
"""Test that argument validation errors are generic when include_detailed_errors=False in streaming mode."""
@ai_function(name="typed_function")
@tool(name="typed_function", approval_mode="never_require")
def typed_func(arg1: int) -> str: # Expects int, not str
return f"Got {arg1}"
@@ -2149,18 +2151,17 @@ async def test_streaming_argument_validation_error_without_detailed_errors(chat_
async def test_streaming_multiple_function_calls_parallel_execution(chat_client_base: ChatClientProtocol):
"""Test that multiple function calls are executed in parallel in streaming mode."""
import asyncio
exec_order = []
@ai_function(name="func1")
@tool(name="func1", approval_mode="never_require")
async def func1(arg1: str) -> str:
exec_order.append("func1_start")
await asyncio.sleep(0.01) # Small delay
exec_order.append("func1_end")
return f"Result1 {arg1}"
@ai_function(name="func2")
@tool(name="func2", approval_mode="never_require")
async def func2(arg1: str) -> str:
exec_order.append("func2_start")
await asyncio.sleep(0.01) # Small delay
@@ -2202,7 +2203,7 @@ async def test_streaming_approval_requests_in_assistant_message(chat_client_base
"""Approval requests should be added to assistant updates in streaming mode."""
exec_counter = 0
@ai_function(name="test_func", approval_mode="always_require")
@tool(name="test_func", approval_mode="always_require")
def func_with_approval(arg1: str) -> str:
nonlocal exec_counter
exec_counter += 1
@@ -2238,7 +2239,7 @@ async def test_streaming_error_recovery_resets_counter(chat_client_base: ChatCli
call_count = 0
@ai_function(name="sometimes_fails")
@tool(name="sometimes_fails", approval_mode="never_require")
def sometimes_fails(arg1: str) -> str:
nonlocal call_count
call_count += 1
@@ -2306,7 +2307,7 @@ async def test_terminate_loop_single_function_call(chat_client_base: ChatClientP
"""Test that terminate_loop=True exits the function calling loop after single function call."""
exec_counter = 0
@ai_function(name="test_function")
@tool(name="test_function", approval_mode="never_require")
def ai_func(arg1: str) -> str:
nonlocal exec_counter
exec_counter += 1
@@ -2367,13 +2368,13 @@ async def test_terminate_loop_multiple_function_calls_one_terminates(chat_client
normal_call_count = 0
terminating_call_count = 0
@ai_function(name="normal_function")
@tool(name="normal_function", approval_mode="never_require")
def normal_func(arg1: str) -> str:
nonlocal normal_call_count
normal_call_count += 1
return f"Normal {arg1}"
@ai_function(name="terminating_function")
@tool(name="terminating_function", approval_mode="never_require")
def terminating_func(arg1: str) -> str:
nonlocal terminating_call_count
terminating_call_count += 1
@@ -2423,7 +2424,7 @@ async def test_terminate_loop_streaming_single_function_call(chat_client_base: C
"""Test that terminate_loop=True exits the streaming function calling loop."""
exec_counter = 0
@ai_function(name="test_function")
@tool(name="test_function", approval_mode="never_require")
def ai_func(arg1: str) -> str:
nonlocal exec_counter
exec_counter += 1
python/packages/core/tests/core/test_kwargs_propagation_to_ai_function.py
+13 -13
View File
@@ -1,6 +1,6 @@
# Copyright (c) Microsoft. All rights reserved.
"""Tests for kwargs propagation from get_response() to @ai_function tools."""
"""Tests for kwargs propagation from get_response() to @tool functions."""
from typing import Any
@@ -9,19 +9,19 @@ from agent_framework import (
ChatResponse,
ChatResponseUpdate,
Content,
ai_function,
tool,
)
from agent_framework._tools import _handle_function_calls_response, _handle_function_calls_streaming_response
class TestKwargsPropagationToAIFunction:
"""Test cases for kwargs flowing from get_response() to @ai_function tools."""
class TestKwargsPropagationToFunctionTool:
"""Test cases for kwargs flowing from get_response() to @tool functions."""
async def test_kwargs_propagate_to_ai_function_with_kwargs(self) -> None:
"""Test that kwargs passed to get_response() are available in @ai_function **kwargs."""
async def test_kwargs_propagate_to_tool_with_kwargs(self) -> None:
"""Test that kwargs passed to get_response() are available in @tool **kwargs."""
captured_kwargs: dict[str, Any] = {}
@ai_function
@tool(approval_mode="never_require")
def capture_kwargs_tool(x: int, **kwargs: Any) -> str:
"""A tool that captures kwargs for testing."""
captured_kwargs.update(kwargs)
@@ -75,10 +75,10 @@ class TestKwargsPropagationToAIFunction:
# Verify result
assert result.messages[-1].text == "Done!"
async def test_kwargs_not_forwarded_to_ai_function_without_kwargs(self) -> None:
"""Test that kwargs are NOT forwarded to @ai_function that doesn't accept **kwargs."""
async def test_kwargs_not_forwarded_to_tool_without_kwargs(self) -> None:
"""Test that kwargs are NOT forwarded to @tool that doesn't accept **kwargs."""
@ai_function
@tool(approval_mode="never_require")
def simple_tool(x: int) -> str:
"""A simple tool without **kwargs."""
# This should not receive any extra kwargs
@@ -120,7 +120,7 @@ class TestKwargsPropagationToAIFunction:
"""Test that kwargs don't leak between different function call invocations."""
invocation_kwargs: list[dict[str, Any]] = []
@ai_function
@tool(approval_mode="never_require")
def tracking_tool(name: str, **kwargs: Any) -> str:
"""A tool that tracks kwargs from each invocation."""
invocation_kwargs.append(dict(kwargs))
@@ -170,10 +170,10 @@ class TestKwargsPropagationToAIFunction:
assert result.messages[-1].text == "All done!"
async def test_streaming_response_kwargs_propagation(self) -> None:
"""Test that kwargs propagate to @ai_function in streaming mode."""
"""Test that kwargs propagate to @tool in streaming mode."""
captured_kwargs: dict[str, Any] = {}
@ai_function
@tool(approval_mode="never_require")
def streaming_capture_tool(value: str, **kwargs: Any) -> str:
"""A tool that captures kwargs during streaming."""
captured_kwargs.update(kwargs)
python/packages/core/tests/core/test_mcp.py
+6 -6
View File
@@ -1228,7 +1228,7 @@ async def test_streamable_http_integration():
if not url.startswith("http"):
pytest.skip("LOCAL_MCP_URL is not an HTTP URL")
tool = MCPStreamableHTTPTool(name="integration_test", url=url)
tool = MCPStreamableHTTPTool(name="integration_test", url=url, approval_mode="never_require")
async with tool:
# Test that we can connect and load tools
@@ -1260,7 +1260,7 @@ async def test_mcp_connection_reset_integration():
"""
url = os.environ.get("LOCAL_MCP_URL")
tool = MCPStreamableHTTPTool(name="integration_test", url=url)
tool = MCPStreamableHTTPTool(name="integration_test", url=url, approval_mode="never_require")
async with tool:
# Verify initial connection
@@ -1620,7 +1620,7 @@ def test_mcp_websocket_tool_get_mcp_client_with_kwargs():
async def test_mcp_tool_deduplication():
"""Test that MCP tools are not duplicated in MCPTool"""
from agent_framework._mcp import MCPTool
from agent_framework._tools import AIFunction
from agent_framework._tools import FunctionTool
# Create MCPStreamableHTTPTool instance
tool = MCPTool(name="test_mcp_tool")
@@ -1629,12 +1629,12 @@ async def test_mcp_tool_deduplication():
tool._functions = []
# Add initial functions
func1 = AIFunction(
func1 = FunctionTool(
func=lambda x: f"Result: {x}",
name="analyze_content",
description="Analyzes content",
)
func2 = AIFunction(
func2 = FunctionTool(
func=lambda x: f"Extract: {x}",
name="extract_info",
description="Extracts information",
@@ -1662,7 +1662,7 @@ async def test_mcp_tool_deduplication():
if tool_name in existing_names:
continue # Skip duplicates
new_func = AIFunction(func=lambda x: f"Process: {x}", name=tool_name, description=description)
new_func = FunctionTool(func=lambda x: f"Process: {x}", name=tool_name, description=description)
tool._functions.append(new_func)
existing_names.add(tool_name)
added_count += 1
python/packages/core/tests/core/test_middleware.py
+15 -15
View File
@@ -28,7 +28,7 @@ from agent_framework._middleware import (
FunctionMiddleware,
FunctionMiddlewarePipeline,
)
from agent_framework._tools import AIFunction
from agent_framework._tools import FunctionTool
class TestAgentRunContext:
@@ -73,7 +73,7 @@ class TestAgentRunContext:
class TestFunctionInvocationContext:
"""Test cases for FunctionInvocationContext."""
def test_init_with_defaults(self, mock_function: AIFunction[Any, Any]) -> None:
def test_init_with_defaults(self, mock_function: FunctionTool[Any, Any]) -> None:
"""Test FunctionInvocationContext initialization with default values."""
arguments = FunctionTestArgs(name="test")
context = FunctionInvocationContext(function=mock_function, arguments=arguments)
@@ -82,7 +82,7 @@ class TestFunctionInvocationContext:
assert context.arguments == arguments
assert context.metadata == {}
def test_init_with_custom_metadata(self, mock_function: AIFunction[Any, Any]) -> None:
def test_init_with_custom_metadata(self, mock_function: FunctionTool[Any, Any]) -> None:
"""Test FunctionInvocationContext initialization with custom metadata."""
arguments = FunctionTestArgs(name="test")
metadata = {"key": "value"}
@@ -419,7 +419,7 @@ class TestFunctionMiddlewarePipeline:
await next(context)
context.terminate = True
async def test_execute_with_pre_next_termination(self, mock_function: AIFunction[Any, Any]) -> None:
async def test_execute_with_pre_next_termination(self, mock_function: FunctionTool[Any, Any]) -> None:
"""Test pipeline execution with termination before next()."""
middleware = self.PreNextTerminateFunctionMiddleware()
pipeline = FunctionMiddlewarePipeline([middleware])
@@ -438,7 +438,7 @@ class TestFunctionMiddlewarePipeline:
# Handler should not be called when terminated before next()
assert execution_order == []
async def test_execute_with_post_next_termination(self, mock_function: AIFunction[Any, Any]) -> None:
async def test_execute_with_post_next_termination(self, mock_function: FunctionTool[Any, Any]) -> None:
"""Test pipeline execution with termination after next()."""
middleware = self.PostNextTerminateFunctionMiddleware()
pipeline = FunctionMiddlewarePipeline([middleware])
@@ -477,7 +477,7 @@ class TestFunctionMiddlewarePipeline:
pipeline = FunctionMiddlewarePipeline([test_middleware])
assert pipeline.has_middlewares
async def test_execute_no_middleware(self, mock_function: AIFunction[Any, Any]) -> None:
async def test_execute_no_middleware(self, mock_function: FunctionTool[Any, Any]) -> None:
"""Test pipeline execution with no middleware."""
pipeline = FunctionMiddlewarePipeline()
arguments = FunctionTestArgs(name="test")
@@ -491,7 +491,7 @@ class TestFunctionMiddlewarePipeline:
result = await pipeline.execute(mock_function, arguments, context, final_handler)
assert result == expected_result
async def test_execute_with_middleware(self, mock_function: AIFunction[Any, Any]) -> None:
async def test_execute_with_middleware(self, mock_function: FunctionTool[Any, Any]) -> None:
"""Test pipeline execution with middleware."""
execution_order: list[str] = []
@@ -778,7 +778,7 @@ class TestClassBasedMiddleware:
assert context.metadata["after"] is True
assert metadata_updates == ["before", "handler", "after"]
async def test_function_middleware_execution(self, mock_function: AIFunction[Any, Any]) -> None:
async def test_function_middleware_execution(self, mock_function: FunctionTool[Any, Any]) -> None:
"""Test class-based function middleware execution."""
metadata_updates: list[str] = []
@@ -840,7 +840,7 @@ class TestFunctionBasedMiddleware:
assert context.metadata["function_middleware"] is True
assert execution_order == ["function_before", "handler", "function_after"]
async def test_function_function_middleware(self, mock_function: AIFunction[Any, Any]) -> None:
async def test_function_function_middleware(self, mock_function: FunctionTool[Any, Any]) -> None:
"""Test function-based function middleware."""
execution_order: list[str] = []
@@ -902,7 +902,7 @@ class TestMixedMiddleware:
assert result is not None
assert execution_order == ["class_before", "function_before", "handler", "function_after", "class_after"]
async def test_mixed_function_middleware(self, mock_function: AIFunction[Any, Any]) -> None:
async def test_mixed_function_middleware(self, mock_function: FunctionTool[Any, Any]) -> None:
"""Test mixed class and function-based function middleware."""
execution_order: list[str] = []
@@ -1022,7 +1022,7 @@ class TestMultipleMiddlewareOrdering:
]
assert execution_order == expected_order
async def test_function_middleware_execution_order(self, mock_function: AIFunction[Any, Any]) -> None:
async def test_function_middleware_execution_order(self, mock_function: FunctionTool[Any, Any]) -> None:
"""Test that multiple function middleware execute in registration order."""
execution_order: list[str] = []
@@ -1150,7 +1150,7 @@ class TestContextContentValidation:
result = await pipeline.execute(mock_agent, messages, context, final_handler)
assert result is not None
async def test_function_context_validation(self, mock_function: AIFunction[Any, Any]) -> None:
async def test_function_context_validation(self, mock_function: FunctionTool[Any, Any]) -> None:
"""Test that function context contains expected data."""
class ContextValidationMiddleware(FunctionMiddleware):
@@ -1498,7 +1498,7 @@ class TestMiddlewareExecutionControl:
assert not handler_called
assert context.result is None
async def test_function_middleware_no_next_no_execution(self, mock_function: AIFunction[Any, Any]) -> None:
async def test_function_middleware_no_next_no_execution(self, mock_function: FunctionTool[Any, Any]) -> None:
"""Test that when function middleware doesn't call next(), no execution happens."""
class FunctionTestArgs(BaseModel):
@@ -1672,9 +1672,9 @@ def mock_agent() -> AgentProtocol:
@pytest.fixture
def mock_function() -> AIFunction[Any, Any]:
def mock_function() -> FunctionTool[Any, Any]:
"""Mock function for testing."""
function = MagicMock(spec=AIFunction[Any, Any])
function = MagicMock(spec=FunctionTool[Any, Any])
function.name = "test_function"
return function
python/packages/core/tests/core/test_middleware_context_result.py
+7 -7
View File
@@ -24,7 +24,7 @@ from agent_framework._middleware import (
FunctionMiddleware,
FunctionMiddlewarePipeline,
)
from agent_framework._tools import AIFunction
from agent_framework._tools import FunctionTool
from .conftest import MockChatClient
@@ -103,7 +103,7 @@ class TestResultOverrideMiddleware:
assert updates[0].text == "overridden"
assert updates[1].text == " stream"
async def test_function_middleware_result_override(self, mock_function: AIFunction[Any, Any]) -> None:
async def test_function_middleware_result_override(self, mock_function: FunctionTool[Any, Any]) -> None:
"""Test that function middleware can override result."""
override_result = "overridden function result"
@@ -260,7 +260,7 @@ class TestResultOverrideMiddleware:
assert execute_result.messages[0].text == "executed response"
assert handler_called
async def test_function_middleware_conditional_no_next(self, mock_function: AIFunction[Any, Any]) -> None:
async def test_function_middleware_conditional_no_next(self, mock_function: FunctionTool[Any, Any]) -> None:
"""Test that when function middleware conditionally doesn't call next(), no execution happens."""
class ConditionalNoNextFunctionMiddleware(FunctionMiddleware):
@@ -345,7 +345,7 @@ class TestResultObservability:
assert observed_responses[0].messages[0].text == "executed response"
assert result == observed_responses[0]
async def test_function_middleware_result_observability(self, mock_function: AIFunction[Any, Any]) -> None:
async def test_function_middleware_result_observability(self, mock_function: FunctionTool[Any, Any]) -> None:
"""Test that middleware can observe function result after execution."""
observed_results: list[str] = []
@@ -414,7 +414,7 @@ class TestResultObservability:
assert result is not None
assert result.messages[0].text == "modified after execution"
async def test_function_middleware_post_execution_override(self, mock_function: AIFunction[Any, Any]) -> None:
async def test_function_middleware_post_execution_override(self, mock_function: FunctionTool[Any, Any]) -> None:
"""Test that middleware can override function result after observing execution."""
class PostExecutionOverrideMiddleware(FunctionMiddleware):
@@ -456,8 +456,8 @@ def mock_agent() -> AgentProtocol:
@pytest.fixture
def mock_function() -> AIFunction[Any, Any]:
def mock_function() -> FunctionTool[Any, Any]:
"""Mock function for testing."""
function = MagicMock(spec=AIFunction[Any, Any])
function = MagicMock(spec=FunctionTool[Any, Any])
function.name = "test_function"
return function
python/packages/core/tests/core/test_middleware_with_agent.py
+42 -7
View File
@@ -14,6 +14,7 @@ from agent_framework import (
ChatResponse,
ChatResponseUpdate,
Content,
FunctionTool,
Role,
agent_middleware,
chat_middleware,
@@ -214,9 +215,13 @@ class TestChatAgentFunctionBasedMiddleware:
execution_order.append("function_called")
return "test_result"
test_function_tool = FunctionTool(
func=test_function, name="test_function", description="Test function", approval_mode="never_require"
)
# Create ChatAgent with function middleware and test function
middleware = PreTerminationFunctionMiddleware()
agent = ChatAgent(chat_client=chat_client, middleware=[middleware], tools=[test_function])
agent = ChatAgent(chat_client=chat_client, middleware=[middleware], tools=[test_function_tool])
# Execute the agent
await agent.run(messages)
@@ -271,9 +276,13 @@ class TestChatAgentFunctionBasedMiddleware:
execution_order.append("function_called")
return "test_result"
test_function_tool = FunctionTool(
func=test_function, name="test_function", description="Test function", approval_mode="never_require"
)
# Create ChatAgent with function middleware and test function
middleware = PostTerminationFunctionMiddleware()
agent = ChatAgent(chat_client=chat_client, middleware=[middleware], tools=[test_function])
agent = ChatAgent(chat_client=chat_client, middleware=[middleware], tools=[test_function_tool])
# Execute the agent
response = await agent.run(messages)
@@ -518,11 +527,19 @@ class TestChatAgentMultipleMiddlewareOrdering:
# region Tool Functions for Testing
def sample_tool_function(location: str) -> str:
def _sample_tool_function_impl(location: str) -> str:
"""A simple tool function for middleware testing."""
return f"Weather in {location}: sunny"
sample_tool_function = FunctionTool(
func=_sample_tool_function_impl,
name="sample_tool_function",
description="A simple tool function for middleware testing.",
approval_mode="never_require",
)
# region ChatAgent Function Middleware Tests with Tools
@@ -1157,6 +1174,10 @@ class TestRunLevelMiddleware:
execution_log.append("tool_executed")
return f"Tool response: {message}"
custom_tool_wrapped = FunctionTool(
func=custom_tool, name="custom_tool", description="Custom tool", approval_mode="never_require"
)
# Set up mock to return a function call first, then a regular response
function_call_response = ChatResponse(
messages=[
@@ -1179,7 +1200,7 @@ class TestRunLevelMiddleware:
agent = ChatAgent(
chat_client=chat_client,
middleware=[AgentLevelAgentMiddleware(), AgentLevelFunctionMiddleware()],
tools=[custom_tool],
tools=[custom_tool_wrapped],
)
# Execute with run-level middleware
@@ -1246,6 +1267,10 @@ class TestMiddlewareDecoratorLogic:
execution_order.append("tool_executed")
return f"Tool response: {message}"
custom_tool_wrapped = FunctionTool(
func=custom_tool, name="custom_tool", description="Custom tool", approval_mode="never_require"
)
# Set up mock to return a function call first, then a regular response
function_call_response = ChatResponse(
messages=[
@@ -1268,7 +1293,7 @@ class TestMiddlewareDecoratorLogic:
agent = ChatAgent(
chat_client=chat_client,
middleware=[matching_agent_middleware, matching_function_middleware],
tools=[custom_tool],
tools=[custom_tool_wrapped],
)
response = await agent.run([ChatMessage(role=Role.USER, text="test")])
@@ -1313,6 +1338,10 @@ class TestMiddlewareDecoratorLogic:
execution_order.append("tool_executed")
return f"Tool response: {message}"
custom_tool_wrapped = FunctionTool(
func=custom_tool, name="custom_tool", description="Custom tool", approval_mode="never_require"
)
# Set up mock to return a function call first, then a regular response
function_call_response = ChatResponse(
messages=[
@@ -1333,7 +1362,9 @@ class TestMiddlewareDecoratorLogic:
# Should work - relies on decorator
agent = ChatAgent(
chat_client=chat_client, middleware=[decorator_only_agent, decorator_only_function], tools=[custom_tool]
chat_client=chat_client,
middleware=[decorator_only_agent, decorator_only_function],
tools=[custom_tool_wrapped],
)
response = await agent.run([ChatMessage(role=Role.USER, text="test")])
@@ -1363,6 +1394,10 @@ class TestMiddlewareDecoratorLogic:
execution_order.append("tool_executed")
return f"Tool response: {message}"
custom_tool_wrapped = FunctionTool(
func=custom_tool, name="custom_tool", description="Custom tool", approval_mode="never_require"
)
# Set up mock to return a function call first, then a regular response
function_call_response = ChatResponse(
messages=[
@@ -1383,7 +1418,7 @@ class TestMiddlewareDecoratorLogic:
# Should work - relies on type annotations
agent = ChatAgent(
chat_client=chat_client, middleware=[type_only_agent, type_only_function], tools=[custom_tool]
chat_client=chat_client, middleware=[type_only_agent, type_only_function], tools=[custom_tool_wrapped]
)
response = await agent.run([ChatMessage(role=Role.USER, text="test")])
python/packages/core/tests/core/test_middleware_with_chat.py
+17 -2
View File
@@ -11,6 +11,7 @@ from agent_framework import (
ChatResponse,
Content,
FunctionInvocationContext,
FunctionTool,
Role,
chat_middleware,
function_middleware,
@@ -337,6 +338,13 @@ class TestChatMiddleware:
"""Get weather for a location."""
return f"Weather in {location}: sunny"
sample_tool_wrapped = FunctionTool(
func=sample_tool,
name="sample_tool",
description="Get weather for a location",
approval_mode="never_require",
)
# Create function-invocation enabled chat client
chat_client = use_chat_middleware(use_function_invocation(MockBaseChatClient))()
@@ -366,7 +374,7 @@ class TestChatMiddleware:
# Execute the chat client directly with tools - this should trigger function invocation and middleware
messages = [ChatMessage(role=Role.USER, text="What's the weather in San Francisco?")]
response = await chat_client.get_response(messages, options={"tools": [sample_tool]})
response = await chat_client.get_response(messages, options={"tools": [sample_tool_wrapped]})
# Verify response
assert response is not None
@@ -396,6 +404,13 @@ class TestChatMiddleware:
"""Get weather for a location."""
return f"Weather in {location}: sunny"
sample_tool_wrapped = FunctionTool(
func=sample_tool,
name="sample_tool",
description="Get weather for a location",
approval_mode="never_require",
)
# Create function-invocation enabled chat client
chat_client = use_function_invocation(MockBaseChatClient)()
@@ -423,7 +438,7 @@ class TestChatMiddleware:
# Execute the chat client directly with run-level middleware and tools
messages = [ChatMessage(role=Role.USER, text="What's the weather in New York?")]
response = await chat_client.get_response(
messages, options={"tools": [sample_tool]}, middleware=[run_level_function_middleware]
messages, options={"tools": [sample_tool_wrapped]}, middleware=[run_level_function_middleware]
)
# Verify response
python/packages/core/tests/core/test_observability.py
+2 -2
View File
@@ -21,8 +21,8 @@ from agent_framework import (
ChatResponseUpdate,
Role,
UsageDetails,
ai_function,
prepend_agent_framework_to_user_agent,
tool,
)
from agent_framework.exceptions import AgentInitializationError, ChatClientInitializationError
from agent_framework.observability import (
@@ -606,7 +606,7 @@ async def test_function_call_with_error_handling(span_exporter: InMemorySpanExpo
"""Test that function call errors are properly captured in telemetry."""
# Create a function that raises an error using the decorator
@ai_function(name="failing_function", description="A function that fails")
@tool(name="failing_function", description="A function that fails")
async def failing_function(param: str) -> str:
raise ValueError("Function execution failed")
python/packages/core/tests/core/test_tools.py
+92 -89
View File
@@ -8,13 +8,13 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanE
from pydantic import BaseModel, ValidationError
from agent_framework import (
AIFunction,
Content,
FunctionTool,
HostedCodeInterpreterTool,
HostedImageGenerationTool,
HostedMCPTool,
ToolProtocol,
ai_function,
tool,
)
from agent_framework._tools import (
_build_pydantic_model_from_json_schema,
@@ -24,19 +24,19 @@ from agent_framework._tools import (
from agent_framework.exceptions import ToolException
from agent_framework.observability import OtelAttr
# region AIFunction and ai_function decorator tests
# region FunctionTool and tool decorator tests
def test_ai_function_decorator():
"""Test the ai_function decorator."""
def test_tool_decorator():
"""Test the tool decorator."""
@ai_function(name="test_tool", description="A test tool")
@tool(name="test_tool", description="A test tool")
def test_tool(x: int, y: int) -> int:
"""A simple function that adds two numbers."""
return x + y
assert isinstance(test_tool, ToolProtocol)
assert isinstance(test_tool, AIFunction)
assert isinstance(test_tool, FunctionTool)
assert test_tool.name == "test_tool"
assert test_tool.description == "A test tool"
assert test_tool.parameters() == {
@@ -48,16 +48,16 @@ def test_ai_function_decorator():
assert test_tool(1, 2) == 3
def test_ai_function_decorator_without_args():
"""Test the ai_function decorator."""
def test_tool_decorator_without_args():
"""Test the tool decorator."""
@ai_function
@tool
def test_tool(x: int, y: int) -> int:
"""A simple function that adds two numbers."""
return x + y
assert isinstance(test_tool, ToolProtocol)
assert isinstance(test_tool, AIFunction)
assert isinstance(test_tool, FunctionTool)
assert test_tool.name == "test_tool"
assert test_tool.description == "A simple function that adds two numbers."
assert test_tool.parameters() == {
@@ -67,18 +67,19 @@ def test_ai_function_decorator_without_args():
"type": "object",
}
assert test_tool(1, 2) == 3
assert test_tool.approval_mode == "never_require"
def test_ai_function_without_args():
"""Test the ai_function decorator."""
def test_tool_without_args():
"""Test the tool decorator."""
@ai_function
@tool
def test_tool() -> int:
"""A simple function that adds two numbers."""
return 1 + 2
assert isinstance(test_tool, ToolProtocol)
assert isinstance(test_tool, AIFunction)
assert isinstance(test_tool, FunctionTool)
assert test_tool.name == "test_tool"
assert test_tool.description == "A simple function that adds two numbers."
assert test_tool.parameters() == {
@@ -89,16 +90,16 @@ def test_ai_function_without_args():
assert test_tool() == 3
async def test_ai_function_decorator_with_async():
"""Test the ai_function decorator with an async function."""
async def test_tool_decorator_with_async():
"""Test the tool decorator with an async function."""
@ai_function(name="async_test_tool", description="An async test tool")
@tool(name="async_test_tool", description="An async test tool")
async def async_test_tool(x: int, y: int) -> int:
"""An async function that adds two numbers."""
return x + y
assert isinstance(async_test_tool, ToolProtocol)
assert isinstance(async_test_tool, AIFunction)
assert isinstance(async_test_tool, FunctionTool)
assert async_test_tool.name == "async_test_tool"
assert async_test_tool.description == "An async test tool"
assert async_test_tool.parameters() == {
@@ -110,11 +111,11 @@ async def test_ai_function_decorator_with_async():
assert (await async_test_tool(1, 2)) == 3
def test_ai_function_decorator_in_class():
"""Test the ai_function decorator."""
def test_tool_decorator_in_class():
"""Test the tool decorator."""
class my_tools:
@ai_function(name="test_tool", description="A test tool")
@tool(name="test_tool", description="A test tool")
def test_tool(self, x: int, y: int) -> int:
"""A simple function that adds two numbers."""
return x + y
@@ -122,7 +123,7 @@ def test_ai_function_decorator_in_class():
test_tool = my_tools().test_tool
assert isinstance(test_tool, ToolProtocol)
assert isinstance(test_tool, AIFunction)
assert isinstance(test_tool, FunctionTool)
assert test_tool.name == "test_tool"
assert test_tool.description == "A test tool"
assert test_tool.parameters() == {
@@ -134,15 +135,15 @@ def test_ai_function_decorator_in_class():
assert test_tool(1, 2) == 3
def test_ai_function_with_literal_type_parameter():
"""Test ai_function decorator with Literal type parameter (issue #2891)."""
def test_tool_with_literal_type_parameter():
"""Test tool decorator with Literal type parameter (issue #2891)."""
@ai_function
@tool
def search_flows(category: Literal["Data", "Security", "Network"], issue: str) -> str:
"""Search flows by category."""
return f"{category}: {issue}"
assert isinstance(search_flows, AIFunction)
assert isinstance(search_flows, FunctionTool)
schema = search_flows.parameters()
assert schema == {
"properties": {
@@ -157,18 +158,18 @@ def test_ai_function_with_literal_type_parameter():
assert search_flows("Data", "test issue") == "Data: test issue"
def test_ai_function_with_literal_type_in_class_method():
"""Test ai_function decorator with Literal type parameter in a class method (issue #2891)."""
def test_tool_with_literal_type_in_class_method():
"""Test tool decorator with Literal type parameter in a class method (issue #2891)."""
class MyTools:
@ai_function
@tool
def search_flows(self, category: Literal["Data", "Security", "Network"], issue: str) -> str:
"""Search flows by category."""
return f"{category}: {issue}"
tools = MyTools()
search_tool = tools.search_flows
assert isinstance(search_tool, AIFunction)
assert isinstance(search_tool, FunctionTool)
schema = search_tool.parameters()
assert schema == {
"properties": {
@@ -183,15 +184,15 @@ def test_ai_function_with_literal_type_in_class_method():
assert search_tool("Security", "test issue") == "Security: test issue"
def test_ai_function_with_literal_int_type():
"""Test ai_function decorator with Literal int type parameter."""
def test_tool_with_literal_int_type():
"""Test tool decorator with Literal int type parameter."""
@ai_function
@tool
def set_priority(priority: Literal[1, 2, 3], task: str) -> str:
"""Set priority for a task."""
return f"Priority {priority}: {task}"
assert isinstance(set_priority, AIFunction)
assert isinstance(set_priority, FunctionTool)
schema = set_priority.parameters()
assert schema == {
"properties": {
@@ -205,10 +206,10 @@ def test_ai_function_with_literal_int_type():
assert set_priority(1, "important task") == "Priority 1: important task"
def test_ai_function_with_literal_and_annotated():
"""Test ai_function decorator with Literal type combined with Annotated for description."""
def test_tool_with_literal_and_annotated():
"""Test tool decorator with Literal type combined with Annotated for description."""
@ai_function
@tool
def categorize(
category: Annotated[Literal["A", "B", "C"], "The category to assign"],
name: str,
@@ -216,14 +217,14 @@ def test_ai_function_with_literal_and_annotated():
"""Categorize an item."""
return f"{category}: {name}"
assert isinstance(categorize, AIFunction)
assert isinstance(categorize, FunctionTool)
schema = categorize.parameters()
# Literal type inside Annotated should preserve enum values
assert schema["properties"]["category"]["enum"] == ["A", "B", "C"]
assert categorize("A", "test") == "A: test"
async def test_ai_function_decorator_shared_state():
async def test_tool_decorator_shared_state():
"""Test that decorated methods maintain shared state across multiple calls and tool usage."""
class StatefulCounter:
@@ -233,20 +234,20 @@ async def test_ai_function_decorator_shared_state():
self.counter = initial_value
self.operation_log: list[str] = []
@ai_function(name="increment", description="Increment the counter")
@tool(name="increment", description="Increment the counter")
def increment(self, amount: int) -> str:
"""Increment the counter by the given amount."""
self.counter += amount
self.operation_log.append(f"increment({amount})")
return f"Counter incremented by {amount}. New value: {self.counter}"
@ai_function(name="get_value", description="Get the current counter value")
@tool(name="get_value", description="Get the current counter value")
def get_value(self) -> str:
"""Get the current counter value."""
self.operation_log.append("get_value()")
return f"Current counter value: {self.counter}"
@ai_function(name="multiply", description="Multiply the counter")
@tool(name="multiply", description="Multiply the counter")
def multiply(self, factor: int) -> str:
"""Multiply the counter by the given factor."""
self.counter *= factor
@@ -261,10 +262,10 @@ async def test_ai_function_decorator_shared_state():
get_value_tool = counter_instance.get_value
multiply_tool = counter_instance.multiply
# Verify they are AIFunction instances
assert isinstance(increment_tool, AIFunction)
assert isinstance(get_value_tool, AIFunction)
assert isinstance(multiply_tool, AIFunction)
# Verify they are FunctionTool instances
assert isinstance(increment_tool, FunctionTool)
assert isinstance(get_value_tool, FunctionTool)
assert isinstance(multiply_tool, FunctionTool)
# Tool 1 (increment) is used
result1 = increment_tool(5)
@@ -329,10 +330,10 @@ async def test_ai_function_decorator_shared_state():
assert counter_instance.counter == 60
async def test_ai_function_invoke_telemetry_enabled(span_exporter: InMemorySpanExporter):
"""Test the ai_function invoke method with telemetry enabled."""
async def test_tool_invoke_telemetry_enabled(span_exporter: InMemorySpanExporter):
"""Test the tool invoke method with telemetry enabled."""
@ai_function(
@tool(
name="telemetry_test_tool",
description="A test tool for telemetry",
)
@@ -373,10 +374,10 @@ async def test_ai_function_invoke_telemetry_enabled(span_exporter: InMemorySpanE
@pytest.mark.parametrize("enable_sensitive_data", [False], indirect=True)
async def test_ai_function_invoke_telemetry_sensitive_disabled(span_exporter: InMemorySpanExporter):
"""Test the ai_function invoke method with telemetry enabled."""
async def test_tool_invoke_telemetry_sensitive_disabled(span_exporter: InMemorySpanExporter):
"""Test the tool invoke method with telemetry enabled."""
@ai_function(
@tool(
name="telemetry_test_tool",
description="A test tool for telemetry",
)
@@ -416,10 +417,10 @@ async def test_ai_function_invoke_telemetry_sensitive_disabled(span_exporter: In
assert attributes[OtelAttr.TOOL_CALL_ID] == "test_call_id"
async def test_ai_function_invoke_ignores_additional_kwargs() -> None:
"""Ensure ai_function tools drop unknown kwargs when invoked with validated arguments."""
async def test_tool_invoke_ignores_additional_kwargs() -> None:
"""Ensure tools drop unknown kwargs when invoked with validated arguments."""
@ai_function
@tool
async def simple_tool(message: str) -> str:
"""Echo tool."""
return message.upper()
@@ -436,10 +437,10 @@ async def test_ai_function_invoke_ignores_additional_kwargs() -> None:
assert result == "HELLO WORLD"
async def test_ai_function_invoke_telemetry_with_pydantic_args(span_exporter: InMemorySpanExporter):
"""Test the ai_function invoke method with Pydantic model arguments."""
async def test_tool_invoke_telemetry_with_pydantic_args(span_exporter: InMemorySpanExporter):
"""Test the tool invoke method with Pydantic model arguments."""
@ai_function(
@tool(
name="pydantic_test_tool",
description="A test tool with Pydantic args",
)
@@ -470,10 +471,10 @@ async def test_ai_function_invoke_telemetry_with_pydantic_args(span_exporter: In
assert span.attributes[OtelAttr.TOOL_ARGUMENTS] == '{"x":5,"y":10}'
async def test_ai_function_invoke_telemetry_with_exception(span_exporter: InMemorySpanExporter):
"""Test the ai_function invoke method with telemetry when an exception occurs."""
async def test_tool_invoke_telemetry_with_exception(span_exporter: InMemorySpanExporter):
"""Test the tool invoke method with telemetry when an exception occurs."""
@ai_function(
@tool(
name="exception_test_tool",
description="A test tool that raises an exception",
)
@@ -507,10 +508,10 @@ async def test_ai_function_invoke_telemetry_with_exception(span_exporter: InMemo
assert attributes[OtelAttr.ERROR_TYPE] == ValueError.__name__
async def test_ai_function_invoke_telemetry_async_function(span_exporter: InMemorySpanExporter):
"""Test the ai_function invoke method with telemetry on async function."""
async def test_tool_invoke_telemetry_async_function(span_exporter: InMemorySpanExporter):
"""Test the tool invoke method with telemetry on async function."""
@ai_function(
@tool(
name="async_telemetry_test",
description="An async test tool for telemetry",
)
@@ -544,10 +545,10 @@ async def test_ai_function_invoke_telemetry_async_function(span_exporter: InMemo
assert attributes[OtelAttr.MEASUREMENT_FUNCTION_TAG_NAME] == "async_telemetry_test"
async def test_ai_function_invoke_invalid_pydantic_args():
"""Test the ai_function invoke method with invalid Pydantic model arguments."""
async def test_tool_invoke_invalid_pydantic_args():
"""Test the tool invoke method with invalid Pydantic model arguments."""
@ai_function(name="invalid_args_test", description="A test tool for invalid args")
@tool(name="invalid_args_test", description="A test tool for invalid args")
def invalid_args_test(x: int, y: int) -> int:
"""A function for testing invalid Pydantic args."""
return x + y
@@ -564,20 +565,18 @@ async def test_ai_function_invoke_invalid_pydantic_args():
await invalid_args_test.invoke(arguments=wrong_args)
def test_ai_function_serialization():
"""Test AIFunction serialization and deserialization."""
def test_tool_serialization():
"""Test FunctionTool serialization and deserialization."""
def serialize_test(x: int, y: int) -> int:
"""A function for testing serialization."""
return x - y
serialize_test_ai_function = ai_function(name="serialize_test", description="A test tool for serialization")(
serialize_test
)
serialize_test_tool = tool(name="serialize_test", description="A test tool for serialization")(serialize_test)
# Serialize to dict
tool_dict = serialize_test_ai_function.to_dict()
assert tool_dict["type"] == "ai_function"
tool_dict = serialize_test_tool.to_dict()
assert tool_dict["type"] == "function_tool"
assert tool_dict["name"] == "serialize_test"
assert tool_dict["description"] == "A test tool for serialization"
assert tool_dict["input_model"] == {
@@ -588,21 +587,21 @@ def test_ai_function_serialization():
}
# Deserialize from dict
restored_tool = AIFunction.from_dict(tool_dict, dependencies={"ai_function": {"func": serialize_test}})
assert isinstance(restored_tool, AIFunction)
restored_tool = FunctionTool.from_dict(tool_dict, dependencies={"function_tool": {"func": serialize_test}})
assert isinstance(restored_tool, FunctionTool)
assert restored_tool.name == "serialize_test"
assert restored_tool.description == "A test tool for serialization"
assert restored_tool.parameters() == serialize_test_ai_function.parameters()
assert restored_tool.parameters() == serialize_test_tool.parameters()
assert restored_tool(10, 4) == 6
# Deserialize from dict with instance name
restored_tool_2 = AIFunction.from_dict(
tool_dict, dependencies={"ai_function": {"name:serialize_test": {"func": serialize_test}}}
restored_tool_2 = FunctionTool.from_dict(
tool_dict, dependencies={"function_tool": {"name:serialize_test": {"func": serialize_test}}}
)
assert isinstance(restored_tool_2, AIFunction)
assert isinstance(restored_tool_2, FunctionTool)
assert restored_tool_2.name == "serialize_test"
assert restored_tool_2.description == "A test tool for serialization"
assert restored_tool_2.parameters() == serialize_test_ai_function.parameters()
assert restored_tool_2.parameters() == serialize_test_tool.parameters()
assert restored_tool_2(10, 4) == 6
@@ -979,13 +978,17 @@ def mock_chat_client():
return MockChatClient()
@ai_function(name="no_approval_tool", description="Tool that doesn't require approval")
@tool(
name="no_approval_tool",
description="Tool that doesn't require approval",
approval_mode="never_require",
)
def no_approval_tool(x: int) -> int:
"""A tool that doesn't require approval."""
return x * 2
@ai_function(
@tool(
name="requires_approval_tool",
description="Tool that requires approval",
approval_mode="always_require",
@@ -1451,10 +1454,10 @@ async def test_streaming_two_functions_mixed_approval():
assert all(c.type == "function_approval_request" for c in updates[2].contents)
async def test_ai_function_with_kwargs_injection():
"""Test that ai_function correctly handles kwargs injection and hides them from schema."""
async def test_tool_with_kwargs_injection():
"""Test that tool correctly handles kwargs injection and hides them from schema."""
@ai_function
@tool
def tool_with_kwargs(x: int, **kwargs: Any) -> str:
"""A tool that accepts kwargs."""
user_id = kwargs.get("user_id", "unknown")
@@ -1647,11 +1650,11 @@ def test_build_pydantic_model_from_json_schema_array_of_objects_issue():
# CRITICAL: Validate using the same methods that actual chat clients use
# This is what would actually be sent to the LLM
# Create an AIFunction wrapper to access the client-facing APIs
# Create a FunctionTool wrapper to access the client-facing APIs
def dummy_func(**kwargs):
return kwargs
test_func = AIFunction(
test_func = FunctionTool(
func=dummy_func,
name="create_sales_order",
description="Create a sales order",
python/packages/core/tests/core/test_types.py
+6 -6
View File
@@ -24,10 +24,10 @@ from agent_framework import (
ToolMode,
ToolProtocol,
UsageDetails,
ai_function,
detect_media_type_from_base64,
merge_chat_options,
prepare_function_call_results,
tool,
)
from agent_framework.exceptions import ContentError
@@ -51,10 +51,10 @@ def ai_tool() -> ToolProtocol:
@fixture
def ai_function_tool() -> ToolProtocol:
def tool_tool() -> ToolProtocol:
"""Returns a executable ToolProtocol."""
@ai_function
@tool
def simple_function(x: int, y: int) -> int:
"""A simple function that adds two numbers."""
return x + y
@@ -1017,13 +1017,13 @@ def test_chat_options_tool_choice_validation():
validate_tool_mode({"mode": "auto", "required_function_name": "should_not_be_here"})
def test_chat_options_merge(ai_function_tool, ai_tool) -> None:
def test_chat_options_merge(tool_tool, ai_tool) -> None:
"""Test merge_chat_options utility function."""
from agent_framework import merge_chat_options
options1: ChatOptions = {
"model_id": "gpt-4o",
"tools": [ai_function_tool],
"tools": [tool_tool],
"logit_bias": {"x": 1},
"metadata": {"a": "b"},
}
@@ -1034,7 +1034,7 @@ def test_chat_options_merge(ai_function_tool, ai_tool) -> None:
options3 = merge_chat_options(options1, options2)
assert options3.get("model_id") == "gpt-4.1"
assert options3.get("tools") == [ai_function_tool, ai_tool] # tools are combined
assert options3.get("tools") == [tool_tool, ai_tool] # tools are combined
assert options3.get("logit_bias") == {"x": 1} # base value preserved
assert options3.get("metadata") == {"a": "b"} # base value preserved
python/packages/core/tests/openai/test_assistant_provider.py
+12 -11
View File
@@ -8,7 +8,7 @@ import pytest
from openai.types.beta.assistant import Assistant
from pydantic import BaseModel, Field
from agent_framework import ChatAgent, HostedCodeInterpreterTool, HostedFileSearchTool, ai_function, normalize_tools
from agent_framework import ChatAgent, HostedCodeInterpreterTool, HostedFileSearchTool, normalize_tools, tool
from agent_framework.exceptions import ServiceInitializationError
from agent_framework.openai import OpenAIAssistantProvider
from agent_framework.openai._shared import from_assistant_tools, to_assistant_tools
@@ -244,11 +244,11 @@ class TestOpenAIAssistantProviderCreateAgent:
assert call_kwargs["tools"][0]["type"] == "function"
assert call_kwargs["tools"][0]["function"]["name"] == "get_weather"
async def test_create_agent_with_ai_function(self, mock_async_openai: MagicMock) -> None:
"""Test assistant creation with AIFunction."""
async def test_create_agent_with_tool(self, mock_async_openai: MagicMock) -> None:
"""Test assistant creation with FunctionTool."""
provider = OpenAIAssistantProvider(mock_async_openai)
@ai_function
@tool
def my_function(x: int) -> int:
"""Double a number."""
return x * 2
@@ -537,10 +537,10 @@ class TestOpenAIAssistantProviderAsAgent:
class TestToolConversion:
"""Tests for tool conversion utilities (shared functions)."""
def test_to_assistant_tools_ai_function(self) -> None:
"""Test AIFunction conversion to API format."""
def test_to_assistant_tools_tool(self) -> None:
"""Test FunctionTool conversion to API format."""
@ai_function
@tool
def test_func(x: int) -> int:
"""Test function."""
return x
@@ -555,7 +555,7 @@ class TestToolConversion:
def test_to_assistant_tools_callable(self) -> None:
"""Test raw callable conversion via normalize_tools."""
# normalize_tools converts callables to AIFunction
# normalize_tools converts callables to FunctionTool
normalized = normalize_tools([get_weather])
api_tools = to_assistant_tools(normalized)
@@ -666,12 +666,12 @@ class TestToolValidation:
# Should not raise
provider._validate_function_tools(assistant_tools, None) # type: ignore[reportPrivateUsage]
def test_validate_with_ai_function(self, mock_async_openai: MagicMock) -> None:
"""Test validation with AIFunction."""
def test_validate_with_tool(self, mock_async_openai: MagicMock) -> None:
"""Test validation with FunctionTool."""
provider = OpenAIAssistantProvider(mock_async_openai)
assistant_tools = [create_function_tool("get_weather")]
wrapped = ai_function(get_weather)
wrapped = tool(get_weather)
# Should not raise
provider._validate_function_tools(assistant_tools, [wrapped]) # type: ignore[reportPrivateUsage]
@@ -789,6 +789,7 @@ class TestOpenAIAssistantProviderIntegration:
"""Integration test with function tools."""
provider = OpenAIAssistantProvider()
@tool(approval_mode="never_require")
def get_current_time() -> str:
"""Get the current time."""
from datetime import datetime
python/packages/core/tests/openai/test_openai_assistants_client.py
+5 -4
View File
@@ -23,7 +23,7 @@ from agent_framework import (
HostedCodeInterpreterTool,
HostedFileSearchTool,
Role,
ai_function,
tool,
)
from agent_framework.exceptions import ServiceInitializationError
from agent_framework.openai import OpenAIAssistantsClient
@@ -709,13 +709,13 @@ def test_prepare_options_basic(mock_async_openai: MagicMock) -> None:
assert tool_results is None
def test_prepare_options_with_ai_function_tool(mock_async_openai: MagicMock) -> None:
"""Test _prepare_options with AIFunction tool."""
def test_prepare_options_with_tool_tool(mock_async_openai: MagicMock) -> None:
"""Test _prepare_options with a FunctionTool."""
chat_client = create_test_openai_assistants_client(mock_async_openai)
# Create a simple function for testing and decorate it
@ai_function
@tool(approval_mode="never_require")
def test_function(query: str) -> str:
"""A test function."""
return f"Result for {query}"
@@ -998,6 +998,7 @@ def test_update_agent_name_and_description_none(mock_async_openai: MagicMock) ->
assert chat_client.assistant_name is None
@tool(approval_mode="never_require")
def get_weather(
location: Annotated[str, Field(description="The location to get the weather for.")],
) -> str:
python/packages/core/tests/openai/test_openai_chat_client.py
+4 -4
View File
@@ -19,8 +19,8 @@ from agent_framework import (
Content,
HostedWebSearchTool,
ToolProtocol,
ai_function,
prepare_function_call_results,
tool,
)
from agent_framework.exceptions import ServiceInitializationError, ServiceResponseException
from agent_framework.openai import OpenAIChatClient
@@ -175,7 +175,7 @@ def test_unsupported_tool_handling(openai_unit_test_env: dict[str, str]) -> None
"""Test that unsupported tool types are handled correctly."""
client = OpenAIChatClient()
# Create a mock ToolProtocol that's not an AIFunction
# Create a mock ToolProtocol that's not a FunctionTool
unsupported_tool = MagicMock(spec=ToolProtocol)
unsupported_tool.__class__.__name__ = "UnsupportedAITool"
@@ -189,7 +189,7 @@ def test_unsupported_tool_handling(openai_unit_test_env: dict[str, str]) -> None
assert result["tools"] == [dict_tool]
@ai_function
@tool(approval_mode="never_require")
def get_story_text() -> str:
"""Returns a story about Emily and David."""
return (
@@ -200,7 +200,7 @@ def get_story_text() -> str:
)
@ai_function
@tool(approval_mode="never_require")
def get_weather(location: str) -> str:
"""Get the current weather for a location."""
return f"The weather in {location} is sunny and 72°F."
python/packages/core/tests/openai/test_openai_responses_client.py
+3 -3
View File
@@ -40,7 +40,7 @@ from agent_framework import (
HostedMCPTool,
HostedWebSearchTool,
Role,
ai_function,
tool,
)
from agent_framework.exceptions import (
ServiceInitializationError,
@@ -96,7 +96,7 @@ async def delete_vector_store(client: OpenAIResponsesClient, file_id: str, vecto
await client.client.files.delete(file_id=file_id)
@ai_function
@tool(approval_mode="never_require")
async def get_weather(location: Annotated[str, "The location as a city name"]) -> str:
"""Get the current weather in a given location."""
# Implementation of the tool to get weather
@@ -642,7 +642,7 @@ def test_response_content_creation_with_function_call() -> None:
assert function_call.arguments == '{"location": "Seattle"}'
def test_prepare_content_for_openai_function_approval_response() -> None:
def test_prepare_content_for_opentool_approval_response() -> None:
"""Test _prepare_content_for_openai with function approval response content."""
client = OpenAIResponsesClient(model_id="test-model", api_key="test-key")
python/packages/core/tests/workflow/test_agent_executor_tool_calls.py
+2 -2
View File
@@ -25,8 +25,8 @@ from agent_framework import (
WorkflowBuilder,
WorkflowContext,
WorkflowOutputEvent,
ai_function,
executor,
tool,
use_function_invocation,
)
@@ -132,7 +132,7 @@ async def test_agent_executor_emits_tool_calls_in_streaming_mode() -> None:
assert "sunny" in events[3].data.contents[0].text
@ai_function(approval_mode="always_require")
@tool(approval_mode="always_require")
def mock_tool_requiring_approval(query: str) -> str:
"""Mock tool that requires approval before execution."""
return f"Executed tool with query: {query}"
python/packages/core/tests/workflow/test_workflow_kwargs.py
+2 -2
View File
@@ -20,7 +20,7 @@ from agent_framework import (
SequentialBuilder,
WorkflowRunState,
WorkflowStatusEvent,
ai_function,
tool,
)
from agent_framework._workflows._const import WORKFLOW_RUN_KWARGS_KEY
@@ -28,7 +28,7 @@ from agent_framework._workflows._const import WORKFLOW_RUN_KWARGS_KEY
_received_kwargs: list[dict[str, Any]] = []
@ai_function
@tool(approval_mode="never_require")
def tool_with_kwargs(
action: Annotated[str, "The action to perform"],
**kwargs: Any,