mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-31 08:13:11 +03:00
fix(usage): record provider stream timing
This commit is contained in:
@@ -1259,6 +1259,53 @@ class LLMProvider(ABC):
|
|||||||
"""Call chat_stream() and convert unexpected exceptions to error responses."""
|
"""Call chat_stream() and convert unexpected exceptions to error responses."""
|
||||||
started_at_ms = time.time_ns() // 1_000_000
|
started_at_ms = time.time_ns() // 1_000_000
|
||||||
started_at_ns = time.monotonic_ns()
|
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:
|
try:
|
||||||
provider_context = kwargs.pop("provider_context", None)
|
provider_context = kwargs.pop("provider_context", None)
|
||||||
if isinstance(provider_context, ProviderCallContext):
|
if isinstance(provider_context, ProviderCallContext):
|
||||||
@@ -1284,7 +1331,7 @@ class LLMProvider(ABC):
|
|||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
response = LLMResponse(content=f"Error calling LLM: {exc}", finish_reason="error")
|
response = LLMResponse(content=f"Error calling LLM: {exc}", finish_reason="error")
|
||||||
return self._observe_llm_call(
|
return self._observe_llm_call(
|
||||||
response,
|
_attach_stream_timing(response),
|
||||||
kwargs,
|
kwargs,
|
||||||
started_at_ms=started_at_ms,
|
started_at_ms=started_at_ms,
|
||||||
started_at_ns=started_at_ns,
|
started_at_ns=started_at_ns,
|
||||||
|
|||||||
@@ -43,6 +43,26 @@ class _BlockingProvider(LLMProvider):
|
|||||||
return "test-model"
|
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
|
@pytest.mark.asyncio
|
||||||
async def test_observer_receives_every_retry_attempt() -> None:
|
async def test_observer_receives_every_retry_attempt() -> None:
|
||||||
provider = _SequenceProvider(
|
provider = _SequenceProvider(
|
||||||
@@ -114,6 +134,40 @@ async def test_observer_failure_never_breaks_provider_call() -> None:
|
|||||||
assert response.content == "ok"
|
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.asyncio
|
||||||
@pytest.mark.parametrize("stream", [False, True])
|
@pytest.mark.parametrize("stream", [False, True])
|
||||||
async def test_observer_records_cancelled_provider_attempt(stream: bool) -> None:
|
async def test_observer_records_cancelled_provider_attempt(stream: bool) -> None:
|
||||||
|
|||||||
Reference in New Issue
Block a user