From 674a17f6375e648b0c5889e898ec6606a6820955 Mon Sep 17 00:00:00 2001 From: patcher99 Date: Thu, 14 Mar 2024 22:19:32 +0530 Subject: [PATCH] add finish_reason = event.choices[0].finish_reason --- src/dokumetry/async_mistral.py | 4 +++- src/dokumetry/mistral.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/src/dokumetry/async_mistral.py b/src/dokumetry/async_mistral.py index 5970f7b..eb341fc 100644 --- a/src/dokumetry/async_mistral.py +++ b/src/dokumetry/async_mistral.py @@ -107,6 +107,7 @@ async def stream_generator(): prompt_tokens = event.usage.prompt_tokens completion_tokens = event.usage.completion_tokens total_tokens = event.usage.total_tokens + finish_reason = event.choices[0].finish_reason yield event end_time = time.time() duration = end_time - start_time @@ -142,7 +143,8 @@ async def stream_generator(): "response": accumulated_content, "promptTokens": prompt_tokens, "completionTokens": completion_tokens, - "totalTokens": total_tokens + "totalTokens": total_tokens, + "finishReason": finish_reason } send_data(data, doku_url, api_key) diff --git a/src/dokumetry/mistral.py b/src/dokumetry/mistral.py index 213f3e6..71ac5ca 100644 --- a/src/dokumetry/mistral.py +++ b/src/dokumetry/mistral.py @@ -107,6 +107,7 @@ def stream_generator(): prompt_tokens = event.usage.prompt_tokens completion_tokens = event.usage.completion_tokens total_tokens = event.usage.total_tokens + finish_reason = event.choices[0].finish_reason yield event end_time = time.time() duration = end_time - start_time @@ -142,7 +143,8 @@ def stream_generator(): "response": accumulated_content, "promptTokens": prompt_tokens, "completionTokens": completion_tokens, - "totalTokens": total_tokens + "totalTokens": total_tokens, + "finishReason": finish_reason } send_data(data, doku_url, api_key)