igbot-base 0.0.19__py3-none-any.whl → 0.0.20__py3-none-any.whl

igbot_base/models.py

@@ -1,12 +1,16 @@
 from enum import Enum
 from openai import OpenAI
 
+from tokenizer import BaseTokenizer, OpenAiTokenizer
+
 
 class ModelInfo:
 
-    def __init__(self, name, client):
+    def __init__(self, name, client, tokenizer: BaseTokenizer, max_tokens):
         self.__name = name
         self.__client = client
+        self.__tokenizer = tokenizer
+        self.__max_tokens = max_tokens
 
     def get_name(self):
         return self.__name
@@ -14,9 +18,19 @@ class ModelInfo:
     def get_client(self):
         return self.__client()
 
+    def get_tokenizer(self) -> BaseTokenizer:
+        return self.__tokenizer
+
+    def get_max_tokens(self):
+        return self.__max_tokens
+
 
 class Model(Enum):
-    OLLAMA_3_2_LOCAL = ModelInfo("llama3.2", lambda: OpenAI(base_url="http://localhost:11434/v1", api_key='ollama'))
-    OPENAI_GPT_4o_MINI = ModelInfo("gpt-4o-mini", lambda: OpenAI())
-    OPENAI_GPT_4o = ModelInfo("gpt-4o", lambda: OpenAI())
-    OPENAI_GPT_4o_MINI_JSON = ModelInfo("gpt-4o-mini-2024-07-18", lambda: OpenAI())
+    OLLAMA_3_2_LOCAL = ModelInfo("llama3.2", lambda: OpenAI(base_url="http://localhost:11434/v1", api_key='ollama'),
+                                 OpenAiTokenizer("gpt-4o"), 128_000)
+    OPENAI_GPT_4o_MINI = ModelInfo("gpt-4o-mini", lambda: OpenAI(),
+                                   OpenAiTokenizer("gpt-4o"), 128_000)
+    OPENAI_GPT_4o = ModelInfo("gpt-4o", lambda: OpenAI(),
+                              OpenAiTokenizer("gpt-4o"), 128_000)
+    OPENAI_GPT_4o_MINI_JSON = ModelInfo("gpt-4o-mini-2024-07-18", lambda: OpenAI(),
+                                        OpenAiTokenizer("gpt-4o"), 128_000)

igbot_base/tokenizer.py

@@ -0,0 +1,16 @@
+from abc import ABC, abstractmethod
+import tiktoken
+
+
+class BaseTokenizer(ABC):
+    @abstractmethod
+    def count_tokens(self, text: str):
+        pass
+
+
+class OpenAiTokenizer(BaseTokenizer):
+    def __init__(self, model_name):
+        self.__tokenizer = tiktoken.encoding_for_model(model_name)
+
+    def count_tokens(self, text: str):
+        return len(self.__tokenizer.encode(text))

{igbot_base-0.0.19.dist-info → igbot_base-0.0.20.dist-info}/METADATA

@@ -1,8 +1,9 @@
 Metadata-Version: 2.4
 Name: igbot_base
-Version: 0.0.19
+Version: 0.0.20
 Summary: Base classes for igbot
 Author-email: Igor Kopeć <igor.kopec95@gmail.com>
 License: LGPL-3.0-or-later
 Requires-Python: >=3.12
 Requires-Dist: openai
+Requires-Dist: tiktoken

{igbot_base-0.0.19.dist-info → igbot_base-0.0.20.dist-info}/RECORD

@@ -6,14 +6,15 @@ igbot_base/exception_handler.py,sha256=Er5LbNt2ga2-e0GSzkJw66E-NS1dN_2msTRyyac1q
 igbot_base/llm.py,sha256=0zul2Q3aiCpdhcLbTcpdSPkLdX7Dvwhhiv2V_KH7YRo,2157
 igbot_base/llmmemory.py,sha256=I6eSBPW1DK8UNiHhayon-DLbSbWSinO_EwM5S8C3tHA,1090
 igbot_base/logging_adapter.py,sha256=kMuAaQPCba2luk0-WPxAbPNJaBL1yZsaff35ltKWgww,256
-igbot_base/models.py,sha256=15bitUZy4Z4JhZbLd_SMx5pD7m8YuZiIq-wYGumf0ws,622
+igbot_base/models.py,sha256=n3mxd9-WjVuLbTG3bWDVRauc8YhhBZew4_0Is9oQQgk,1215
 igbot_base/persistable_memory.py,sha256=867Z4l_awD23QSTFDK3uN9TEttKfDpxx9g19c7ULG1Y,167
 igbot_base/prompt_template.py,sha256=q7bnU2Rp2Cvm4ixEmKHCyuolTPan6CHfYVAomT5kp10,1573
 igbot_base/response_formats.py,sha256=Sp679VfkcyxsR19ddFEHTOM7Ox0u2WBhdji_Y9pVc2M,94
 igbot_base/retriever.py,sha256=YHBHxiuPttsIAtjwwIz0fa9q2uFtsg0WoMRCorCt1PM,635
+igbot_base/tokenizer.py,sha256=VfCBZeI8DLJalguUc1e9XdeCmK8VAgyMHxssv0D5Y1M,385
 igbot_base/tool.py,sha256=ghjuU3cK_VElHGStUl3IYf9dlRDOGM_aQucUQNUH3mU,431
 igbot_base/vectorstore.py,sha256=f-S9DmOuehC0rtVAPGz22Guo7ddb1jVw7BygzFJM624,928
-igbot_base-0.0.19.dist-info/METADATA,sha256=hoqGiJOM-oEgXfXMtmdsH5N22MVDJICpcZT35CUMASY,211
-igbot_base-0.0.19.dist-info/WHEEL,sha256=lTU6B6eIfYoiQJTZNc-fyaR6BpL6ehTzU3xGYxn2n8k,91
-igbot_base-0.0.19.dist-info/top_level.txt,sha256=RGpEN0pLsnfNLoLfpw1D_nLSyZwi3ggSaDq07o4ZEAI,11
-igbot_base-0.0.19.dist-info/RECORD,,
+igbot_base-0.0.20.dist-info/METADATA,sha256=q_M4Ul2sCNwaeI2vwC7725Nmd9rdsVpns5D5pMT-tC8,235
+igbot_base-0.0.20.dist-info/WHEEL,sha256=ck4Vq1_RXyvS4Jt6SI0Vz6fyVs4GWg7AINwpsaGEgPE,91
+igbot_base-0.0.20.dist-info/top_level.txt,sha256=RGpEN0pLsnfNLoLfpw1D_nLSyZwi3ggSaDq07o4ZEAI,11
+igbot_base-0.0.20.dist-info/RECORD,,

{igbot_base-0.0.19.dist-info → igbot_base-0.0.20.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (78.1.1)
+Generator: setuptools (80.0.0)
 Root-Is-Purelib: true
 Tag: py3-none-any