Use custom format_output_for_agent override for tool output

When a `BaseTool` subclass overrides `format_output_for_agent`, route
the agent-facing text through it instead of the default JSON/`str()`
serialization. The structured tool wrapper now delegates to the original
tool via `_original_tool`, so a tool can present Markdown or any custom
representation to the agent while `tool.run(...)` still returns the raw
Python value.
This commit is contained in:
Vinicius Brasil
2026-06-18 21:41:34 -07:00
parent 8c35dedfb5
commit 9b8ecc7df5
6 changed files with 155 additions and 0 deletions

View File

@@ -52,6 +52,10 @@ def _infer_output_schema_from_callable(
def _format_tool_output_for_agent(tool: Any, raw_result: Any) -> str:
original_tool = getattr(tool, "_original_tool", None)
if original_tool is not None:
return cast(str, original_tool.format_output_for_agent(raw_result))
output_schema = getattr(tool, "output_schema", None)
if output_schema is None:
return str(raw_result)

View File

@@ -584,6 +584,30 @@ class TestToolOutputSchema:
"score": 0.8,
}
def test_custom_agent_output_formatter_carries_over_to_structured_tool(
self,
) -> None:
class MarkdownSearchTool(BaseTool):
name: str = "markdown_search"
description: str = "Search for information"
output_schema: type[BaseModel] = SearchOutput
def _run(self, query: str) -> SearchOutput:
return SearchOutput(query=query, score=0.8)
def format_output_for_agent(self, raw_result: object) -> str:
result = self.output_schema.model_validate(raw_result)
return f"### Search result\n\n- Query: `{result.query}`\n- Score: {result.score}"
structured = MarkdownSearchTool().to_structured_tool()
raw_result = structured.invoke({"query": "crew"})
assert raw_result == SearchOutput(query="crew", score=0.8)
assert structured.format_output_for_agent(raw_result) == (
"### Search result\n\n- Query: `crew`\n- Score: 0.8"
)
# Async arun() Schema Validation Tests

View File

@@ -1079,6 +1079,51 @@ class TestExecuteSingleNativeToolCall:
"score": 0.9,
}
def test_custom_agent_output_formatter_is_used_from_structured_tool(
self,
) -> None:
clear_before_tool_call_hooks()
clear_after_tool_call_hooks()
class SearchOutput(BaseModel):
query: str
score: float
class MarkdownSearchTool(BaseTool):
name: str = "markdown_search"
description: str = "Search for a query"
output_schema: type[BaseModel] = SearchOutput
def _run(self, query: str) -> SearchOutput:
return SearchOutput(query=query, score=0.9)
def format_output_for_agent(self, raw_result: Any) -> str:
result = self.output_schema.model_validate(raw_result)
return f"### {result.query}\n\nScore: **{result.score}**"
tool = MarkdownSearchTool()
tool_call = MagicMock()
tool_call.id = "call_1"
tool_call.function.name = "markdown_search"
tool_call.function.arguments = '{"query": "crew"}'
result = execute_single_native_tool_call(
tool_call,
available_functions={"markdown_search": tool._run},
original_tools=[],
structured_tools=[tool.to_structured_tool()],
tools_handler=None,
agent=None,
task=None,
crew=None,
event_source=MagicMock(),
printer=None,
verbose=False,
)
assert result.result == "### crew\n\nScore: **0.9**"
assert result.tool_message["content"] == "### crew\n\nScore: **0.9**"
def test_after_hook_includes_raw_tool_result_for_typed_output(self) -> None:
clear_after_tool_call_hooks()