Merge pull request #2517 from arc53/better-retry

Better retry

application/llm/base.py

@@ -136,32 +136,22 @@ class BaseLLM(ABC):
             return args_dict
         return {k: v for k, v in args_dict.items() if v is not None}
 
-    @staticmethod
-    def _is_non_retriable_client_error(exc: BaseException) -> bool:
-        """4xx errors mean the request itself is malformed — retrying with
-        a different model fails identically and doubles the work. Only
-        transient/5xx/connection errors should trigger fallback."""
-        try:
-            from google.genai.errors import ClientError as _GenaiClientError
-
-            if isinstance(exc, _GenaiClientError):
-                return True
-        except ImportError:
-            pass
-        for attr in ("status_code", "code", "http_status"):
-            v = getattr(exc, attr, None)
-            if isinstance(v, int) and 400 <= v < 500:
-                return True
-        resp = getattr(exc, "response", None)
-        v = getattr(resp, "status_code", None)
-        return isinstance(v, int) and 400 <= v < 500
-
     def _execute_with_fallback(
         self, method_name: str, decorators: list, *args, **kwargs
     ):
         """
         Execute method with fallback support.
 
+        Any error raised by the primary model triggers a single attempt on
+        the fallback model, when one is configured. There is no error
+        classification: 5xx/transient failures are obviously recoverable on a
+        different model, and for client-side (4xx) errors — including rate
+        limits (429) and provider-specific payload rejections — the one extra
+        attempt is cheap insurance, since a second provider often accepts what
+        the first refused. ``GeneratorExit``/cancellation are ``BaseException``
+        subclasses and so bypass this handler (no fallback on client
+        disconnect), which is intentional.
+
         Args:
             method_name: Name of the raw method ('_raw_gen' or '_raw_gen_stream')
             decorators: List of decorators to apply
@@ -185,12 +175,6 @@ class BaseLLM(ABC):
         try:
             return decorated_method()
         except Exception as e:
-            if self._is_non_retriable_client_error(e):
-                logger.error(
-                    f"Primary LLM failed with non-retriable client error; "
-                    f"skipping fallback: {str(e)}"
-                )
-                raise
             if not self.fallback_llm:
                 logger.error(f"Primary LLM failed and no fallback configured: {str(e)}")
                 raise
@@ -200,6 +184,18 @@ class BaseLLM(ABC):
                 f"{fallback.model_id}. Error: {str(e)}"
             )
 
+            # Mirror the streaming path: emit the fallback's own start event so
+            # dashboards attribute the response to the backup provider, not the
+            # failed primary.
+            fallback._emit_gen_start_log(
+                fallback.model_id,
+                kwargs.get("messages"),
+                kwargs.get("tools"),
+                bool(
+                    kwargs.get("_usage_attachments")
+                    or kwargs.get("attachments")
+                ),
+            )
             # Apply decorators to fallback's raw method directly — calling
             # fallback.gen() would re-enter the orchestrator and recurse via
             # fallback.fallback_llm.
@@ -210,13 +206,7 @@ class BaseLLM(ABC):
             try:
                 return fallback_method(fallback, *args, **fallback_kwargs)
             except Exception as e2:
-                if self._is_non_retriable_client_error(e2):
-                    logger.error(
-                        f"Fallback LLM failed with non-retriable client "
-                        f"error; giving up: {str(e2)}"
-                    )
-                else:
-                    logger.error(f"Fallback LLM also failed; giving up: {str(e2)}")
+                logger.error(f"Fallback LLM also failed; giving up: {str(e2)}")
                 raise
 
     def _stream_with_fallback(
@@ -233,12 +223,6 @@ class BaseLLM(ABC):
         try:
             yield from decorated_method()
         except Exception as e:
-            if self._is_non_retriable_client_error(e):
-                logger.error(
-                    f"Primary LLM failed mid-stream with non-retriable client "
-                    f"error; skipping fallback: {str(e)}"
-                )
-                raise
             if not self.fallback_llm:
                 logger.error(
                     f"Primary LLM failed and no fallback configured: {str(e)}"
@@ -270,18 +254,18 @@ class BaseLLM(ABC):
             try:
                 yield from fallback_method(fallback, *args, **fallback_kwargs)
             except Exception as e2:
-                if self._is_non_retriable_client_error(e2):
-                    logger.error(
-                        f"Fallback LLM failed mid-stream with non-retriable "
-                        f"client error; giving up: {str(e2)}"
-                    )
-                else:
-                    logger.error(
-                        f"Fallback LLM also failed mid-stream; giving up: {str(e2)}"
-                    )
+                logger.error(
+                    f"Fallback LLM also failed mid-stream; giving up: {str(e2)}"
+                )
                 raise
 
     def gen(self, model, messages, stream=False, tools=None, *args, **kwargs):
+        # Mirror gen_stream: emit the start event before the decorators run so
+        # ``_usage_attachments`` is still in kwargs (the gen decorators pop it).
+        has_attachments = bool(
+            kwargs.get("_usage_attachments") or kwargs.get("attachments")
+        )
+        self._emit_gen_start_log(model, messages, tools, has_attachments)
         decorators = [gen_token_usage, gen_cache]
         return self._execute_with_fallback(
             "_raw_gen",
@@ -294,6 +278,53 @@ class BaseLLM(ABC):
             **kwargs,
         )
 
+    def _emit_gen_start_log(self, model, messages, tools, has_attachments):
+        # Non-streaming counterpart to ``_emit_stream_start_log``. Emitted by
+        # ``gen()`` before the call — and again for the fallback provider in
+        # ``_execute_with_fallback`` — so non-streaming invocations are
+        # observable from the first log line, not just streaming ones. A
+        # distinct event name keeps non-stream calls out of stream dashboards.
+        logging.info(
+            "llm_gen_start",
+            extra={
+                "model": model,
+                "provider": self.provider_name,
+                "message_count": len(messages) if messages is not None else 0,
+                "has_attachments": bool(has_attachments),
+                "has_tools": bool(tools),
+            },
+        )
+
+    def _emit_gen_finished_log(
+        self,
+        model,
+        *,
+        prompt_tokens,
+        completion_tokens,
+        latency_ms,
+        cached_tokens=None,
+        error=None,
+    ):
+        # Non-streaming counterpart to ``_emit_stream_finished_log``. Paired
+        # with ``llm_gen_start`` so cost dashboards can join start/finish for
+        # non-streaming calls just as they do for streams. Token counts are
+        # client-side estimates from ``gen_token_usage``; ``status`` is
+        # ``"error"`` when the call raised. A distinct event name keeps
+        # non-stream calls out of stream dashboards.
+        extra = {
+            "model": model,
+            "provider": self.provider_name,
+            "prompt_tokens": int(prompt_tokens),
+            "completion_tokens": int(completion_tokens),
+            "latency_ms": int(latency_ms),
+            "status": "error" if error is not None else "ok",
+        }
+        if cached_tokens is not None:
+            extra["cached_tokens"] = int(cached_tokens)
+        if error is not None:
+            extra["error_class"] = type(error).__name__
+        logging.info("llm_gen_finished", extra=extra)
+
     def _emit_stream_start_log(self, model, messages, tools, has_attachments):
         # Stamped with ``self.provider_name`` so dashboards can group calls
         # by vendor; the fallback path emits its own copy on the fallback

application/usage.py

@@ -147,6 +147,11 @@ def gen_token_usage(func):
     paths that introspect it (e.g., logging, response payloads). DB
     persistence happens here for every call so primary streams,
     side-channel LLMs, and no-save flows all produce rows uniformly.
+
+    Mirrors ``stream_token_usage``: persistence and the
+    ``llm_gen_finished`` log fire from a ``finally`` block, so a failed
+    call still records the prompt tokens it consumed and emits a
+    ``status="error"`` finish event.
     """
     def wrapper(self, model, messages, stream, tools, **kwargs):
         usage_attachments = kwargs.pop("_usage_attachments", None)
@@ -157,12 +162,31 @@ def gen_token_usage(func):
             usage_attachments=usage_attachments,
             **kwargs,
         )
-        result = func(self, model, messages, stream, tools, **kwargs)
-        call_usage["generated_tokens"] += _count_tokens(result)
-        self.token_usage["prompt_tokens"] += call_usage["prompt_tokens"]
-        self.token_usage["generated_tokens"] += call_usage["generated_tokens"]
-        _persist_call_usage(self, call_usage)
-        return result
+        started_at = time.monotonic()
+        error: BaseException | None = None
+        try:
+            result = func(self, model, messages, stream, tools, **kwargs)
+            call_usage["generated_tokens"] += _count_tokens(result)
+            return result
+        except Exception as exc:
+            error = exc
+            raise
+        finally:
+            self.token_usage["prompt_tokens"] += call_usage["prompt_tokens"]
+            self.token_usage["generated_tokens"] += call_usage["generated_tokens"]
+            _persist_call_usage(self, call_usage)
+            emit = getattr(self, "_emit_gen_finished_log", None)
+            if callable(emit):
+                try:
+                    emit(
+                        model,
+                        prompt_tokens=call_usage["prompt_tokens"],
+                        completion_tokens=call_usage["generated_tokens"],
+                        latency_ms=int((time.monotonic() - started_at) * 1000),
+                        error=error,
+                    )
+                except Exception:
+                    logger.exception("Failed to emit llm_gen_finished")
 
     return wrapper
 

tests/llm/test_base.py

@@ -81,6 +81,160 @@ class TestGenMethods:
         result = llm.gen(model="m", messages=[{"role": "user", "content": "hi"}])
         assert result == "hello"
 
+    @patch("application.llm.base.gen_cache", lambda f: f)
+    @patch("application.llm.base.gen_token_usage", lambda f: f)
+    def test_gen_emits_llm_gen_start_event(self, caplog):
+        # Non-streaming counterpart to the llm_stream_start event: gen() must
+        # log before the model is queried so every call is observable.
+        import logging as _logging
+
+        class FakeProvider(StubLLM):
+            provider_name = "fake-provider"
+
+        llm = FakeProvider(raw_gen_return="hi")
+        with caplog.at_level(_logging.INFO, logger="root"):
+            llm.gen(
+                model="m1",
+                messages=[
+                    {"role": "user", "content": "hi"},
+                    {"role": "assistant", "content": "hey"},
+                ],
+                tools=[{"name": "t"}],
+                _usage_attachments=[{"path": "/tmp/a.png"}],
+            )
+
+        starts = [r for r in caplog.records if r.message == "llm_gen_start"]
+        assert len(starts) == 1
+        evt = starts[0]
+        assert evt.model == "m1"
+        assert evt.provider == "fake-provider"
+        assert evt.message_count == 2
+        assert evt.has_attachments is True
+        assert evt.has_tools is True
+
+    @patch("application.llm.base.gen_cache", lambda f: f)
+    @patch("application.llm.base.gen_token_usage", lambda f: f)
+    def test_gen_emits_event_without_attachments_or_tools(self, caplog):
+        import logging as _logging
+
+        llm = StubLLM(raw_gen_return="hi")
+        with caplog.at_level(_logging.INFO, logger="root"):
+            llm.gen(model="m1", messages=[])
+
+        evt = next(r for r in caplog.records if r.message == "llm_gen_start")
+        assert evt.message_count == 0
+        assert evt.has_attachments is False
+        assert evt.has_tools is False
+        # BaseLLM default — concrete providers always override.
+        assert evt.provider == "unknown"
+
+    @patch("application.llm.base.gen_cache", lambda f: f)
+    @patch("application.llm.base.gen_token_usage", lambda f: f)
+    def test_gen_fallback_emits_gen_start_for_fallback_provider(self, caplog):
+        # The fallback raw path bypasses gen(), so _execute_with_fallback must
+        # emit a second llm_gen_start tagged with the backup vendor/model —
+        # otherwise dashboards record only the failed primary.
+        import logging as _logging
+
+        class PrimaryProvider(FailingLLM):
+            provider_name = "primary-vendor"
+
+        class FallbackProvider(FallbackLLM):
+            provider_name = "fallback-vendor"
+
+        primary = PrimaryProvider()
+        primary._fallback_llm = FallbackProvider(model_id="backup-model-id")
+
+        with caplog.at_level(_logging.INFO, logger="root"):
+            result = primary.gen(
+                model="primary-model",
+                messages=[{"role": "user", "content": "hi"}],
+            )
+
+        assert result == "fallback_result"
+        starts = [r for r in caplog.records if r.message == "llm_gen_start"]
+        assert len(starts) == 2
+        assert starts[0].provider == "primary-vendor"
+        assert starts[0].model == "primary-model"
+        assert starts[1].provider == "fallback-vendor"
+        assert starts[1].model == "backup-model-id"
+
+    @patch("application.llm.base.gen_cache", lambda f: f)
+    def test_gen_emits_llm_gen_finished_on_success(self, caplog):
+        # Real gen_token_usage (only gen_cache patched) so the emit-from-
+        # finally path runs. user_api_key=None makes _persist_call_usage
+        # short-circuit before any DB access — same trick as the stream test.
+        import logging as _logging
+
+        class FakeProvider(StubLLM):
+            provider_name = "fake-provider"
+
+        llm = FakeProvider(raw_gen_return="alpha beta")
+        llm.user_api_key = None
+        with caplog.at_level(_logging.INFO, logger="root"):
+            llm.gen(model="m1", messages=[{"role": "user", "content": "hi"}])
+
+        finished = [r for r in caplog.records if r.message == "llm_gen_finished"]
+        assert len(finished) == 1
+        evt = finished[0]
+        assert evt.model == "m1"
+        assert evt.provider == "fake-provider"
+        assert evt.status == "ok"
+        assert isinstance(evt.prompt_tokens, int) and evt.prompt_tokens >= 0
+        assert isinstance(evt.completion_tokens, int) and evt.completion_tokens > 0
+        assert isinstance(evt.latency_ms, int) and evt.latency_ms >= 0
+        # cached_tokens / error_class are intentionally absent on success.
+        assert not hasattr(evt, "cached_tokens")
+        assert not hasattr(evt, "error_class")
+
+    @patch("application.llm.base.gen_cache", lambda f: f)
+    def test_gen_emits_llm_gen_finished_on_error(self, caplog):
+        import logging as _logging
+
+        class FakeProvider(BaseLLM):
+            provider_name = "fake-provider"
+
+            def _raw_gen(self, baseself, model, messages, stream=False, tools=None, **kw):
+                raise RuntimeError("gen_boom")
+
+            def _raw_gen_stream(self, baseself, model, messages, stream=True, tools=None, **kw):
+                yield "x"
+
+        llm = FakeProvider()
+        llm.user_api_key = None
+        with caplog.at_level(_logging.INFO, logger="root"), pytest.raises(RuntimeError):
+            llm.gen(model="m1", messages=[{"role": "user", "content": "hi"}])
+
+        finished = [r for r in caplog.records if r.message == "llm_gen_finished"]
+        assert len(finished) == 1
+        evt = finished[0]
+        assert evt.status == "error"
+        assert evt.error_class == "RuntimeError"
+        # Prompt tokens are still recorded — the request was sent and billed.
+        assert evt.prompt_tokens > 0
+
+    @patch("application.llm.base.gen_cache", lambda f: f)
+    def test_gen_finished_event_paired_with_gen_start(self, caplog):
+        # The two events form a pair the cost dashboards join on; verify they
+        # come in order and from the same provider/model.
+        import logging as _logging
+
+        class FakeProvider(StubLLM):
+            provider_name = "fake-provider"
+
+        llm = FakeProvider(raw_gen_return="x")
+        llm.user_api_key = None
+        with caplog.at_level(_logging.INFO, logger="root"):
+            llm.gen(model="m1", messages=[])
+
+        records = [
+            r for r in caplog.records
+            if r.message in ("llm_gen_start", "llm_gen_finished")
+        ]
+        assert [r.message for r in records] == ["llm_gen_start", "llm_gen_finished"]
+        assert records[0].model == records[1].model == "m1"
+        assert records[0].provider == records[1].provider == "fake-provider"
+
     @patch("application.llm.base.stream_cache", lambda f: f)
     @patch("application.llm.base.stream_token_usage", lambda f: f)
     def test_gen_stream_yields_results(self):
@@ -367,54 +521,39 @@ class _ClientErrorLLM(BaseLLM):
 
 
 @pytest.mark.unit
-class TestNonRetriableClientError:
+class TestClientErrorFallback:
+    """Any primary error triggers a single fallback attempt.
 
-    def test_helper_detects_4xx_status_code(self):
-        assert BaseLLM._is_non_retriable_client_error(_StatusError(400))
-        assert BaseLLM._is_non_retriable_client_error(_StatusError(404))
-        assert BaseLLM._is_non_retriable_client_error(_StatusError(429))
-
-    def test_helper_passes_5xx_through(self):
-        # 5xx and connection errors should still trigger fallback.
-        assert not BaseLLM._is_non_retriable_client_error(_StatusError(500))
-        assert not BaseLLM._is_non_retriable_client_error(_StatusError(503))
-        assert not BaseLLM._is_non_retriable_client_error(RuntimeError("oops"))
-
-    def test_helper_detects_genai_client_error(self):
-        try:
-            from google.genai.errors import ClientError
-        except ImportError:
-            pytest.skip("google-genai not installed")
-        # ClientError(code, response_json, response=None)
-        exc = ClientError(400, {"error": {"message": "bad", "code": 400}}, None)
-        assert BaseLLM._is_non_retriable_client_error(exc)
-
-    def test_helper_detects_response_status_code(self):
-        exc = RuntimeError("wrapped")
-        exc.response = type("R", (), {"status_code": 401})()
-        assert BaseLLM._is_non_retriable_client_error(exc)
+    There is no error-class gate: 4xx (including 429 rate limits and
+    provider-specific payload rejections) and Gemini ``ClientError`` now
+    fall back just like 5xx/transient failures. Client disconnect
+    (``GeneratorExit``) is a ``BaseException`` and still bypasses fallback.
+    """
 
+    # 4xx that used to be force-skipped; each should now reach the fallback.
+    @pytest.mark.parametrize("status_code", [400, 401, 403, 404, 413, 422, 429])
     @patch("application.llm.base.gen_cache", lambda f: f)
     @patch("application.llm.base.gen_token_usage", lambda f: f)
-    def test_4xx_skips_fallback(self):
+    def test_4xx_now_falls_back(self, status_code):
         fallback = FallbackLLM(model_id="fallback-model")
-        llm = _ClientErrorLLM(status_code=400)
+        llm = _ClientErrorLLM(status_code=status_code)
         llm._fallback_llm = fallback
 
-        with pytest.raises(_StatusError):
-            llm.gen(model="m", messages=[])
-        assert not fallback.gen_called
+        result = llm.gen(model="m", messages=[])
+        assert result == "fallback_result"
+        assert fallback.gen_called
 
+    @pytest.mark.parametrize("status_code", [400, 429])
     @patch("application.llm.base.stream_cache", lambda f: f)
     @patch("application.llm.base.stream_token_usage", lambda f: f)
-    def test_4xx_skips_stream_fallback(self):
+    def test_4xx_now_falls_back_stream(self, status_code):
         fallback = FallbackLLM(model_id="fallback-model")
-        llm = _ClientErrorLLM(status_code=400)
+        llm = _ClientErrorLLM(status_code=status_code)
         llm._fallback_llm = fallback
 
-        with pytest.raises(_StatusError):
-            list(llm.gen_stream(model="m", messages=[]))
-        assert not fallback.gen_stream_called
+        result = list(llm.gen_stream(model="m", messages=[]))
+        assert "fallback_chunk" in result
+        assert fallback.gen_stream_called
 
     @patch("application.llm.base.gen_cache", lambda f: f)
     @patch("application.llm.base.gen_token_usage", lambda f: f)
@@ -427,6 +566,50 @@ class TestNonRetriableClientError:
         assert result == "fallback_result"
         assert fallback.gen_called
 
+    @patch("application.llm.base.gen_cache", lambda f: f)
+    @patch("application.llm.base.gen_token_usage", lambda f: f)
+    def test_genai_client_error_now_falls_back(self):
+        """A Gemini ClientError (4xx) is no longer force-skipped."""
+        try:
+            from google.genai.errors import ClientError
+        except ImportError:
+            pytest.skip("google-genai not installed")
+
+        class _GenaiErrorLLM(BaseLLM):
+            def _raw_gen(self, baseself, model, messages, stream=False, tools=None, **kw):
+                raise ClientError(400, {"error": {"message": "bad", "code": 400}}, None)
+
+            def _raw_gen_stream(self, baseself, model, messages, stream=True, tools=None, **kw):
+                raise ClientError(400, {"error": {"message": "bad", "code": 400}}, None)
+                yield  # unreachable; makes this a generator
+
+        fallback = FallbackLLM(model_id="fallback-model")
+        llm = _GenaiErrorLLM()
+        llm._fallback_llm = fallback
+
+        result = llm.gen(model="m", messages=[])
+        assert result == "fallback_result"
+        assert fallback.gen_called
+
+    @patch("application.llm.base.gen_cache", lambda f: f)
+    @patch("application.llm.base.gen_token_usage", lambda f: f)
+    def test_fallback_failure_propagates(self):
+        """When the fallback also fails, its error propagates to the caller."""
+
+        class _FailingFallback(BaseLLM):
+            def _raw_gen(self, baseself, model, messages, stream=False, tools=None, **kw):
+                raise _StatusError(400, "fallback_failed")
+
+            def _raw_gen_stream(self, baseself, model, messages, stream=True, tools=None, **kw):
+                raise _StatusError(400, "fallback_failed")
+                yield  # unreachable; makes this a generator
+
+        llm = _ClientErrorLLM(status_code=429)
+        llm._fallback_llm = _FailingFallback(model_id="fb")
+
+        with pytest.raises(_StatusError, match="fallback_failed"):
+            llm.gen(model="m", messages=[])
+
 
 # ---------------------------------------------------------------------------
 # fallback_llm property: backup model resolution