mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-15 17:30:13 +08:00
perf: update FileReadTool description to mention image, PDF and docx support, and enhance modality checking in tool result case (#7506)
* feat: update FileReadTool description to mention image and PDF support Add explicit mention of image (OCR) and PDF (text extraction) support to the FileReadTool description for better discoverability. * feat: update FileReadTool description to include support for docx and epub files; change base64 decoding to utf-8 * feat: enhance ToolLoopAgentRunner to support image and audio modalities; add context sanitization logic
This commit is contained in:
Soulter
@@ -14,6 +14,7 @@ sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")
|
||||
from astrbot.core.agent.agent import Agent
|
||||
from astrbot.core.agent.handoff import HandoffTool
|
||||
from astrbot.core.agent.hooks import BaseAgentRunHooks
|
||||
from astrbot.core.agent.message import ImageURLPart, Message, TextPart
|
||||
from astrbot.core.agent.run_context import ContextWrapper
|
||||
from astrbot.core.agent.runners.tool_loop_agent_runner import ToolLoopAgentRunner
|
||||
from astrbot.core.agent.tool import FunctionTool, ToolSet
|
||||
@@ -156,6 +157,25 @@ class MockErrProvider(MockProvider):
|
||||
)
|
||||
|
||||
|
||||
class CapturingProvider(MockProvider):
|
||||
def __init__(self, modalities: list[str]):
|
||||
super().__init__()
|
||||
self.provider_config["modalities"] = modalities
|
||||
self.received_contexts = []
|
||||
self.received_func_tools = []
|
||||
self.should_call_tools = False
|
||||
|
||||
async def text_chat(self, **kwargs) -> LLMResponse:
|
||||
self.call_count += 1
|
||||
self.received_contexts.append(kwargs.get("contexts"))
|
||||
self.received_func_tools.append(kwargs.get("func_tool"))
|
||||
return LLMResponse(
|
||||
role="assistant",
|
||||
completion_text="final",
|
||||
usage=TokenUsage(input_other=10, output=5),
|
||||
)
|
||||
|
||||
|
||||
class MockEmptyOutputThenSuccessProvider(MockProvider):
|
||||
def __init__(self, failures_before_success: int = 1):
|
||||
super().__init__()
|
||||
@@ -615,6 +635,99 @@ async def test_tool_result_includes_all_calltoolresult_content(
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runner_replaces_runtime_image_context_before_provider_call(
|
||||
runner, provider_request, mock_hooks
|
||||
):
|
||||
provider = CapturingProvider(modalities=["tool_use"])
|
||||
|
||||
await runner.reset(
|
||||
provider=provider,
|
||||
request=provider_request,
|
||||
run_context=ContextWrapper(context=None),
|
||||
tool_executor=MockToolExecutor,
|
||||
agent_hooks=mock_hooks,
|
||||
streaming=False,
|
||||
)
|
||||
|
||||
runner.run_context.messages.append(
|
||||
Message(
|
||||
role="user",
|
||||
content=[
|
||||
TextPart(text="Review this image"),
|
||||
ImageURLPart(
|
||||
image_url=ImageURLPart.ImageURL(
|
||||
url="data:image/png;base64,dGVzdA=="
|
||||
)
|
||||
),
|
||||
],
|
||||
)
|
||||
)
|
||||
|
||||
async for _ in runner.step_until_done(1):
|
||||
pass
|
||||
|
||||
assert provider.received_contexts
|
||||
sent_context = provider.received_contexts[0]
|
||||
assert sent_context[-1]["content"] == [
|
||||
{"type": "text", "text": "Review this image"},
|
||||
{"type": "text", "text": "[Image]"},
|
||||
]
|
||||
assert len(runner.run_context.messages[-2].content) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runner_builds_placeholder_for_unsupported_request_image(
|
||||
runner, mock_hooks, tool_set
|
||||
):
|
||||
provider = CapturingProvider(modalities=["tool_use"])
|
||||
request = ProviderRequest(
|
||||
prompt="Describe it",
|
||||
image_urls=["/path/that/should/not/be/read.jpg"],
|
||||
func_tool=tool_set,
|
||||
contexts=[],
|
||||
)
|
||||
|
||||
await runner.reset(
|
||||
provider=provider,
|
||||
request=request,
|
||||
run_context=ContextWrapper(context=None),
|
||||
tool_executor=MockToolExecutor,
|
||||
agent_hooks=mock_hooks,
|
||||
streaming=False,
|
||||
)
|
||||
|
||||
async for _ in runner.step_until_done(1):
|
||||
pass
|
||||
|
||||
sent_context = provider.received_contexts[0]
|
||||
assert sent_context[-1]["content"] == [
|
||||
{"type": "text", "text": "Describe it"},
|
||||
{"type": "text", "text": "[Image]"},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runner_clears_tools_for_provider_without_tool_use(
|
||||
runner, provider_request, mock_hooks, mock_tool_executor
|
||||
):
|
||||
provider = CapturingProvider(modalities=["text"])
|
||||
|
||||
await runner.reset(
|
||||
provider=provider,
|
||||
request=provider_request,
|
||||
run_context=ContextWrapper(context=None),
|
||||
tool_executor=mock_tool_executor,
|
||||
agent_hooks=mock_hooks,
|
||||
streaming=False,
|
||||
)
|
||||
|
||||
async for _ in runner.step_until_done(1):
|
||||
pass
|
||||
|
||||
assert provider.received_func_tools == [None]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_same_tool_consecutive_results_include_escalating_guidance(
|
||||
runner, mock_tool_executor, mock_hooks
|
||||
|
||||
@@ -713,144 +713,6 @@ class TestDecorateLlmRequest:
|
||||
assert req.prompt == "Hello"
|
||||
|
||||
|
||||
class TestModalitiesFix:
|
||||
"""Tests for _modalities_fix function."""
|
||||
|
||||
def test_modalities_fix_image_not_supported(self, mock_provider):
|
||||
"""Test modality fix when image is not supported."""
|
||||
module = ama
|
||||
mock_provider.provider_config = {"modalities": ["text"]}
|
||||
req = ProviderRequest(prompt="Hello", image_urls=["/path/to/image.jpg"])
|
||||
|
||||
module._modalities_fix(mock_provider, req)
|
||||
|
||||
assert "[Image]" in req.prompt
|
||||
assert req.image_urls == []
|
||||
|
||||
def test_modalities_fix_tool_not_supported(self, mock_provider):
|
||||
"""Test modality fix when tool is not supported."""
|
||||
module = ama
|
||||
mock_provider.provider_config = {"modalities": ["text", "image"]}
|
||||
req = ProviderRequest(prompt="Hello")
|
||||
req.func_tool = ToolSet()
|
||||
req.func_tool.add_tool(
|
||||
FunctionTool(
|
||||
name="dummy_tool",
|
||||
description="dummy",
|
||||
parameters={"type": "object", "properties": {}},
|
||||
)
|
||||
)
|
||||
|
||||
module._modalities_fix(mock_provider, req)
|
||||
|
||||
assert req.func_tool is None
|
||||
|
||||
def test_modalities_fix_all_supported(self, mock_provider):
|
||||
"""Test modality fix when all features are supported."""
|
||||
module = ama
|
||||
mock_provider.provider_config = {"modalities": ["image", "tool_use"]}
|
||||
tool_set = ToolSet()
|
||||
tool_set.add_tool(
|
||||
FunctionTool(
|
||||
name="dummy_tool",
|
||||
description="dummy",
|
||||
parameters={"type": "object", "properties": {}},
|
||||
)
|
||||
)
|
||||
req = ProviderRequest(
|
||||
prompt="Hello",
|
||||
image_urls=["/path/to/image.jpg"],
|
||||
func_tool=tool_set,
|
||||
)
|
||||
|
||||
module._modalities_fix(mock_provider, req)
|
||||
|
||||
assert req.prompt == "Hello"
|
||||
assert len(req.image_urls) == 1
|
||||
assert req.func_tool is not None
|
||||
|
||||
|
||||
class TestSanitizeContextByModalities:
|
||||
"""Tests for _sanitize_context_by_modalities function."""
|
||||
|
||||
def test_sanitize_no_op(self, mock_provider):
|
||||
"""Test sanitize when disabled or modalities support everything."""
|
||||
module = ama
|
||||
config = module.MainAgentBuildConfig(
|
||||
tool_call_timeout=60, sanitize_context_by_modalities=False
|
||||
)
|
||||
mock_provider.provider_config = {"modalities": ["image", "tool_use"]}
|
||||
req = ProviderRequest(contexts=[{"role": "user", "content": "Hello"}])
|
||||
|
||||
module._sanitize_context_by_modalities(config, mock_provider, req)
|
||||
|
||||
assert len(req.contexts) == 1
|
||||
|
||||
def test_sanitize_removes_tool_messages(self, mock_provider):
|
||||
"""Test sanitize removes tool messages when tool_use not supported."""
|
||||
module = ama
|
||||
config = module.MainAgentBuildConfig(
|
||||
tool_call_timeout=60, sanitize_context_by_modalities=True
|
||||
)
|
||||
mock_provider.provider_config = {"modalities": ["image"]}
|
||||
req = ProviderRequest(
|
||||
contexts=[
|
||||
{"role": "user", "content": "Hello"},
|
||||
{"role": "tool", "content": "Tool result"},
|
||||
]
|
||||
)
|
||||
|
||||
module._sanitize_context_by_modalities(config, mock_provider, req)
|
||||
|
||||
assert len(req.contexts) == 1
|
||||
assert req.contexts[0]["role"] == "user"
|
||||
|
||||
def test_sanitize_removes_tool_calls(self, mock_provider):
|
||||
"""Test sanitize removes tool_calls from assistant messages."""
|
||||
module = ama
|
||||
config = module.MainAgentBuildConfig(
|
||||
tool_call_timeout=60, sanitize_context_by_modalities=True
|
||||
)
|
||||
mock_provider.provider_config = {"modalities": ["image"]}
|
||||
req = ProviderRequest(
|
||||
contexts=[
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Response",
|
||||
"tool_calls": [{"name": "tool"}],
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
module._sanitize_context_by_modalities(config, mock_provider, req)
|
||||
|
||||
assert "tool_calls" not in req.contexts[0]
|
||||
|
||||
def test_sanitize_removes_image_blocks(self, mock_provider):
|
||||
"""Test sanitize removes image blocks when image not supported."""
|
||||
module = ama
|
||||
config = module.MainAgentBuildConfig(
|
||||
tool_call_timeout=60, sanitize_context_by_modalities=True
|
||||
)
|
||||
mock_provider.provider_config = {"modalities": ["tool_use"]}
|
||||
req = ProviderRequest(
|
||||
contexts=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Hello"},
|
||||
{"type": "image_url", "url": "image.jpg"},
|
||||
],
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
module._sanitize_context_by_modalities(config, mock_provider, req)
|
||||
|
||||
assert len(req.contexts[0]["content"]) == 1
|
||||
assert req.contexts[0]["content"][0]["type"] == "text"
|
||||
|
||||
|
||||
class TestPluginToolFix:
|
||||
"""Tests for _plugin_tool_fix function."""
|
||||
|
||||
|
||||
Reference in New Issue
Block a user