From 3892bc322401f9090a51f00b811940414b5ff649 Mon Sep 17 00:00:00 2001 From: Ryan Westfall Date: Fri, 31 Jul 2026 09:29:14 -0500 Subject: [PATCH] Persist Ollama token usage from streamed LLM responses (#16) Streaming dropped provider usage (StrOutputParser + None extract path), so PromptMetric tokens_in/out stayed null. Capture generation_info while streaming and stop double-closing websockets. --- llm_be/chat_backend/consumers.py | 26 ++++++--- llm_be/chat_backend/consumers_graph.py | 25 ++++++--- .../services/data_analysis_service.py | 2 +- llm_be/chat_backend/services/llm_service.py | 3 +- llm_be/chat_backend/services/rag_services.py | 2 +- llm_be/chat_backend/tests/test_utils.py | 48 ++++++++++++++++ llm_be/chat_backend/utils.py | 55 ++++++++++++++++++- 7 files changed, 139 insertions(+), 22 deletions(-) diff --git a/llm_be/chat_backend/consumers.py b/llm_be/chat_backend/consumers.py index 48ae5ab..7f97115 100644 --- a/llm_be/chat_backend/consumers.py +++ b/llm_be/chat_backend/consumers.py @@ -28,6 +28,8 @@ from .services.moderation_classifier import moderation_classifier, ModerationLab from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType from .services.data_analysis_service import AsyncDataAnalysisService from .utils import ( + TokenUsageCollector, + aiter_text_chunks, extract_token_usage, has_usable_user_prompt, is_heartbeat_payload, @@ -257,7 +259,9 @@ class ChatConsumerAgain(AsyncWebsocketConsumer): await self.accept() async def disconnect(self, close_code): - await self.close() + # Connection already closing — do not call self.close() again + # (triggers ASGI 'websocket.close' after close completed). + pass async def send_json_message(self, data_str): """ @@ -509,26 +513,30 @@ class ChatConsumerAgain(AsyncWebsocketConsumer): response_generator_or_dict = await generate_response_step(step2) full_response = "" - + tokens_in = tokens_out = None + if isinstance(response_generator_or_dict, dict): # It's an error or simple message content = response_generator_or_dict.get("content", "") await self.send_json_message(json.dumps(response_generator_or_dict)) full_response = content + tokens_in, tokens_out = extract_token_usage( + response_generator_or_dict + ) else: - # It's an async generator - async for chunk in response_generator_or_dict: + # Stream raw LLM chunks so final Ollama generation_info + # (prompt_eval_count / eval_count) is not stripped. + usage = TokenUsageCollector() + async for chunk in aiter_text_chunks( + response_generator_or_dict, usage + ): full_response += chunk await self.send_json_message(chunk) + tokens_in, tokens_out = usage.pair await self.send("END_OF_THE_STREAM_ENDER_GAME_42") await save_generated_message(conversation_id, full_response) - tokens_in, tokens_out = extract_token_usage( - response_generator_or_dict - if isinstance(response_generator_or_dict, dict) - else None - ) await finish_prompt_metric( prompt_metric, len(full_response), diff --git a/llm_be/chat_backend/consumers_graph.py b/llm_be/chat_backend/consumers_graph.py index 8e3cc0f..2f6102c 100644 --- a/llm_be/chat_backend/consumers_graph.py +++ b/llm_be/chat_backend/consumers_graph.py @@ -23,6 +23,8 @@ from .services.moderation_classifier import moderation_classifier, ModerationLab from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType from .services.data_analysis_service import AsyncDataAnalysisService from .utils import ( + TokenUsageCollector, + aiter_text_chunks, extract_token_usage, has_usable_user_prompt, is_heartbeat_payload, @@ -323,7 +325,9 @@ class ChatConsumerGraph(AsyncWebsocketConsumer): await self.accept() async def disconnect(self, close_code): - await self.close() + # Connection already closing — do not call self.close() again + # (triggers ASGI 'websocket.close' after close completed). + pass async def send_json_message(self, data_str): try: @@ -453,24 +457,29 @@ class ChatConsumerGraph(AsyncWebsocketConsumer): await self.send("START_OF_THE_STREAM_ENDER_GAME_42") full_response = "" - + tokens_in = tokens_out = None + if isinstance(response_generator_or_dict, dict): content = response_generator_or_dict.get("content", "") await self.send_json_message(json.dumps(response_generator_or_dict)) full_response = content + tokens_in, tokens_out = extract_token_usage( + response_generator_or_dict + ) else: - async for chunk in response_generator_or_dict: + # Stream raw LLM chunks so final Ollama generation_info + # (prompt_eval_count / eval_count) is not stripped. + usage = TokenUsageCollector() + async for chunk in aiter_text_chunks( + response_generator_or_dict, usage + ): full_response += chunk await self.send_json_message(chunk) + tokens_in, tokens_out = usage.pair await self.send("END_OF_THE_STREAM_ENDER_GAME_42") await save_generated_message(conversation_id, full_response) - tokens_in, tokens_out = extract_token_usage( - response_generator_or_dict - if isinstance(response_generator_or_dict, dict) - else None - ) await finish_prompt_metric( prompt_metric, len(full_response), diff --git a/llm_be/chat_backend/services/data_analysis_service.py b/llm_be/chat_backend/services/data_analysis_service.py index 4da2926..4766491 100644 --- a/llm_be/chat_backend/services/data_analysis_service.py +++ b/llm_be/chat_backend/services/data_analysis_service.py @@ -55,7 +55,7 @@ Answer:""" } | self.prompt | self.llm - | self.output_parser + # No StrOutputParser: keep Ollama generation_info token counts. ) def _get_dataframe_summary(self, df: pd.DataFrame) -> str: diff --git a/llm_be/chat_backend/services/llm_service.py b/llm_be/chat_backend/services/llm_service.py index 220f013..95e80c0 100644 --- a/llm_be/chat_backend/services/llm_service.py +++ b/llm_be/chat_backend/services/llm_service.py @@ -118,7 +118,8 @@ class AsyncLLMService(LLMService): } | self.prompt | self.llm - | self.output_parser + # No StrOutputParser: Ollama puts prompt_eval_count/eval_count on the + # final GenerationChunk.generation_info; the parser would drop it. ) async def _format_history(self, conversation: list) -> str: diff --git a/llm_be/chat_backend/services/rag_services.py b/llm_be/chat_backend/services/rag_services.py index 74dafc2..94bd645 100644 --- a/llm_be/chat_backend/services/rag_services.py +++ b/llm_be/chat_backend/services/rag_services.py @@ -328,7 +328,7 @@ class AsyncRAGService(RAGService): } | self.prompt | self.llm - | StrOutputParser() + # No StrOutputParser: keep Ollama generation_info token counts. ) async def _format_history(self, conversation: Conversation) -> str: diff --git a/llm_be/chat_backend/tests/test_utils.py b/llm_be/chat_backend/tests/test_utils.py index a26f7cf..0f5a9cd 100644 --- a/llm_be/chat_backend/tests/test_utils.py +++ b/llm_be/chat_backend/tests/test_utils.py @@ -12,6 +12,9 @@ from chat_backend.ollama_config import ( ollama_model, ) from chat_backend.utils import ( + TokenUsageCollector, + aiter_text_chunks, + chunk_text, extract_token_usage, has_usable_user_prompt, is_heartbeat_payload, @@ -65,6 +68,51 @@ class ExtractTokenUsageTestCase(SimpleTestCase): (10, 20), ) + def test_reads_generation_info_attribute(self): + class Chunk: + generation_info = { + "done": True, + "prompt_eval_count": 22, + "eval_count": 55, + } + + self.assertEqual(extract_token_usage(Chunk()), (22, 55)) + + def test_chunk_text_from_generation_chunk(self): + class Chunk: + text = "hello" + generation_info = {"prompt_eval_count": 1, "eval_count": 2} + + self.assertEqual(chunk_text(Chunk()), "hello") + self.assertEqual(chunk_text("plain"), "plain") + self.assertEqual(chunk_text(None), "") + + async def test_aiter_text_chunks_collects_final_usage(self): + class Chunk: + def __init__(self, text, info=None): + self.text = text + self.generation_info = info or {} + + async def stream(): + yield Chunk("Hel") + yield Chunk("lo", {"prompt_eval_count": 11, "eval_count": 3}) + + usage = TokenUsageCollector() + texts = [t async for t in aiter_text_chunks(stream(), usage)] + + self.assertEqual("".join(texts), "Hello") + self.assertEqual(usage.pair, (11, 3)) + + async def test_aiter_text_chunks_without_usage_stays_null(self): + async def stream(): + yield "only-text" + + usage = TokenUsageCollector() + texts = [t async for t in aiter_text_chunks(stream(), usage)] + + self.assertEqual(texts, ["only-text"]) + self.assertEqual(usage.pair, (None, None)) + class LastDayOfMonthTestCase(SimpleTestCase): @parameterized.expand( diff --git a/llm_be/chat_backend/utils.py b/llm_be/chat_backend/utils.py index 687b1cf..8d51a6b 100644 --- a/llm_be/chat_backend/utils.py +++ b/llm_be/chat_backend/utils.py @@ -50,7 +50,8 @@ 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``), or any object exposing those attributes. + ``response_metadata``), an Ollama ``GenerationChunk`` (``generation_info``), + or any object exposing those attributes. """ if source is None: return None @@ -60,7 +61,7 @@ def _as_usage_mapping(source): if isinstance(nested, dict): return nested return source - for attr in ("usage_metadata", "response_metadata"): + for attr in ("usage_metadata", "response_metadata", "generation_info"): nested = getattr(source, attr, None) if isinstance(nested, dict): mapping = _as_usage_mapping(nested) @@ -79,3 +80,53 @@ def extract_token_usage(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 -- 2.54.0