Add count tokens param in base class

kagisearch · Feb 28, 2024 · 9ffa28c · 9ffa28c
1 parent cfde824
commit 9ffa28c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/llms/providers/base_provider.py b/llms/providers/base_provider.py
@@ -44,12 +44,12 @@ def compute_cost(self, prompt_tokens: int, completion_tokens: int) -> float:
         cost = round(cost, 5)
         return cost
 
-    def count_tokens(self):
+    def count_tokens(self, content):
         raise NotImplementedError(
             f"Count tokens is currently not supported with {self.__name__}"
         )
 
-    def complete(self):
+    def complete(self, *args, **kwargs):
         raise NotImplementedError
 
     async def acomplete(self):