merge all token tracker (#111)

TengHu · Dec 30, 2023 · 07290db · 07290db
1 parent 111ad3a
commit 07290db
Show file tree

Hide file tree

Showing 7 changed files with 4 additions and 135 deletions.
diff --git a/actionweaver/llms/azure/chat.py b/actionweaver/llms/azure/chat.py
@@ -16,10 +16,10 @@
 
 from actionweaver.actions.action import Action, ActionHandlers
 from actionweaver.llms.azure.functions import Functions
-from actionweaver.llms.azure.tokens import TokenUsageTracker
 from actionweaver.telemetry import traceable
 from actionweaver.utils import DEFAULT_ACTION_SCOPE
 from actionweaver.utils.stream import get_first_element_and_iterator, merge_dicts
+from actionweaver.utils.tokens import TokenUsageTracker
 
 # TODO: support AsyncAzureOpenAI
 

diff --git a/actionweaver/llms/general/tokens.py b/actionweaver/llms/general/tokens.py
diff --git a/actionweaver/llms/openai/functions/chat.py b/actionweaver/llms/openai/functions/chat.py
@@ -14,9 +14,9 @@
 
 from actionweaver.actions.action import Action, ActionHandlers
 from actionweaver.llms.openai.functions.functions import Functions
-from actionweaver.llms.openai.functions.tokens import TokenUsageTracker
 from actionweaver.utils import DEFAULT_ACTION_SCOPE
 from actionweaver.utils.stream import get_first_element_and_iterator, merge_dicts
+from actionweaver.utils.tokens import TokenUsageTracker
 
 
 class OpenAIChatCompletionException(Exception):
@@ -27,9 +27,7 @@ class OpenAIChatCompletion:
     def __init__(self, model, token_usage_tracker=None, logger=None):
         self.model = model
         self.logger = logger or logging.getLogger(__name__)
-        self.token_usage_tracker = token_usage_tracker or TokenUsageTracker(
-            logger=logger
-        )
+        self.token_usage_tracker = token_usage_tracker or TokenUsageTracker()
         self.client = OpenAI()
 
         print(

diff --git a/actionweaver/llms/openai/functions/tokens.py b/actionweaver/llms/openai/functions/tokens.py
diff --git a/actionweaver/llms/openai/tools/chat.py b/actionweaver/llms/openai/tools/chat.py
@@ -14,11 +14,11 @@
 )
 
 from actionweaver.actions.action import Action, ActionHandlers
-from actionweaver.llms.openai.tools.tokens import TokenUsageTracker
 from actionweaver.llms.openai.tools.tools import Tools
 from actionweaver.telemetry import traceable
 from actionweaver.utils import DEFAULT_ACTION_SCOPE
 from actionweaver.utils.stream import get_first_element_and_iterator, merge_dicts
+from actionweaver.utils.tokens import TokenUsageTracker
 
 
 class OpenAIChatCompletionException(Exception):

diff --git a/actionweaver/llms/openai/tools/tokens.py b/actionweaver/llms/openai/tools/tokens.py
diff --git a/actionweaver/llms/azure/tokens.py → actionweaver/utils/tokens.py b/actionweaver/llms/azure/tokens.py → actionweaver/utils/tokens.py
@@ -1,6 +1,4 @@
 import collections
-import logging
-import time
 from typing import Dict