From 28085f6f048f2e8b25e28f6bb556879be5c8a025 Mon Sep 17 00:00:00 2001
From: CelineNi2 <144465886+CelineNi2@users.noreply.github.com>
Date: Thu, 16 Oct 2025 01:48:48 +0200
Subject: [PATCH] Add messages regarding the use of token during query (#147)

* Add messages regarding the use of token during query

* fix: apply ruff format
---
 packages/leann-core/src/leann/chat.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/packages/leann-core/src/leann/chat.py b/packages/leann-core/src/leann/chat.py
index 1f870bf..9d874bf 100644
--- a/packages/leann-core/src/leann/chat.py
+++ b/packages/leann-core/src/leann/chat.py
@@ -834,6 +834,11 @@ class OpenAIChat(LLMInterface):
 
         try:
             response = self.client.chat.completions.create(**params)
+            print(
+                f"Total tokens = {response.usage.total_tokens}, prompt tokens = {response.usage.prompt_tokens}, completion tokens = {response.usage.completion_tokens}"
+            )
+            if response.choices[0].finish_reason == "length":
+                print("The query is exceeding the maximum allowed number of tokens")
             return response.choices[0].message.content.strip()
         except Exception as e:
             logger.error(f"Error communicating with OpenAI: {e}")