westfarn
cc45ae5808
## Summary - Fixes token tracking for [#16](#16): streaming chat never persisted `PromptMetric.tokens_in` / `tokens_out` (admin + account usage showed `—`). - Drop `StrOutputParser` on async LLM/RAG/data-analysis chains so Ollama `generation_info` (`prompt_eval_count` / `eval_count`) survives; collect usage while streaming via `TokenUsageCollector`. - Stop calling `self.close()` in `disconnect` (fixes Grafana `Unexpected ASGI message 'websocket.close'`). ## Test plan - [x] Unit tests: `test_utils`, consumers, LLM/RAG/data-analysis services, finance quotas - [ ] Deploy / local: send a chat prompt, confirm admin Prompt Metrics shows Tokens In/Out - [ ] Reload Account usage card — in/out no longer `—` for new turns - [ ] Confirm WS disconnect no longer raises double-close in logsReviewed-on: #38
133 lines
4.4 KiB
Python
133 lines
4.4 KiB
Python
import datetime
|
|
|
|
|
|
def is_heartbeat_payload(data) -> bool:
|
|
"""True for app-level WS keepalive frames (see FE buildHeartbeatPayload)."""
|
|
return isinstance(data, dict) and data.get("type") == "ping"
|
|
|
|
|
|
def normalize_user_message(message):
|
|
"""Return stripped message text, or None if missing/blank."""
|
|
if message is None:
|
|
return None
|
|
if not isinstance(message, str):
|
|
message = str(message)
|
|
stripped = message.strip()
|
|
return stripped or None
|
|
|
|
|
|
def has_usable_user_prompt(message, file=None) -> bool:
|
|
"""Reject empty/whitespace chat text. ``file`` kept for call-site clarity."""
|
|
return normalize_user_message(message) is not None
|
|
|
|
|
|
def last_day_of_month(any_day):
|
|
# The day 28 exists in every month. 4 days later, it's always next month
|
|
next_month = any_day.replace(day=28) + datetime.timedelta(days=4)
|
|
# subtracting the number of the current day brings us back one month
|
|
return next_month - datetime.timedelta(days=next_month.day)
|
|
|
|
|
|
# Keys different providers use for input/output token counts. We only ever read
|
|
# real usage the provider reports; we never estimate, so absence maps to None.
|
|
_TOKENS_IN_KEYS = ("input_tokens", "prompt_tokens", "prompt_eval_count")
|
|
_TOKENS_OUT_KEYS = ("output_tokens", "completion_tokens", "eval_count")
|
|
|
|
|
|
def _first_int(mapping, keys):
|
|
for key in keys:
|
|
value = mapping.get(key)
|
|
if isinstance(value, bool):
|
|
continue
|
|
if isinstance(value, int):
|
|
return value
|
|
if isinstance(value, float) and value.is_integer():
|
|
return int(value)
|
|
return None
|
|
|
|
|
|
def _as_usage_mapping(source):
|
|
"""Best-effort pull of a usage dict out of a provider response.
|
|
|
|
Accepts a raw dict, a LangChain message (``usage_metadata`` /
|
|
``response_metadata``), an Ollama ``GenerationChunk`` (``generation_info``),
|
|
or any object exposing those attributes.
|
|
"""
|
|
if source is None:
|
|
return None
|
|
if isinstance(source, dict):
|
|
for nested_key in ("usage_metadata", "usage", "token_usage"):
|
|
nested = source.get(nested_key)
|
|
if isinstance(nested, dict):
|
|
return nested
|
|
return source
|
|
for attr in ("usage_metadata", "response_metadata", "generation_info"):
|
|
nested = getattr(source, attr, None)
|
|
if isinstance(nested, dict):
|
|
mapping = _as_usage_mapping(nested)
|
|
if mapping:
|
|
return mapping
|
|
return None
|
|
|
|
|
|
def extract_token_usage(source):
|
|
"""Return ``(tokens_in, tokens_out)`` from a provider usage payload.
|
|
|
|
Values are only returned when the provider actually reports them; anything
|
|
missing comes back as ``None`` so callers never persist estimated counts.
|
|
"""
|
|
mapping = _as_usage_mapping(source)
|
|
if not mapping:
|
|
return None, None
|
|
return _first_int(mapping, _TOKENS_IN_KEYS), _first_int(mapping, _TOKENS_OUT_KEYS)
|
|
|
|
|
|
def chunk_text(chunk) -> str:
|
|
"""Pull display text out of a stream chunk (str, GenerationChunk, message)."""
|
|
if chunk is None:
|
|
return ""
|
|
if isinstance(chunk, str):
|
|
return chunk
|
|
text = getattr(chunk, "text", None)
|
|
if isinstance(text, str) and text:
|
|
return text
|
|
content = getattr(chunk, "content", None)
|
|
if isinstance(content, str):
|
|
return content
|
|
return ""
|
|
|
|
|
|
class TokenUsageCollector:
|
|
"""Accumulate provider-reported token counts while streaming LLM chunks."""
|
|
|
|
def __init__(self):
|
|
self.tokens_in = None
|
|
self.tokens_out = None
|
|
|
|
def observe(self, source) -> None:
|
|
tin, tout = extract_token_usage(source)
|
|
if tin is not None:
|
|
self.tokens_in = tin
|
|
if tout is not None:
|
|
self.tokens_out = tout
|
|
|
|
@property
|
|
def pair(self):
|
|
return self.tokens_in, self.tokens_out
|
|
|
|
|
|
async def aiter_text_chunks(stream, collector: TokenUsageCollector | None = None):
|
|
"""Yield text from a provider stream, optionally capturing token usage.
|
|
|
|
Ollama reports ``prompt_eval_count`` / ``eval_count`` on the final
|
|
``GenerationChunk.generation_info`` when ``done`` is true. ``StrOutputParser``
|
|
strips that metadata, so callers must stream the raw LLM chain and use this
|
|
helper (or equivalent) to persist real usage.
|
|
"""
|
|
async for chunk in stream:
|
|
if collector is not None:
|
|
collector.observe(chunk)
|
|
text = chunk_text(chunk)
|
|
if text:
|
|
yield text
|