From 7bb54e0f107864d8979f56550f0cebb9f7212b84 Mon Sep 17 00:00:00 2001 From: Alex Date: Sun, 31 May 2026 16:43:40 +0200 Subject: [PATCH 1/3] Fix: change model retry strategy --- application/llm/base.py | 62 ++++++-------------------- tests/llm/test_base.py | 99 ++++++++++++++++++++++++++--------------- 2 files changed, 78 insertions(+), 83 deletions(-) diff --git a/application/llm/base.py b/application/llm/base.py index f6448d50..98fb98a5 100644 --- a/application/llm/base.py +++ b/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 @@ -210,13 +194,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 +211,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,15 +242,9 @@ 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): diff --git a/tests/llm/test_base.py b/tests/llm/test_base.py index 6fbe0f96..d19a967d 100644 --- a/tests/llm/test_base.py +++ b/tests/llm/test_base.py @@ -367,54 +367,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 +412,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 From 85a2811e442038f80148a3354171260880350186 Mon Sep 17 00:00:00 2001 From: Alex Date: Sun, 31 May 2026 17:42:42 +0200 Subject: [PATCH 2/3] feat: log gen starts --- application/llm/base.py | 35 ++++++++++++++++++ tests/llm/test_base.py | 78 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 113 insertions(+) diff --git a/application/llm/base.py b/application/llm/base.py index 98fb98a5..c2f51011 100644 --- a/application/llm/base.py +++ b/application/llm/base.py @@ -184,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. @@ -248,6 +260,12 @@ class BaseLLM(ABC): 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", @@ -260,6 +278,23 @@ 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_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 diff --git a/tests/llm/test_base.py b/tests/llm/test_base.py index d19a967d..3e5c1e11 100644 --- a/tests/llm/test_base.py +++ b/tests/llm/test_base.py @@ -81,6 +81,84 @@ 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.stream_cache", lambda f: f) @patch("application.llm.base.stream_token_usage", lambda f: f) def test_gen_stream_yields_results(self): From badd3830b79c87145820a1ff49723a8c6d8f56e7 Mon Sep 17 00:00:00 2001 From: Alex Date: Sun, 31 May 2026 18:00:51 +0200 Subject: [PATCH 3/3] feat: log gen end --- application/llm/base.py | 30 ++++++++++++++++ application/usage.py | 36 +++++++++++++++---- tests/llm/test_base.py | 76 +++++++++++++++++++++++++++++++++++++++++ 3 files changed, 136 insertions(+), 6 deletions(-) diff --git a/application/llm/base.py b/application/llm/base.py index c2f51011..b84aa91e 100644 --- a/application/llm/base.py +++ b/application/llm/base.py @@ -295,6 +295,36 @@ class BaseLLM(ABC): }, ) + 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 diff --git a/application/usage.py b/application/usage.py index 61518c8e..9ae875fa 100644 --- a/application/usage.py +++ b/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 diff --git a/tests/llm/test_base.py b/tests/llm/test_base.py index 3e5c1e11..330a6fca 100644 --- a/tests/llm/test_base.py +++ b/tests/llm/test_base.py @@ -159,6 +159,82 @@ class TestGenMethods: 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):