Improvements to token counting for images and type annotations (#1244)

* Improve token counting

* Update tests

* Adding pillow dependency to backend

* Update pillow-types

* Add source in comment

* Comments and types

* Adding comment

Author: pamelafox

app/backend/approaches/chatapproach.py

@@ -110,7 +110,7 @@ def get_messages_from_history(
 
         total_token_count = 0
         for existing_message in message_builder.messages:
-            total_token_count += message_builder.count_tokens_for_message(dict(existing_message))  # type: ignore
+            total_token_count += message_builder.count_tokens_for_message(existing_message)
 
         newest_to_oldest = list(reversed(history[:-1]))
         for message in newest_to_oldest:

app/backend/core/imageshelper.py

@@ -1,8 +1,12 @@
 import base64
+import math
 import os
+import re
+from io import BytesIO
 from typing import Optional
 
 from azure.storage.blob.aio import ContainerClient
+from PIL import Image
 from typing_extensions import Literal, Required, TypedDict
 
 from approaches.approach import Document
@@ -34,3 +38,51 @@ async def fetch_image(blob_container_client: ContainerClient, result: Document)
         else:
             return None
     return None
+
+
+def get_image_dims(image_uri: str) -> tuple[int, int]:
+    # From https://github.com/openai/openai-cookbook/pull/881/files
+    if re.match(r"data:image\/\w+;base64", image_uri):
+        image_uri = re.sub(r"data:image\/\w+;base64,", "", image_uri)
+        image = Image.open(BytesIO(base64.b64decode(image_uri)))
+        return image.size
+    else:
+        raise ValueError("Image must be a base64 string.")
+
+
+def calculate_image_token_cost(image_uri: str, detail: str = "auto") -> 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
+
+    if detail == "auto":
+        # assume high detail for now
+        detail = "high"
+
+    if detail == "low":
+        # Low detail images have a fixed cost
+        return LOW_DETAIL_COST
+    elif detail == "high":
+        # Calculate token cost for high detail images
+        width, height = get_image_dims(image_uri)
+        # Check if resizing is needed to fit within a 2048 x 2048 square
+        if max(width, height) > 2048:
+            # Resize dimensions to fit within a 2048 x 2048 square
+            ratio = 2048 / max(width, height)
+            width = int(width * ratio)
+            height = int(height * ratio)
+        # Further scale down to 768px on the shortest side
+        if min(width, height) > 768:
+            ratio = 768 / min(width, height)
+            width = int(width * ratio)
+            height = int(height * ratio)
+        # 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
+    else:
+        # Invalid detail_option
+        raise ValueError("Invalid value for detail parameter. Use 'low' or 'high'.")

app/backend/core/messagebuilder.py

@@ -1,4 +1,5 @@
 import unicodedata
+from collections.abc import Mapping
 from typing import List, Union
 
 from openai.types.chat import (
@@ -52,7 +53,7 @@ def insert_message(self, role: str, content: Union[str, List[ChatCompletionConte
             raise ValueError(f"Invalid role: {role}")
         self.messages.insert(index, message)
 
-    def count_tokens_for_message(self, message: dict[str, str]):
+    def count_tokens_for_message(self, message: Mapping[str, object]):
         return num_tokens_from_messages(message, self.model)
 
     def normalize_content(self, content: Union[str, List[ChatCompletionContentPartParam]]):

app/backend/core/modelhelper.py

@@ -1,7 +1,11 @@
 from __future__ import annotations
 
+from collections.abc import Mapping
+
 import tiktoken
 
+from .imageshelper import calculate_image_token_cost
+
 MODELS_2_TOKEN_LIMITS = {
     "gpt-35-turbo": 4000,
     "gpt-3.5-turbo": 4000,
@@ -22,11 +26,11 @@ def get_token_limit(model_id: str) -> int:
     return MODELS_2_TOKEN_LIMITS[model_id]
 
 
-def num_tokens_from_messages(message: dict[str, str], model: str) -> int:
+def num_tokens_from_messages(message: Mapping[str, object], model: str) -> int:
     """
     Calculate the number of tokens required to encode a message.
     Args:
-        message (dict): The message to encode, represented as a dictionary.
+        message (Mapping): The message to encode, in a dictionary-like object.
         model (str): The name of the model to use for encoding.
     Returns:
         int: The total number of tokens required to encode the message.
@@ -39,14 +43,19 @@ def num_tokens_from_messages(message: dict[str, str], model: str) -> int:
 
     encoding = tiktoken.encoding_for_model(get_oai_chatmodel_tiktok(model))
     num_tokens = 2  # For "role" and "content" keys
-    for key, value in message.items():
+    for value in message.values():
         if isinstance(value, list):
-            for v in value:
-                # TODO: Update token count for images https://github.com/openai/openai-cookbook/pull/881/files
-                if isinstance(v, str):
-                    num_tokens += len(encoding.encode(v))
-        else:
+            # For GPT-4-vision support, based on https://github.com/openai/openai-cookbook/pull/881/files
+            for item in value:
+                num_tokens += len(encoding.encode(item["type"]))
+                if item["type"] == "text":
+                    num_tokens += len(encoding.encode(item["text"]))
+                elif item["type"] == "image_url":
+                    num_tokens += calculate_image_token_cost(item["image_url"]["url"], item["image_url"]["detail"])
+        elif isinstance(value, str):
             num_tokens += len(encoding.encode(value))
+        else:
+            raise ValueError(f"Could not encode unsupported message value type: {type(value)}")
     return num_tokens
 
 

app/backend/requirements.in

@@ -17,3 +17,5 @@ msal
 azure-keyvault-secrets
 cryptography
 python-jose[cryptography]
+Pillow
+types-Pillow

app/backend/requirements.txt

@@ -257,6 +257,8 @@ pandas==2.2.0
     # via openai
 pandas-stubs==2.1.4.231227
     # via openai
+pillow==10.2.0
+    # via -r requirements.in
 portalocker==2.8.2
     # via msal-extensions
 priority==2.0.0
@@ -315,6 +317,8 @@ tiktoken==0.5.2
     # via -r requirements.in
 tqdm==4.66.1
     # via openai
+types-pillow==10.2.0.20240206
+    # via -r requirements.in
 types-pytz==2023.4.0.20240130
     # via pandas-stubs
 typing-extensions==4.9.0

scripts/requirements.txt

@@ -174,7 +174,7 @@ tiktoken==0.5.2
     # via -r requirements.in
 tqdm==4.66.1
     # via openai
-types-pillow==10.2.0.20240125
+types-pillow==10.2.0.20240206
     # via -r requirements.in
 types-pytz==2023.4.0.20240130
     # via pandas-stubs

tests/image_large.png

@@ -28,7 +28,7 @@ def __init__(self):
         )
 
     async def readall(self):
-        return b"\x89PNG\x50\x4e\x47\x0d\x0a\x1a\x0a\x00\x00\x00\x0d\x49\x48\x44\x52\x00\x00\x00\x01\x00\x00\x00\x01\x01\x00\x00\x00\x00\x37\x6e\xf9\x24\x00\x00\x00\x0a\x49\x44\x41\x54\x78\x9c\x63\x00\x01\x00\x00\x05\x00\x01\x0d\x0d\x2d\xba\x1b\x00\x00\x00\x00\x49\x45\x4e\x44\xae\x42\x60\x82"
+        return b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x06\x00\x00\x00\x1f\x15\xc4\x89\x00\x00\x00\rIDATx\xdac\xfc\xcf\xf0\xbf\x1e\x00\x06\x83\x02\x7f\x94\xad\xd0\xeb\x00\x00\x00\x00IEND\xaeB`\x82"
 
 
 class MockKeyVaultSecret:

tests/mocks.py

@@ -6,7 +6,7 @@
                     "images": [
                         {
                             "detail": "auto",
-                            "url": "data:image/png;base64,iVBOR1BORw0KGgoAAAANSUhEUgAAAAEAAAABAQAAAAA3bvkkAAAACklEQVR4nGMAAQAABQABDQ0tuhsAAAAASUVORK5CYII="
+                            "url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z/C/HgAGgwJ/lK3Q6wAAAABJRU5ErkJggg=="
                         }
                     ],
                     "text": [
@@ -46,7 +46,7 @@
                     {
                         "description": [
                             "{'role': 'system', 'content': \"You are an intelligent assistant helping analyze the Annual Financial Report of Contoso Ltd., The documents contain text, graphs, tables and images. Each image source has the file name in the top left corner of the image with coordinates (10,10) pixels and is in the format SourceFileName:<file_name> Each text source starts in a new line and has the file name followed by colon and the actual information Always include the source name from the image or text for each fact you use in the response in the format: [filename] Answer the following question using only the data provided in the sources below. For tabular information return it as an html table. Do not return markdown format. The text and image source can be the same file name, don't use the image title when citing the image source, only use the file name as mentioned If you cannot answer using the sources below, say you don't know. Return just the answer without any input texts \"}",
-                            "{'role': 'user', 'content': [{'text': 'Are interest rates high?', 'type': 'text'}, {'text': 'Financial Market Analysis Report 2023-6.png: 3</td><td>1</td></tr></table> Financial markets are interconnected, with movements in one segment often influencing others. This section examines the correlations between stock indices, cryptocurrency prices, and commodity prices, revealing how changes in one market can have ripple effects across the financial ecosystem.Impact of Macroeconomic Factors Impact of Interest Rates, Inflation, and GDP Growth on Financial Markets 5 4 3 2 1 0 -1 2018 2019 -2 -3 -4 -5 2020 2021 2022 2023 Macroeconomic factors such as interest rates, inflation, and GDP growth play a pivotal role in shaping financial markets. This section analyzes how these factors have influenced stock, cryptocurrency, and commodity markets over recent years, providing insights into the complex relationship between the economy and financial market performance. -Interest Rates % -Inflation Data % GDP Growth % :unselected: :unselected:Future Predictions and Trends Relative Growth Trends for S&P 500, Bitcoin, and Oil Prices (2024 Indexed to 100) 2028 Based on historical data, current trends, and economic indicators, this section presents predictions ', 'type': 'text'}, {'image_url': {'url': 'data:image/png;base64,iVBOR1BORw0KGgoAAAANSUhEUgAAAAEAAAABAQAAAAA3bvkkAAAACklEQVR4nGMAAQAABQABDQ0tuhsAAAAASUVORK5CYII=', 'detail': 'auto'}, 'type': 'image_url'}]}"
+                            "{'role': 'user', 'content': [{'text': 'Are interest rates high?', 'type': 'text'}, {'text': 'Financial Market Analysis Report 2023-6.png: 3</td><td>1</td></tr></table> Financial markets are interconnected, with movements in one segment often influencing others. This section examines the correlations between stock indices, cryptocurrency prices, and commodity prices, revealing how changes in one market can have ripple effects across the financial ecosystem.Impact of Macroeconomic Factors Impact of Interest Rates, Inflation, and GDP Growth on Financial Markets 5 4 3 2 1 0 -1 2018 2019 -2 -3 -4 -5 2020 2021 2022 2023 Macroeconomic factors such as interest rates, inflation, and GDP growth play a pivotal role in shaping financial markets. This section analyzes how these factors have influenced stock, cryptocurrency, and commodity markets over recent years, providing insights into the complex relationship between the economy and financial market performance. -Interest Rates % -Inflation Data % GDP Growth % :unselected: :unselected:Future Predictions and Trends Relative Growth Trends for S&P 500, Bitcoin, and Oil Prices (2024 Indexed to 100) 2028 Based on historical data, current trends, and economic indicators, this section presents predictions ', 'type': 'text'}, {'image_url': {'url': 'data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z/C/HgAGgwJ/lK3Q6wAAAABJRU5ErkJggg==', 'detail': 'auto'}, 'type': 'image_url'}]}"
                         ],
                         "props": null,
                         "title": "Prompt"