unique_toolkit 0.7.7__py3-none-any.whl → 1.23.0__py3-none-any.whl

unique_toolkit/agentic/short_term_memory_manager/persistent_short_term_memory_manager.py

@@ -0,0 +1,141 @@
+import base64
+import zlib
+from logging import getLogger
+from typing import Generic, Type, TypeVar
+
+from pydantic import BaseModel
+
+from unique_toolkit.agentic.tools.utils.execution.execution import SafeTaskExecutor
+from unique_toolkit.short_term_memory.schemas import ShortTermMemory
+from unique_toolkit.short_term_memory.service import ShortTermMemoryService
+
+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/agentic/thinking_manager/thinking_manager.py

@@ -0,0 +1,103 @@
+from logging import Logger
+
+from pydantic import BaseModel, Field
+
+from unique_toolkit.agentic.tools.tool_progress_reporter import (
+    ToolProgressReporter,
+)
+from unique_toolkit.chat.service import ChatService
+from unique_toolkit.language_model.schemas import (
+    LanguageModelStreamResponse,
+)
+
+
+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/agentic/tools/__init__.py

@@ -0,0 +1 @@
+"""Tools module for the Unique Toolkit."""

unique_toolkit/agentic/tools/a2a/__init__.py

@@ -0,0 +1,36 @@
+from unique_toolkit.agentic.tools.a2a.config import ExtendedSubAgentToolConfig
+from unique_toolkit.agentic.tools.a2a.evaluation import (
+    SubAgentEvaluationService,
+    SubAgentEvaluationServiceConfig,
+    SubAgentEvaluationSpec,
+)
+from unique_toolkit.agentic.tools.a2a.manager import A2AManager
+from unique_toolkit.agentic.tools.a2a.postprocessing import (
+    SubAgentDisplaySpec,
+    SubAgentReferencesPostprocessor,
+    SubAgentResponsesDisplayPostprocessor,
+    SubAgentResponsesPostprocessorConfig,
+)
+from unique_toolkit.agentic.tools.a2a.prompts import (
+    REFERENCING_INSTRUCTIONS_FOR_SYSTEM_PROMPT,
+    REFERENCING_INSTRUCTIONS_FOR_USER_PROMPT,
+)
+from unique_toolkit.agentic.tools.a2a.response_watcher import SubAgentResponseWatcher
+from unique_toolkit.agentic.tools.a2a.tool import SubAgentTool, SubAgentToolConfig
+
+__all__ = [
+    "SubAgentToolConfig",
+    "SubAgentTool",
+    "SubAgentResponsesDisplayPostprocessor",
+    "SubAgentResponsesPostprocessorConfig",
+    "SubAgentDisplaySpec",
+    "A2AManager",
+    "ExtendedSubAgentToolConfig",
+    "SubAgentEvaluationServiceConfig",
+    "SubAgentEvaluationService",
+    "REFERENCING_INSTRUCTIONS_FOR_SYSTEM_PROMPT",
+    "REFERENCING_INSTRUCTIONS_FOR_USER_PROMPT",
+    "SubAgentResponseWatcher",
+    "SubAgentReferencesPostprocessor",
+    "SubAgentEvaluationSpec",
+]

unique_toolkit/agentic/tools/a2a/config.py

@@ -0,0 +1,17 @@
+from pydantic import Field
+
+from unique_toolkit.agentic.tools.a2a.evaluation import SubAgentEvaluationConfig
+from unique_toolkit.agentic.tools.a2a.postprocessing import SubAgentDisplayConfig
+from unique_toolkit.agentic.tools.a2a.tool import SubAgentToolConfig
+
+
+# SubAgentToolConfig with display and evaluation configs
+class ExtendedSubAgentToolConfig(SubAgentToolConfig):
+    response_display_config: SubAgentDisplayConfig = Field(
+        default_factory=SubAgentDisplayConfig,
+        description="Configuration for how to display the sub-agent response.",
+    )
+    evaluation_config: SubAgentEvaluationConfig = Field(
+        default_factory=SubAgentEvaluationConfig,
+        description="Configuration for handling assessments of the sub-agent response.",
+    )

unique_toolkit/agentic/tools/a2a/evaluation/__init__.py

@@ -0,0 +1,15 @@
+from unique_toolkit.agentic.tools.a2a.evaluation.config import (
+    SubAgentEvaluationConfig,
+    SubAgentEvaluationServiceConfig,
+)
+from unique_toolkit.agentic.tools.a2a.evaluation.evaluator import (
+    SubAgentEvaluationService,
+    SubAgentEvaluationSpec,
+)
+
+__all__ = [
+    "SubAgentEvaluationService",
+    "SubAgentEvaluationServiceConfig",
+    "SubAgentEvaluationConfig",
+    "SubAgentEvaluationSpec",
+]

unique_toolkit/agentic/tools/a2a/evaluation/_utils.py

@@ -0,0 +1,66 @@
+import logging
+
+import unique_sdk
+
+from unique_toolkit.chat.schemas import (
+    ChatMessageAssessmentLabel,
+    ChatMessageAssessmentStatus,
+)
+
+logger = logging.getLogger(__name__)
+
+_ASSESSMENT_LABEL_COMPARISON_DICT: dict[str, int] = {
+    ChatMessageAssessmentLabel.RED: 0,
+    ChatMessageAssessmentLabel.YELLOW: 1,
+    ChatMessageAssessmentLabel.GREEN: 2,
+}
+
+
+def sort_assessments(
+    assessments: list[unique_sdk.Space.Assessment],
+) -> list[unique_sdk.Space.Assessment]:
+    return sorted(
+        assessments,
+        key=lambda x: _ASSESSMENT_LABEL_COMPARISON_DICT[x["label"]],  # type: ignore (should be checked before sorting)
+    )
+
+
+def get_worst_label(
+    *labels: str,
+) -> str:
+    return min(
+        labels,
+        key=lambda x: _ASSESSMENT_LABEL_COMPARISON_DICT[x],
+    )
+
+
+def get_valid_assessments(
+    assessments: list[unique_sdk.Space.Assessment],
+    display_name: str,
+    sequence_number: int,
+) -> list[unique_sdk.Space.Assessment]:
+    valid_assessments = []
+    for assessment in assessments:
+        if (
+            assessment["label"] is None
+            or assessment["label"] not in ChatMessageAssessmentLabel
+        ):
+            logger.warning(
+                "Unkown assistant label %s for assistant %s (sequence number: %s) will be ignored",
+                assessment["label"],
+                display_name,
+                sequence_number,
+            )
+            continue
+        if assessment["status"] != ChatMessageAssessmentStatus.DONE:
+            logger.warning(
+                "Assessment %s for assistant %s (sequence number: %s) is not done (status: %s) will be ignored",
+                assessment["label"],
+                display_name,
+                sequence_number,
+                assessment["status"],
+            )
+            continue
+        valid_assessments.append(assessment)
+
+    return valid_assessments

unique_toolkit/agentic/tools/a2a/evaluation/config.py

@@ -0,0 +1,55 @@
+from pathlib import Path
+
+from pydantic import AliasChoices, BaseModel, Field
+
+from unique_toolkit._common.pydantic_helpers import get_configuration_dict
+from unique_toolkit._common.validators import LMI, get_LMI_default_field
+from unique_toolkit.chat.schemas import (
+    ChatMessageAssessmentType,
+)
+from unique_toolkit.language_model.default_language_model import DEFAULT_GPT_4o
+
+DEFAULT_EVALUATION_SYSTEM_MESSAGE_TEMPLATE = """
+You are a through and precise summarization model.
+You will receive a list of "assessments" of one or more agent(s) response(s).
+Your task is to give a brief summary (1-10 sentences) of the received assessments, following the following guidelines:
+1. You must NOT in ANY case state a fact that is not stated in the given assessments.
+2. You must focus first and foremost on the failing assessments, labeled `RED` below.
+3. You must mention each agent's name when summarizing its list of assessments.
+4. You must NOT use any markdown formatting in your response as this will FAIL to render in the chat frontend.
+""".strip()
+
+with open(Path(__file__).parent / "summarization_user_message.j2", "r") as file:
+    DEFAULT_SUMMARIZATION_USER_MESSAGE_TEMPLATE = file.read().strip()
+
+
+class SubAgentEvaluationServiceConfig(BaseModel):
+    model_config = get_configuration_dict()
+
+    assessment_type: ChatMessageAssessmentType = Field(
+        default=ChatMessageAssessmentType.COMPLIANCE,
+        description="The type of assessment to use in the display.",
+    )
+    summarization_model: LMI = get_LMI_default_field(DEFAULT_GPT_4o)
+    summarization_system_message: str = Field(
+        default=DEFAULT_EVALUATION_SYSTEM_MESSAGE_TEMPLATE,
+        description="The system message template for the summarization model.",
+    )
+    summarization_user_message_template: str = Field(
+        default=DEFAULT_SUMMARIZATION_USER_MESSAGE_TEMPLATE,
+        description="The user message template for the summarization model.",
+    )
+
+
+class SubAgentEvaluationConfig(BaseModel):
+    model_config = get_configuration_dict()
+
+    include_evaluation: bool = Field(
+        default=True,
+        description="Whether to include the evaluation in the response.",
+        validation_alias=AliasChoices(
+            "includeEvaluation",
+            "displayEvalution",  # typo in old config name
+            "display_evalution",
+        ),
+    )