add finish_reason = event.choices[0].finish_reason

openlit · Mar 14, 2024 · 674a17f · 674a17f
1 parent d1db1eb
commit 674a17f
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git 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
@@ -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)