Merge pull request #17 from pamelafox/gpt-4o-mini

Add gpt4o mini support

Author: pamelafox

.github/workflows/python.yaml

@@ -22,7 +22,7 @@ jobs:
             python-version: ${{ matrix.python_version }}
         - name: Install dependencies
           run: |
-            python -m pip install --upgrade pip
+            python3 -m pip install --upgrade pip
             python3 -m pip install -e '.[dev]'
         - name: Lint with ruff
           run: ruff check .
@@ -32,4 +32,4 @@ jobs:
           run: |
             python3 -m pytest -s -vv --cov --cov-fail-under=97
         - name: Run type checks
-          run: mypy .
+          run: python3 -m mypy .

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 All notable changes to this project will be documented in this file.
 
+## [0.1.9] - Aug 7, 2024
+
+- Add gpt-4o-mini support, by adding a 33.3x multiplier to the token cost.
+
 ## [0.1.8] - Aug 3, 2024
 
 - Fix the type for the tool_choice param to be inclusive of "auto" and other options.

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "openai-messages-token-helper"
 description = "A helper library for estimating tokens used by messages sent through OpenAI Chat Completions API."
-version = "0.1.8"
+version = "0.1.9"
 authors = [{name = "Pamela Fox"}]
 requires-python = ">=3.9"
 readme = "README.md"

src/openai_messages_token_helper/images_helper.py

@@ -1,7 +1,9 @@
 import base64
 import math
 import re
+from fractions import Fraction
 from io import BytesIO
+from typing import Optional
 
 from PIL import Image
 
@@ -16,20 +18,23 @@ def get_image_dims(image_uri: str) -> tuple[int, int]:
         raise ValueError("Image must be a base64 string.")
 
 
-def count_tokens_for_image(image_uri: str, detail: str = "auto") -> int:
+def count_tokens_for_image(image_uri: str, detail: str = "auto", model: Optional[str] = None) -> int:
     # From https://github.com/openai/openai-cookbook/pull/881/files
     # Based on https://platform.openai.com/docs/guides/vision
-    LOW_DETAIL_COST = 85
-    HIGH_DETAIL_COST_PER_TILE = 170
-    ADDITIONAL_COST = 85
+    multiplier = Fraction(1, 1)
+    if model == "gpt-4o-mini":
+        multiplier = Fraction(100, 3)
+    COST_PER_TILE = 85 * multiplier
+    LOW_DETAIL_COST = COST_PER_TILE
+    HIGH_DETAIL_COST_PER_TILE = COST_PER_TILE * 2
 
     if detail == "auto":
         # assume high detail for now
         detail = "high"
 
     if detail == "low":
         # Low detail images have a fixed cost
-        return LOW_DETAIL_COST
+        return int(LOW_DETAIL_COST)
     elif detail == "high":
         # Calculate token cost for high detail images
         width, height = get_image_dims(image_uri)
@@ -47,8 +52,8 @@ def count_tokens_for_image(image_uri: str, detail: str = "auto") -> int:
         # Calculate the number of 512px squares
         num_squares = math.ceil(width / 512) * math.ceil(height / 512)
         # Calculate the total token cost
-        total_cost = num_squares * HIGH_DETAIL_COST_PER_TILE + ADDITIONAL_COST
-        return total_cost
+        total_cost = num_squares * HIGH_DETAIL_COST_PER_TILE + COST_PER_TILE
+        return math.ceil(total_cost)
     else:
         # Invalid detail_option
         raise ValueError("Invalid value for detail parameter. Use 'low' or 'high'.")

src/openai_messages_token_helper/message_builder.py

@@ -126,7 +126,7 @@ def build_messages(
         if tool_calls is not None and not isinstance(tool_calls, Iterable):
             raise ValueError("tool_calls must be a list of tool calls")
         message_builder.insert_message(
-            shot["role"], shot.get("content"), tool_calls=tool_calls, tool_call_id=tool_call_id
+            shot["role"], shot.get("content"), tool_calls=tool_calls, tool_call_id=tool_call_id  # type: ignore[arg-type]
         )
 
     append_index = len(few_shots)
@@ -149,6 +149,6 @@ def build_messages(
 
         if message["role"] is None or message["content"] is None:
             raise ValueError("Few-shot messages must have both role and content")
-        message_builder.insert_message(message["role"], message["content"], index=append_index)
+        message_builder.insert_message(message["role"], message["content"], index=append_index)  # type: ignore[arg-type]
         total_token_count += potential_message_count
     return message_builder.all_messages

src/openai_messages_token_helper/model_helper.py

@@ -22,6 +22,7 @@
     "gpt-4-32k": 32000,
     "gpt-4v": 128000,
     "gpt-4o": 128000,
+    "gpt-4o-mini": 128000,
 }
 
 
@@ -106,7 +107,7 @@ def count_tokens_for_message(model: str, message: ChatCompletionMessageParam, de
                 if item["type"] == "text":
                     num_tokens += len(encoding.encode(item["text"]))
                 elif item["type"] == "image_url":
-                    num_tokens += count_tokens_for_image(item["image_url"]["url"], item["image_url"]["detail"])
+                    num_tokens += count_tokens_for_image(item["image_url"]["url"], item["image_url"]["detail"], model)
         elif isinstance(value, str):
             num_tokens += len(encoding.encode(value))
         else:

tests/image_messages.py

@@ -146,5 +146,4 @@
     assistant_message_perf,
     assistant_message_perf_short,
     assistant_message_dresscode,
-    text_and_image_message,
 ]

tests/messages.py

@@ -18,6 +18,7 @@ def large_image():
 
 def test_count_tokens_for_image(small_image, large_image):
     assert count_tokens_for_image(small_image, "low") == 85
+    assert count_tokens_for_image(small_image, "low", "gpt-4o-mini") == 2833
     assert count_tokens_for_image(small_image, "high") == 255
     assert count_tokens_for_image(small_image) == 255
     assert count_tokens_for_image(large_image, "low") == 85

tests/test_imageshelper.py

@@ -9,14 +9,14 @@
 from openai_messages_token_helper import build_messages, count_tokens_for_message
 
 from .functions import search_sources_toolchoice_auto
+from .image_messages import text_and_tiny_image_message
 from .messages import (
     assistant_message_dresscode,
     assistant_message_perf,
     assistant_message_perf_short,
     system_message_long,
     system_message_short,
     system_message_unicode,
-    text_and_image_message,
     user_message,
     user_message_dresscode,
     user_message_perf,
@@ -35,9 +35,9 @@ def test_messagebuilder_imagemessage():
     messages = build_messages(
         "gpt-35-turbo",
         system_message_short["message"]["content"],
-        new_user_content=text_and_image_message["message"]["content"],
+        new_user_content=text_and_tiny_image_message["message"]["content"],
     )
-    assert messages == [system_message_short["message"], text_and_image_message["message"]]
+    assert messages == [system_message_short["message"], text_and_tiny_image_message["message"]]
 
 
 def test_messagebuilder_append():