diff --git a/docsgpt/alembic/versions/0035_token_usage_latency.py b/docsgpt/alembic/versions/0035_token_usage_latency.py new file mode 100644 index 00000000..b06d3203 --- /dev/null +++ b/docsgpt/alembic/versions/0035_token_usage_latency.py @@ -0,0 +1,42 @@ +"""0035 token usage latency — how long the call took, and time to first token. + +Nothing in the schema recorded how long an LLM call took. The wrappers in +``docsgpt/usage.py`` already measured it for the ``llm_gen_finished`` / +``llm_stream_finished`` log lines and then threw the number away, so an +operator could see what an instance spent but never how slow it was. + +Two columns, both nullable because both are genuinely unknown for some rows: + +* ``duration_ms`` — wall-clock for the whole call. +* ``ttft_ms`` — time to the first streamed chunk. NULL on non-streaming calls + and on streams that failed before yielding anything, which is the honest + reading: "no first token", not "instant". + +NULL also means "written before this migration", so historical rows never +drag a p50 towards zero. + +Idempotent both ways. + +Revision ID: 0035_token_usage_latency +Revises: 0034_auth_events_actor_target +""" + +from typing import Sequence, Union + +from alembic import op + + +revision: str = "0035_token_usage_latency" +down_revision: Union[str, None] = "0034_auth_events_actor_target" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.execute("ALTER TABLE token_usage ADD COLUMN IF NOT EXISTS duration_ms INTEGER;") + op.execute("ALTER TABLE token_usage ADD COLUMN IF NOT EXISTS ttft_ms INTEGER;") + + +def downgrade() -> None: + op.execute("ALTER TABLE token_usage DROP COLUMN IF EXISTS ttft_ms;") + op.execute("ALTER TABLE token_usage DROP COLUMN IF EXISTS duration_ms;") diff --git a/docsgpt/storage/db/models.py b/docsgpt/storage/db/models.py index b3bd9a1f..edab66a5 100644 --- a/docsgpt/storage/db/models.py +++ b/docsgpt/storage/db/models.py @@ -250,6 +250,11 @@ token_usage_table = Table( # Added in ``0033_quotas``. USD cost of the call at write time; 0 for # unpriced and bring-your-own models. Column("cost", Numeric(12, 8), nullable=False, server_default="0"), + # Added in ``0035_token_usage_latency``. Wall-clock for the call, and time + # to the first streamed chunk. ``ttft_ms`` is NULL for non-streaming calls + # and for streams that failed before yielding -- "no first token", not 0. + Column("duration_ms", Integer), + Column("ttft_ms", Integer), ) user_logs_table = Table( diff --git a/docsgpt/storage/db/repositories/token_usage.py b/docsgpt/storage/db/repositories/token_usage.py index fb256b55..ec380b85 100644 --- a/docsgpt/storage/db/repositories/token_usage.py +++ b/docsgpt/storage/db/repositories/token_usage.py @@ -38,6 +38,8 @@ class TokenUsageRepository: cached_tokens: Optional[int] = None, cache_write_tokens: Optional[int] = None, cost: float = 0.0, + duration_ms: Optional[int] = None, + ttft_ms: Optional[int] = None, ) -> None: # Attribution guard: the ``token_usage_attribution_chk`` CHECK # constraint requires at least one of ``user_id`` / ``api_key`` @@ -64,6 +66,7 @@ class TokenUsageRepository: user_id, api_key, agent_id, prompt_tokens, generated_tokens, cached_tokens, cache_write_tokens, cost, + duration_ms, ttft_ms, source, request_id, model_id, timestamp ) VALUES ( @@ -71,6 +74,7 @@ class TokenUsageRepository: CAST(:agent_id AS uuid), :prompt_tokens, :generated_tokens, :cached_tokens, :cache_write_tokens, :cost, + :duration_ms, :ttft_ms, :source, :request_id, :model_id, COALESCE(:timestamp, now()) ) """ @@ -84,6 +88,8 @@ class TokenUsageRepository: "cached_tokens": cached_tokens, "cache_write_tokens": cache_write_tokens, "cost": cost, + "duration_ms": duration_ms, + "ttft_ms": ttft_ms, "source": source, "request_id": request_id, "model_id": model_id, diff --git a/docsgpt/usage.py b/docsgpt/usage.py index 82fc3eb8..17adc24a 100644 --- a/docsgpt/usage.py +++ b/docsgpt/usage.py @@ -93,7 +93,7 @@ def _count_prompt_tokens(messages, tools=None, usage_attachments=None, **kwargs) return prompt_tokens -def _persist_call_usage(llm, call_usage): +def _persist_call_usage(llm, call_usage, *, duration_ms=None, ttft_ms=None): """Write one ``token_usage`` row per LLM call. Always-on; no flag. Source defaults to ``agent_stream`` and can be overridden per @@ -101,6 +101,13 @@ def _persist_call_usage(llm, call_usage): title / compression / rag_condense / fallback). A ``_request_id`` stamped on the LLM lets ``count_in_range`` deduplicate the multiple rows produced by a single multi-tool agent run. + + Args: + llm: The LLM instance the call ran on. + call_usage: The call's token counts. + duration_ms: Wall-clock for the call, measured by the wrapper. + ttft_ms: Time to the first streamed chunk; None for a non-streaming + call and for a stream that failed before yielding anything. """ if call_usage["prompt_tokens"] == 0 and call_usage["generated_tokens"] == 0: return @@ -149,6 +156,8 @@ def _persist_call_usage(llm, call_usage): ), request_id=getattr(llm, "_request_id", None), model_id=model_id, + duration_ms=duration_ms, + ttft_ms=ttft_ms, ) except Exception: logger.exception("token_usage persist failed") @@ -258,10 +267,12 @@ def gen_token_usage(func): error = exc raise finally: + duration_ms = int((time.monotonic() - started_at) * 1000) call_usage = _prefer_provider_usage(self, call_usage) self.token_usage["prompt_tokens"] += call_usage["prompt_tokens"] self.token_usage["generated_tokens"] += call_usage["generated_tokens"] - _persist_call_usage(self, call_usage) + # A non-streaming call has no first-token moment; ttft stays NULL. + _persist_call_usage(self, call_usage, duration_ms=duration_ms) emit = getattr(self, "_emit_gen_finished_log", None) if callable(emit): try: @@ -269,7 +280,7 @@ def gen_token_usage(func): model, prompt_tokens=call_usage["prompt_tokens"], completion_tokens=call_usage["generated_tokens"], - latency_ms=int((time.monotonic() - started_at) * 1000), + latency_ms=duration_ms, cached_tokens=call_usage.get("cached_tokens"), cache_write_tokens=call_usage.get("cache_write_tokens"), error=error, @@ -293,10 +304,13 @@ def stream_token_usage(func): ) batch = [] started_at = time.monotonic() + first_chunk_at: float | None = None error: BaseException | None = None try: result = func(self, model, messages, stream, tools, **kwargs) for r in result: + if first_chunk_at is None: + first_chunk_at = time.monotonic() batch.append(r) yield r except Exception as exc: @@ -306,12 +320,22 @@ def stream_token_usage(func): error = exc raise finally: + duration_ms = int((time.monotonic() - started_at) * 1000) + # NULL, not 0, when the stream failed before yielding: "no first + # token" must not read as an instant one in a p50. + ttft_ms = ( + int((first_chunk_at - started_at) * 1000) + if first_chunk_at is not None + else None + ) for line in batch: call_usage["generated_tokens"] += _count_tokens(line) call_usage = _prefer_provider_usage(self, call_usage) self.token_usage["prompt_tokens"] += call_usage["prompt_tokens"] self.token_usage["generated_tokens"] += call_usage["generated_tokens"] - _persist_call_usage(self, call_usage) + _persist_call_usage( + self, call_usage, duration_ms=duration_ms, ttft_ms=ttft_ms + ) emit = getattr(self, "_emit_stream_finished_log", None) if callable(emit): try: @@ -319,7 +343,7 @@ def stream_token_usage(func): model, prompt_tokens=call_usage["prompt_tokens"], completion_tokens=call_usage["generated_tokens"], - latency_ms=int((time.monotonic() - started_at) * 1000), + latency_ms=duration_ms, cached_tokens=call_usage.get("cached_tokens"), cache_write_tokens=call_usage.get("cache_write_tokens"), error=error, diff --git a/tests/storage/db/test_migration_0035.py b/tests/storage/db/test_migration_0035.py new file mode 100644 index 00000000..4fe13f9a --- /dev/null +++ b/tests/storage/db/test_migration_0035.py @@ -0,0 +1,85 @@ +"""Migration round-trip test for 0035_token_usage_latency.""" + +from __future__ import annotations + +import os +import subprocess +import sys +from pathlib import Path + +import pytest +from sqlalchemy import text + + +pytestmark = pytest.mark.integration + + +def _alembic_ini() -> Path: + return Path(__file__).resolve().parents[3] / "docsgpt" / "alembic.ini" + + +def _run_alembic(url: str, *args: str) -> None: + subprocess.check_call( + [sys.executable, "-m", "alembic", "-c", str(_alembic_ini()), *args], + timeout=60, + env={**os.environ, "POSTGRES_URI": url}, + ) + + +def _alembic_version(conn) -> str: + return conn.execute(text("SELECT version_num FROM alembic_version")).scalar() + + +def _column_exists(conn, table: str, column: str) -> bool: + row = conn.execute( + text( + "SELECT 1 FROM information_schema.columns " + "WHERE table_name = :t AND column_name = :c AND table_schema = 'public'" + ), + {"t": table, "c": column}, + ).fetchone() + return row is not None + + +_0035 = "0035_token_usage_latency" +_0034 = "0034_auth_events_actor_target" + + +class TestMigration0035RoundTrip: + def test_head_has_latency_columns(self, pg_engine): + with pg_engine.connect() as conn: + assert _alembic_version(conn) >= _0035 + assert _column_exists(conn, "token_usage", "duration_ms") + assert _column_exists(conn, "token_usage", "ttft_ms") + + def test_downgrade_drops_then_upgrade_restores(self, pg_engine): + url = pg_engine.url.render_as_string(hide_password=False) + _run_alembic(url, "downgrade", _0034) + with pg_engine.connect() as conn: + assert _alembic_version(conn) == _0034 + assert not _column_exists(conn, "token_usage", "duration_ms") + assert not _column_exists(conn, "token_usage", "ttft_ms") + _run_alembic(url, "upgrade", "head") + with pg_engine.connect() as conn: + assert _column_exists(conn, "token_usage", "duration_ms") + + def test_existing_rows_read_null_latency(self, pg_engine): + """Rows written before 0035 must read NULL, not 0 — a 0ms call is a lie.""" + url = pg_engine.url.render_as_string(hide_password=False) + _run_alembic(url, "downgrade", _0034) + with pg_engine.begin() as conn: + conn.execute( + text( + "INSERT INTO token_usage (user_id, prompt_tokens, generated_tokens) " + "VALUES ('u-mig35', 10, 1)" + ) + ) + _run_alembic(url, "upgrade", "head") + with pg_engine.connect() as conn: + row = conn.execute( + text( + "SELECT duration_ms, ttft_ms FROM token_usage " + "WHERE user_id = 'u-mig35'" + ) + ).fetchone() + assert tuple(row) == (None, None) diff --git a/tests/test_usage_latency.py b/tests/test_usage_latency.py new file mode 100644 index 00000000..bf54a4ff --- /dev/null +++ b/tests/test_usage_latency.py @@ -0,0 +1,94 @@ +"""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