From 8c6c1baa309dc90dac3ede73aaa68b91c86ee1c4 Mon Sep 17 00:00:00 2001
From: yym68686 <yym68686@outlook.com>
Date: Thu, 29 Aug 2024 04:55:23 +0800
Subject: [PATCH] Fix the bug in Gemini API token calculation error.

---
 setup.py                        | 2 +-
 src/ModelMerge/models/genimi.py | 5 +++--
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/setup.py b/setup.py
index 916d791..72bfd1b 100644
--- a/setup.py
+++ b/setup.py
@@ -4,7 +4,7 @@
 
 setup(
     name="modelmerge",
-    version="0.11.18",
+    version="0.11.19",
     description="modelmerge is a multi-large language model API aggregator.",
     long_description=Path.open(Path("README.md"), encoding="utf-8").read(),
     long_description_content_type="text/markdown",
diff --git a/src/ModelMerge/models/genimi.py b/src/ModelMerge/models/genimi.py
index 257e33e..ad0a8a2 100644
--- a/src/ModelMerge/models/genimi.py
+++ b/src/ModelMerge/models/genimi.py
@@ -305,6 +305,7 @@ async def ask_stream_async(
         function_full_response: str = "{"
         need_function_call = False
         revicing_function_call = False
+        total_tokens = 0
         try:
             async with self.aclient.stream(
                 "post",
@@ -317,7 +318,7 @@ async def ask_stream_async(
                     async for line in response.aiter_lines():
                         if not line:
                             continue
-                        print(line)
+                        # print(line)
                         if line and '\"text\": \"' in line:
                             content = line.split('\"text\": \"')[1][:-1]
                             content = "\n".join(content.split("\\n"))
@@ -368,4 +369,4 @@ async def ask_stream_async(
             async for chunk in self.ask_stream_async(function_response, response_role, convo_id=convo_id, function_name=function_call_name, total_tokens=total_tokens, model=model, function_arguments=function_call, api_key=kwargs.get('api_key', self.api_key), plugins=kwargs.get("plugins", PLUGINS)):
                 yield chunk
         else:
-            self.add_to_conversation([{"text": full_response}], response_role, convo_id=convo_id, pass_history=pass_history)
\ No newline at end of file
+            self.add_to_conversation([{"text": full_response}], response_role, convo_id=convo_id, total_tokens=total_tokens, pass_history=pass_history)
\ No newline at end of file