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):