From 105ef95262b1feb5aa015b0f6d14d6c9e6ea4e9b Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Mon, 9 Feb 2026 19:28:13 +0000 Subject: [PATCH] Fix mistral metrics --- .../model/text_completion/mistral/llm.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/trustgraph-flow/trustgraph/model/text_completion/mistral/llm.py b/trustgraph-flow/trustgraph/model/text_completion/mistral/llm.py index 7952b1df..fab41ecd 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/mistral/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/mistral/llm.py @@ -156,6 +156,9 @@ class Processor(LlmService): response_format={"type": "text"} ) + total_input_tokens = 0 + total_output_tokens = 0 + for chunk in stream: if chunk.data.choices and chunk.data.choices[0].delta.content: yield LlmChunk( @@ -166,11 +169,16 @@ class Processor(LlmService): is_final=False ) - # Send final chunk + # Capture usage data when available (typically in final chunk) + if chunk.data.usage: + total_input_tokens = chunk.data.usage.prompt_tokens + total_output_tokens = chunk.data.usage.completion_tokens + + # Send final chunk with token counts yield LlmChunk( text="", - in_token=None, - out_token=None, + in_token=total_input_tokens, + out_token=total_output_tokens, model=model_name, is_final=True )