mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-14 08:17:34 +00:00
feat(usage): persist LLM call latency
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.
This commit is contained in:
1 parent
e30631a060
commit
b17f40a993
6 files changed
+261
-5
No files matched your search
+29
-5
@@ -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,
|
||||
|
||||
Reference in new issue
Block a user