mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
[BREAKING] Python: clean up kwargs across agents, chat clients, tools, and sessions (#4581)
* Python: clean up kwargs across agents, chat clients, tools, and sessions (#3642) Audit and refactor public **kwargs usage across core agents, chat clients, tools, sessions, and provider packages per the migration strategy codified in CODING_STANDARD.md. Key changes: - Add explicit runtime buckets: function_invocation_kwargs and client_kwargs on RawAgent.run() and chat client get_response() layers. - Refactor FunctionTool to prefer explicit ctx: FunctionInvocationContext injection; legacy **kwargs tools still work via _forward_runtime_kwargs. - Refactor Agent.as_tool() to use direct JSON schema, always-streaming wrapper, approval_mode parameter, and UserInputRequiredException propagation (integrates PR #4568 behavior). - Remove implicit session bleeding into FunctionInvocationContext; tools that need a session must receive it via function_invocation_kwargs. - Lower chat-client layers after FunctionInvocationLayer accept only compatibility **kwargs (client_kwargs flattened, function_invocation_kwargs ignored). - Add layered docstring composition from Raw... implementations via _docstrings.py helper. - Clean up provider constructors to use explicit additional_properties. - Deprecation warnings on legacy direct kwargs paths. - Update samples, tests, and typing across all 23 packages. Resolves #3642 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * clarified docstring * feedback fixes * Add unit tests for _docstrings.py build/apply helpers Tests cover: no docstring source, no extra kwargs, appending to existing Keyword Args section, inserting after Args, inserting in plain docstrings, multiline descriptions, ordering, and apply_layered_docstring. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Add test for propagate_session TypeError on non-AgentSession values Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Add tests for multi-content and empty UserInputRequiredException propagation Cover the branching logic in _try_execute_function_calls for: - Multiple user_input_request items in a single exception (extra_user_input_contents path) - Empty contents list (fallback function_result path) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Add tests for DurableAIAgent.get_session forwarding service_session_id Verifies get_session correctly forwards service_session_id and session_id to the executor's get_new_session, replacing the removed kwargs test. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Simplify ag-ui test stub to read session from client_kwargs only Remove dual-mode detection (client_kwargs vs raw kwargs fallback) from the test mock. Session is now read exclusively from client_kwargs, matching the settled public calling convention. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * updated create and get sessions in durable * fixed docstrings * fix test * updated session handling * updated from main * updated tests --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
Eduard van Valkenburg
committed by
GitHub
Unverified
GitHub
parent
b7990908fe
commit
a4b9539b62
@@ -1,6 +1,7 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import contextlib
|
||||
import inspect
|
||||
from collections.abc import AsyncIterable, MutableSequence
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
@@ -31,6 +32,7 @@ from agent_framework import (
|
||||
)
|
||||
from agent_framework._agents import _get_tool_name, _merge_options, _sanitize_agent_name
|
||||
from agent_framework._mcp import MCPTool, _build_prefixed_mcp_name, _normalize_mcp_name
|
||||
from agent_framework._middleware import FunctionInvocationContext
|
||||
|
||||
|
||||
class _FixedTokenizer:
|
||||
@@ -101,6 +103,30 @@ def test_chat_client_agent_type(client: SupportsChatGetResponse) -> None:
|
||||
assert isinstance(chat_client_agent, SupportsAgentRun)
|
||||
|
||||
|
||||
def test_agent_init_docstring_surfaces_raw_agent_constructor_docs() -> None:
|
||||
docstring = inspect.getdoc(Agent.__init__)
|
||||
|
||||
assert docstring is not None
|
||||
assert "client: The chat client to use for the agent." in docstring
|
||||
assert "middleware: List of middleware to intercept agent and function invocations." in docstring
|
||||
|
||||
|
||||
def test_agent_run_docstring_surfaces_raw_agent_runtime_docs() -> None:
|
||||
docstring = inspect.getdoc(Agent.run)
|
||||
|
||||
assert docstring is not None
|
||||
assert "Run the agent with the given messages and options." in docstring
|
||||
assert "function_invocation_kwargs: Keyword arguments forwarded to tool invocation." in docstring
|
||||
assert "middleware: Optional per-run agent, chat, and function middleware." in docstring
|
||||
|
||||
|
||||
def test_agent_run_is_defined_on_agent_class() -> None:
|
||||
signature = inspect.signature(Agent.run)
|
||||
|
||||
assert Agent.run.__qualname__ == "Agent.run"
|
||||
assert "middleware" in signature.parameters
|
||||
|
||||
|
||||
async def test_chat_client_agent_init(client: SupportsChatGetResponse) -> None:
|
||||
agent_id = str(uuid4())
|
||||
agent = Agent(client=client, id=agent_id, description="Test")
|
||||
@@ -121,6 +147,13 @@ async def test_chat_client_agent_init_with_name(
|
||||
assert agent.description == "Test"
|
||||
|
||||
|
||||
def test_agent_init_warns_for_direct_additional_properties(client: SupportsChatGetResponse) -> None:
|
||||
with pytest.warns(DeprecationWarning, match="additional_properties"):
|
||||
agent = Agent(client=client, legacy_key="legacy-value")
|
||||
|
||||
assert agent.additional_properties["legacy_key"] == "legacy-value"
|
||||
|
||||
|
||||
async def test_chat_client_agent_run(client: SupportsChatGetResponse) -> None:
|
||||
agent = Agent(client=client)
|
||||
|
||||
@@ -253,33 +286,38 @@ async def test_prepare_session_does_not_mutate_agent_chat_options(
|
||||
assert len(agent.default_options["tools"]) == 1
|
||||
|
||||
|
||||
async def test_prepare_run_context_keeps_compaction_overrides_out_of_kwargs(
|
||||
async def test_prepare_run_context_handles_function_kwargs(
|
||||
chat_client_base: SupportsChatGetResponse,
|
||||
) -> None:
|
||||
strategy = SlidingWindowStrategy(keep_last_groups=2)
|
||||
tokenizer = _FixedTokenizer(13)
|
||||
agent = Agent(client=chat_client_base)
|
||||
session = agent.create_session()
|
||||
|
||||
ctx = await agent._prepare_run_context( # type: ignore[reportPrivateUsage]
|
||||
messages=[Message(role="user", text="Hello")],
|
||||
session=None,
|
||||
messages="Hello",
|
||||
session=session,
|
||||
tools=None,
|
||||
options=None,
|
||||
compaction_strategy=strategy,
|
||||
tokenizer=tokenizer,
|
||||
kwargs={"custom_flag": True},
|
||||
options={
|
||||
"temperature": 0.4,
|
||||
"additional_function_arguments": {"from_options": "options-value"},
|
||||
},
|
||||
compaction_strategy=None,
|
||||
tokenizer=None,
|
||||
legacy_kwargs={"legacy_key": "legacy-value"},
|
||||
function_invocation_kwargs={"runtime_key": "runtime-value"},
|
||||
client_kwargs={"client_key": "client-value"},
|
||||
)
|
||||
|
||||
assert ctx["compaction_strategy"] is strategy
|
||||
assert ctx["tokenizer"] is tokenizer
|
||||
assert ctx["filtered_kwargs"].get("custom_flag") is True
|
||||
assert "compaction_strategy" not in ctx["filtered_kwargs"]
|
||||
assert "tokenizer" not in ctx["filtered_kwargs"]
|
||||
assert ctx["chat_options"]["temperature"] == 0.4
|
||||
assert "additional_function_arguments" not in ctx["chat_options"]
|
||||
assert ctx["function_invocation_kwargs"]["from_options"] == "options-value"
|
||||
assert ctx["function_invocation_kwargs"]["legacy_key"] == "legacy-value"
|
||||
assert ctx["function_invocation_kwargs"]["runtime_key"] == "runtime-value"
|
||||
assert "session" not in ctx["function_invocation_kwargs"]
|
||||
assert ctx["client_kwargs"]["client_key"] == "client-value"
|
||||
assert ctx["client_kwargs"]["session"] is session
|
||||
|
||||
|
||||
async def test_chat_client_agent_run_with_session(
|
||||
chat_client_base: SupportsChatGetResponse,
|
||||
) -> None:
|
||||
async def test_chat_client_agent_run_with_session(chat_client_base: SupportsChatGetResponse) -> None:
|
||||
mock_response = ChatResponse(
|
||||
messages=[Message(role="assistant", contents=[Content.from_text("test response")])],
|
||||
conversation_id="123",
|
||||
@@ -720,8 +758,9 @@ async def test_chat_agent_as_tool_basic(client: SupportsChatGetResponse) -> None
|
||||
|
||||
assert tool.name == "TestAgent"
|
||||
assert tool.description == "Test agent for as_tool"
|
||||
assert tool.approval_mode == "never_require"
|
||||
assert hasattr(tool, "func")
|
||||
assert hasattr(tool, "input_model")
|
||||
assert tool.input_model is None
|
||||
|
||||
|
||||
async def test_chat_agent_as_tool_custom_parameters(
|
||||
@@ -735,13 +774,15 @@ async def test_chat_agent_as_tool_custom_parameters(
|
||||
description="Custom description",
|
||||
arg_name="query",
|
||||
arg_description="Custom input description",
|
||||
approval_mode="always_require",
|
||||
)
|
||||
|
||||
assert tool.name == "CustomTool"
|
||||
assert tool.description == "Custom description"
|
||||
assert tool.approval_mode == "always_require"
|
||||
|
||||
# Check that the input model has the custom field name
|
||||
schema = tool.input_model.model_json_schema()
|
||||
schema = tool.parameters()
|
||||
assert "query" in schema["properties"]
|
||||
assert schema["properties"]["query"]["description"] == "Custom input description"
|
||||
|
||||
@@ -760,7 +801,7 @@ async def test_chat_agent_as_tool_defaults(client: SupportsChatGetResponse) -> N
|
||||
assert tool.description == "" # Should default to empty string
|
||||
|
||||
# Check default input field
|
||||
schema = tool.input_model.model_json_schema()
|
||||
schema = tool.parameters()
|
||||
assert "task" in schema["properties"]
|
||||
assert "Task for TestAgent" in schema["properties"]["task"]["description"]
|
||||
|
||||
@@ -783,12 +824,12 @@ async def test_chat_agent_as_tool_function_execution(
|
||||
tool = agent.as_tool()
|
||||
|
||||
# Test function execution
|
||||
result = await tool.invoke(arguments=tool.input_model(task="Hello"))
|
||||
result = await tool.invoke(arguments={"task": "Hello"})
|
||||
|
||||
# Should return the agent's response text as a list of Content items
|
||||
assert isinstance(result, list)
|
||||
assert len(result) == 1
|
||||
assert result[0].text == "test response" # From mock chat client
|
||||
assert result[0].text == "test streaming response another update" # From mock streaming client
|
||||
|
||||
|
||||
async def test_chat_agent_as_tool_with_stream_callback(
|
||||
@@ -806,7 +847,7 @@ async def test_chat_agent_as_tool_with_stream_callback(
|
||||
tool = agent.as_tool(stream_callback=stream_callback)
|
||||
|
||||
# Execute the tool
|
||||
result = await tool.invoke(arguments=tool.input_model(task="Hello"))
|
||||
result = await tool.invoke(arguments={"task": "Hello"})
|
||||
|
||||
# Should have collected streaming updates
|
||||
assert len(collected_updates) > 0
|
||||
@@ -826,9 +867,9 @@ async def test_chat_agent_as_tool_with_custom_arg_name(
|
||||
tool = agent.as_tool(arg_name="prompt", arg_description="Custom prompt input")
|
||||
|
||||
# Test that the custom argument name works
|
||||
result = await tool.invoke(arguments=tool.input_model(prompt="Test prompt"))
|
||||
result = await tool.invoke(arguments={"prompt": "Test prompt"})
|
||||
assert isinstance(result, list)
|
||||
assert result[0].text == "test response"
|
||||
assert result[0].text == "test streaming response another update"
|
||||
|
||||
|
||||
async def test_chat_agent_as_tool_with_async_stream_callback(
|
||||
@@ -846,7 +887,7 @@ async def test_chat_agent_as_tool_with_async_stream_callback(
|
||||
tool = agent.as_tool(stream_callback=async_stream_callback)
|
||||
|
||||
# Execute the tool
|
||||
result = await tool.invoke(arguments=tool.input_model(task="Hello"))
|
||||
result = await tool.invoke(arguments={"task": "Hello"})
|
||||
|
||||
# Should have collected streaming updates
|
||||
assert len(collected_updates) > 0
|
||||
@@ -877,17 +918,14 @@ async def test_chat_agent_as_tool_name_sanitization(
|
||||
assert tool.name == expected_tool_name, f"Expected {expected_tool_name}, got {tool.name} for input {agent_name}"
|
||||
|
||||
|
||||
async def test_chat_agent_as_tool_propagate_session_true(
|
||||
client: SupportsChatGetResponse,
|
||||
) -> None:
|
||||
"""Test that propagate_session=True forwards the parent's session to the sub-agent."""
|
||||
async def test_chat_agent_as_tool_propagate_session_true(client: SupportsChatGetResponse) -> None:
|
||||
"""Test that propagate_session=True forwards the session to the sub-agent."""
|
||||
agent = Agent(client=client, name="SubAgent", description="Sub agent")
|
||||
tool = agent.as_tool(propagate_session=True)
|
||||
|
||||
parent_session = AgentSession(session_id="parent-session-123")
|
||||
parent_session.state["shared_key"] = "shared_value"
|
||||
|
||||
# Spy on the agent's run method to capture the session argument
|
||||
original_run = agent.run
|
||||
captured_session = None
|
||||
|
||||
@@ -898,16 +936,20 @@ async def test_chat_agent_as_tool_propagate_session_true(
|
||||
|
||||
agent.run = capturing_run # type: ignore[assignment, method-assign]
|
||||
|
||||
await tool.invoke(arguments=tool.input_model(task="Hello"), session=parent_session)
|
||||
await tool.invoke(
|
||||
context=FunctionInvocationContext(
|
||||
function=tool,
|
||||
arguments={"task": "Hello"},
|
||||
session=parent_session,
|
||||
)
|
||||
)
|
||||
|
||||
assert captured_session is parent_session
|
||||
assert captured_session.session_id == "parent-session-123"
|
||||
assert captured_session.state["shared_key"] == "shared_value"
|
||||
|
||||
|
||||
async def test_chat_agent_as_tool_propagate_session_false_by_default(
|
||||
client: SupportsChatGetResponse,
|
||||
) -> None:
|
||||
async def test_chat_agent_as_tool_propagate_session_false_by_default(client: SupportsChatGetResponse) -> None:
|
||||
"""Test that propagate_session defaults to False and does not forward the session."""
|
||||
agent = Agent(client=client, name="SubAgent", description="Sub agent")
|
||||
tool = agent.as_tool() # default: propagate_session=False
|
||||
@@ -924,22 +966,25 @@ async def test_chat_agent_as_tool_propagate_session_false_by_default(
|
||||
|
||||
agent.run = capturing_run # type: ignore[assignment, method-assign]
|
||||
|
||||
await tool.invoke(arguments=tool.input_model(task="Hello"), session=parent_session)
|
||||
await tool.invoke(
|
||||
context=FunctionInvocationContext(
|
||||
function=tool,
|
||||
arguments={"task": "Hello"},
|
||||
session=parent_session,
|
||||
)
|
||||
)
|
||||
|
||||
assert captured_session is None
|
||||
|
||||
|
||||
async def test_chat_agent_as_tool_propagate_session_shares_state(
|
||||
client: SupportsChatGetResponse,
|
||||
) -> None:
|
||||
"""Test that shared session allows the sub-agent to read and write parent's state."""
|
||||
async def test_chat_agent_as_tool_propagate_session_shares_state(client: SupportsChatGetResponse) -> None:
|
||||
"""Test that a propagated session allows the sub-agent to read and write parent state."""
|
||||
agent = Agent(client=client, name="SubAgent", description="Sub agent")
|
||||
tool = agent.as_tool(propagate_session=True)
|
||||
|
||||
parent_session = AgentSession(session_id="shared-session")
|
||||
parent_session.state["counter"] = 0
|
||||
|
||||
# The sub-agent receives the same session object, so mutations are shared
|
||||
original_run = agent.run
|
||||
captured_session = None
|
||||
|
||||
@@ -952,9 +997,14 @@ async def test_chat_agent_as_tool_propagate_session_shares_state(
|
||||
|
||||
agent.run = capturing_run # type: ignore[assignment, method-assign]
|
||||
|
||||
await tool.invoke(arguments=tool.input_model(task="Hello"), session=parent_session)
|
||||
await tool.invoke(
|
||||
context=FunctionInvocationContext(
|
||||
function=tool,
|
||||
arguments={"task": "Hello"},
|
||||
session=parent_session,
|
||||
)
|
||||
)
|
||||
|
||||
# The parent's state should reflect the sub-agent's mutation
|
||||
assert parent_session.state["counter"] == 1
|
||||
|
||||
|
||||
@@ -1131,7 +1181,7 @@ async def test_agent_run_accepts_prefixed_mcp_tools(chat_client_base: Any) -> No
|
||||
|
||||
|
||||
async def test_agent_tool_receives_session_in_kwargs(chat_client_base: Any) -> None:
|
||||
"""Verify tool execution receives 'session' inside **kwargs when function is called by client."""
|
||||
"""Verify legacy **kwargs tools receive the session when agent.run() is called with one."""
|
||||
|
||||
captured: dict[str, Any] = {}
|
||||
|
||||
@@ -1142,7 +1192,6 @@ async def test_agent_tool_receives_session_in_kwargs(chat_client_base: Any) -> N
|
||||
captured["has_state"] = session.state is not None if isinstance(session, AgentSession) else False
|
||||
return f"echo: {text}"
|
||||
|
||||
# Make the base client emit a function call for our tool
|
||||
chat_client_base.run_responses = [
|
||||
ChatResponse(
|
||||
messages=Message(
|
||||
@@ -1162,17 +1211,52 @@ async def test_agent_tool_receives_session_in_kwargs(chat_client_base: Any) -> N
|
||||
agent = Agent(client=chat_client_base, tools=[echo_session_info])
|
||||
session = agent.create_session()
|
||||
|
||||
result = await agent.run(
|
||||
"hello",
|
||||
session=session,
|
||||
options={"additional_function_arguments": {"session": session}},
|
||||
)
|
||||
result = await agent.run("hello", session=session)
|
||||
|
||||
assert result.text == "done"
|
||||
assert captured.get("has_session") is True
|
||||
assert captured.get("has_state") is True
|
||||
|
||||
|
||||
async def test_agent_tool_receives_explicit_session_via_function_invocation_context_kwargs(
|
||||
chat_client_base: Any,
|
||||
) -> None:
|
||||
"""Verify ctx-based tools receive the session via FunctionInvocationContext.session."""
|
||||
|
||||
captured: dict[str, Any] = {}
|
||||
|
||||
@tool(name="capture_session_context", approval_mode="never_require")
|
||||
def capture_session_context(text: str, ctx: FunctionInvocationContext) -> str:
|
||||
captured["session"] = ctx.session
|
||||
captured["has_state"] = ctx.session.state is not None if isinstance(ctx.session, AgentSession) else False
|
||||
return f"echo: {text}"
|
||||
|
||||
chat_client_base.run_responses = [
|
||||
ChatResponse(
|
||||
messages=Message(
|
||||
role="assistant",
|
||||
contents=[
|
||||
Content.from_function_call(
|
||||
call_id="1",
|
||||
name="capture_session_context",
|
||||
arguments='{"text": "hello"}',
|
||||
)
|
||||
],
|
||||
)
|
||||
),
|
||||
ChatResponse(messages=Message(role="assistant", text="done")),
|
||||
]
|
||||
|
||||
agent = Agent(client=chat_client_base, tools=[capture_session_context])
|
||||
session = agent.create_session()
|
||||
|
||||
result = await agent.run("hello", session=session)
|
||||
|
||||
assert result.text == "done"
|
||||
assert captured["session"] is session
|
||||
assert captured["has_state"] is True
|
||||
|
||||
|
||||
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."""
|
||||
|
||||
@@ -1859,4 +1943,26 @@ async def test_stores_by_default_with_store_false_in_default_options_injects_inm
|
||||
assert any(isinstance(p, InMemoryHistoryProvider) for p in agent.context_providers)
|
||||
|
||||
|
||||
# endregion
|
||||
# region as_tool user_input_request propagation
|
||||
|
||||
|
||||
async def test_as_tool_raises_on_user_input_request(client: SupportsChatGetResponse) -> None:
|
||||
"""Test that as_tool raises when the wrapped sub-agent requests user input."""
|
||||
from agent_framework.exceptions import UserInputRequiredException
|
||||
|
||||
consent_content = Content.from_oauth_consent_request(
|
||||
consent_link="https://login.microsoftonline.com/consent",
|
||||
)
|
||||
client.streaming_responses = [ # type: ignore[attr-defined]
|
||||
[ChatResponseUpdate(contents=[consent_content], role="assistant")],
|
||||
]
|
||||
|
||||
agent = Agent(client=client, name="OAuthAgent", description="Agent requiring consent")
|
||||
agent_tool = agent.as_tool()
|
||||
|
||||
with raises(UserInputRequiredException) as exc_info:
|
||||
await agent_tool.invoke(arguments={"task": "Do something"})
|
||||
|
||||
assert len(exc_info.value.contents) == 1
|
||||
assert exc_info.value.contents[0].type == "oauth_consent_request"
|
||||
assert exc_info.value.contents[0].consent_link == "https://login.microsoftonline.com/consent"
|
||||
|
||||
@@ -6,7 +6,7 @@ from collections.abc import Awaitable, Callable
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import Agent, ChatResponse, Content, Message, agent_middleware
|
||||
from agent_framework._middleware import AgentContext
|
||||
from agent_framework._middleware import AgentContext, FunctionInvocationContext
|
||||
|
||||
from .conftest import MockChatClient
|
||||
|
||||
@@ -14,14 +14,28 @@ from .conftest import MockChatClient
|
||||
class TestAsToolKwargsPropagation:
|
||||
"""Test cases for kwargs propagation through as_tool() delegation."""
|
||||
|
||||
@staticmethod
|
||||
def _build_context(
|
||||
tool: Any,
|
||||
*,
|
||||
task: str,
|
||||
runtime_kwargs: dict[str, Any] | None = None,
|
||||
) -> FunctionInvocationContext:
|
||||
return FunctionInvocationContext(
|
||||
function=tool,
|
||||
arguments={"task": task},
|
||||
kwargs=runtime_kwargs,
|
||||
)
|
||||
|
||||
async def test_as_tool_forwards_runtime_kwargs(self, client: MockChatClient) -> None:
|
||||
"""Test that runtime kwargs are forwarded through as_tool() to sub-agent."""
|
||||
"""Test that runtime kwargs are forwarded through as_tool() to sub-agent tools."""
|
||||
captured_kwargs: dict[str, Any] = {}
|
||||
captured_function_invocation_kwargs: dict[str, Any] = {}
|
||||
|
||||
@agent_middleware
|
||||
async def capture_middleware(context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
|
||||
# Capture kwargs passed to the sub-agent
|
||||
captured_kwargs.update(context.kwargs)
|
||||
captured_function_invocation_kwargs.update(context.function_invocation_kwargs)
|
||||
await call_next()
|
||||
|
||||
# Setup mock response
|
||||
@@ -39,29 +53,31 @@ class TestAsToolKwargsPropagation:
|
||||
# Create tool from sub-agent
|
||||
tool = sub_agent.as_tool(name="delegate", arg_name="task")
|
||||
|
||||
# Directly invoke the tool with kwargs (simulating what happens during agent execution)
|
||||
# Directly invoke the tool with explicit runtime context (simulating agent execution).
|
||||
_ = await tool.invoke(
|
||||
arguments=tool.input_model(task="Test delegation"),
|
||||
api_token="secret-xyz-123",
|
||||
user_id="user-456",
|
||||
session_id="session-789",
|
||||
context=self._build_context(
|
||||
tool,
|
||||
task="Test delegation",
|
||||
runtime_kwargs={
|
||||
"api_token": "secret-xyz-123",
|
||||
"user_id": "user-456",
|
||||
"session_id": "session-789",
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
# Verify kwargs were forwarded to sub-agent
|
||||
assert "api_token" in captured_kwargs, f"Expected 'api_token' in {captured_kwargs}"
|
||||
assert captured_kwargs["api_token"] == "secret-xyz-123"
|
||||
assert "user_id" in captured_kwargs
|
||||
assert captured_kwargs["user_id"] == "user-456"
|
||||
assert "session_id" in captured_kwargs
|
||||
assert captured_kwargs["session_id"] == "session-789"
|
||||
assert captured_kwargs == {}
|
||||
assert captured_function_invocation_kwargs["api_token"] == "secret-xyz-123"
|
||||
assert captured_function_invocation_kwargs["user_id"] == "user-456"
|
||||
assert captured_function_invocation_kwargs["session_id"] == "session-789"
|
||||
|
||||
async def test_as_tool_excludes_arg_name_from_forwarded_kwargs(self, client: MockChatClient) -> None:
|
||||
"""Test that the arg_name parameter is not forwarded as a kwarg."""
|
||||
captured_kwargs: dict[str, Any] = {}
|
||||
async def test_as_tool_forwards_context_kwargs_verbatim(self, client: MockChatClient) -> None:
|
||||
"""Test that runtime kwargs are forwarded exactly from FunctionInvocationContext.kwargs."""
|
||||
captured_function_invocation_kwargs: dict[str, Any] = {}
|
||||
|
||||
@agent_middleware
|
||||
async def capture_middleware(context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
|
||||
captured_kwargs.update(context.kwargs)
|
||||
captured_function_invocation_kwargs.update(context.function_invocation_kwargs)
|
||||
await call_next()
|
||||
|
||||
# Setup mock response
|
||||
@@ -79,25 +95,26 @@ class TestAsToolKwargsPropagation:
|
||||
|
||||
# Invoke tool with both the arg_name field and additional kwargs
|
||||
await tool.invoke(
|
||||
arguments=tool.input_model(custom_task="Test task"),
|
||||
api_token="token-123",
|
||||
custom_task="should_be_excluded", # This should be filtered out
|
||||
context=FunctionInvocationContext(
|
||||
function=tool,
|
||||
arguments={"custom_task": "Test task"},
|
||||
kwargs={
|
||||
"api_token": "token-123",
|
||||
"custom_task": "should_be_excluded",
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
# The arg_name ("custom_task") should NOT be in the forwarded kwargs
|
||||
assert "custom_task" not in captured_kwargs
|
||||
# But other kwargs should be present
|
||||
assert "api_token" in captured_kwargs
|
||||
assert captured_kwargs["api_token"] == "token-123"
|
||||
assert captured_function_invocation_kwargs["custom_task"] == "should_be_excluded"
|
||||
assert captured_function_invocation_kwargs["api_token"] == "token-123"
|
||||
|
||||
async def test_as_tool_nested_delegation_propagates_kwargs(self, client: MockChatClient) -> None:
|
||||
"""Test that kwargs propagate through multiple levels of delegation (A → B → C)."""
|
||||
captured_kwargs_list: list[dict[str, Any]] = []
|
||||
"""Test that runtime kwargs propagate through multiple levels of delegation (A -> B -> C)."""
|
||||
captured_function_invocation_kwargs_list: list[dict[str, Any]] = []
|
||||
|
||||
@agent_middleware
|
||||
async def capture_middleware(context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
|
||||
# Capture kwargs at each level
|
||||
captured_kwargs_list.append(dict(context.kwargs))
|
||||
captured_function_invocation_kwargs_list.append(dict(context.function_invocation_kwargs))
|
||||
await call_next()
|
||||
|
||||
# Setup mock responses to trigger nested tool invocation: B calls tool C, then completes.
|
||||
@@ -140,24 +157,29 @@ class TestAsToolKwargsPropagation:
|
||||
|
||||
# Invoke tool B with kwargs - should propagate to both B and C
|
||||
await tool_b.invoke(
|
||||
arguments=tool_b.input_model(task="Test cascade"),
|
||||
trace_id="trace-abc-123",
|
||||
tenant_id="tenant-xyz",
|
||||
options={"additional_function_arguments": {"trace_id": "trace-abc-123", "tenant_id": "tenant-xyz"}},
|
||||
context=self._build_context(
|
||||
tool_b,
|
||||
task="Test cascade",
|
||||
runtime_kwargs={
|
||||
"trace_id": "trace-abc-123",
|
||||
"tenant_id": "tenant-xyz",
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
# Verify kwargs were forwarded to the first agent invocation.
|
||||
assert len(captured_kwargs_list) >= 1
|
||||
assert captured_kwargs_list[0].get("trace_id") == "trace-abc-123"
|
||||
assert captured_kwargs_list[0].get("tenant_id") == "tenant-xyz"
|
||||
assert len(captured_function_invocation_kwargs_list) >= 1
|
||||
assert captured_function_invocation_kwargs_list[0].get("trace_id") == "trace-abc-123"
|
||||
assert captured_function_invocation_kwargs_list[0].get("tenant_id") == "tenant-xyz"
|
||||
|
||||
async def test_as_tool_streaming_mode_forwards_kwargs(self, client: MockChatClient) -> None:
|
||||
"""Test that kwargs are forwarded in streaming mode."""
|
||||
"""Test that runtime kwargs are forwarded in streaming mode."""
|
||||
captured_kwargs: dict[str, Any] = {}
|
||||
captured_function_invocation_kwargs: dict[str, Any] = {}
|
||||
|
||||
@agent_middleware
|
||||
async def capture_middleware(context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
|
||||
captured_kwargs.update(context.kwargs)
|
||||
captured_function_invocation_kwargs.update(context.function_invocation_kwargs)
|
||||
await call_next()
|
||||
|
||||
# Setup mock streaming responses
|
||||
@@ -182,13 +204,15 @@ class TestAsToolKwargsPropagation:
|
||||
|
||||
# Invoke tool with kwargs while streaming callback is active
|
||||
await tool.invoke(
|
||||
arguments=tool.input_model(task="Test streaming"),
|
||||
api_key="streaming-key-999",
|
||||
context=self._build_context(
|
||||
tool,
|
||||
task="Test streaming",
|
||||
runtime_kwargs={"api_key": "streaming-key-999"},
|
||||
),
|
||||
)
|
||||
|
||||
# Verify kwargs were forwarded even in streaming mode
|
||||
assert "api_key" in captured_kwargs
|
||||
assert captured_kwargs["api_key"] == "streaming-key-999"
|
||||
assert captured_kwargs == {}
|
||||
assert captured_function_invocation_kwargs["api_key"] == "streaming-key-999"
|
||||
assert len(captured_updates) == 1
|
||||
|
||||
async def test_as_tool_empty_kwargs_still_works(self, client: MockChatClient) -> None:
|
||||
@@ -206,18 +230,20 @@ class TestAsToolKwargsPropagation:
|
||||
tool = sub_agent.as_tool()
|
||||
|
||||
# Invoke without any extra kwargs - should work without errors
|
||||
result = await tool.invoke(arguments=tool.input_model(task="Simple task"))
|
||||
result = await tool.invoke(arguments={"task": "Simple task"})
|
||||
|
||||
# Verify tool executed successfully
|
||||
assert result is not None
|
||||
|
||||
async def test_as_tool_kwargs_with_chat_options(self, client: MockChatClient) -> None:
|
||||
"""Test that kwargs including chat_options are properly forwarded."""
|
||||
"""Test that runtime kwargs are forwarded only via function_invocation_kwargs."""
|
||||
captured_kwargs: dict[str, Any] = {}
|
||||
captured_function_invocation_kwargs: dict[str, Any] = {}
|
||||
|
||||
@agent_middleware
|
||||
async def capture_middleware(context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
|
||||
captured_kwargs.update(context.kwargs)
|
||||
captured_function_invocation_kwargs.update(context.function_invocation_kwargs)
|
||||
await call_next()
|
||||
|
||||
# Setup mock response
|
||||
@@ -235,24 +261,26 @@ class TestAsToolKwargsPropagation:
|
||||
|
||||
# Invoke with various kwargs
|
||||
await tool.invoke(
|
||||
arguments=tool.input_model(task="Test with options"),
|
||||
temperature=0.8,
|
||||
max_tokens=500,
|
||||
custom_param="custom_value",
|
||||
context=self._build_context(
|
||||
tool,
|
||||
task="Test with options",
|
||||
runtime_kwargs={
|
||||
"temperature": 0.8,
|
||||
"max_tokens": 500,
|
||||
"custom_param": "custom_value",
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
# Verify all kwargs were forwarded
|
||||
assert "temperature" in captured_kwargs
|
||||
assert captured_kwargs["temperature"] == 0.8
|
||||
assert "max_tokens" in captured_kwargs
|
||||
assert captured_kwargs["max_tokens"] == 500
|
||||
assert "custom_param" in captured_kwargs
|
||||
assert captured_kwargs["custom_param"] == "custom_value"
|
||||
assert captured_kwargs == {}
|
||||
assert captured_function_invocation_kwargs["temperature"] == 0.8
|
||||
assert captured_function_invocation_kwargs["max_tokens"] == 500
|
||||
assert captured_function_invocation_kwargs["custom_param"] == "custom_value"
|
||||
|
||||
async def test_as_tool_kwargs_isolated_per_invocation(self, client: MockChatClient) -> None:
|
||||
"""Test that kwargs are isolated per invocation and don't leak between calls."""
|
||||
first_call_kwargs: dict[str, Any] = {}
|
||||
second_call_kwargs: dict[str, Any] = {}
|
||||
"""Test that runtime kwargs are isolated per invocation and don't leak between calls."""
|
||||
first_call_function_invocation_kwargs: dict[str, Any] = {}
|
||||
second_call_function_invocation_kwargs: dict[str, Any] = {}
|
||||
call_count = 0
|
||||
|
||||
@agent_middleware
|
||||
@@ -260,9 +288,9 @@ class TestAsToolKwargsPropagation:
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if call_count == 1:
|
||||
first_call_kwargs.update(context.kwargs)
|
||||
first_call_function_invocation_kwargs.update(context.function_invocation_kwargs)
|
||||
elif call_count == 2:
|
||||
second_call_kwargs.update(context.kwargs)
|
||||
second_call_function_invocation_kwargs.update(context.function_invocation_kwargs)
|
||||
await call_next()
|
||||
|
||||
# Setup mock responses for both calls
|
||||
@@ -281,33 +309,35 @@ class TestAsToolKwargsPropagation:
|
||||
|
||||
# First call with specific kwargs
|
||||
await tool.invoke(
|
||||
arguments=tool.input_model(task="First task"),
|
||||
session_id="session-1",
|
||||
api_token="token-1",
|
||||
context=self._build_context(
|
||||
tool,
|
||||
task="First task",
|
||||
runtime_kwargs={"session_id": "session-1", "api_token": "token-1"},
|
||||
),
|
||||
)
|
||||
|
||||
# Second call with different kwargs
|
||||
await tool.invoke(
|
||||
arguments=tool.input_model(task="Second task"),
|
||||
session_id="session-2",
|
||||
api_token="token-2",
|
||||
context=self._build_context(
|
||||
tool,
|
||||
task="Second task",
|
||||
runtime_kwargs={"session_id": "session-2", "api_token": "token-2"},
|
||||
),
|
||||
)
|
||||
|
||||
# Verify first call had its own kwargs
|
||||
assert first_call_kwargs.get("session_id") == "session-1"
|
||||
assert first_call_kwargs.get("api_token") == "token-1"
|
||||
assert first_call_function_invocation_kwargs.get("session_id") == "session-1"
|
||||
assert first_call_function_invocation_kwargs.get("api_token") == "token-1"
|
||||
|
||||
# Verify second call had its own kwargs (not leaked from first)
|
||||
assert second_call_kwargs.get("session_id") == "session-2"
|
||||
assert second_call_kwargs.get("api_token") == "token-2"
|
||||
assert second_call_function_invocation_kwargs.get("session_id") == "session-2"
|
||||
assert second_call_function_invocation_kwargs.get("api_token") == "token-2"
|
||||
|
||||
async def test_as_tool_excludes_conversation_id_from_forwarded_kwargs(self, client: MockChatClient) -> None:
|
||||
"""Test that conversation_id is not forwarded to sub-agent."""
|
||||
captured_kwargs: dict[str, Any] = {}
|
||||
async def test_as_tool_forwards_conversation_id_from_context_kwargs(self, client: MockChatClient) -> None:
|
||||
"""Test that conversation_id is forwarded when explicitly present in runtime context kwargs."""
|
||||
captured_function_invocation_kwargs: dict[str, Any] = {}
|
||||
|
||||
@agent_middleware
|
||||
async def capture_middleware(context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
|
||||
captured_kwargs.update(context.kwargs)
|
||||
captured_function_invocation_kwargs.update(context.function_invocation_kwargs)
|
||||
await call_next()
|
||||
|
||||
# Setup mock response
|
||||
@@ -325,17 +355,17 @@ class TestAsToolKwargsPropagation:
|
||||
|
||||
# Invoke tool with conversation_id in kwargs (simulating parent's conversation state)
|
||||
await tool.invoke(
|
||||
arguments=tool.input_model(task="Test delegation"),
|
||||
conversation_id="conv-parent-456",
|
||||
api_token="secret-xyz-123",
|
||||
user_id="user-456",
|
||||
context=self._build_context(
|
||||
tool,
|
||||
task="Test delegation",
|
||||
runtime_kwargs={
|
||||
"conversation_id": "conv-parent-456",
|
||||
"api_token": "secret-xyz-123",
|
||||
"user_id": "user-456",
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
# Verify conversation_id was NOT forwarded to sub-agent
|
||||
assert "conversation_id" not in captured_kwargs, (
|
||||
f"conversation_id should not be forwarded, but got: {captured_kwargs}"
|
||||
)
|
||||
|
||||
# Verify other kwargs were still forwarded
|
||||
assert captured_kwargs.get("api_token") == "secret-xyz-123"
|
||||
assert captured_kwargs.get("user_id") == "user-456"
|
||||
assert captured_function_invocation_kwargs.get("conversation_id") == "conv-parent-456"
|
||||
assert captured_function_invocation_kwargs.get("api_token") == "secret-xyz-123"
|
||||
assert captured_function_invocation_kwargs.get("user_id") == "user-456"
|
||||
|
||||
@@ -1,9 +1,12 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
|
||||
import inspect
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from agent_framework import (
|
||||
GROUP_ANNOTATION_KEY,
|
||||
GROUP_TOKEN_COUNT_KEY,
|
||||
@@ -50,6 +53,60 @@ def test_base_client(chat_client_base: SupportsChatGetResponse):
|
||||
assert isinstance(chat_client_base, SupportsChatGetResponse)
|
||||
|
||||
|
||||
def test_base_client_warns_for_direct_additional_properties(chat_client_base: SupportsChatGetResponse) -> None:
|
||||
with pytest.warns(DeprecationWarning, match="additional_properties"):
|
||||
client = type(chat_client_base)(legacy_key="legacy-value")
|
||||
|
||||
assert client.additional_properties["legacy_key"] == "legacy-value"
|
||||
|
||||
|
||||
def test_base_client_as_agent_uses_explicit_additional_properties(chat_client_base: SupportsChatGetResponse) -> None:
|
||||
agent = chat_client_base.as_agent(additional_properties={"team": "core"})
|
||||
|
||||
assert agent.additional_properties == {"team": "core"}
|
||||
|
||||
|
||||
def test_openai_chat_client_get_response_docstring_surfaces_layered_runtime_docs() -> None:
|
||||
from agent_framework.openai import OpenAIChatClient
|
||||
|
||||
docstring = inspect.getdoc(OpenAIChatClient.get_response)
|
||||
|
||||
assert docstring is not None
|
||||
assert "Get a response from a chat client." in docstring
|
||||
assert "function_invocation_kwargs" in docstring
|
||||
assert "function_middleware: Optional per-call function middleware." in docstring
|
||||
assert "middleware: Optional per-call chat and function middleware." in docstring
|
||||
|
||||
|
||||
def test_openai_chat_client_get_response_is_defined_on_openai_class() -> None:
|
||||
from agent_framework.openai import OpenAIChatClient
|
||||
|
||||
signature = inspect.signature(OpenAIChatClient.get_response)
|
||||
|
||||
assert OpenAIChatClient.get_response.__qualname__ == "OpenAIChatClient.get_response"
|
||||
assert "function_middleware" in signature.parameters
|
||||
assert "middleware" in signature.parameters
|
||||
|
||||
|
||||
async def test_base_client_get_response_uses_explicit_client_kwargs(chat_client_base: SupportsChatGetResponse) -> None:
|
||||
async def fake_inner_get_response(**kwargs):
|
||||
assert kwargs["trace_id"] == "trace-123"
|
||||
assert "function_invocation_kwargs" not in kwargs
|
||||
return ChatResponse(messages=[Message(role="assistant", text="ok")])
|
||||
|
||||
with patch.object(
|
||||
chat_client_base,
|
||||
"_inner_get_response",
|
||||
side_effect=fake_inner_get_response,
|
||||
) as mock_inner_get_response:
|
||||
await chat_client_base.get_response(
|
||||
[Message(role="user", text="hello")],
|
||||
function_invocation_kwargs={"tool_request_id": "tool-123"},
|
||||
client_kwargs={"trace_id": "trace-123"},
|
||||
)
|
||||
mock_inner_get_response.assert_called_once()
|
||||
|
||||
|
||||
async def test_base_client_get_response(chat_client_base: SupportsChatGetResponse):
|
||||
response = await chat_client_base.get_response([Message(role="user", text="Hello")])
|
||||
assert response.messages[0].role == "assistant"
|
||||
|
||||
@@ -0,0 +1,175 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from agent_framework._docstrings import apply_layered_docstring, build_layered_docstring
|
||||
|
||||
# -- Helpers: stub functions with various docstring shapes --
|
||||
|
||||
|
||||
def _source_with_full_docstring(x: int) -> int:
|
||||
"""Do something useful.
|
||||
|
||||
Args:
|
||||
x: The input value.
|
||||
|
||||
Keyword Args:
|
||||
timeout: Max seconds to wait.
|
||||
|
||||
Returns:
|
||||
The computed result.
|
||||
"""
|
||||
return x
|
||||
|
||||
|
||||
def _source_with_args_only(x: int) -> int:
|
||||
"""Do something useful.
|
||||
|
||||
Args:
|
||||
x: The input value.
|
||||
|
||||
Returns:
|
||||
The computed result.
|
||||
"""
|
||||
return x
|
||||
|
||||
|
||||
def _source_no_sections() -> None:
|
||||
"""A plain summary with no Google-style sections."""
|
||||
|
||||
|
||||
def _source_no_docstring() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def _target_stub() -> None:
|
||||
pass
|
||||
|
||||
|
||||
# -- build_layered_docstring tests --
|
||||
|
||||
|
||||
def test_build_returns_none_when_source_has_no_docstring() -> None:
|
||||
result = build_layered_docstring(_source_no_docstring)
|
||||
assert result is None
|
||||
|
||||
|
||||
def test_build_returns_original_when_no_extra_kwargs() -> None:
|
||||
result = build_layered_docstring(_source_with_full_docstring)
|
||||
assert result is not None
|
||||
assert "Do something useful." in result
|
||||
assert "Keyword Args:" in result
|
||||
|
||||
|
||||
def test_build_returns_original_when_extra_kwargs_empty() -> None:
|
||||
result = build_layered_docstring(_source_with_full_docstring, extra_keyword_args={})
|
||||
assert result is not None
|
||||
assert result == build_layered_docstring(_source_with_full_docstring)
|
||||
|
||||
|
||||
def test_build_appends_to_existing_keyword_args_section() -> None:
|
||||
result = build_layered_docstring(
|
||||
_source_with_full_docstring,
|
||||
extra_keyword_args={"retries": "Number of retries."},
|
||||
)
|
||||
assert result is not None
|
||||
assert "timeout: Max seconds to wait." in result
|
||||
assert "retries: Number of retries." in result
|
||||
# Both should be under Keyword Args
|
||||
lines = result.splitlines()
|
||||
kw_index = next(i for i, line in enumerate(lines) if line == "Keyword Args:")
|
||||
ret_index = next(i for i, line in enumerate(lines) if line == "Returns:")
|
||||
retries_index = next(i for i, line in enumerate(lines) if "retries:" in line)
|
||||
assert kw_index < retries_index < ret_index
|
||||
|
||||
|
||||
def test_build_inserts_keyword_args_after_args_section() -> None:
|
||||
result = build_layered_docstring(
|
||||
_source_with_args_only,
|
||||
extra_keyword_args={"verbose": "Enable verbose output."},
|
||||
)
|
||||
assert result is not None
|
||||
assert "Keyword Args:" in result
|
||||
assert "verbose: Enable verbose output." in result
|
||||
lines = result.splitlines()
|
||||
args_index = next(i for i, line in enumerate(lines) if line == "Args:")
|
||||
kw_index = next(i for i, line in enumerate(lines) if line == "Keyword Args:")
|
||||
ret_index = next(i for i, line in enumerate(lines) if line == "Returns:")
|
||||
assert args_index < kw_index < ret_index
|
||||
|
||||
|
||||
def test_build_inserts_keyword_args_in_docstring_with_no_sections() -> None:
|
||||
result = build_layered_docstring(
|
||||
_source_no_sections,
|
||||
extra_keyword_args={"debug": "Enable debug mode."},
|
||||
)
|
||||
assert result is not None
|
||||
assert "A plain summary" in result
|
||||
assert "Keyword Args:" in result
|
||||
assert "debug: Enable debug mode." in result
|
||||
|
||||
|
||||
def test_build_handles_multiline_descriptions() -> None:
|
||||
result = build_layered_docstring(
|
||||
_source_with_args_only,
|
||||
extra_keyword_args={
|
||||
"config": "The configuration object.\nMust be a valid mapping.\nDefaults to empty.",
|
||||
},
|
||||
)
|
||||
assert result is not None
|
||||
lines = result.splitlines()
|
||||
config_line = next(line for line in lines if "config:" in line)
|
||||
assert "The configuration object." in config_line
|
||||
# Continuation lines should be indented
|
||||
config_idx = lines.index(config_line)
|
||||
assert "Must be a valid mapping." in lines[config_idx + 1]
|
||||
assert "Defaults to empty." in lines[config_idx + 2]
|
||||
|
||||
|
||||
def test_build_preserves_multiple_extra_kwargs_order() -> None:
|
||||
result = build_layered_docstring(
|
||||
_source_with_args_only,
|
||||
extra_keyword_args={
|
||||
"alpha": "First.",
|
||||
"beta": "Second.",
|
||||
"gamma": "Third.",
|
||||
},
|
||||
)
|
||||
assert result is not None
|
||||
lines = result.splitlines()
|
||||
alpha_idx = next(i for i, line in enumerate(lines) if "alpha:" in line)
|
||||
beta_idx = next(i for i, line in enumerate(lines) if "beta:" in line)
|
||||
gamma_idx = next(i for i, line in enumerate(lines) if "gamma:" in line)
|
||||
assert alpha_idx < beta_idx < gamma_idx
|
||||
|
||||
|
||||
# -- apply_layered_docstring tests --
|
||||
|
||||
|
||||
def test_apply_sets_docstring_on_target() -> None:
|
||||
def target() -> None:
|
||||
pass
|
||||
|
||||
apply_layered_docstring(target, _source_with_full_docstring)
|
||||
assert target.__doc__ is not None
|
||||
assert "Do something useful." in target.__doc__
|
||||
|
||||
|
||||
def test_apply_with_extra_kwargs() -> None:
|
||||
def target() -> None:
|
||||
pass
|
||||
|
||||
apply_layered_docstring(
|
||||
target,
|
||||
_source_with_args_only,
|
||||
extra_keyword_args={"flag": "A boolean flag."},
|
||||
)
|
||||
assert target.__doc__ is not None
|
||||
assert "flag: A boolean flag." in target.__doc__
|
||||
assert "Keyword Args:" in target.__doc__
|
||||
|
||||
|
||||
def test_apply_sets_none_when_source_has_no_docstring() -> None:
|
||||
def target() -> None:
|
||||
"""Original."""
|
||||
|
||||
apply_layered_docstring(target, _source_no_docstring)
|
||||
assert target.__doc__ is None
|
||||
@@ -4,6 +4,8 @@ from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
|
||||
import pytest
|
||||
|
||||
from agent_framework import (
|
||||
BaseEmbeddingClient,
|
||||
Embedding,
|
||||
@@ -63,6 +65,11 @@ def test_base_additional_properties_custom() -> None:
|
||||
assert client.additional_properties == {"key": "value"}
|
||||
|
||||
|
||||
def test_base_embedding_client_rejects_unknown_kwargs() -> None:
|
||||
with pytest.raises(TypeError):
|
||||
MockEmbeddingClient(legacy_key="value") # type: ignore[call-arg]
|
||||
|
||||
|
||||
# --- SupportsGetEmbeddings protocol tests ---
|
||||
|
||||
|
||||
|
||||
@@ -3651,3 +3651,131 @@ class TestUpdateConversationId:
|
||||
|
||||
|
||||
# endregion
|
||||
async def test_user_input_request_propagates_through_as_tool(chat_client_base: SupportsChatGetResponse):
|
||||
"""Test that user_input_request content from a sub-agent wrapped as a tool propagates to the parent response."""
|
||||
from agent_framework.exceptions import UserInputRequiredException
|
||||
|
||||
@tool(name="delegate_agent", approval_mode="never_require")
|
||||
def delegate_tool(task: str) -> str:
|
||||
del task
|
||||
raise UserInputRequiredException(
|
||||
contents=[
|
||||
Content.from_oauth_consent_request(
|
||||
consent_link="https://login.microsoftonline.com/consent",
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
chat_client_base.run_responses = [
|
||||
ChatResponse(
|
||||
messages=Message(
|
||||
role="assistant",
|
||||
contents=[
|
||||
Content.from_function_call(call_id="1", name="delegate_agent", arguments='{"task": "do it"}'),
|
||||
],
|
||||
)
|
||||
)
|
||||
]
|
||||
|
||||
response = await chat_client_base.get_response(
|
||||
[Message(role="user", text="delegate this")],
|
||||
options={"tool_choice": "auto", "tools": [delegate_tool]},
|
||||
)
|
||||
|
||||
user_requests = [
|
||||
content
|
||||
for msg in response.messages
|
||||
for content in msg.contents
|
||||
if isinstance(content, Content) and content.user_input_request
|
||||
]
|
||||
assert len(user_requests) == 1
|
||||
assert user_requests[0].type == "oauth_consent_request"
|
||||
assert user_requests[0].consent_link == "https://login.microsoftonline.com/consent"
|
||||
assert user_requests[0].user_input_request is True
|
||||
|
||||
|
||||
async def test_user_input_request_multiple_contents_propagate(chat_client_base: SupportsChatGetResponse):
|
||||
"""Test that multiple user_input_request items in a single exception all propagate to the parent response."""
|
||||
from agent_framework.exceptions import UserInputRequiredException
|
||||
|
||||
@tool(name="multi_request_tool", approval_mode="never_require")
|
||||
def multi_request(task: str) -> str:
|
||||
del task
|
||||
raise UserInputRequiredException(
|
||||
contents=[
|
||||
Content.from_oauth_consent_request(
|
||||
consent_link="https://example.com/consent1",
|
||||
),
|
||||
Content.from_oauth_consent_request(
|
||||
consent_link="https://example.com/consent2",
|
||||
),
|
||||
Content.from_oauth_consent_request(
|
||||
consent_link="https://example.com/consent3",
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
chat_client_base.run_responses = [
|
||||
ChatResponse(
|
||||
messages=Message(
|
||||
role="assistant",
|
||||
contents=[
|
||||
Content.from_function_call(call_id="1", name="multi_request_tool", arguments='{"task": "do it"}'),
|
||||
],
|
||||
)
|
||||
)
|
||||
]
|
||||
|
||||
response = await chat_client_base.get_response(
|
||||
[Message(role="user", text="do something")],
|
||||
options={"tool_choice": "auto", "tools": [multi_request]},
|
||||
)
|
||||
|
||||
user_requests = [
|
||||
content
|
||||
for msg in response.messages
|
||||
for content in msg.contents
|
||||
if isinstance(content, Content) and content.user_input_request
|
||||
]
|
||||
assert len(user_requests) == 3
|
||||
consent_links = {r.consent_link for r in user_requests}
|
||||
assert consent_links == {
|
||||
"https://example.com/consent1",
|
||||
"https://example.com/consent2",
|
||||
"https://example.com/consent3",
|
||||
}
|
||||
|
||||
|
||||
async def test_user_input_request_empty_contents_returns_fallback(chat_client_base: SupportsChatGetResponse):
|
||||
"""Test that UserInputRequiredException with empty contents produces a fallback function_result."""
|
||||
from agent_framework.exceptions import UserInputRequiredException
|
||||
|
||||
@tool(name="empty_request_tool", approval_mode="never_require")
|
||||
def empty_request(task: str) -> str:
|
||||
del task
|
||||
raise UserInputRequiredException(contents=[])
|
||||
|
||||
chat_client_base.run_responses = [
|
||||
ChatResponse(
|
||||
messages=Message(
|
||||
role="assistant",
|
||||
contents=[
|
||||
Content.from_function_call(call_id="1", name="empty_request_tool", arguments='{"task": "do it"}'),
|
||||
],
|
||||
)
|
||||
),
|
||||
ChatResponse(messages=Message(role="assistant", text="handled")),
|
||||
]
|
||||
|
||||
response = await chat_client_base.get_response(
|
||||
[Message(role="user", text="do something")],
|
||||
options={"tool_choice": "auto", "tools": [empty_request]},
|
||||
)
|
||||
|
||||
# With empty contents, the handler returns a function_result with an error message
|
||||
# and the loop continues to the next chat response.
|
||||
function_results = [
|
||||
content for msg in response.messages for content in msg.contents if content.type == "function_result"
|
||||
]
|
||||
assert len(function_results) >= 1
|
||||
assert any("user input" in (fr.result or "").lower() for fr in function_results)
|
||||
|
||||
@@ -6,11 +6,13 @@ from collections.abc import AsyncIterable, Awaitable, MutableSequence, Sequence
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import (
|
||||
Agent,
|
||||
BaseChatClient,
|
||||
ChatMiddlewareLayer,
|
||||
ChatResponse,
|
||||
ChatResponseUpdate,
|
||||
Content,
|
||||
FunctionInvocationContext,
|
||||
FunctionInvocationLayer,
|
||||
Message,
|
||||
ResponseStream,
|
||||
@@ -97,6 +99,7 @@ class TestKwargsPropagationToFunctionTool:
|
||||
|
||||
async def test_kwargs_propagate_to_tool_with_kwargs(self) -> None:
|
||||
"""Test that kwargs passed to get_response() are available in @tool **kwargs."""
|
||||
# TODO(Copilot): Remove this legacy coverage once runtime ``**kwargs`` tool injection is removed.
|
||||
captured_kwargs: dict[str, Any] = {}
|
||||
|
||||
@tool(approval_mode="never_require")
|
||||
@@ -149,6 +152,7 @@ class TestKwargsPropagationToFunctionTool:
|
||||
|
||||
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."""
|
||||
# TODO(Copilot): Remove this legacy coverage once runtime ``**kwargs`` tool injection is removed.
|
||||
|
||||
@tool(approval_mode="never_require")
|
||||
def simple_tool(x: int) -> str:
|
||||
@@ -185,6 +189,7 @@ class TestKwargsPropagationToFunctionTool:
|
||||
|
||||
async def test_kwargs_isolated_between_function_calls(self) -> None:
|
||||
"""Test that kwargs are consistent across multiple function call invocations."""
|
||||
# TODO(Copilot): Remove this legacy coverage once runtime ``**kwargs`` tool injection is removed.
|
||||
invocation_kwargs: list[dict[str, Any]] = []
|
||||
|
||||
@tool(approval_mode="never_require")
|
||||
@@ -235,6 +240,7 @@ class TestKwargsPropagationToFunctionTool:
|
||||
|
||||
async def test_streaming_response_kwargs_propagation(self) -> None:
|
||||
"""Test that kwargs propagate to @tool in streaming mode."""
|
||||
# TODO(Copilot): Remove this legacy coverage once runtime ``**kwargs`` tool injection is removed.
|
||||
captured_kwargs: dict[str, Any] = {}
|
||||
|
||||
@tool(approval_mode="never_require")
|
||||
@@ -287,3 +293,59 @@ class TestKwargsPropagationToFunctionTool:
|
||||
assert "streaming_session" in captured_kwargs, f"Expected 'streaming_session' in {captured_kwargs}"
|
||||
assert captured_kwargs["streaming_session"] == "session-xyz"
|
||||
assert captured_kwargs["correlation_id"] == "corr-123"
|
||||
|
||||
async def test_agent_run_injects_function_invocation_context(self) -> None:
|
||||
"""Test that Agent.run injects FunctionInvocationContext for ctx-based tools."""
|
||||
captured_context_kwargs: dict[str, Any] = {}
|
||||
captured_client_kwargs: dict[str, Any] = {}
|
||||
captured_options: dict[str, Any] = {}
|
||||
|
||||
@tool(approval_mode="never_require")
|
||||
def capture_context_tool(x: int, ctx: FunctionInvocationContext) -> str:
|
||||
captured_context_kwargs.update(ctx.kwargs)
|
||||
return f"result: x={x}"
|
||||
|
||||
class CapturingFunctionInvokingMockClient(FunctionInvokingMockClient):
|
||||
async def _get_non_streaming_response(
|
||||
self,
|
||||
*,
|
||||
messages: MutableSequence[Message],
|
||||
options: dict[str, Any],
|
||||
**kwargs: Any,
|
||||
) -> ChatResponse:
|
||||
captured_options.update(options)
|
||||
captured_client_kwargs.update(kwargs)
|
||||
return await super()._get_non_streaming_response(messages=messages, options=options, **kwargs)
|
||||
|
||||
client = CapturingFunctionInvokingMockClient()
|
||||
client.run_responses = [
|
||||
ChatResponse(
|
||||
messages=[
|
||||
Message(
|
||||
role="assistant",
|
||||
contents=[
|
||||
Content.from_function_call(
|
||||
call_id="call_1",
|
||||
name="capture_context_tool",
|
||||
arguments='{"x": 42}',
|
||||
)
|
||||
],
|
||||
)
|
||||
]
|
||||
),
|
||||
ChatResponse(messages=[Message(role="assistant", text="Done!")]),
|
||||
]
|
||||
|
||||
agent = Agent(client=client, tools=[capture_context_tool])
|
||||
result = await agent.run(
|
||||
[Message(role="user", text="Test")],
|
||||
function_invocation_kwargs={"tool_request_id": "tool-123"},
|
||||
client_kwargs={"client_request_id": "client-456"},
|
||||
)
|
||||
|
||||
assert captured_context_kwargs["tool_request_id"] == "tool-123"
|
||||
assert "client_request_id" not in captured_context_kwargs
|
||||
assert captured_client_kwargs["client_request_id"] == "client-456"
|
||||
assert "tool_request_id" not in captured_client_kwargs
|
||||
assert "additional_function_arguments" not in captured_options
|
||||
assert result.messages[-1].text == "Done!"
|
||||
|
||||
@@ -192,10 +192,10 @@ class ConcreteHistoryProvider(BaseHistoryProvider):
|
||||
self.stored: list[Message] = []
|
||||
self._stored_messages = stored_messages or []
|
||||
|
||||
async def get_messages(self, session_id: str | None, **kwargs) -> list[Message]:
|
||||
async def get_messages(self, session_id: str | None, *, state=None, **kwargs) -> list[Message]:
|
||||
return list(self._stored_messages)
|
||||
|
||||
async def save_messages(self, session_id: str | None, messages: Sequence[Message], **kwargs) -> None:
|
||||
async def save_messages(self, session_id: str | None, messages: Sequence[Message], *, state=None, **kwargs) -> None:
|
||||
self.stored.extend(messages)
|
||||
|
||||
|
||||
|
||||
@@ -12,6 +12,7 @@ from agent_framework import (
|
||||
FunctionTool,
|
||||
tool,
|
||||
)
|
||||
from agent_framework._middleware import FunctionInvocationContext
|
||||
from agent_framework._tools import (
|
||||
_parse_annotation,
|
||||
_parse_inputs,
|
||||
@@ -952,6 +953,128 @@ async def test_ai_function_with_kwargs_injection():
|
||||
assert result_default[0].text == "x=10, user=unknown"
|
||||
|
||||
|
||||
async def test_ai_function_with_explicit_invocation_context():
|
||||
"""Test that invoke() can receive runtime kwargs via FunctionInvocationContext."""
|
||||
|
||||
@tool
|
||||
def tool_with_context(x: int, ctx: FunctionInvocationContext) -> str:
|
||||
"""A tool that accepts runtime context injection."""
|
||||
user_id = ctx.kwargs.get("user_id", "unknown")
|
||||
return f"x={x}, user={user_id}"
|
||||
|
||||
assert tool_with_context.parameters() == {
|
||||
"properties": {"x": {"title": "X", "type": "integer"}},
|
||||
"required": ["x"],
|
||||
"title": "tool_with_context_input",
|
||||
"type": "object",
|
||||
}
|
||||
|
||||
context = FunctionInvocationContext(
|
||||
function=tool_with_context,
|
||||
arguments=tool_with_context.input_model(x=7),
|
||||
kwargs={"user_id": "ctx-user"},
|
||||
)
|
||||
|
||||
result = await tool_with_context.invoke(context=context)
|
||||
|
||||
assert result[0].text == "x=7, user=ctx-user"
|
||||
|
||||
|
||||
async def test_ai_function_with_typed_context_parameter_using_custom_name():
|
||||
"""Test that typed context injection works for names other than ctx."""
|
||||
|
||||
@tool
|
||||
def tool_with_runtime_context(x: int, runtime: FunctionInvocationContext) -> str:
|
||||
"""A tool that uses a custom context parameter name."""
|
||||
user_id = runtime.kwargs.get("user_id", "unknown")
|
||||
return f"x={x}, user={user_id}"
|
||||
|
||||
assert tool_with_runtime_context.parameters() == {
|
||||
"properties": {"x": {"title": "X", "type": "integer"}},
|
||||
"required": ["x"],
|
||||
"title": "tool_with_runtime_context_input",
|
||||
"type": "object",
|
||||
}
|
||||
|
||||
context = FunctionInvocationContext(
|
||||
function=tool_with_runtime_context,
|
||||
arguments=tool_with_runtime_context.input_model(x=8),
|
||||
kwargs={"user_id": "runtime-user"},
|
||||
)
|
||||
|
||||
result = await tool_with_runtime_context.invoke(context=context)
|
||||
|
||||
assert result[0].text == "x=8, user=runtime-user"
|
||||
|
||||
|
||||
async def test_ai_function_with_explicit_schema_and_untyped_ctx():
|
||||
"""Test that explicit schemas allow an untyped ctx parameter."""
|
||||
|
||||
class ToolInput(BaseModel):
|
||||
x: int
|
||||
|
||||
@tool(schema=ToolInput)
|
||||
def tool_with_schema(x, ctx) -> str:
|
||||
"""A tool with explicit schema and implicit ctx injection."""
|
||||
return f"x={x}, user={ctx.kwargs.get('user_id', 'unknown')}"
|
||||
|
||||
context = FunctionInvocationContext(
|
||||
function=tool_with_schema,
|
||||
arguments=ToolInput(x=9),
|
||||
kwargs={"user_id": "schema-user"},
|
||||
)
|
||||
|
||||
result = await tool_with_schema.invoke(context=context)
|
||||
|
||||
assert result[0].text == "x=9, user=schema-user"
|
||||
|
||||
|
||||
async def test_ai_function_with_explicit_schema_and_typed_ctx():
|
||||
"""Test that explicit schemas also work with typed context injection."""
|
||||
|
||||
class ToolInput(BaseModel):
|
||||
x: int
|
||||
|
||||
@tool(schema=ToolInput)
|
||||
def tool_with_schema(x: int, runtime: FunctionInvocationContext) -> str:
|
||||
"""A tool with explicit schema and typed context injection."""
|
||||
return f"x={x}, user={runtime.kwargs.get('user_id', 'unknown')}"
|
||||
|
||||
context = FunctionInvocationContext(
|
||||
function=tool_with_schema,
|
||||
arguments=ToolInput(x=11),
|
||||
kwargs={"user_id": "typed-schema-user"},
|
||||
)
|
||||
|
||||
result = await tool_with_schema.invoke(context=context)
|
||||
|
||||
assert tool_with_schema.parameters() == ToolInput.model_json_schema()
|
||||
assert result[0].text == "x=11, user=typed-schema-user"
|
||||
|
||||
|
||||
def test_ai_function_with_multiple_typed_context_parameters_fails():
|
||||
"""Test that tools reject multiple typed FunctionInvocationContext parameters."""
|
||||
|
||||
with pytest.raises(ValueError, match="multiple FunctionInvocationContext parameters"):
|
||||
|
||||
@tool
|
||||
def invalid_tool(ctx_one: FunctionInvocationContext, ctx_two: FunctionInvocationContext) -> str:
|
||||
return f"{ctx_one.kwargs}-{ctx_two.kwargs}"
|
||||
|
||||
|
||||
def test_ai_function_with_ctx_and_typed_context_parameter_fails():
|
||||
"""Test that explicit-schema tools reject both implicit ctx and typed context parameters."""
|
||||
|
||||
class ToolInput(BaseModel):
|
||||
x: int
|
||||
|
||||
with pytest.raises(ValueError, match="multiple FunctionInvocationContext parameters"):
|
||||
|
||||
@tool(schema=ToolInput)
|
||||
def invalid_tool(x, ctx, runtime: FunctionInvocationContext) -> str:
|
||||
return f"{x}-{ctx.kwargs}-{runtime.kwargs}"
|
||||
|
||||
|
||||
# region _parse_annotation tests
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user