unique_toolkit 0.8.13__py3-none-any.whl → 0.8.15__py3-none-any.whl

unique_toolkit/_common/default_language_model.py

@@ -0,0 +1,6 @@
+from unique_toolkit.language_model.infos import LanguageModelName
+
+DEFAULT_GPT_35_TURBO = LanguageModelName.AZURE_GPT_35_TURBO_0125
+DEFAULT_GPT_4o = LanguageModelName.AZURE_GPT_4o_2024_1120
+DEFAULT_GPT_4o_STRUCTURED_OUTPUT = LanguageModelName.AZURE_GPT_4o_2024_0806
+DEFAULT_GPT_4o_MINI = LanguageModelName.AZURE_GPT_4o_MINI_2024_0718

unique_toolkit/_common/token/image_token_counting.py

@@ -0,0 +1,67 @@
+import base64
+import math
+import re
+from enum import Enum
+from io import BytesIO
+
+from PIL import Image
+
+
+class DetailLevel(Enum):
+    LOW = "low"
+    HIGH = "high"
+
+
+# https://platform.openai.com/docs/guides/vision/calculating-costs#calculating-costs
+def calculate_image_tokens(width, height, detail: DetailLevel):
+    """
+    Calculate the token cost of an image based on its dimensions and detail level.
+    NOTE: While we followed the documentation provided by openai to calculate image token cost, in practice,
+    we notice that this function overestimate the number of tokens consumed by the model.
+
+    Parameters:
+    - width (int): The width of the image in pixels.
+    - height (int): The height of the image in pixels.
+    - detail (str): The detail level, either "low" or "high".
+
+    Returns:
+    - int: The token cost of the image.
+    """
+    # Base cost for low detail
+    if detail == DetailLevel.LOW:
+        return 85
+
+    # Scaling for high detail
+    # Scale down to fit within 2048x2048 square
+    max_long_dim = 2048
+    long_dim = max(width, height)
+    if long_dim > max_long_dim:
+        scale_factor = long_dim / max_long_dim
+        width = int(width / scale_factor)
+        height = int(height / scale_factor)
+
+    # Scale down the shortest side to 768
+    max_short_dim = 768
+    short_dim = min(width, height)
+    if short_dim > max_short_dim:
+        scale_factor = short_dim / max_short_dim
+        width = int(width / scale_factor)
+        height = int(height / scale_factor)
+
+    # Step 3: Calculate the number of 512x512 tiles
+    tiles = math.ceil(width / 512) * math.ceil(height / 512)
+    # Step 4: Compute token cost
+    token_cost = (tiles * 170) + 85
+    return token_cost
+
+
+def calculate_image_tokens_from_base64(base64_string: str):
+    base64_string = remove_base64_header(base64_string)
+    image = Image.open(BytesIO(base64.b64decode(base64_string)))
+    # DETAIL LEVEL HIGH IS THE DEFAULT TO BE ON THE SAFE SIDE
+    return calculate_image_tokens(image.width, image.height, DetailLevel.HIGH)
+
+
+def remove_base64_header(base64_string: str):
+    header_pattern = r"^data:image/\w+;base64,"
+    return re.sub(header_pattern, "", base64_string)

unique_toolkit/_common/token/token_counting.py

@@ -0,0 +1,196 @@
+# Original source
+# https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb
+
+import json
+from typing import Any, Callable
+
+from pydantic import BaseModel
+from unique_toolkit.language_model import (
+    LanguageModelMessage,
+    LanguageModelMessages,
+    LanguageModelName,
+)
+
+from _common.utils.token.image_token_counting import (
+    calculate_image_tokens_from_base64,
+)
+
+
+class SpecialToolCallingTokens(BaseModel):
+    func_init: int = 0
+    prop_init: int = 0
+    prop_key: int = 0
+    enum_init: int = 0
+    enum_item: int = 0
+    func_end: int = 0
+
+
+def get_special_token(model: LanguageModelName) -> SpecialToolCallingTokens:
+    special_token = SpecialToolCallingTokens()
+
+    match model:
+        case (
+            LanguageModelName.AZURE_GPT_4o_2024_0513
+            | LanguageModelName.AZURE_GPT_4o_2024_0806
+            | LanguageModelName.AZURE_GPT_4o_MINI_2024_0718
+            | LanguageModelName.AZURE_GPT_4o_2024_1120
+        ):
+            special_token.func_init = 7
+            special_token.prop_init = 3
+            special_token.prop_key = 3
+            special_token.enum_init = -3
+            special_token.enum_item = 3
+            special_token.func_end = 12
+
+        case (
+            LanguageModelName.AZURE_GPT_35_TURBO_0125
+            | LanguageModelName.AZURE_GPT_4_0613
+            | LanguageModelName.AZURE_GPT_4_32K_0613
+            | LanguageModelName.AZURE_GPT_4_TURBO_2024_0409
+        ):
+            special_token.func_init = 10
+            special_token.prop_init = 3
+            special_token.prop_key = 3
+            special_token.enum_init = -3
+            special_token.enum_item = 3
+            special_token.func_end = 12
+
+        case _:
+            raise NotImplementedError(
+                f"""num_tokens_for_tools() is not implemented for model {model}."""
+            )
+    return special_token
+
+
+def num_tokens_per_messages(
+    messages: list[dict[str, str]], encode: Callable[[str], list[int]]
+) -> list[int]:
+    """Return the number of tokens used by a list of messages."""
+
+    num_token_per_message = []
+    for message in messages:
+        num_tokens = 3  # extra_tokens_per_message
+        for key, value in message.items():
+            if isinstance(value, str):
+                num_tokens += len(encode(value))
+            elif isinstance(value, list):
+                # NOTE: The result returned by the function below is not 100% accurate.
+                num_tokens += handle_message_with_images(value, encode)
+            if key == "name":
+                num_tokens += 1  # extra_tokens_per_name
+
+        num_token_per_message.append(num_tokens)
+
+    return num_token_per_message
+
+
+def num_tokens_from_messages(
+    messages: list[dict[str, str]], encode: Callable[[str], list[int]]
+) -> int:
+    """Return the number of tokens used by a list of messages."""
+
+    num_tokens_per_message = num_tokens_per_messages(messages, encode)
+    num_tokens = sum(num_tokens_per_message) + 3
+
+    return num_tokens
+
+
+def num_tokens_for_tools(
+    functions: list[dict[str, Any]],
+    special_token: SpecialToolCallingTokens,
+    encode: Callable[[str], list[int]],
+):
+    def num_token_function_enum(
+        properties: dict[str, Any], encode: Callable[[str], list[int]]
+    ):
+        enum_token_count = 0
+        enum_token_count += special_token.enum_init
+        for item in properties[key]["enum"]:
+            enum_token_count += special_token.enum_item
+            enum_token_count += len(encode(item))
+
+        return enum_token_count
+
+    func_token_count = 0
+    if len(functions) > 0:
+        for func in functions:
+            func_token_count += special_token.func_init
+            function = func.get("function", {})
+            func_token_count += len(
+                encode(
+                    function.get("name", "")
+                    + ":"
+                    + function.get("description", "").rstrip(".").rstrip()
+                )
+            )
+            if len(function.get("parameters", {}).get("properties", "")) > 0:
+                properties = function.get("parameters", {}).get(
+                    "properties", ""
+                )
+                func_token_count += special_token.prop_init
+
+                for key in list(properties.keys()):
+                    func_token_count += special_token.prop_key
+
+                    if "enum" in properties[key].keys():
+                        func_token_count += num_token_function_enum(
+                            properties, encode
+                        )
+
+                    func_token_count += len(
+                        encode(
+                            f"{key}:{properties[key]['type']}:{properties[key]['description'].rstrip('.').rstrip()}"
+                        )
+                    )
+
+        func_token_count += special_token.func_end
+
+    return func_token_count
+
+
+def handle_message_with_images(
+    message: list[dict], encode: Callable[[str], list[int]]
+):
+    token_count = 0
+    for item in message:
+        if item.get("type") == "image_url":
+            image_url = item.get("imageUrl", {}).get("url")
+            if image_url:
+                token_count += calculate_image_tokens_from_base64(image_url)
+        elif item.get("type") == "text":
+            token_count += len(encode(item.get("text", "")))
+    return token_count
+
+
+def messages_to_openai_messages(
+    messages: LanguageModelMessages | list[LanguageModelMessage],
+):
+    if isinstance(messages, list):
+        messages = LanguageModelMessages(messages)
+
+    return [
+        {
+            k: v
+            for k, v in m.items()
+            if (k in ["content", "role"] and v is not None)
+        }
+        for m in json.loads(messages.model_dump_json())
+    ]
+
+
+def num_tokens_per_language_model_message(
+    messages: LanguageModelMessages | list[LanguageModelMessage],
+    encode: Callable[[str], list[int]],
+) -> list[int]:
+    return num_tokens_per_messages(
+        messages=messages_to_openai_messages(messages), encode=encode
+    )
+
+
+def num_token_for_language_model_messages(
+    messages: LanguageModelMessages | list[LanguageModelMessage],
+    encode: Callable[[str], list[int]],
+) -> int:
+    return num_tokens_from_messages(
+        messages_to_openai_messages(messages), encode
+    )

unique_toolkit/history_manager/history_construction_with_contents.py

@@ -0,0 +1,307 @@
+import base64
+import mimetypes
+
+from datetime import datetime
+from enum import StrEnum
+
+import numpy as np
+import tiktoken
+
+from pydantic import RootModel
+
+from _common.token.token_counting import num_tokens_per_language_model_message
+from chat.service import ChatService
+from content.service import ContentService
+from language_model.schemas import LanguageModelMessages
+from unique_toolkit.app import ChatEventUserMessage
+from unique_toolkit.chat.schemas import ChatMessage
+from unique_toolkit.chat.schemas import ChatMessageRole as ChatRole
+from unique_toolkit.content.schemas import Content
+from unique_toolkit.language_model import LanguageModelMessageRole as LLMRole
+from unique_toolkit.language_model.infos import EncoderName
+
+
+
+# TODO: Test this once it moves into the unique toolkit
+
+map_chat_llm_message_role = {
+    ChatRole.USER: LLMRole.USER,
+    ChatRole.ASSISTANT: LLMRole.ASSISTANT,
+}
+
+
+class ImageMimeType(StrEnum):
+    JPEG = "image/jpeg"
+    PNG = "image/png"
+    GIF = "image/gif"
+    BMP = "image/bmp"
+    WEBP = "image/webp"
+    TIFF = "image/tiff"
+    SVG = "image/svg+xml"
+
+
+class FileMimeType(StrEnum):
+    PDF = "application/pdf"
+    DOCX = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
+    DOC = "application/msword"
+    XLSX = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
+    XLS = "application/vnd.ms-excel"
+    PPTX = "application/vnd.openxmlformats-officedocument.presentationml.presentation"
+    CSV = "text/csv"
+    HTML = "text/html"
+    MD = "text/markdown"
+    TXT = "text/plain"
+
+
+class ChatMessageWithContents(ChatMessage):
+    contents: list[Content] = []
+
+
+class ChatHistoryWithContent(RootModel):
+    root: list[ChatMessageWithContents]
+
+    @classmethod
+    def from_chat_history_and_contents(
+        cls,
+        chat_history: list[ChatMessage],
+        chat_contents: list[Content],
+    ):
+        combined = chat_contents + chat_history
+        combined.sort(key=lambda x: x.created_at or datetime.min)
+
+        grouped_elements = []
+        content_container = []
+
+        # Content is collected and added to the next chat message
+        for c in combined:
+            if isinstance(c, ChatMessage):
+                grouped_elements.append(
+                    ChatMessageWithContents(
+                        contents=content_container.copy(),
+                        **c.model_dump(),
+                    ),
+                )
+                content_container.clear()
+            else:
+                content_container.append(c)
+
+        return cls(root=grouped_elements)
+
+    def __iter__(self):
+        return iter(self.root)
+
+    def __getitem__(self, item):
+        return self.root[item]
+
+
+def is_image_content(filename: str) -> bool:
+    mimetype, _ = mimetypes.guess_type(filename)
+
+    if not mimetype:
+        return False
+
+    return mimetype in ImageMimeType.__members__.values()
+
+
+def is_file_content(filename: str) -> bool:
+    mimetype, _ = mimetypes.guess_type(filename)
+
+    if not mimetype:
+        return False
+
+    return mimetype in FileMimeType.__members__.values()
+
+
+def get_chat_history_with_contents(
+    user_message: ChatEventUserMessage,
+    chat_id: str,
+    chat_history: list[ChatMessage],
+    content_service: ContentService,
+) -> ChatHistoryWithContent:
+    last_user_message = ChatMessage(
+        id=user_message.id,
+        chat_id=chat_id,
+        text=user_message.text,
+        originalText=user_message.original_text,
+        role=ChatRole.USER,
+        gpt_request=None,
+        created_at=datetime.fromisoformat(user_message.created_at),
+    )
+    if len(chat_history) > 0 and last_user_message.id == chat_history[-1].id:
+        pass
+    else:
+        chat_history.append(last_user_message)
+
+    chat_contents = content_service.search_contents(
+        where={
+            "ownerId": {
+                "equals": chat_id,
+            },
+        },
+    )
+
+    return ChatHistoryWithContent.from_chat_history_and_contents(
+        chat_history,
+        chat_contents,
+    )
+
+
+def download_encoded_images(
+    contents: list[Content],
+    content_service: ContentService,
+    chat_id: str,
+) -> list[str]:
+    base64_encoded_images = []
+    for im in contents:
+        if is_image_content(im.key):
+            try:
+                file_bytes = content_service.download_content_to_bytes(
+                    content_id=im.id,
+                    chat_id=chat_id,
+                )
+
+                mime_type, _ = mimetypes.guess_type(im.key)
+                encoded_string = base64.b64encode(file_bytes).decode("utf-8")
+                image_string = f"data:{mime_type};base64," + encoded_string
+                base64_encoded_images.append(image_string)
+            except Exception as e:
+                print(e)
+    return base64_encoded_images
+
+
+class FileContentSerialization(StrEnum):
+    NONE = "none"
+    FILE_NAME = "file_name"
+
+
+class ImageContentInclusion(StrEnum):
+    NONE = "none"
+    ALL = "all"
+
+
+def file_content_serialization(
+    file_contents: list[Content],
+    file_content_serialization: FileContentSerialization,
+) -> str:
+    match file_content_serialization:
+        case FileContentSerialization.NONE:
+            return ""
+        case FileContentSerialization.FILE_NAME:
+            file_names = [
+                f"- Uploaded file: {f.key} at {f.created_at}"
+                for f in file_contents
+            ]
+            return "\n".join(
+                [
+                    "Files Uploaded to Chat can be accessed by internal search tool if available:\n",
+                ]
+                + file_names,
+            )
+
+
+def get_full_history_with_contents(
+    user_message: ChatEventUserMessage,
+    chat_id: str,
+    chat_service: ChatService,
+    content_service: ContentService,
+    include_images: ImageContentInclusion = ImageContentInclusion.ALL,
+    file_content_serialization_type: FileContentSerialization = FileContentSerialization.FILE_NAME,
+) -> LanguageModelMessages:
+    grouped_elements = get_chat_history_with_contents(
+        user_message=user_message,
+        chat_id=chat_id,
+        chat_history=chat_service.get_full_history(),
+        content_service=content_service,
+    )
+
+    builder = LanguageModelMessages([]).builder()
+    for c in grouped_elements:
+        # LanguageModelUserMessage has not field original content
+        text = c.original_content if c.original_content else c.content
+        if text is None:
+            if c.role == ChatRole.USER:
+                raise ValueError(
+                    "Content or original_content of LanguageModelMessages should exist.",
+                )
+            text = ""
+
+        if len(c.contents) > 0:
+            file_contents = [
+                co for co in c.contents if is_file_content(co.key)
+            ]
+            image_contents = [
+                co for co in c.contents if is_image_content(co.key)
+            ]
+
+            content = (
+                text
+                + "\n\n"
+                + file_content_serialization(
+                    file_contents,
+                    file_content_serialization_type,
+                )
+            )
+            content = content.strip()
+
+            if include_images and len(image_contents) > 0:
+                builder.image_message_append(
+                    content=content,
+                    images=download_encoded_images(
+                        contents=image_contents,
+                        content_service=content_service,
+                        chat_id=chat_id,
+                    ),
+                    role=map_chat_llm_message_role[c.role],
+                )
+            else:
+                builder.message_append(
+                    role=map_chat_llm_message_role[c.role],
+                    content=content,
+                )
+        else:
+            builder.message_append(
+                role=map_chat_llm_message_role[c.role],
+                content=text,
+            )
+    return builder.build()
+
+
+def get_full_history_as_llm_messages(
+    chat_service: ChatService,
+) -> LanguageModelMessages:
+    chat_history = chat_service.get_full_history()
+
+    map_chat_llm_message_role = {
+        ChatRole.USER: LLMRole.USER,
+        ChatRole.ASSISTANT: LLMRole.ASSISTANT,
+    }
+
+    builder = LanguageModelMessages([]).builder()
+    for c in chat_history:
+        builder.message_append(
+            role=map_chat_llm_message_role[c.role],
+            content=c.content or "",
+        )
+    return builder.build()
+
+
+
+def limit_to_token_window(
+    messages: LanguageModelMessages,
+    token_limit: int,
+    encoding_name: EncoderName = EncoderName.O200K_BASE,
+) -> LanguageModelMessages:
+    encoder = tiktoken.get_encoding(encoding_name)
+    token_per_message_reversed = num_tokens_per_language_model_message(
+        messages,
+        encode=encoder.encode,
+    )
+
+    to_take: list[bool] = (
+        np.cumsum(token_per_message_reversed) < token_limit
+    ).tolist()
+    to_take.reverse()
+
+    return LanguageModelMessages(
+        root=[m for m, tt in zip(messages, to_take, strict=False) if tt],
+    )