unique_toolkit 0.8.11__py3-none-any.whl → 0.8.13__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/service.py

@@ -6,6 +6,7 @@ from typing_extensions import deprecated
 
 from unique_toolkit._common.validate_required_values import validate_required_values
 from unique_toolkit.app.schemas import BaseEvent, ChatEvent, Event
+from unique_toolkit.app.unique_settings import UniqueSettings
 from unique_toolkit.content.schemas import ContentChunk
 from unique_toolkit.language_model.constants import (
     DEFAULT_COMPLETE_TEMPERATURE,
@@ -89,6 +90,16 @@ class LanguageModelService:
         """
         return cls(company_id=event.company_id, user_id=event.user_id)
 
+    @classmethod
+    def from_settings(cls, settings: UniqueSettings):
+        """
+        Initialize the LanguageModelService with a settings object.
+        """
+        return cls(
+            company_id=settings.auth.company_id.get_secret_value(),
+            user_id=settings.auth.user_id.get_secret_value(),
+        )
+
     @property
     @deprecated(
         "The event property is deprecated and will be removed in a future version."

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)

unique_toolkit/thinking_manager/thinking_manager.py

@@ -0,0 +1,102 @@
+from logging import Logger
+from pydantic import BaseModel, Field
+
+from unique_toolkit.chat.service import ChatService
+from unique_toolkit.language_model.schemas import (
+    LanguageModelAssistantMessage,
+    LanguageModelStreamResponse,
+)
+from unique_toolkit.tools.tool_progress_reporter import (
+    ToolProgressReporter,
+)
+
+
+class ThinkingManagerConfig(BaseModel):
+    thinking_steps_display: bool = Field(
+        default=True, description="Whether to display thinking steps in the chat."
+    )
+
+
+class ThinkingManager:
+    """
+    Manages the display and tracking of thinking steps during response generation.
+
+    This class is responsible for:
+    - Tracking and formatting thinking steps as part of the response process.
+    - Updating the tool progress reporter with the latest thinking step information.
+    - Managing the display of thinking steps in the assistant's response.
+    - Closing and finalizing the thinking steps section when the process is complete.
+
+    Key Features:
+    - Thinking Step Tracking: Maintains a sequential log of thinking steps with step numbers.
+    - Configurable Display: Supports enabling or disabling the display of thinking steps based on configuration.
+    - Integration with Tool Progress: Updates the tool progress reporter to reflect the current thinking state.
+    - Dynamic Response Updates: Modifies the assistant's response to include or finalize thinking steps.
+    - Flexible Formatting: Formats thinking steps in a structured and user-friendly HTML-like format.
+
+    The ThinkingManager enhances transparency and user understanding by providing a clear view of the assistant's reasoning process.
+    """
+    def __init__(
+        self,
+        logger: Logger,
+        config: ThinkingManagerConfig,
+        tool_progress_reporter: ToolProgressReporter,
+        chat_service: ChatService,
+    ):
+        self._chat_service = chat_service
+        self._config = config
+        self._thinking_steps = ""
+        self._thinking_step_number = 1
+        self._tool_progress_reporter = tool_progress_reporter
+
+    def thinking_is_displayed(self) -> bool:
+        return self._config.thinking_steps_display
+
+    def update_tool_progress_reporter(self, loop_response: LanguageModelStreamResponse):
+        if self._config.thinking_steps_display and (
+            not loop_response.message.text
+            == self._tool_progress_reporter._progress_start_text
+        ):
+            self._tool_progress_reporter.tool_statuses = {}
+            self._tool_progress_reporter._progress_start_text = (
+                loop_response.message.text
+            )
+
+    def update_start_text(
+        self, start_text: str, loop_response: LanguageModelStreamResponse
+    ) -> str:
+        if not self._config.thinking_steps_display:
+            return start_text
+        if not loop_response.message.original_text:
+            return start_text
+        if loop_response.message.original_text == "":
+            return start_text
+
+        update_message = loop_response.message.original_text
+
+        if start_text == "":
+            self._thinking_steps = f"\n<i><b>Step 1:</b>\n{update_message}</i>\n"
+            start_text = f"""<details open>\n<summary><b>Thinking steps</b></summary>\n{self._thinking_steps}\n</details>\n\n---\n\n"""
+        else:
+            self._thinking_steps += f"\n\n<i><b>Step {self._thinking_step_number}:</b>\n{update_message}</i>\n\n"
+            start_text = f"""<details open>\n<summary><b>Thinking steps</b></summary>\n<i>{self._thinking_steps}\n\n</i>\n</details>\n\n---\n\n"""
+
+        self._thinking_step_number += 1
+        return start_text
+
+    def close_thinking_steps(self, loop_response: LanguageModelStreamResponse):
+        if not self._config.thinking_steps_display:
+            return
+        if not self._thinking_steps:
+            return
+        if not loop_response.message.text:
+            return
+        if not loop_response.message.text.startswith("<details open>"):
+            return
+
+        loop_response.message.text = loop_response.message.text.replace(
+            "<details open>", "<details>"
+        )
+
+        self._chat_service.modify_assistant_message(content=loop_response.message.text)
+        return

unique_toolkit/tools/schemas.py

@@ -134,5 +134,4 @@ class ToolPrompts(BaseModel):
     display_name: str
     tool_description: str
     tool_format_information_for_system_prompt: str
-    tool_format_information_for_system_prompt: str
     input_model: dict[str, Any]

unique_toolkit/tools/tool.py

@@ -18,7 +18,7 @@ from unique_toolkit.language_model.schemas import (
 from unique_toolkit.language_model.service import LanguageModelService
 
 
-from unique_toolkit.evaluators.schemas import EvaluationMetricName
+from unique_toolkit.evals.schemas import EvaluationMetricName
 from unique_toolkit.tools.agent_chunks_handler import AgentChunksHandler
 from unique_toolkit.tools.config import ToolBuildConfig, ToolSelectionPolicy
 from unique_toolkit.tools.schemas import BaseToolConfig, ToolCallResponse, ToolPrompts