Persist Ollama token usage from streamed LLM responses (#16) (#38)
Deploy Beta / unit-tests (push) Successful in 10s
Unit Tests / test (push) Successful in 10s
Deploy Beta / docker (push) Successful in 19s
Deploy Beta / deploy-beta (push) Successful in 48s

## 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
This commit was merged in pull request #38.
This commit is contained in:
2026-07-31 10:46:57 -07:00
parent 841c0962d9
commit cc45ae5808
7 changed files with 139 additions and 22 deletions
@@ -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:
+2 -1
View File
@@ -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:
+1 -1
View File
@@ -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: