diff --git a/nanobot/providers/base.py b/nanobot/providers/base.py index a794c1adb..23a43b32d 100644 --- a/nanobot/providers/base.py +++ b/nanobot/providers/base.py @@ -1259,6 +1259,53 @@ class LLMProvider(ABC): """Call chat_stream() and convert unexpected exceptions to error responses.""" started_at_ms = time.time_ns() // 1_000_000 started_at_ns = time.monotonic_ns() + first_output_at_ns: int | None = None + + def _mark_output(delta: str) -> None: + nonlocal first_output_at_ns + if delta and first_output_at_ns is None: + first_output_at_ns = time.monotonic_ns() + + if self._llm_call_observer is not None: + content_callback = kwargs.get("on_content_delta") + if callable(content_callback): + typed_content_callback = cast( + Callable[[str], Awaitable[None]], + content_callback, + ) + + async def _timed_content_delta(delta: str) -> None: + _mark_output(delta) + await typed_content_callback(delta) + + kwargs["on_content_delta"] = _timed_content_delta + + thinking_callback = kwargs.get("on_thinking_delta") + if callable(thinking_callback): + typed_thinking_callback = cast( + Callable[[str], Awaitable[None]], + thinking_callback, + ) + + async def _timed_thinking_delta(delta: str) -> None: + _mark_output(delta) + await typed_thinking_callback(delta) + + kwargs["on_thinking_delta"] = _timed_thinking_delta + + def _attach_stream_timing(response: LLMResponse) -> LLMResponse: + if first_output_at_ns is None: + return response + finished_at_ns = time.monotonic_ns() + if response.ttft_ms is None: + response.ttft_ms = max(0, round((first_output_at_ns - started_at_ns) / 1_000_000)) + if response.generation_ms is None: + response.generation_ms = max( + 1, + round((finished_at_ns - first_output_at_ns) / 1_000_000), + ) + return response + try: provider_context = kwargs.pop("provider_context", None) if isinstance(provider_context, ProviderCallContext): @@ -1284,7 +1331,7 @@ class LLMProvider(ABC): except Exception as exc: response = LLMResponse(content=f"Error calling LLM: {exc}", finish_reason="error") return self._observe_llm_call( - response, + _attach_stream_timing(response), kwargs, started_at_ms=started_at_ms, started_at_ns=started_at_ns, diff --git a/tests/providers/test_llm_usage_observer.py b/tests/providers/test_llm_usage_observer.py index 6c22ca5ac..0e5cbc097 100644 --- a/tests/providers/test_llm_usage_observer.py +++ b/tests/providers/test_llm_usage_observer.py @@ -43,6 +43,26 @@ class _BlockingProvider(LLMProvider): return "test-model" +class _StreamingProvider(LLMProvider): + async def chat(self, **_kwargs: object) -> LLMResponse: + raise AssertionError("streaming path expected") + + async def chat_stream(self, **kwargs: object) -> LLMResponse: + on_thinking_delta = kwargs.get("on_thinking_delta") + if callable(on_thinking_delta): + await on_thinking_delta("thinking") + on_content_delta = kwargs.get("on_content_delta") + if callable(on_content_delta): + await on_content_delta("ok") + return LLMResponse( + content="ok", + usage=LLMUsage.reported(input_tokens=12, output_tokens=2), + ) + + def get_default_model(self) -> str: + return "test-model" + + @pytest.mark.asyncio async def test_observer_receives_every_retry_attempt() -> None: provider = _SequenceProvider( @@ -114,6 +134,40 @@ async def test_observer_failure_never_breaks_provider_call() -> None: assert response.content == "ok" +@pytest.mark.asyncio +async def test_stream_observer_records_physical_attempt_timing(monkeypatch) -> None: + provider = _StreamingProvider(provider_name="streaming-provider") + events: list[LLMCallRecord] = [] + provider.set_llm_call_observer(events.append) + monotonic_values = iter( + [ + 1_000_000_000, + 1_005_000_000, + 1_012_000_000, + 1_013_000_000, + ] + ) + monkeypatch.setattr( + "nanobot.providers.base.time.monotonic_ns", + lambda: next(monotonic_values), + ) + + response = await provider.chat_stream_with_retry( + messages=[{"role": "user", "content": "hello"}], + on_content_delta=lambda _delta: asyncio.sleep(0), + on_thinking_delta=lambda _delta: asyncio.sleep(0), + ) + + assert len(events) == 1 + usage = events[0].usage + assert usage is not None + assert usage.ttft_ms == 5 + assert usage.generation_ms == 7 + assert usage.timed_requests == 1 + assert usage.measured_output_tokens == 2 + assert response.usage == usage + + @pytest.mark.asyncio @pytest.mark.parametrize("stream", [False, True]) async def test_observer_records_cancelled_provider_attempt(stream: bool) -> None: