unique_toolkit 0.7.9__py3-none-any.whl → 1.33.3__py3-none-any.whl

unique_toolkit/_common/token/token_counting.py

@@ -0,0 +1,204 @@
+# 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._common.token.image_token_counting import (
+    calculate_image_tokens_from_base64,
+)
+from unique_toolkit.language_model import (
+    LanguageModelMessage,
+    LanguageModelMessages,
+    LanguageModelName,
+)
+
+
+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 key == "content":
+                if message.get("role") == "tool":
+                    """
+                    We have observed a general difference in the way tool response messages are handled.
+                    Specifically, if we take a list of tool responses and artificially transform them into user messages (content field stays the same)
+                    the token consumption goes does drastically, this seems to scale with the number of tokens in the tool responses.
+
+                    the json.dumps() method was found by trial and error. It will give a conservative estimate, but seems to be close to the token count
+                    returned from the openai call.
+                    """
+                    num_tokens += len(encode(json.dumps(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)
+                else:
+                    num_tokens += len(encode(value))
+            elif isinstance(value, str):
+                num_tokens += len(encode(value))
+
+            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", "name"] and v is not None
+            )  # Ignore tool_calls for now
+        }
+        for m in messages.model_dump(mode="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/_common/utils/__init__.py

@@ -0,0 +1 @@
+# Utils package for common utilities

unique_toolkit/_common/utils/files.py

@@ -0,0 +1,43 @@
+import mimetypes
+from enum import StrEnum
+
+
+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 ImageMimeType(StrEnum):
+    JPEG = "image/jpeg"
+    PNG = "image/png"
+    GIF = "image/gif"
+    BMP = "image/bmp"
+    WEBP = "image/webp"
+    TIFF = "image/tiff"
+    SVG = "image/svg+xml"
+
+
+def is_file_content(filename: str) -> bool:
+    mimetype, _ = mimetypes.guess_type(filename)
+
+    if not mimetype:
+        return False
+
+    return mimetype in FileMimeType.__members__.values()
+
+
+def is_image_content(filename: str) -> bool:
+    mimetype, _ = mimetypes.guess_type(filename)
+
+    if not mimetype:
+        return False
+
+    return mimetype in ImageMimeType.__members__.values()

unique_toolkit/_common/utils/image/encode.py

@@ -0,0 +1,25 @@
+import base64
+import io
+
+from PIL.ImageFile import ImageFile
+
+
+def image_to_base64(image: ImageFile) -> str:
+    # Convert to RGB if needed
+    img = image
+    if image.mode != "RGB":
+        img = image.convert("RGB")
+
+    # Create BytesIO object to store compressed image
+    img_byte_arr = io.BytesIO()
+
+    # Save with compression
+    img.save(img_byte_arr, format="JPEG", quality=85, optimize=True)
+    img_byte_arr.seek(0)
+
+    # Encode compressed image
+    encoded_string = base64.b64encode(img_byte_arr.getvalue())
+    image_string = encoded_string.decode("utf-8")
+
+    image_string = "data:image/jpeg;base64," + image_string
+    return image_string

unique_toolkit/_common/utils/jinja/helpers.py

@@ -0,0 +1,10 @@
+from pathlib import Path
+
+
+def get_parent_dir(file_path: str) -> Path:
+    return Path(file_path).parent
+
+
+def load_template(parent_dir: Path, template_name: str) -> str:
+    with open(parent_dir / template_name, "r") as file:
+        return file.read().strip()

unique_toolkit/_common/utils/jinja/render.py

@@ -0,0 +1,18 @@
+from typing import Any
+
+from jinja2 import Template
+
+from unique_toolkit._common.utils.jinja.schema import Jinja2PromptParams
+
+
+def render_template(
+    template: str, params: Jinja2PromptParams | dict[str, Any] | None = None, **kwargs
+) -> str:
+    params = params or {}
+
+    if isinstance(params, Jinja2PromptParams):
+        params = params.model_dump(exclude_none=True, mode="json")
+
+    params.update(kwargs)
+
+    return Template(template, lstrip_blocks=True).render(**params)

unique_toolkit/_common/utils/jinja/schema.py

@@ -0,0 +1,65 @@
+from datetime import date, datetime
+from typing import Annotated, Any
+
+from jinja2 import Template
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    SerializerFunctionWrapHandler,
+    WrapSerializer,
+)
+
+from unique_toolkit.agentic.tools.tool import Tool
+
+
+class Jinja2PromptParams(BaseModel):
+    model_config = ConfigDict(str_strip_whitespace=True)
+
+    def render_template(self, template: str) -> str:
+        params = self.model_dump(exclude_none=True, mode="json")
+
+        return Template(template, lstrip_blocks=True).render(**params)
+
+
+class ToolPromptParams(Jinja2PromptParams):
+    name: str
+    tool_description_for_system_prompt: str = ""
+    tool_format_information_for_system_prompt: str = ""
+    tool_format_reminder_for_user_prompt: str = ""
+
+    @classmethod
+    def from_tool(cls, tool: Tool) -> "ToolPromptParams":
+        return cls(
+            name=tool.name,
+            tool_description_for_system_prompt=tool.tool_description_for_system_prompt(),
+            tool_format_information_for_system_prompt=tool.tool_format_information_for_system_prompt(),
+            tool_format_reminder_for_user_prompt=tool.tool_format_reminder_for_user_prompt(),
+        )
+
+
+def serialize_iso8601_date(v: Any, handler: SerializerFunctionWrapHandler) -> str:
+    if isinstance(v, date):
+        return v.isoformat()
+    return handler(v)
+
+
+ISO8601Date = Annotated[
+    date,
+    WrapSerializer(serialize_iso8601_date, return_type=str),
+]
+
+
+class AgentSystemPromptParams(Jinja2PromptParams):
+    info_cutoff_at: ISO8601Date | None
+    current_date: ISO8601Date = Field(default_factory=lambda: datetime.now().date())
+    tools: list[ToolPromptParams]
+    used_tools: list[ToolPromptParams]
+    add_citation_appendix: bool = True
+    max_tools_per_iteration: int
+    max_loop_iterations: int
+    current_iteration: int
+
+
+class AgentUserPromptParams(Jinja2PromptParams):
+    user_prompt: str

unique_toolkit/_common/utils/jinja/utils.py

@@ -0,0 +1,80 @@
+from jinja2 import Environment
+from jinja2.nodes import Const, Getattr, Getitem, Name
+from pydantic import BaseModel
+
+
+class TemplateValidationResult(BaseModel):
+    is_valid: bool
+    missing_placeholders: list[str]
+    optional_placeholders: list[str]
+    unexpected_placeholders: list[str]
+
+
+def _get_nested_variables(node):
+    """Recursively extract all variable references from a Jinja2 AST node."""
+    variables = set()
+
+    if isinstance(node, Name):
+        variables.add(node.name)
+    elif isinstance(node, (Getattr, Getitem)):
+        # For nested attributes like example.category
+        if isinstance(node.node, Name):
+            if isinstance(node, Getattr):
+                variables.add(f"{node.node.name}.{node.attr}")
+            else:  # Getitem
+                if isinstance(node.arg, Const):
+                    variables.add(f"{node.node.name}.{node.arg.value}")
+                else:
+                    # For dynamic indices, just use the base variable
+                    variables.add(node.node.name)
+        # Recursively process nested nodes
+        variables.update(_get_nested_variables(node.node))
+
+    # Process child nodes
+    for child in node.iter_child_nodes():
+        variables.update(_get_nested_variables(child))
+
+    return variables
+
+
+def validate_template_placeholders(
+    template_content: str,
+    required_placeholders: set[str],
+    optional_placeholders: set[str],
+) -> TemplateValidationResult:
+    """
+    Validates that all required placeholders in the template are present.
+    Handles both top-level and nested variables (e.g. example.category).
+
+    Args:
+        template_content (str): The content of the Jinja template
+        required_placeholders (set[str]): Set of required placeholder names
+        optional_placeholders (set[str]): Set of optional placeholder names
+
+    Returns:
+        TemplateValidationResult: A result object containing validation information
+    """
+    # Create a Jinja environment
+    env = Environment()
+
+    # Parse the template and get all variables including nested ones
+    ast = env.parse(template_content)
+    template_vars = _get_nested_variables(ast)
+
+    # Check for missing required placeholders
+    missing_placeholders = required_placeholders - template_vars
+
+    # Check for optional placeholders present
+    present_optional = optional_placeholders & template_vars
+
+    # Check for any unexpected placeholders
+    unexpected_placeholders = template_vars - (
+        required_placeholders | optional_placeholders
+    )
+
+    return TemplateValidationResult(
+        is_valid=len(missing_placeholders) == 0,
+        missing_placeholders=sorted(list(missing_placeholders)),
+        optional_placeholders=sorted(list(present_optional)),
+        unexpected_placeholders=sorted(list(unexpected_placeholders)),
+    )

unique_toolkit/_common/utils/structured_output/__init__.py

@@ -0,0 +1 @@
+# Structured output utilities package

unique_toolkit/_common/utils/structured_output/schema.py

@@ -0,0 +1,5 @@
+from pydantic import BaseModel, ConfigDict
+
+
+class StructuredOutputModel(BaseModel):
+    model_config = ConfigDict(extra="forbid")

unique_toolkit/_common/utils/write_configuration.py

@@ -0,0 +1,51 @@
+import json
+from pathlib import Path
+
+from humps import kebabize, pascalize
+from pydantic import BaseModel
+
+
+def write_module_configuration_schema(
+    service_folderpath: Path,
+    write_folderpath: Path,
+    config: BaseModel,
+    sub_name: str = "",
+):
+    filename_prefix = pascalize(service_folderpath.name)
+
+    filepath = (
+        write_folderpath
+        / f"{filename_prefix}{f'-{sub_name}' if sub_name else ''}Schema.json"
+    )
+
+    with open(filepath, "w") as f:
+        json.dump(config.model_json_schema(by_alias=True), f, indent=4)
+
+
+def write_service_configuration(
+    service_folderpath: Path,
+    write_folderpath: Path,
+    config: BaseModel,
+    sub_name: str = "",
+):
+    filename_prefix = kebabize(service_folderpath.name)
+
+    filepath = (
+        write_folderpath
+        / f"{filename_prefix}{f'-{sub_name}' if sub_name else ''}-configuration-schema.json"
+    )
+
+    with open(filepath, "w") as f:
+        json.dump(config.model_json_schema(by_alias=True), f, indent=4)
+    filepath = (
+        write_folderpath
+        / f"{filename_prefix}{f'-{sub_name}' if sub_name else ''}-default-configuration.json"
+    )
+
+    # We exclude language_model_info as it is infered from language_model_name
+    with open(filepath, "w") as f:
+        f.write(
+            config.model_dump_json(
+                by_alias=True, indent=4, exclude=set(["language_model_info"])
+            )
+        )

unique_toolkit/_common/validators.py

@@ -1,8 +1,105 @@
-from unique_toolkit.language_model import LanguageModel, LanguageModelName
+import logging
+from typing import Annotated, Any
 
+from pydantic import BeforeValidator, Field, PlainSerializer, ValidationInfo
+from pydantic.fields import FieldInfo
+
+from unique_toolkit.language_model import LanguageModelName
+from unique_toolkit.language_model.infos import (
+    LanguageModelInfo,
+    LanguageModelProvider,
+)
+
+logger = logging.getLogger(__name__)
+
+# TODO @klcd: Inform on deprecation of str as input
+LMI = Annotated[
+    LanguageModelInfo,
+    BeforeValidator(
+        lambda v: validate_and_init_language_model_info(v),
+        json_schema_input_type=LanguageModelName
+        | Annotated[
+            str,
+            Field(
+                title="Language Model String",
+            ),
+        ]
+        | LanguageModelInfo,
+    ),
+    PlainSerializer(
+        lambda v: serialize_lmi(v),
+        when_used="json",
+        return_type=str | LanguageModelInfo,
+    ),
+]
+
+
+def get_LMI_default_field(llm_name: LanguageModelName, **kwargs) -> Any:
+    return Field(
+        default=LanguageModelInfo.from_name(llm_name),
+        json_schema_extra={"default": llm_name},
+        **kwargs,
+    )
+
+
+def serialize_lmi(model: LanguageModelInfo) -> str | LanguageModelInfo:
+    if model.provider == LanguageModelProvider.CUSTOM:
+        return model
+
+    return model.name
+
+
+def validate_and_init_language_model_info(
+    v: str | LanguageModelName | LanguageModelInfo,
+) -> LanguageModelInfo:
+    """Validate and initialize a LanguageModelInfo object.
+
+    Args:
+        v: The input value to validate and initialize.
+
+    Returns:
+        LanguageModelInfo: The validated and initialized LanguageModelInfo object.
+
+    """
+    if isinstance(v, LanguageModelName):
+        return LanguageModelInfo.from_name(v)
+    if isinstance(v, str):
+        if v in [name.value for name in LanguageModelName]:
+            return LanguageModelInfo.from_name(LanguageModelName(v))
+
+        return LanguageModelInfo(
+            name=v,
+            version="custom",
+            provider=LanguageModelProvider.CUSTOM,
+        )
+
+    return v
+
+
+def ClipInt(*, min_value: int, max_value: int) -> tuple[BeforeValidator, FieldInfo]:
+    def _validator(value: Any, info: ValidationInfo) -> Any:
+        if not isinstance(value, int):
+            value = int(value)
+
+        field_name = info.field_name
+        if value < min_value:
+            logger.warning(
+                "Field %s is below the allowed minimum of %s. It will be set to %s.",
+                field_name,
+                min_value,
+                min_value,
+            )
+            return min_value
+
+        if value > max_value:
+            logger.warning(
+                "Field %s is above the allowed maximum of %s. It will be set to %s.",
+                field_name,
+                max_value,
+                max_value,
+            )
+            return max_value
 
-def validate_and_init_language_model(value: LanguageModelName | LanguageModel | str):
-    if isinstance(value, LanguageModel):
         return value
 
-    return LanguageModel(value)
+    return (BeforeValidator(_validator), Field(ge=min_value, le=max_value))

unique_toolkit/agentic/__init__.py

@@ -0,0 +1 @@
+

unique_toolkit/agentic/debug_info_manager/debug_info_manager.py

@@ -0,0 +1,28 @@
+from typing import Any
+
+from unique_toolkit.agentic.tools.schemas import ToolCallResponse
+
+
+class DebugInfoManager:
+    def __init__(self):
+        self.debug_info = {"tools": []}
+
+    def extract_tool_debug_info(
+        self,
+        tool_call_responses: list[ToolCallResponse],
+        loop_iteration_index: int | None = None,
+    ):
+        for tool_call_response in tool_call_responses:
+            tool_info = {
+                "name": tool_call_response.name,
+                "info": tool_call_response.debug_info,
+            }
+            if loop_iteration_index is not None:
+                tool_info["info"]["loop_iteration"] = loop_iteration_index
+            self.debug_info["tools"].append(tool_info)
+
+    def add(self, key: str, value: Any) -> None:
+        self.debug_info = self.debug_info | {key: value}
+
+    def get(self) -> dict[str, Any]:
+        return self.debug_info