Change token counter implementation & support embedding token counting (defaults to tiktoken) (#92)

* Change token counter implementation & support embedding token counting (defaults to tiktoken)

* Refactor code

* Add a todo

Author: gaya3-zipstack

pdm.lock

@@ -48,7 +48,7 @@ dependencies = [
     "llama-index-llms-vertex==0.2.2",
     "llama-index-llms-replicate==0.1.3",
     "llama-index-llms-ollama==0.2.2",
-    "llama-index-llms-bedrock==0.1.12",
+    "llama-index-llms-bedrock==0.1.13",
     # For Llama Parse X2Text
     "llama-parse==0.4.9",
     # OCR

pyproject.toml

@@ -1,4 +1,4 @@
-__version__ = "0.45.2"
+__version__ = "0.46.0"
 
 
 def get_sdk_version():

src/unstract/sdk/__init__.py

@@ -2,6 +2,7 @@
 
 from deprecated import deprecated
 from llama_index.core.base.embeddings.base import Embedding
+from llama_index.core.callbacks import CallbackManager as LlamaIndexCallbackManager
 from llama_index.core.embeddings import BaseEmbedding
 
 from unstract.sdk.adapter import ToolAdapter
@@ -104,6 +105,17 @@ def get_class_name(self) -> str:
         """
         return self._embedding_instance.class_name()
 
+    def get_callback_manager(self) -> LlamaIndexCallbackManager:
+        """Gets the llama-index callback manager set on the model.
+
+        Args:
+            NA
+
+            Returns:
+                llama-index callback manager
+        """
+        return self._embedding_instance.callback_manager
+
     @deprecated("Use Embedding instead of ToolEmbedding")
     def get_embedding_length(self, embedding: BaseEmbedding) -> int:
         return self._get_embedding_length()

src/unstract/sdk/embedding.py

@@ -289,7 +289,9 @@ def index(
             try:
                 if chunk_size == 0:
                     parser = SimpleNodeParser.from_defaults(
-                        chunk_size=len(documents[0].text) + 10, chunk_overlap=0
+                        chunk_size=len(documents[0].text) + 10,
+                        chunk_overlap=0,
+                        callback_manager=embedding.get_callback_manager(),
                     )
                     nodes = parser.get_nodes_from_documents(
                         documents, show_progress=True
@@ -301,7 +303,9 @@ def index(
                 else:
                     storage_context = vector_db.get_storage_context()
                     parser = SimpleNodeParser.from_defaults(
-                        chunk_size=chunk_size, chunk_overlap=chunk_overlap
+                        chunk_size=chunk_size,
+                        chunk_overlap=chunk_overlap,
+                        callback_manager=embedding.get_callback_manager(),
                     )
                     self.tool.stream_log("Adding nodes to vector db...")
                     # TODO: Phase 2:
@@ -320,6 +324,7 @@ def index(
                         show_progress=True,
                         embed_model=embedding,
                         node_parser=parser,
+                        callback_manager=embedding.get_callback_manager(),
                     )
             except Exception as e:
                 self.tool.stream_log(

src/unstract/sdk/index.py

@@ -1,24 +1,10 @@
-from typing import Any
+from typing import Any, Union
 
 from llama_index.core.callbacks.schema import EventPayload
-from llama_index.core.utilities.token_counting import TokenCounter
-from openai.types import CompletionUsage
-from openai.types.chat import ChatCompletion
+from llama_index.core.llms import ChatResponse, CompletionResponse
 
 
 class Constants:
-    KEY_USAGE = "usage"
-    KEY_USAGE_METADATA = "usage_metadata"
-    KEY_EVAL_COUNT = "eval_count"
-    KEY_PROMPT_EVAL_COUNT = "prompt_eval_count"
-    KEY_RAW_RESPONSE = "_raw_response"
-    KEY_TEXT_TOKEN_COUNT = "inputTextTokenCount"
-    KEY_TOKEN_COUNT = "tokenCount"
-    KEY_RESULTS = "results"
-    INPUT_TOKENS = "input_tokens"
-    OUTPUT_TOKENS = "output_tokens"
-    PROMPT_TOKENS = "prompt_tokens"
-    COMPLETION_TOKENS = "completion_tokens"
     DEFAULT_TOKEN_COUNT = 0
 
 
@@ -35,69 +21,25 @@ def __init__(self, input_tokens, output_tokens):
             self.prompt_llm_token_count + self.completion_llm_token_count
         )
 
+    # TODO: Add unit test cases for the following function
+    #  for ease of manintenance
     @staticmethod
-    def get_llm_token_counts(payload: dict[str, Any]) -> TokenCounter:
+    def get_llm_token_counts(payload: dict[str, Any]):
         prompt_tokens = Constants.DEFAULT_TOKEN_COUNT
         completion_tokens = Constants.DEFAULT_TOKEN_COUNT
         if EventPayload.PROMPT in payload:
-            completion_raw = payload.get(EventPayload.COMPLETION).raw
-            if completion_raw:
-                # For Open AI models, token count is part of ChatCompletion
-                if isinstance(completion_raw, ChatCompletion):
-                    if hasattr(completion_raw, Constants.KEY_USAGE):
-                        token_counts: dict[
-                            str, int
-                        ] = TokenCounter._get_prompt_completion_tokens(completion_raw)
-                        prompt_tokens = token_counts[Constants.PROMPT_TOKENS]
-                        completion_tokens = token_counts[Constants.COMPLETION_TOKENS]
-                # For other models
-                elif isinstance(completion_raw, dict):
-                    # For Gemini models
-                    if completion_raw.get(Constants.KEY_RAW_RESPONSE):
-                        if hasattr(
-                            completion_raw.get(Constants.KEY_RAW_RESPONSE),
-                            Constants.KEY_USAGE_METADATA,
-                        ):
-                            usage = completion_raw.get(
-                                Constants.KEY_RAW_RESPONSE
-                            ).usage_metadata
-                            prompt_tokens = usage.prompt_token_count
-                            completion_tokens = usage.candidates_token_count
-                    elif completion_raw.get(Constants.KEY_USAGE):
-                        token_counts: dict[
-                            str, int
-                        ] = TokenCounter._get_prompt_completion_tokens(completion_raw)
-                        prompt_tokens = token_counts[Constants.PROMPT_TOKENS]
-                        completion_tokens = token_counts[Constants.COMPLETION_TOKENS]
-                    # For Bedrock models
-                    elif Constants.KEY_TEXT_TOKEN_COUNT in completion_raw:
-                        prompt_tokens = completion_raw[Constants.KEY_TEXT_TOKEN_COUNT]
-                        if Constants.KEY_RESULTS in completion_raw:
-                            result_list: list = completion_raw[Constants.KEY_RESULTS]
-                            if len(result_list) > 0:
-                                result: dict = result_list[0]
-                                if Constants.KEY_TOKEN_COUNT in result:
-                                    completion_tokens = result.get(
-                                        Constants.KEY_TOKEN_COUNT
-                                    )
-                    else:
-                        if completion_raw.get(Constants.KEY_PROMPT_EVAL_COUNT):
-                            prompt_tokens = completion_raw.get(
-                                Constants.KEY_PROMPT_EVAL_COUNT
-                            )
-                        if completion_raw.get(Constants.KEY_EVAL_COUNT):
-                            completion_tokens = completion_raw.get(
-                                Constants.KEY_EVAL_COUNT
-                            )
-        # For Anthropic models
+            response = payload.get(EventPayload.COMPLETION)
+            (
+                prompt_tokens,
+                completion_tokens,
+            ) = TokenCounter._get_tokens_from_response(response)
         elif EventPayload.MESSAGES in payload:
-            response_raw = payload.get(EventPayload.RESPONSE).raw
-            if response_raw:
-                token_counts: dict[
-                    str, int
-                ] = TokenCounter._get_prompt_completion_tokens(response_raw)
-                prompt_tokens = token_counts[Constants.PROMPT_TOKENS]
-                completion_tokens = token_counts[Constants.COMPLETION_TOKENS]
+            response = payload.get(EventPayload.RESPONSE)
+            if response:
+                (
+                    prompt_tokens,
+                    completion_tokens,
+                ) = TokenCounter._get_tokens_from_response(response)
 
         token_counter = TokenCounter(
             input_tokens=prompt_tokens,
@@ -106,33 +48,72 @@ def get_llm_token_counts(payload: dict[str, Any]) -> TokenCounter:
         return token_counter
 
     @staticmethod
-    def _get_prompt_completion_tokens(response) -> dict[str, int]:
-        usage = None
-        prompt_tokens = Constants.DEFAULT_TOKEN_COUNT
-        completion_tokens = Constants.DEFAULT_TOKEN_COUNT
-        # For OpenAI models,response is an obj of CompletionUsage
-        if (
-            isinstance(response, ChatCompletion)
-            and hasattr(response, Constants.KEY_USAGE)
-            and isinstance(response.usage, CompletionUsage)
+    def _get_tokens_from_response(
+        response: Union[CompletionResponse, ChatResponse, dict]
+    ) -> tuple[int, int]:
+        """Get the token counts from a raw response."""
+        prompt_tokens, completion_tokens = 0, 0
+        if isinstance(response, CompletionResponse) or isinstance(
+            response, ChatResponse
         ):
-            usage = response.usage
-        # For LLM models other than OpenAI, response is a dict
-        elif isinstance(response, dict) and Constants.KEY_USAGE in response:
-            usage = response.get(Constants.KEY_USAGE)
+            raw_response = response.raw
+            if not isinstance(raw_response, dict):
+                raw_response = dict(raw_response)
+
+            usage = raw_response.get("usage", None)
+        if usage is None:
+            if (
+                hasattr(response, "additional_kwargs")
+                and "prompt_tokens" in response.additional_kwargs
+            ):
+                usage = response.additional_kwargs
+            elif hasattr(response, "raw"):
+                completion_raw = response.raw
+                if ("_raw_response" in completion_raw) and hasattr(
+                    completion_raw["_raw_response"], "usage_metadata"
+                ):
+                    usage = completion_raw["_raw_response"].usage_metadata
+                    prompt_tokens = usage.prompt_token_count
+                    completion_tokens = usage.candidates_token_count
+                    return prompt_tokens, completion_tokens
+                elif "inputTextTokenCount" in completion_raw:
+                    prompt_tokens = completion_raw["inputTextTokenCount"]
+                    if "results" in completion_raw:
+                        result_list: list = completion_raw["results"]
+                        if len(result_list) > 0:
+                            result: dict = result_list[0]
+                            if "tokenCount" in result:
+                                completion_tokens = result.get("tokenCount", 0)
+                    return prompt_tokens, completion_tokens
+                else:
+                    usage = response.raw
+            else:
+                usage = response
+
+        if not isinstance(usage, dict):
+            usage = usage.model_dump()
+
+        possible_input_keys = (
+            "prompt_tokens",
+            "input_tokens",
+            "prompt_eval_count",
+        )
+        possible_output_keys = (
+            "completion_tokens",
+            "output_tokens",
+            "eval_count",
+        )
 
-        if usage:
-            if hasattr(usage, Constants.INPUT_TOKENS):
-                prompt_tokens = usage.input_tokens
-            elif hasattr(usage, Constants.PROMPT_TOKENS):
-                prompt_tokens = usage.prompt_tokens
+        prompt_tokens = 0
+        for input_key in possible_input_keys:
+            if input_key in usage:
+                prompt_tokens = int(usage[input_key])
+                break
 
-            if hasattr(usage, Constants.OUTPUT_TOKENS):
-                completion_tokens = usage.output_tokens
-            elif hasattr(usage, Constants.COMPLETION_TOKENS):
-                completion_tokens = usage.completion_tokens
+        completion_tokens = 0
+        for output_key in possible_output_keys:
+            if output_key in usage:
+                completion_tokens = int(usage[output_key])
+                break
 
-        token_counts: dict[str, int] = dict()
-        token_counts[Constants.PROMPT_TOKENS] = prompt_tokens
-        token_counts[Constants.COMPLETION_TOKENS] = completion_tokens
-        return token_counts
+        return prompt_tokens, completion_tokens

src/unstract/sdk/utils/token_counter.py

@@ -124,6 +124,7 @@ def get_vector_store_index_from_storage_context(
         documents: Sequence[Document],
         storage_context: Optional[StorageContext] = None,
         show_progress: bool = False,
+        callback_manager=None,
         **kwargs,
     ) -> IndexType:
         if not self._embedding_instance:
@@ -135,6 +136,7 @@ def get_vector_store_index_from_storage_context(
             show_progress=show_progress,
             embed_model=self._embedding_instance,
             node_parser=parser,
+            callback_manager=callback_manager,
         )
 
     def get_vector_store_index(self, **kwargs: Any) -> VectorStoreIndex:
@@ -143,7 +145,7 @@ def get_vector_store_index(self, **kwargs: Any) -> VectorStoreIndex:
         return VectorStoreIndex.from_vector_store(
             vector_store=self._vector_db_instance,
             embed_model=self._embedding_instance,
-            kwargs=kwargs,
+            callback_manager=kwargs.get("callback_manager"),
         )
 
     def get_storage_context(self) -> StorageContext: