unique_toolkit 0.8.12__py3-none-any.whl → 0.8.14__py3-none-any.whl

unique_toolkit/history_manager/utils.py

@@ -0,0 +1,174 @@
+import json
+import logging
+from copy import deepcopy
+
+from unique_toolkit.content.schemas import ContentChunk, ContentMetadata
+from unique_toolkit.language_model.schemas import (
+    LanguageModelAssistantMessage,
+    LanguageModelMessage,
+    LanguageModelToolMessage,
+)
+from unique_toolkit.tools.schemas import Source
+from unique_toolkit.tools.utils.source_handling.schema import (
+    SourceFormatConfig,
+)
+
+
+logger = logging.getLogger(__name__)
+
+
+def convert_tool_interactions_to_content_messages(
+    loop_history: list[LanguageModelMessage],
+) -> list[LanguageModelMessage]:
+    new_loop_history = []
+    copy_loop_history = deepcopy(loop_history)
+
+    for message in copy_loop_history:
+        if isinstance(message, LanguageModelAssistantMessage) and message.tool_calls:
+            new_loop_history.append(_convert_tool_call_to_content(message))
+
+        elif isinstance(message, LanguageModelToolMessage):
+            new_loop_history.append(_convert_tool_call_response_to_content(message))
+        else:
+            new_loop_history.append(message)
+
+    return new_loop_history
+
+
+def _convert_tool_call_to_content(
+    assistant_message: LanguageModelAssistantMessage,
+) -> LanguageModelAssistantMessage:
+    assert assistant_message.tool_calls is not None
+    new_content = "The assistant requested the following tool_call:"
+    for tool_call in assistant_message.tool_calls:
+        new_content += (
+            f"\n\n- {tool_call.function.name}: {tool_call.function.arguments}"
+        )
+    assistant_message.tool_calls = None
+    assistant_message.content = new_content
+
+    return assistant_message
+
+
+def _convert_tool_call_response_to_content(
+    tool_message: LanguageModelToolMessage,
+) -> LanguageModelAssistantMessage:
+    new_content = f"The assistant received the following tool_call_response: {tool_message.name}, {tool_message.content}"
+    assistant_message = LanguageModelAssistantMessage(
+        content=new_content,
+    )
+    return assistant_message
+
+
+def transform_chunks_to_string(
+    content_chunks: list[ContentChunk],
+    max_source_number: int,
+    cfg: SourceFormatConfig | None,
+    full_sources_serialize_dump: bool = False,
+) -> str:
+    """Transform content chunks into a string of sources.
+
+    Args:
+        content_chunks (list[ContentChunk]): The content chunks to transform
+        max_source_number (int): The maximum source number to use
+        feature_full_sources (bool, optional): Whether to include the full source object. Defaults to False which is the old format.
+
+    Returns:
+        str: String for the tool call response
+    """
+    if not content_chunks:
+        return "No relevant sources found."
+    if full_sources_serialize_dump:
+        sources = [
+            Source(
+                source_number=max_source_number + i,
+                key=chunk.key,
+                id=chunk.id,
+                order=chunk.order,
+                content=chunk.text,
+                chunk_id=chunk.chunk_id,
+                metadata=(
+                    _format_metadata(chunk.metadata, cfg) or None
+                    if chunk.metadata
+                    else None
+                ),
+                url=chunk.url,
+            ).model_dump(
+                exclude_none=True,
+                exclude_defaults=True,
+                by_alias=True,
+            )
+            for i, chunk in enumerate(content_chunks)
+        ]
+    else:
+        sources = [
+            {
+                "source_number": max_source_number + i,
+                "content": chunk.text,
+                **(
+                    {"metadata": meta}
+                    if (
+                        meta := _format_metadata(chunk.metadata, cfg)
+                    )  # only add when not empty
+                    else {}
+                ),
+            }
+            for i, chunk in enumerate(content_chunks)
+        ]
+    return json.dumps(sources)
+
+
+def load_sources_from_string(
+    source_string: str,
+) -> list[Source] | None:
+    """Transform JSON string from language model tool message in the tool call response into Source objects"""
+
+    try:
+        # First, try parsing as JSON (new format)
+        sources_data = json.loads(source_string)
+        return [Source.model_validate(source) for source in sources_data]
+    except (json.JSONDecodeError, ValueError):
+        logger.warning("Failed to parse source string")
+        return None
+
+
+def _format_metadata(
+    metadata: ContentMetadata | None,
+    cfg: SourceFormatConfig | None,
+) -> str:
+    """
+    Build the concatenated tag string from the chunk's metadata
+    and the templates found in cfg.sections.
+    Example result:
+      "<|topic|>GenAI<|/topic|>\n<|date|>This document is from: 2025-04-29<|/date|>\n"
+    """
+    if metadata is None:
+        return ""
+
+    if cfg is None or not cfg.sections:
+        # If no configuration is provided, return empty string
+        return ""
+
+    meta_dict = metadata.model_dump(exclude_none=True, by_alias=True)
+    sections = cfg.sections
+
+    parts: list[str] = []
+    for key, template in sections.items():
+        if key in meta_dict:
+            parts.append(template.format(meta_dict[key]))
+
+    return "".join(parts)
+
+
+### In case we do not want any formatting of metadata we could use this function
+# def _filtered_metadata(
+#     meta: ContentMetadata | None,
+#     cfg: SourceFormatConfig,
+# ) -> dict[str, str] | None:
+#     if meta is None:
+#         return None
+
+#     allowed = set(cfg.sections)
+#     raw = meta.model_dump(exclude_none=True, by_alias=True)
+#     pruned = {k: v for k, v in raw.items() if k in allowed}
+#     return pruned or None

unique_toolkit/language_model/infos.py

@@ -44,6 +44,10 @@ class LanguageModelName(StrEnum):
     GEMINI_2_5_PRO = "litellm:gemini-2-5-pro"
     GEMINI_2_5_PRO_EXP_0325 = "litellm:gemini-2-5-pro-exp-03-25"
     GEMINI_2_5_PRO_PREVIEW_0605 = "litellm:gemini-2-5-pro-preview-06-05"
+    LITELLM_OPENAI_GPT_5 = "litellm:openai-gpt-5"
+    LITELLM_OPENAI_GPT_5_MINI = "litellm:openai-gpt-5-mini"
+    LITELLM_OPENAI_GPT_5_NANO = "litellm:openai-gpt-5-nano"
+    LITELLM_OPENAI_GPT_5_CHAT = "litellm:openai-gpt-5-chat"
 
 
 class EncoderName(StrEnum):
@@ -71,6 +75,10 @@ def get_encoder_name(model_name: LanguageModelName) -> EncoderName:
             | LMN.AZURE_GPT_5_MINI_2025_0807
             | LMN.AZURE_GPT_5_NANO_2025_0807
             | LMN.AZURE_GPT_5_CHAT_2025_0807
+            | LMN.LITELLM_OPENAI_GPT_5
+            | LMN.LITELLM_OPENAI_GPT_5_MINI
+            | LMN.LITELLM_OPENAI_GPT_5_NANO
+            | LMN.LITELLM_OPENAI_GPT_5_CHAT
         ):
             return EncoderName.O200K_BASE
         case _:
@@ -215,8 +223,10 @@ class LanguageModelInfo(BaseModel):
                     capabilities=[
                         ModelCapabilities.FUNCTION_CALLING,
                         ModelCapabilities.STREAMING,
+                        ModelCapabilities.REASONING,
                         ModelCapabilities.VISION,
                         ModelCapabilities.STRUCTURED_OUTPUT,
+                        ModelCapabilities.PARALLEL_FUNCTION_CALLING,
                     ],
                     token_limits=LanguageModelTokenLimits(
                         token_limit_input=272000, token_limit_output=128000
@@ -241,8 +251,10 @@ class LanguageModelInfo(BaseModel):
                     capabilities=[
                         ModelCapabilities.FUNCTION_CALLING,
                         ModelCapabilities.STREAMING,
+                        ModelCapabilities.REASONING,
                         ModelCapabilities.VISION,
                         ModelCapabilities.STRUCTURED_OUTPUT,
+                        ModelCapabilities.PARALLEL_FUNCTION_CALLING,
                     ],
                     token_limits=LanguageModelTokenLimits(
                         token_limit_input=272000, token_limit_output=128000
@@ -264,6 +276,10 @@ class LanguageModelInfo(BaseModel):
                     provider=LanguageModelProvider.AZURE,
                     version="2025-08-07",
                     encoder_name=EncoderName.O200K_BASE,
+                    capabilities=[
+                        ModelCapabilities.STREAMING,
+                        ModelCapabilities.VISION,
+                    ],
                     token_limits=LanguageModelTokenLimits(
                         token_limit_input=128000, token_limit_output=16384
                     ),
@@ -757,6 +773,109 @@ class LanguageModelInfo(BaseModel):
                     info_cutoff_at=date(2025, 1, day=1),
                     published_at=date(2025, 6, 5),
                 )
+            case LanguageModelName.LITELLM_OPENAI_GPT_5:
+                return cls(
+                    name=model_name,
+                    provider=LanguageModelProvider.LITELLM,
+                    version="gpt-5",
+                    encoder_name=EncoderName.O200K_BASE,
+                    capabilities=[
+                        ModelCapabilities.FUNCTION_CALLING,
+                        ModelCapabilities.STREAMING,
+                        ModelCapabilities.REASONING,
+                        ModelCapabilities.VISION,
+                        ModelCapabilities.STRUCTURED_OUTPUT,
+                        ModelCapabilities.PARALLEL_FUNCTION_CALLING,
+                    ],
+                    token_limits=LanguageModelTokenLimits(
+                        token_limit_input=272000, token_limit_output=128000
+                    ),
+                    info_cutoff_at=date(2024, 10, 24),
+                    published_at=date(2025, 8, 7),
+                    deprecated_at=date(2026, 8, 7),
+                    retirement_at=date(2026, 8, 7),
+                    temperature_bounds=TemperatureBounds(
+                        min_temperature=1.0, max_temperature=1.0
+                    ),
+                    default_options={
+                        "reasoning_effort": "minimal",
+                    },
+                )
+            case LanguageModelName.LITELLM_OPENAI_GPT_5_MINI:
+                return cls(
+                    name=model_name,
+                    provider=LanguageModelProvider.LITELLM,
+                    version="gpt-5-mini",
+                    encoder_name=EncoderName.O200K_BASE,
+                    capabilities=[
+                        ModelCapabilities.FUNCTION_CALLING,
+                        ModelCapabilities.STREAMING,
+                        ModelCapabilities.REASONING,
+                        ModelCapabilities.VISION,
+                        ModelCapabilities.STRUCTURED_OUTPUT,
+                        ModelCapabilities.PARALLEL_FUNCTION_CALLING,
+                    ],
+                    token_limits=LanguageModelTokenLimits(
+                        token_limit_input=272000, token_limit_output=128000
+                    ),
+                    info_cutoff_at=date(2024, 6, 24),
+                    published_at=date(2025, 8, 7),
+                    deprecated_at=date(2026, 8, 7),
+                    retirement_at=date(2026, 8, 7),
+                    temperature_bounds=TemperatureBounds(
+                        min_temperature=1.0, max_temperature=1.0
+                    ),
+                    default_options={
+                        "reasoning_effort": "minimal",
+                    },
+                )
+            case LanguageModelName.LITELLM_OPENAI_GPT_5_NANO:
+                return cls(
+                    name=model_name,
+                    provider=LanguageModelProvider.LITELLM,
+                    version="gpt-5-nano",
+                    encoder_name=EncoderName.O200K_BASE,
+                    capabilities=[
+                        ModelCapabilities.FUNCTION_CALLING,
+                        ModelCapabilities.STREAMING,
+                        ModelCapabilities.REASONING,
+                        ModelCapabilities.VISION,
+                        ModelCapabilities.STRUCTURED_OUTPUT,
+                        ModelCapabilities.PARALLEL_FUNCTION_CALLING,
+                    ],
+                    token_limits=LanguageModelTokenLimits(
+                        token_limit_input=272000, token_limit_output=128000
+                    ),
+                    info_cutoff_at=date(2024, 5, 31),
+                    published_at=date(2025, 8, 7),
+                    deprecated_at=date(2026, 8, 7),
+                    retirement_at=date(2026, 8, 7),
+                    temperature_bounds=TemperatureBounds(
+                        min_temperature=1.0, max_temperature=1.0
+                    ),
+                    default_options={
+                        "reasoning_effort": "minimal",
+                    },
+                )
+            case LanguageModelName.LITELLM_OPENAI_GPT_5_CHAT:
+                return cls(
+                    name=model_name,
+                    provider=LanguageModelProvider.LITELLM,
+                    version="gpt-5-chat",
+                    encoder_name=EncoderName.O200K_BASE,
+                    capabilities=[
+                        ModelCapabilities.STREAMING,
+                        ModelCapabilities.VISION,
+                    ],
+                    token_limits=LanguageModelTokenLimits(
+                        token_limit_input=128000, token_limit_output=16384
+                    ),
+                    info_cutoff_at=date(2024, 10, 24),
+                    published_at=date(2025, 8, 7),
+                    deprecated_at=date(2026, 8, 7),
+                    retirement_at=date(2026, 8, 7),
+                )
+            
             case _:
                 if isinstance(model_name, LanguageModelName):
                     raise ValueError(

unique_toolkit/postprocessor/postprocessor_manager.py

@@ -0,0 +1,122 @@
+from abc import ABC
+import asyncio
+from logging import Logger
+
+from unique_toolkit.chat.service import ChatService
+from unique_toolkit.language_model.schemas import (
+    LanguageModelMessage,
+    LanguageModelStreamResponse,
+)
+from unique_toolkit.tools.utils.execution.execution import Result, SafeTaskExecutor
+
+
+class Postprocessor(ABC):
+    def __init__(self, name: str):
+        self.name = name
+
+    def get_name(self) -> str:
+        return self.name
+
+    async def run(self, loop_response: LanguageModelStreamResponse) -> str:
+        raise NotImplementedError("Subclasses must implement this method.")
+
+    async def apply_postprocessing_to_response(
+        self, loop_response: LanguageModelStreamResponse
+    ) -> bool:
+        raise NotImplementedError(
+            "Subclasses must implement this method to apply post-processing to the response."
+        )
+
+    async def remove_from_text(self, text) -> str:
+        raise NotImplementedError(
+            "Subclasses must implement this method to remove post-processing from the message."
+        )
+
+
+class PostprocessorManager:
+    """
+    Manages and executes postprocessors for modifying and refining responses.
+
+    This class is responsible for:
+    - Storing and managing a collection of postprocessor instances.
+    - Executing postprocessors asynchronously to refine loop responses.
+    - Applying modifications to assistant messages based on postprocessor results.
+    - Providing utility methods for text manipulation using postprocessors.
+
+    Key Features:
+    - Postprocessor Management: Allows adding and retrieving postprocessor instances.
+    - Asynchronous Execution: Runs all postprocessors concurrently for efficiency.
+    - Response Modification: Applies postprocessing changes to assistant messages when necessary.
+    - Text Cleanup: Supports removing specific patterns or content from text using postprocessors.
+    - Error Handling: Logs warnings for any postprocessors that fail during execution.
+
+    The PostprocessorManager serves as a centralized system for managing and applying postprocessing logic to enhance response quality and consistency.
+    """
+
+    _postprocessors: list[Postprocessor] = []
+
+    def __init__(
+        self,
+        logger: Logger,
+        chat_service: ChatService,
+    ):
+        self._logger = logger
+        self._chat_service = chat_service
+
+    def add_postprocessor(self, postprocessor: Postprocessor):
+        self._postprocessors.append(postprocessor)
+
+    def get_postprocessors(self, name: str) -> list[Postprocessor]:
+        return self._postprocessors
+
+    async def run_postprocessors(
+        self,
+        loop_response: LanguageModelStreamResponse,
+    ) -> None:
+        task_executor = SafeTaskExecutor(
+            logger=self._logger,
+        )
+
+        tasks = [
+            task_executor.execute_async(
+                self.execute_postprocessors,
+                loop_response=loop_response,
+                postprocessor_instance=postprocessor,
+            )
+            for postprocessor in self._postprocessors
+        ]
+        postprocessor_results = await asyncio.gather(*tasks)
+
+        for i, result in enumerate(postprocessor_results):
+            if not result.success:
+                self._logger.warning(
+                    f"Postprocessor {self._postprocessors[i].get_name()} failed to run."
+                )
+
+        modification_results = [
+            postprocessor.apply_postprocessing_to_response(loop_response)
+            for postprocessor in self._postprocessors
+        ]
+
+        has_been_modified = any(modification_results)
+
+        if has_been_modified:
+            self._chat_service.modify_assistant_message(
+                content=loop_response.message.text,
+                message_id=loop_response.message.id,
+            )
+
+    async def execute_postprocessors(
+        self,
+        loop_response: LanguageModelStreamResponse,
+        postprocessor_instance: Postprocessor,
+    ) -> None:
+        await postprocessor_instance.run(loop_response)
+
+    async def remove_from_text(
+        self,
+        text: str,
+    ) -> str:
+        for postprocessor in self._postprocessors:
+            text = await postprocessor.remove_from_text(text)
+        return text

unique_toolkit/reference_manager/reference_manager.py

@@ -9,6 +9,25 @@ class tool_chunks:
 
 
 class ReferenceManager:
+    """
+    Manages content chunks and references extracted from tool responses.
+
+    This class is responsible for:
+    - Extracting and storing referenceable content chunks from tool responses.
+    - Managing a collection of content chunks and their associated references.
+    - Providing methods to retrieve, replace, and manipulate chunks and references.
+    - Supporting the retrieval of the latest references and their corresponding chunks.
+
+    Key Features:
+    - Chunk Extraction: Extracts content chunks from tool responses and organizes them for reference.
+    - Reference Management: Tracks references to content chunks and allows for easy retrieval.
+    - Latest Reference Access: Provides methods to fetch the most recent references and their associated chunks.
+    - Flexible Chunk Replacement: Allows for replacing the current set of chunks with a new list.
+    - Reference-to-Chunk Mapping: Matches references to their corresponding chunks based on source IDs.
+
+    The ReferenceManager serves as a utility for managing and linking content chunks with references, enabling efficient content tracking and retrieval.
+    """
+
     def __init__(self):
         self._tool_chunks: dict[str, tool_chunks] = {}
         self._chunks: list[ContentChunk] = []

unique_toolkit/short_term_memory/persistent_short_term_memory_manager.py

@@ -0,0 +1,140 @@
+import base64
+import zlib
+from logging import getLogger
+from typing import Generic, Type, TypeVar
+
+from pydantic import BaseModel
+from unique_toolkit.short_term_memory.schemas import ShortTermMemory
+from unique_toolkit.short_term_memory.service import ShortTermMemoryService
+from unique_toolkit.tools.utils.execution.execution import SafeTaskExecutor
+
+
+TSchema = TypeVar("TSchema", bound=BaseModel)
+
+
+logger = getLogger(__name__)
+
+
+def _default_short_term_memory_name(schema: type[BaseModel]) -> str:
+    return f"{schema.__name__}Key"
+
+
+def _compress_data_zlib_base64(data: str) -> str:
+    """Compress data using ZLIB and encode as base64 string."""
+    compressed = zlib.compress(data.encode("utf-8"))
+    return base64.b64encode(compressed).decode("utf-8")
+
+
+def _decompress_data_zlib_base64(compressed_data: str) -> str:
+    """Decompress base64 encoded ZLIB data."""
+    decoded = base64.b64decode(compressed_data.encode("utf-8"))
+    return zlib.decompress(decoded).decode("utf-8")
+
+
+class PersistentShortMemoryManager(Generic[TSchema]):
+    """
+    Manages the storage, retrieval, and processing of short-term memory in a persistent manner.
+
+    This class is responsible for:
+    - Saving and loading short-term memory data, both synchronously and asynchronously.
+    - Compressing and decompressing memory data for efficient storage.
+    - Validating and processing memory data using a predefined schema.
+    - Logging the status of memory operations, such as whether memory was found or saved.
+
+    Key Features:
+    - Persistent Storage: Integrates with a short-term memory service to store and retrieve memory data.
+    - Compression Support: Compresses memory data before saving and decompresses it upon retrieval.
+    - Schema Validation: Ensures memory data adheres to a specified schema for consistency.
+    - Synchronous and Asynchronous Operations: Supports both sync and async methods for flexibility.
+    - Logging and Debugging: Provides detailed logs for memory operations, including success and failure cases.
+
+    The PersistentShortMemoryManager is designed to handle short-term memory efficiently, ensuring data integrity and optimized storage.
+    """
+    def __init__(
+        self,
+        short_term_memory_service: ShortTermMemoryService,
+        short_term_memory_schema: Type[TSchema],
+        short_term_memory_name: str | None = None,
+    ) -> None:
+        self._short_term_memory_name = (
+            short_term_memory_name
+            if short_term_memory_name
+            else _default_short_term_memory_name(short_term_memory_schema)
+        )
+        self._short_term_memory_schema = short_term_memory_schema
+        self._short_term_memory_service = short_term_memory_service
+
+        self._executor = SafeTaskExecutor(
+            log_exceptions=False,
+        )
+
+    def _log_not_found(self) -> None:
+        logger.warning(
+            f"No short term memory found for chat {self._short_term_memory_service.chat_id} and key {self._short_term_memory_name}"
+        )
+
+    def _log_found(self) -> None:
+        logger.debug(
+            f"Short term memory found for chat {self._short_term_memory_service.chat_id} and key {self._short_term_memory_name}"
+        )
+
+    def _find_latest_memory_sync(self) -> ShortTermMemory | None:
+        result = self._executor.execute(
+            self._short_term_memory_service.find_latest_memory,
+            self._short_term_memory_name,
+        )
+
+        self._log_not_found() if not result.success else self._log_found()
+
+        return result.unpack(default=None)
+
+    async def _find_latest_memory_async(self) -> ShortTermMemory | None:
+        result = await self._executor.execute_async(
+            self._short_term_memory_service.find_latest_memory_async,
+            self._short_term_memory_name,
+        )
+
+        self._log_not_found() if not result.success else self._log_found()
+
+        return result.unpack(default=None)
+
+    def save_sync(self, short_term_memory: TSchema) -> None:
+        json_data = short_term_memory.model_dump_json()
+        compressed_data = _compress_data_zlib_base64(json_data)
+        logger.info(
+            f"Saving memory with {len(compressed_data)} characters compressed from {len(json_data)} characters for memory {self._short_term_memory_name}"
+        )
+        self._short_term_memory_service.create_memory(
+            key=self._short_term_memory_name,
+            value=compressed_data,
+        )
+
+    async def save_async(self, short_term_memory: TSchema) -> None:
+        json_data = short_term_memory.model_dump_json()
+        compressed_data = _compress_data_zlib_base64(json_data)
+        logger.info(
+            f"Saving memory with {len(compressed_data)} characters compressed from {len(json_data)} characters for memory {self._short_term_memory_name}"
+        )
+        await self._short_term_memory_service.create_memory_async(
+            key=self._short_term_memory_name,
+            value=compressed_data,
+        )
+
+    def _process_compressed_memory(
+        self, memory: ShortTermMemory | None
+    ) -> TSchema | None:
+        if memory is not None and memory.data is not None:
+            if isinstance(memory.data, str):
+                data = _decompress_data_zlib_base64(memory.data)
+                return self._short_term_memory_schema.model_validate_json(data)
+            elif isinstance(memory.data, dict):
+                return self._short_term_memory_schema.model_validate(memory.data)
+        return None
+
+    def load_sync(self) -> TSchema | None:
+        memory: ShortTermMemory | None = self._find_latest_memory_sync()
+        return self._process_compressed_memory(memory)
+
+    async def load_async(self) -> TSchema | None:
+        memory: ShortTermMemory | None = await self._find_latest_memory_async()
+        return self._process_compressed_memory(memory)