refactor: gemini

Author: wxg0103

apps/setting/models_provider/impl/gemini_model_provider/model/llm.py

@@ -13,7 +13,7 @@
     Tool as GoogleTool,
 )
 from langchain_core.callbacks import CallbackManagerForLLMRun
-from langchain_core.messages import BaseMessage
+from langchain_core.messages import BaseMessage, get_buffer_string
 from langchain_core.outputs import ChatGenerationChunk
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_google_genai._function_utils import _ToolConfigDict, _ToolDict
@@ -22,6 +22,8 @@
 from langchain_google_genai._common import (
     SafetySettingDict,
 )
+
+from common.config.tokenizer_manage_config import TokenizerManage
 from setting.models_provider.base_model_provider import MaxKBBaseModel
 
 
@@ -46,10 +48,18 @@ def get_last_generation_info(self) -> Optional[Dict[str, Any]]:
         return self.__dict__.get('_last_generation_info')
 
     def get_num_tokens_from_messages(self, messages: List[BaseMessage]) -> int:
-        return self.get_last_generation_info().get('input_tokens', 0)
+        try:
+            return self.get_last_generation_info().get('input_tokens', 0)
+        except Exception as e:
+            tokenizer = TokenizerManage.get_tokenizer()
+            return sum([len(tokenizer.encode(get_buffer_string([m]))) for m in messages])
 
     def get_num_tokens(self, text: str) -> int:
-        return self.get_last_generation_info().get('output_tokens', 0)
+        try:
+            return self.get_last_generation_info().get('output_tokens', 0)
+        except Exception as e:
+            tokenizer = TokenizerManage.get_tokenizer()
+            return len(tokenizer.encode(text))
 
     def _stream(
             self,