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

@@ -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
+                    # Stream raw LLM chunks so final Ollama generation_info
-                    async for chunk in response_generator_or_dict:
+                    # (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),

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),

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:

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:

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:

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(

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