Merge pull request #2518 from arc53/backup-fix

fix: backup model response parsing

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)}"

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
 

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

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"