From a5f735575fb162afc8533d430585240e08697619 Mon Sep 17 00:00:00 2001 From: Alex Date: Mon, 1 Jun 2026 11:30:37 +0200 Subject: [PATCH] fix: backup model response parsing --- application/llm/base.py | 10 ++ application/llm/handlers/base.py | 49 +++++- .../test_cross_provider_parse_routing.py | 143 ++++++++++++++++++ tests/llm/test_fallback.py | 57 +++++++ 4 files changed, 255 insertions(+), 4 deletions(-) create mode 100644 tests/llm/handlers/test_cross_provider_parse_routing.py diff --git a/application/llm/base.py b/application/llm/base.py index b84aa91e..f608c0c8 100644 --- a/application/llm/base.py +++ b/application/llm/base.py @@ -38,6 +38,12 @@ class BaseLLM(ABC): # BYOM-resolution scope captured at LLM creation time so backup # / fallback lookups hit the same per-user layer as the primary. self.model_user_id = model_user_id + # Provider whose model actually produced the most recent response. + # Equals ``provider_name`` until a cross-provider fallback swaps the + # responding model mid-call (see ``_stream_with_fallback``); the + # handler layer reads this to parse chunks with the right provider's + # handler instead of the primary's. + self._responding_provider = self.provider_name @property def fallback_llm(self): @@ -172,6 +178,7 @@ class BaseLLM(ABC): decorated_method, method_name, decorators, *args, **kwargs ) + self._responding_provider = self.provider_name try: return decorated_method() except Exception as e: @@ -179,6 +186,7 @@ class BaseLLM(ABC): logger.error(f"Primary LLM failed and no fallback configured: {str(e)}") raise fallback = self.fallback_llm + self._responding_provider = fallback.provider_name logger.warning( f"Primary LLM failed. Falling back to " f"{fallback.model_id}. Error: {str(e)}" @@ -220,6 +228,7 @@ class BaseLLM(ABC): ensures that if the primary LLM fails at any point during streaming (creation or mid-stream), we fall back to the backup model. """ + self._responding_provider = self.provider_name try: yield from decorated_method() except Exception as e: @@ -229,6 +238,7 @@ class BaseLLM(ABC): ) raise fallback = self.fallback_llm + self._responding_provider = fallback.provider_name logger.warning( f"Primary LLM failed mid-stream. Falling back to " f"{fallback.model_id}. Error: {str(e)}" diff --git a/application/llm/handlers/base.py b/application/llm/handlers/base.py index 8466b8dd..79ba8aab 100644 --- a/application/llm/handlers/base.py +++ b/application/llm/handlers/base.py @@ -65,12 +65,53 @@ class LLMHandler(ABC): def __init__(self): self.llm_calls = [] self.tool_calls = [] + # Cache of provider-name -> handler used by ``_parse_for_response`` + # to parse chunks from a model that a cross-provider fallback + # swapped in underneath this handler. + self._parser_by_provider = {} @abstractmethod def parse_response(self, response: Any) -> LLMResponse: """Parse raw LLM response into standardized format.""" pass + def _parse_for_response(self, agent, response: Any) -> "LLMResponse": + """Parse ``response`` with the handler matching the model that + actually produced it. + + ``BaseLLM`` runs model fallback *below* the agent (see + ``BaseLLM._stream_with_fallback``): a Google-primary agent that is + rate-limited can fail over to an OpenAI-compatible backup inside the + same ``gen_stream`` call. This handler was built for the primary + provider, so its ``parse_response`` cannot read the backup's chunk + shape and silently drops tool calls — the agent then stops after the + first text instead of running the tool loop. + + ``parse_response`` is the only provider-specific step that matters + here: both providers' ``_iterate_stream`` and ``create_tool_message`` + are identical, so only this call is routed. The orchestration state + (buffers, ``tool_calls``, ``llm_calls``) stays on ``self``. + """ + provider = getattr(getattr(agent, "llm", None), "_responding_provider", None) + if not isinstance(provider, str): + return self.parse_response(response) + return self._handler_for_provider(provider).parse_response(response) + + def _handler_for_provider(self, provider: str) -> "LLMHandler": + """Resolve (and cache) the handler for ``provider``. Reuses ``self`` + when it already matches that provider, so the common no-fallback path + is unchanged.""" + cached = self._parser_by_provider.get(provider) + if cached is not None: + return cached + from application.llm.handlers.handler_creator import LLMHandlerCreator + + handler = LLMHandlerCreator.create_handler(provider) + if type(handler) is type(self): + handler = self + self._parser_by_provider[provider] = handler + return handler + @abstractmethod def create_tool_message(self, tool_call: ToolCall, result: Any) -> Dict: """Create a tool result message for the conversation history.""" @@ -1044,7 +1085,7 @@ class LLMHandler(ABC): Returns: Final response after processing all tool calls """ - parsed = self.parse_response(response) + parsed = self._parse_for_response(agent, response) self.llm_calls.append(build_stack_data(agent.llm)) iteration = 0 @@ -1094,7 +1135,7 @@ class LLMHandler(ABC): messages=messages, tools=None, ) - parsed = self.parse_response(response) + parsed = self._parse_for_response(agent, response) self.llm_calls.append(build_stack_data(agent.llm)) break @@ -1108,7 +1149,7 @@ class LLMHandler(ABC): messages=messages, tools=agent.tools, ) - parsed = self.parse_response(response) + parsed = self._parse_for_response(agent, response) self.llm_calls.append(build_stack_data(agent.llm)) return parsed.content @@ -1144,7 +1185,7 @@ class LLMHandler(ABC): if isinstance(chunk, str): yield chunk continue - parsed = self.parse_response(chunk) + parsed = self._parse_for_response(agent, chunk) if parsed.reasoning_content: reasoning_buffer += parsed.reasoning_content diff --git a/tests/llm/handlers/test_cross_provider_parse_routing.py b/tests/llm/handlers/test_cross_provider_parse_routing.py new file mode 100644 index 00000000..898daced --- /dev/null +++ b/tests/llm/handlers/test_cross_provider_parse_routing.py @@ -0,0 +1,143 @@ +"""Regression tests for cross-provider fallback tool-call parsing. + +``BaseLLM`` performs model fallback *below* the agent: when a primary +model (e.g. Google) fails, ``BaseLLM._stream_with_fallback`` transparently +fails over to a backup model (e.g. an OpenAI-compatible deployment) inside +the same ``gen_stream`` call. The agent, however, picked its +``LLMHandler`` from the *primary* provider and never switches it. + +Before the fix the Google handler tried to ``parse_response`` the +OpenAI-format chunks the backup emits; an OpenAI streaming ``choice`` has +neither ``candidates`` nor ``function_call``, so the Google parser returned +an empty ``finish_reason="stop"`` response — silently dropping the tool +call. The visible symptom: the agent streams its first text and then stops +instead of running the tool loop. + +These tests pin down that ``parse_response`` follows the model that +actually responded (tracked via ``BaseLLM._responding_provider``). +""" + +from types import SimpleNamespace + +from application.llm.handlers.google import GoogleLLMHandler +from application.llm.handlers.openai import OpenAILLMHandler + + +def _openai_tool_chunk(name, arguments, call_id="call_1", finish_reason="tool_calls"): + """Build an OpenAI-style streaming ``choice`` carrying a tool call. + + Mirrors what ``OpenAILLM._raw_gen_stream`` yields for a tool call: a + ``choice`` object with a ``.delta.tool_calls`` list and a + ``finish_reason``. + """ + fn = SimpleNamespace(name=name, arguments=arguments) + tc = SimpleNamespace(id=call_id, function=fn, index=0, type="function") + delta = SimpleNamespace(content=None, tool_calls=[tc]) + return SimpleNamespace(delta=delta, finish_reason=finish_reason) + + +def _agent_with_responder(provider): + return SimpleNamespace(llm=SimpleNamespace(_responding_provider=provider)) + + +class TestParseRouting: + """``_parse_for_response`` resolves the handler from the responding + provider, not the one the handler was constructed for.""" + + def test_documents_the_bug_google_parser_drops_openai_tool_call(self): + # Baseline: the Google parser cannot read an OpenAI choice. This is + # exactly why a naive (un-routed) parse loses the tool call. + chunk = _openai_tool_chunk("search", '{"q": "x"}') + parsed = GoogleLLMHandler().parse_response(chunk) + assert parsed.tool_calls == [] + assert parsed.finish_reason == "stop" + + def test_google_handler_routes_openai_chunk_to_openai_parser(self): + handler = GoogleLLMHandler() + agent = _agent_with_responder("openai") + chunk = _openai_tool_chunk("search", '{"q": "x"}') + + parsed = handler._parse_for_response(agent, chunk) + + assert parsed.finish_reason == "tool_calls" + assert [tc.name for tc in parsed.tool_calls] == ["search"] + assert parsed.tool_calls[0].arguments == '{"q": "x"}' + + def test_handler_for_provider_reuses_self_when_provider_matches(self): + handler = GoogleLLMHandler() + # Same provider family → reuse self (no behaviour change on the + # common no-fallback path). + assert handler._handler_for_provider("google") is handler + + def test_handler_for_provider_returns_openai_for_openai(self): + handler = GoogleLLMHandler() + resolved = handler._handler_for_provider("openai") + assert isinstance(resolved, OpenAILLMHandler) + assert resolved is not handler + # Cached: same instance on the second lookup. + assert handler._handler_for_provider("openai") is resolved + + def test_non_string_provider_falls_back_to_self_parse(self): + # A Mock/None responder (or any code path that never set the + # attribute) must not trigger routing — parse with self. + handler = OpenAILLMHandler() + agent = SimpleNamespace(llm=SimpleNamespace()) # no _responding_provider + chunk = _openai_tool_chunk("search", "{}") + parsed = handler._parse_for_response(agent, chunk) + assert [tc.name for tc in parsed.tool_calls] == ["search"] + + +class TestStreamingLoopContinuesUnderFallback: + """End-to-end: ``handle_streaming`` must run the tool loop when the + responding model is a different provider than the handler's own.""" + + def _make_agent(self, follow_up_chunks): + agent = SimpleNamespace() + agent.model_id = "primary-model" + agent.tools = [] + agent.context_limit_reached = False + agent._check_context_limit = lambda *a, **k: False + agent._pending_continuation = None + + llm = SimpleNamespace() + llm.model_id = "backup-model" + # Responding model is OpenAI-compatible even though the handler is + # Google — this is the cross-provider fallback condition. + llm._responding_provider = "openai" + llm.gen_stream = lambda **kwargs: iter(follow_up_chunks) + agent.llm = llm + + tool_executor = SimpleNamespace() + tool_executor.check_pause = lambda *a, **k: None + tool_executor._name_to_tool = {} + agent.tool_executor = tool_executor + + executed = [] + + def fake_execute(tools_dict, call): + executed.append(call) + yield {"type": "tool_call", "data": {"status": "pending"}} + return ("tool result", call.id) + + agent._execute_tool_action = fake_execute + agent._executed_calls = executed + return agent + + def test_tool_call_executed_not_dropped(self): + handler = GoogleLLMHandler() + # The backup (OpenAI) streams a bit of text, then a tool call. + stream = [ + "Let me search for relevant provisions", + _openai_tool_chunk("search", '{"q": "hipoteka"}'), + ] + # After the tool runs, the follow-up turn just finishes with text. + agent = self._make_agent(follow_up_chunks=["Here is the answer."]) + + results = list(handler.handle_streaming(agent, iter(stream), {"1": {"name": "t"}}, [])) + + # The tool call survived the cross-provider boundary and executed. + assert len(agent._executed_calls) == 1 + assert agent._executed_calls[0].name == "search" + # The loop continued: first text + follow-up text both streamed. + assert "Let me search for relevant provisions" in results + assert "Here is the answer." in results diff --git a/tests/llm/test_fallback.py b/tests/llm/test_fallback.py index 1dab8af0..f42cb040 100644 --- a/tests/llm/test_fallback.py +++ b/tests/llm/test_fallback.py @@ -754,3 +754,60 @@ class TestLLMCreatorPassesModelUserId: ) assert captured["model_user_id"] == "owner-alice" + + +# Tests — responding-provider tracking (cross-provider fallback handler fix) + + +class _Google(FakeLLM): + provider_name = "google" + + +class _OpenAI(FakeLLM): + provider_name = "openai" + + +@pytest.mark.integration +class TestRespondingProviderTracking: + """The handler that parses a response must follow the model that + actually produced it. ``BaseLLM`` exposes ``_responding_provider`` so + the handler layer can re-route ``parse_response`` after a fallback to a + different-provider model (Google primary -> OpenAI backup), instead of + silently dropping the backup's tool calls.""" + + def test_defaults_to_own_provider_before_any_call(self): + assert _Google()._responding_provider == "google" + + def test_stream_success_keeps_primary_provider(self): + primary = _Google(stream_chunks=["a", "b"]) + list(primary.gen_stream(**CALL_ARGS)) + assert primary._responding_provider == "google" + + def test_stream_fallback_records_backup_provider(self, patch_model_utils): + backup = _OpenAI(stream_chunks=["x"]) + patch_model_utils( + get_provider=lambda m, **_kwargs: "openai", + get_api_key=lambda p: "k", + create_llm=lambda type, **kw: backup, + ) + primary = _Google( + stream_chunks=["a"], fail_at=0, backup_models=["backup-model"] + ) + list(primary.gen_stream(**CALL_ARGS)) + assert primary._responding_provider == "openai" + + def test_gen_success_keeps_primary_provider(self): + primary = _Google(responses=["ok"]) + primary.gen(**CALL_ARGS) + assert primary._responding_provider == "google" + + def test_gen_fallback_records_backup_provider(self, patch_model_utils): + backup = _OpenAI(responses=["backup ok"]) + patch_model_utils( + get_provider=lambda m, **_kwargs: "openai", + get_api_key=lambda p: "k", + create_llm=lambda type, **kw: backup, + ) + primary = _Google(fail_at=0, backup_models=["backup-model"]) + primary.gen(**CALL_ARGS) + assert primary._responding_provider == "openai"