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