Add token count message

Author: johann-petrak

llms_wrapper/llms.py

@@ -14,7 +14,7 @@
 from typing import Optional, Dict, List, Union, Tuple, Callable
 from copy import deepcopy
 
-from litellm import completion, completion_cost
+from litellm import completion, completion_cost, token_counter
 from litellm.utils import get_model_info, get_supported_openai_params, supports_response_schema
 from litellm.utils import supports_function_calling, supports_parallel_function_calling
 from llms_wrapper.utils import dict_except
@@ -402,6 +402,15 @@ def supports_file_upload(self, llmalias: str) -> bool:
         """
         return False
 
+    def count_tokens(self, llmalias: str, messages: List[Dict[str, str]]) -> int:
+        """
+        Count the number of tokens in the given messages. If messages is a string, convert it to a
+        single user message first.
+        """
+        if isinstance(messages, str):
+            messages = [{"role": "user", "content": messages}]
+        return token_counter(model=self.llms[llmalias]["llm"], messages=messages)
+
     def query(
             self,
             llmalias: str,
@@ -604,3 +613,8 @@ def supports_json_schema(self) -> bool:
     def supports_function_calling(self, parallel=False) -> bool:
         return self.llmsobject.supports_function_calling(self.config["alias"], parallel)
 
+    def supports_file_upload(self) -> bool:
+        return self.llmsobject.supports_file_upload(self.config["alias"])
+
+    def count_tokens(self, messages: List[Dict[str, str]]) -> int:
+        return self.llmsobject.count_tokens(self.config["alias"], messages)

llms_wrapper/version.py

@@ -1,3 +1,3 @@
 import importlib.metadata
-__version__ = "0.1.30"
+__version__ = "0.1.31"
 

notebooks/test-api.ipynb

@@ -226,6 +226,7 @@
        " 'api_key_env': 'MY_GEMINI_API_KEY',\n",
        " 'alias': 'gemini/somemodel',\n",
        " '_cost': 0,\n",
+       " '_last_request_time': 0,\n",
        " '_elapsed_time': 0}"
       ]
      },
@@ -321,6 +322,79 @@
     "msg1 = llms.make_messages(\"What is a monoid?\")"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "abac6cf6-7175-4cbb-8527-d24c8847a28b",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "13"
+      ]
+     },
+     "execution_count": 15,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "llms.count_tokens(\"openai/gpt-4.5-preview\", msg1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "a8e72aa6-c1be-4631-86cf-3b3f820c6059",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "llm = llms[\"openai/gpt-4o\"]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "452600c9-ee93-48e3-bbf7-c3ab5eb51ade",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "13"
+      ]
+     },
+     "execution_count": 17,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "llm.count_tokens(msg1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "5a9a8a26-78a6-40ff-b104-24028b1c424c",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "13"
+      ]
+     },
+     "execution_count": 18,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "llm.count_tokens(\"What is a monoid?\")"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 15,
@@ -330,15 +404,15 @@
     {
      "data": {
       "text/plain": [
-       "{'elapsed_time': 21.85033893585205,\n",
-       " 'response': ModelResponse(id='chatcmpl-B84wXy8ULKMSfurBDArzoYPsLaElv', created=1741265325, model='gpt-4.5-preview-2025-02-27', object='chat.completion', system_fingerprint=None, choices=[Choices(finish_reason='stop', index=0, message=Message(content='In abstract algebra, a **monoid** is an algebraic structure consisting of:\\n\\n1. **A set \\\\(M\\\\)**,\\n2. **An associative binary operation**, usually denoted by multiplication (\\\\(\\\\cdot\\\\)), such that \\\\(a \\\\cdot (b \\\\cdot c) = (a \\\\cdot b) \\\\cdot c\\\\) for all \\\\(a, b, c \\\\in M\\\\),\\n3. **An identity element**, typically called unity or identity (often denoted by \\\\(e\\\\)), which satisfies:\\n\\\\[\\ne \\\\cdot a = a \\\\cdot e = a\\n\\\\]\\nfor all \\\\(a \\\\in M\\\\).\\n\\nSimply stated, a monoid is a set equipped with an associative binary operation and a neutral element.\\n\\n**Examples:**\\n\\n- **Natural numbers and addition**: The set \\\\(\\\\mathbb{N}_0 = \\\\{0, 1, 2, \\\\dots\\\\}\\\\) with addition is a monoid. Here, 0 acts as identity element.\\n- **Strings with concatenation**: The set of all strings over an alphabet, combined with concatenation (joining strings end-to-end), forms a monoid. The identity element here is the empty string.\\n- **Multiplication of integers**: The set of integers \\\\(\\\\mathbb{Z}\\\\) is a monoid under multiplication with identity element \\\\(1\\\\).\\n\\n**Why are monoids important?**\\n\\nMonoids appear naturally across many disciplines in mathematics and computer science. In particular:\\n\\n- **Formal language theory and computation**: Monoids arise naturally in studying automata, languages, and formal grammars.\\n- **Category theory**: Monoids can be seen as categories with a single object.\\n- **Parallel and distributed computing**: Monoids are useful for combining, aggregating, and reducing data in parallel computations.\\n\\nIn short, monoids provide an important algebraic abstraction capturing the concepts of \"composable processes\" and \"units of computation\".', role='assistant', tool_calls=None, function_call=None, provider_specific_fields={'refusal': None}))], usage=Usage(completion_tokens=406, prompt_tokens=13, total_tokens=419, completion_tokens_details=CompletionTokensDetailsWrapper(accepted_prediction_tokens=0, audio_tokens=0, reasoning_tokens=0, rejected_prediction_tokens=0, text_tokens=None), prompt_tokens_details=PromptTokensDetailsWrapper(audio_tokens=0, cached_tokens=0, text_tokens=None, image_tokens=None)), service_tier='default'),\n",
+       "{'elapsed_time': 39.25307822227478,\n",
+       " 'response': ModelResponse(id='chatcmpl-BFI6FbX7BtguunCLPpOR1KqsexGuk', created=1742984195, model='gpt-4.5-preview-2025-02-27', object='chat.completion', system_fingerprint=None, choices=[Choices(finish_reason='stop', index=0, message=Message(content='A **monoid** is a fundamental algebraic structure in mathematics. Formally, it\\'s defined as a set equipped with an associative binary operation and an identity element.\\n\\nSpecifically, a monoid is a triple \\\\((M, \\\\cdot, e)\\\\), where:\\n\\n1. **\\\\(M\\\\)** is a set.\\n2. **\\\\(\\\\cdot\\\\)** is a binary operation: \\n   \\\\[\\n   \\\\cdot : M \\\\times M \\\\rightarrow M\\n   \\\\]\\n   satisfying the associativity property:\\n   \\\\[\\n   (a \\\\cdot b) \\\\cdot c = a \\\\cdot (b \\\\cdot c) \\\\quad \\\\text{for every } a, b, c \\\\in M.\\n   \\\\]\\n3. **\\\\(e\\\\)** is a special element in \\\\(M\\\\), called the identity element, such that for every element \\\\(a\\\\) in \\\\(M\\\\):\\n   \\\\[\\n   e \\\\cdot a = a \\\\cdot e = a.\\n   \\\\]\\n\\n**Examples of monoids:**\\n\\n- **Natural numbers or integers** under addition with identity element 0.\\n- **Natural numbers or positive reals** under multiplication with identity element 1.\\n- **Concatenation** of strings with empty string \"\" as identity.\\n- **Lists** with concatenation as the operation and the empty list as identity.\\n\\n**Relation to other algebraic structures:**\\n\\n- A **group** is a monoid in which every element has an inverse.\\n- A **semigroup** is a structure similar to a monoid but does not require an identity element.\\n\\nMonoids arise frequently in algebra, theoretical computer science, category theory, and functional programming, providing a simple yet versatile abstraction for combining or composing entities.', role='assistant', tool_calls=None, function_call=None, provider_specific_fields={'refusal': None, 'annotations': []}))], usage=Usage(completion_tokens=360, prompt_tokens=13, total_tokens=373, completion_tokens_details=CompletionTokensDetailsWrapper(accepted_prediction_tokens=0, audio_tokens=0, reasoning_tokens=0, rejected_prediction_tokens=0, text_tokens=None), prompt_tokens_details=PromptTokensDetailsWrapper(audio_tokens=0, cached_tokens=0, text_tokens=None, image_tokens=None)), service_tier='default'),\n",
        " 'kwargs': {},\n",
-       " 'cost': 0.06187499999999999,\n",
-       " 'n_completion_tokens': 406,\n",
+       " 'cost': 0.05497499999999999,\n",
+       " 'n_completion_tokens': 360,\n",
        " 'n_prompt_tokens': 13,\n",
-       " 'n_total_tokens': 419,\n",
+       " 'n_total_tokens': 373,\n",
        " 'finish_reason': 'stop',\n",
-       " 'answer': 'In abstract algebra, a **monoid** is an algebraic structure consisting of:\\n\\n1. **A set \\\\(M\\\\)**,\\n2. **An associative binary operation**, usually denoted by multiplication (\\\\(\\\\cdot\\\\)), such that \\\\(a \\\\cdot (b \\\\cdot c) = (a \\\\cdot b) \\\\cdot c\\\\) for all \\\\(a, b, c \\\\in M\\\\),\\n3. **An identity element**, typically called unity or identity (often denoted by \\\\(e\\\\)), which satisfies:\\n\\\\[\\ne \\\\cdot a = a \\\\cdot e = a\\n\\\\]\\nfor all \\\\(a \\\\in M\\\\).\\n\\nSimply stated, a monoid is a set equipped with an associative binary operation and a neutral element.\\n\\n**Examples:**\\n\\n- **Natural numbers and addition**: The set \\\\(\\\\mathbb{N}_0 = \\\\{0, 1, 2, \\\\dots\\\\}\\\\) with addition is a monoid. Here, 0 acts as identity element.\\n- **Strings with concatenation**: The set of all strings over an alphabet, combined with concatenation (joining strings end-to-end), forms a monoid. The identity element here is the empty string.\\n- **Multiplication of integers**: The set of integers \\\\(\\\\mathbb{Z}\\\\) is a monoid under multiplication with identity element \\\\(1\\\\).\\n\\n**Why are monoids important?**\\n\\nMonoids appear naturally across many disciplines in mathematics and computer science. In particular:\\n\\n- **Formal language theory and computation**: Monoids arise naturally in studying automata, languages, and formal grammars.\\n- **Category theory**: Monoids can be seen as categories with a single object.\\n- **Parallel and distributed computing**: Monoids are useful for combining, aggregating, and reducing data in parallel computations.\\n\\nIn short, monoids provide an important algebraic abstraction capturing the concepts of \"composable processes\" and \"units of computation\".',\n",
+       " 'answer': 'A **monoid** is a fundamental algebraic structure in mathematics. Formally, it\\'s defined as a set equipped with an associative binary operation and an identity element.\\n\\nSpecifically, a monoid is a triple \\\\((M, \\\\cdot, e)\\\\), where:\\n\\n1. **\\\\(M\\\\)** is a set.\\n2. **\\\\(\\\\cdot\\\\)** is a binary operation: \\n   \\\\[\\n   \\\\cdot : M \\\\times M \\\\rightarrow M\\n   \\\\]\\n   satisfying the associativity property:\\n   \\\\[\\n   (a \\\\cdot b) \\\\cdot c = a \\\\cdot (b \\\\cdot c) \\\\quad \\\\text{for every } a, b, c \\\\in M.\\n   \\\\]\\n3. **\\\\(e\\\\)** is a special element in \\\\(M\\\\), called the identity element, such that for every element \\\\(a\\\\) in \\\\(M\\\\):\\n   \\\\[\\n   e \\\\cdot a = a \\\\cdot e = a.\\n   \\\\]\\n\\n**Examples of monoids:**\\n\\n- **Natural numbers or integers** under addition with identity element 0.\\n- **Natural numbers or positive reals** under multiplication with identity element 1.\\n- **Concatenation** of strings with empty string \"\" as identity.\\n- **Lists** with concatenation as the operation and the empty list as identity.\\n\\n**Relation to other algebraic structures:**\\n\\n- A **group** is a monoid in which every element has an inverse.\\n- A **semigroup** is a structure similar to a monoid but does not require an identity element.\\n\\nMonoids arise frequently in algebra, theoretical computer science, category theory, and functional programming, providing a simple yet versatile abstraction for combining or composing entities.',\n",
        " 'error': '',\n",
        " 'ok': True}"
       ]