From 26e867cc6d60ca02808d3803d2b77fd08fc47eb2 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Wed, 27 May 2026 22:21:55 +0800 Subject: [PATCH] fix(core): route image requests to vision fallback (#8089) --- astrbot/core/astr_main_agent.py | 46 +++++++++- tests/unit/test_astr_main_agent.py | 129 +++++++++++++++++++++++++++++ 2 files changed, 172 insertions(+), 3 deletions(-) diff --git a/astrbot/core/astr_main_agent.py b/astrbot/core/astr_main_agent.py index e522ce545..d848aba40 100644 --- a/astrbot/core/astr_main_agent.py +++ b/astrbot/core/astr_main_agent.py @@ -1192,6 +1192,38 @@ def _get_fallback_chat_providers( return fallbacks +def _provider_supports_modality(provider: Provider, modality: str) -> bool: + modalities = provider.provider_config.get("modalities", None) + return isinstance(modalities, list) and modality in modalities + + +def _select_image_chat_provider( + provider: Provider, + req: ProviderRequest, + fallback_providers: list[Provider], +) -> Provider: + if not req.image_urls or _provider_supports_modality(provider, "image"): + return provider + + provider_id = provider.provider_config.get("id", "") + for fallback_provider in fallback_providers: + if not _provider_supports_modality(fallback_provider, "image"): + continue + fallback_id = fallback_provider.provider_config.get("id", "") + logger.warning( + "Chat provider %s does not support image input, switching this request to fallback provider %s.", + provider_id, + fallback_id, + ) + return fallback_provider + + logger.warning( + "Chat provider %s does not support image input and no image-capable fallback provider is available.", + provider_id, + ) + return provider + + async def build_main_agent( *, event: AstrMessageEvent, @@ -1422,6 +1454,16 @@ async def build_main_agent( ) ) + fallback_providers = _get_fallback_chat_providers( + provider, plugin_context, config.provider_settings + ) + selected_provider = _select_image_chat_provider(provider, req, fallback_providers) + if selected_provider is not provider: + provider = selected_provider + if req.model: + req.model = None + fallback_providers = [p for p in fallback_providers if p is not provider] + if provider.provider_config.get("max_context_tokens", 0) <= 0: model = provider.get_model() if model_info := LLM_METADATAS.get(model): @@ -1474,9 +1516,7 @@ async def build_main_agent( truncate_turns=config.dequeue_context_length, enforce_max_turns=config.max_context_length, tool_schema_mode=config.tool_schema_mode, - fallback_providers=_get_fallback_chat_providers( - provider, plugin_context, config.provider_settings - ), + fallback_providers=fallback_providers, tool_result_overflow_dir=( get_astrbot_system_tmp_path() if req.func_tool and req.func_tool.get_tool("astrbot_file_read_tool") diff --git a/tests/unit/test_astr_main_agent.py b/tests/unit/test_astr_main_agent.py index fdd06d34a..da1c9d966 100644 --- a/tests/unit/test_astr_main_agent.py +++ b/tests/unit/test_astr_main_agent.py @@ -89,6 +89,22 @@ def mock_conversation(): return conv +def test_provider_supports_modality_requires_explicit_list(): + provider = MagicMock(spec=Provider) + + provider.provider_config = {"modalities": ["text", "image"]} + assert ama._provider_supports_modality(provider, "image") + + provider.provider_config = {"modalities": ["text"]} + assert not ama._provider_supports_modality(provider, "image") + + provider.provider_config = {} + assert not ama._provider_supports_modality(provider, "image") + + provider.provider_config = {"modalities": "image"} + assert not ama._provider_supports_modality(provider, "image") + + @pytest.fixture def sample_config(): """Create a sample MainAgentBuildConfig.""" @@ -1069,6 +1085,119 @@ class TestBuildMainAgent: assert result is not None + @pytest.mark.asyncio + async def test_build_main_agent_uses_image_fallback_provider( + self, mock_event, mock_context + ): + """Test image requests use a fallback provider that supports images.""" + module = ama + text_provider = MagicMock(spec=Provider) + text_provider.provider_config = { + "id": "text-provider", + "modalities": ["text", "tool_use"], + "max_context_tokens": 128000, + } + text_provider.get_model.return_value = "text-model" + + image_provider = MagicMock(spec=Provider) + image_provider.provider_config = { + "id": "image-provider", + "modalities": ["text", "image", "tool_use"], + "max_context_tokens": 128000, + } + image_provider.get_model.return_value = "image-model" + + req = ProviderRequest( + prompt="describe this", + image_urls=["/tmp/image.jpg"], + model="text-model", + ) + mock_context.get_provider_by_id.side_effect = lambda provider_id: ( + image_provider if provider_id == "image-provider" else None + ) + mock_context.get_config.return_value = {} + + with ( + patch("astrbot.core.astr_main_agent.AgentRunner") as mock_runner_cls, + patch("astrbot.core.astr_main_agent.AstrAgentContext"), + ): + mock_runner = MagicMock() + mock_runner.reset = AsyncMock() + mock_runner_cls.return_value = mock_runner + + result = await module.build_main_agent( + event=mock_event, + plugin_context=mock_context, + config=module.MainAgentBuildConfig( + tool_call_timeout=60, + llm_safety_mode=False, + computer_use_runtime="none", + add_cron_tools=False, + provider_settings={ + "fallback_chat_models": ["image-provider"], + }, + ), + provider=text_provider, + req=req, + ) + + assert result is not None + assert result.provider is image_provider + assert result.provider_request.image_urls == ["/tmp/image.jpg"] + assert result.provider_request.model is None + assert mock_runner.reset.call_args.kwargs["provider"] is image_provider + assert mock_runner.reset.call_args.kwargs["fallback_providers"] == [] + + @pytest.mark.asyncio + async def test_build_main_agent_keeps_text_provider_without_image_fallback( + self, mock_event, mock_context + ): + """Test image requests fall back to existing sanitizing when no image provider exists.""" + module = ama + text_provider = MagicMock(spec=Provider) + text_provider.provider_config = { + "id": "text-provider", + "modalities": ["text", "tool_use"], + "max_context_tokens": 128000, + } + text_provider.get_model.return_value = "text-model" + + req = ProviderRequest( + prompt="describe this", + image_urls=["/tmp/image.jpg"], + ) + mock_context.get_provider_by_id.return_value = None + mock_context.get_config.return_value = {} + + with ( + patch("astrbot.core.astr_main_agent.AgentRunner") as mock_runner_cls, + patch("astrbot.core.astr_main_agent.AstrAgentContext"), + ): + mock_runner = MagicMock() + mock_runner.reset = AsyncMock() + mock_runner_cls.return_value = mock_runner + + result = await module.build_main_agent( + event=mock_event, + plugin_context=mock_context, + config=module.MainAgentBuildConfig( + tool_call_timeout=60, + llm_safety_mode=False, + computer_use_runtime="none", + add_cron_tools=False, + provider_settings={ + "fallback_chat_models": ["missing-provider"], + }, + ), + provider=text_provider, + req=req, + ) + + assert result is not None + assert result.provider is text_provider + assert result.provider_request.image_urls == ["/tmp/image.jpg"] + assert mock_runner.reset.call_args.kwargs["provider"] is text_provider + @pytest.mark.asyncio async def test_build_main_agent_with_video_attachment( self, mock_event, mock_context, mock_provider