ai_ml_operations/chat_backend
2
0
Fork 0
Code Issues 23 Pull Requests Actions Packages Projects Releases Wiki Activity

Persist Ollama token usage from streamed LLM responses (#16) #38

Merged
westfarn merged 1 commits from fix/persist-stream-token-usage into master 2026-07-31 10:46:57 -07:00
Conversation 0 Commits 1 Files Changed 7
+139 -22
westfarn commented 2026-07-31 07:29:32 -07:00
Owner
Copy Link
Copy Source

Summary

  • Fixes token tracking for #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

  • 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 logs
## Summary - Fixes token tracking for [#16](https://git.aimloperations.com/ai_ml_operations/chat_backend/issues/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 logs
westfarn added 1 commit 2026-07-31 07:29:33 -07:00
Persist Ollama token usage from streamed LLM responses (#16)
CI / test (pull_request) Successful in 10s
Details
Unit Tests / test (pull_request) Successful in 10s
Details
3892bc3224 
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.
westfarn merged commit cc45ae5808 into master 2026-07-31 10:46:57 -07:00
westfarn deleted branch fix/persist-stream-token-usage 2026-07-31 10:46:57 -07:00
Sign in to join this conversation.
Reviewers
No Reviewers
Labels
No items
No labels
Milestone
No items
No Milestone
Projects
Clear projects
No projects
Assignees
Clear assignees
twestfall westfarn (Ryan Westfall)
No Assignees
1 Participants
Notifications
Due Date
No due date set.
Dependencies

No dependencies set.

Reference: ai_ml_operations/chat_backend#38