mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-03 20:12:55 +00:00
arc53-machine
b17f40a993
The wrappers in docsgpt/usage.py already measured how long each call took -- for the llm_gen_finished / llm_stream_finished log lines -- and then threw the number away. Nothing in the schema recorded it, so an operator could see what an instance spent but never how slow it was. Adds token_usage.duration_ms and token_usage.ttft_ms, and threads the measurements the wrappers already take into the insert. The stream wrapper now also stamps the moment of the first yielded chunk. Both columns are nullable, and deliberately so: ttft_ms is NULL for a non-streaming call and for a stream that failed before yielding anything, and every row written before this migration is NULL too. A 0 there would drag a p50 toward an instant first token that never happened.
95 lines
3.1 KiB
Python
95 lines
3.1 KiB
Python
"""Latency is measured by the usage wrappers and persisted with the call.
|
|
|
|
``duration_ms`` and ``ttft_ms`` were already computed for the finish log lines
|
|
and then discarded; these pin that they now reach ``token_usage``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
from docsgpt.usage import gen_token_usage, stream_token_usage
|
|
|
|
|
|
class _LLM:
|
|
"""Minimal stand-in for an LLM instance the wrappers decorate."""
|
|
|
|
def __init__(self):
|
|
self.token_usage = {"prompt_tokens": 0, "generated_tokens": 0}
|
|
self.decoded_token = {"sub": "u1"}
|
|
self.user_api_key = None
|
|
self.agent_id = None
|
|
|
|
|
|
@pytest.fixture
|
|
def persisted():
|
|
calls: list[dict] = []
|
|
|
|
def _capture(llm, call_usage, *, duration_ms=None, ttft_ms=None):
|
|
calls.append({"duration_ms": duration_ms, "ttft_ms": ttft_ms})
|
|
|
|
with patch("docsgpt.usage._persist_call_usage", _capture):
|
|
yield calls
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestNonStreaming:
|
|
def test_records_a_duration_and_no_first_token(self, persisted):
|
|
@gen_token_usage
|
|
def _gen(self, model, messages, stream, tools, **kwargs):
|
|
return "hello"
|
|
|
|
assert _gen(_LLM(), "m", [], False, None) == "hello"
|
|
assert persisted[0]["duration_ms"] >= 0
|
|
# A non-streaming call has no first-token moment.
|
|
assert persisted[0]["ttft_ms"] is None
|
|
|
|
def test_a_failed_call_still_records_its_duration(self, persisted):
|
|
@gen_token_usage
|
|
def _gen(self, model, messages, stream, tools, **kwargs):
|
|
raise RuntimeError("upstream down")
|
|
|
|
with pytest.raises(RuntimeError):
|
|
_gen(_LLM(), "m", [], False, None)
|
|
assert persisted[0]["duration_ms"] >= 0
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestStreaming:
|
|
def test_records_time_to_first_chunk(self, persisted):
|
|
@stream_token_usage
|
|
def _stream(self, model, messages, stream, tools, **kwargs):
|
|
yield "a"
|
|
yield "b"
|
|
|
|
assert list(_stream(_LLM(), "m", [], True, None)) == ["a", "b"]
|
|
row = persisted[0]
|
|
assert row["ttft_ms"] is not None
|
|
# First token cannot land after the call finished.
|
|
assert row["ttft_ms"] <= row["duration_ms"]
|
|
|
|
def test_a_stream_that_never_yields_has_no_first_token(self, persisted):
|
|
"""NULL, not 0 — an instant p50 would be a lie."""
|
|
|
|
@stream_token_usage
|
|
def _stream(self, model, messages, stream, tools, **kwargs):
|
|
raise RuntimeError("refused")
|
|
yield # pragma: no cover - unreachable, makes this a generator
|
|
|
|
with pytest.raises(RuntimeError):
|
|
list(_stream(_LLM(), "m", [], True, None))
|
|
assert persisted[0]["ttft_ms"] is None
|
|
assert persisted[0]["duration_ms"] >= 0
|
|
|
|
def test_a_stream_cut_short_keeps_the_first_token_it_saw(self, persisted):
|
|
@stream_token_usage
|
|
def _stream(self, model, messages, stream, tools, **kwargs):
|
|
yield "a"
|
|
raise RuntimeError("dropped")
|
|
|
|
with pytest.raises(RuntimeError):
|
|
list(_stream(_LLM(), "m", [], True, None))
|
|
assert persisted[0]["ttft_ms"] is not None
|