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

unique_toolkit/agentic/postprocessor/postprocessor_manager.py

@@ -0,0 +1,212 @@
+import asyncio
+from abc import ABC
+from logging import Logger
+
+from unique_toolkit.agentic.tools.utils.execution.execution import SafeTaskExecutor
+from unique_toolkit.chat.service import ChatService
+from unique_toolkit.language_model.schemas import (
+    LanguageModelStreamResponse,
+    ResponsesLanguageModelStreamResponse,
+)
+
+
+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) -> None:
+        raise NotImplementedError("Subclasses must implement this method.")
+
+    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) -> str:
+        raise NotImplementedError(
+            "Subclasses must implement this method to remove post-processing from the message."
+        )
+
+
+class ResponsesApiPostprocessor(ABC):
+    def __init__(self, name: str):
+        self.name = name
+
+    def get_name(self) -> str:
+        return self.name
+
+    async def run(self, loop_response: ResponsesLanguageModelStreamResponse) -> None:
+        raise NotImplementedError("Subclasses must implement this method.")
+
+    def apply_postprocessing_to_response(
+        self, loop_response: ResponsesLanguageModelStreamResponse
+    ) -> bool:
+        raise NotImplementedError(
+            "Subclasses must implement this method to apply post-processing to the response."
+        )
+
+    async def remove_from_text(self, text: str) -> 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.
+    """
+
+    def __init__(
+        self,
+        logger: Logger,
+        chat_service: ChatService,
+    ):
+        self._logger = logger
+        self._chat_service = chat_service
+        self._postprocessors: list[Postprocessor | ResponsesApiPostprocessor] = []
+
+        # Allow to add postprocessors that should be run before or after the others.
+        self._first_postprocessor: Postprocessor | ResponsesApiPostprocessor | None = (
+            None
+        )
+        self._last_postprocessor: Postprocessor | ResponsesApiPostprocessor | None = (
+            None
+        )
+
+    def add_postprocessor(
+        self, postprocessor: Postprocessor | ResponsesApiPostprocessor
+    ):
+        self._postprocessors.append(postprocessor)
+
+    def set_first_postprocessor(
+        self, postprocessor: Postprocessor | ResponsesApiPostprocessor
+    ) -> None:
+        if self._first_postprocessor is not None:
+            raise ValueError("Cannot set first postprocessor if already set.")
+
+        self._first_postprocessor = postprocessor
+
+    def set_last_postprocessor(
+        self, postprocessor: Postprocessor | ResponsesApiPostprocessor
+    ) -> None:
+        if self._last_postprocessor is not None:
+            raise ValueError("Cannot set last postprocessor if already set.")
+
+        self._last_postprocessor = postprocessor
+
+    def get_postprocessors(
+        self, name: str
+    ) -> list[Postprocessor | ResponsesApiPostprocessor]:
+        return self._get_all_postprocessors()
+
+    async def run_postprocessors(
+        self,
+        loop_response: LanguageModelStreamResponse,
+    ) -> None:
+        task_executor = SafeTaskExecutor(
+            logger=self._logger,
+        )
+
+        postprocessors = self._get_valid_postprocessors_for_loop_response(loop_response)
+
+        tasks = [
+            task_executor.execute_async(
+                self.execute_postprocessors,
+                loop_response=loop_response,
+                postprocessor_instance=postprocessor,
+            )
+            for postprocessor in postprocessors
+        ]
+        postprocessor_results = await asyncio.gather(*tasks)
+
+        successful_postprocessors: list[Postprocessor | ResponsesApiPostprocessor] = []
+        for i in range(len(postprocessors)):
+            if postprocessor_results[i].success:
+                successful_postprocessors.append(postprocessors[i])
+
+        modification_results = [
+            postprocessor.apply_postprocessing_to_response(loop_response)  # type: ignore (checked in `get_valid_postprocessors_for_loop_response`)
+            for postprocessor in successful_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,
+                references=loop_response.message.references,
+            )
+
+    async def execute_postprocessors(
+        self,
+        loop_response: LanguageModelStreamResponse,
+        postprocessor_instance: Postprocessor | ResponsesApiPostprocessor,
+    ) -> None:
+        await postprocessor_instance.run(loop_response)  # type: ignore
+
+    async def remove_from_text(
+        self,
+        text: str,
+    ) -> str:
+        for postprocessor in self._get_all_postprocessors():
+            text = await postprocessor.remove_from_text(text)
+        return text
+
+    def _get_all_postprocessors(
+        self,
+    ) -> list[Postprocessor | ResponsesApiPostprocessor]:
+        return [
+            postprocessor
+            for postprocessor in [
+                self._first_postprocessor,
+                *self._postprocessors,
+                self._last_postprocessor,
+            ]
+            if postprocessor is not None
+        ]
+
+    def _get_valid_postprocessors_for_loop_response(
+        self, loop_response: LanguageModelStreamResponse
+    ) -> list[Postprocessor | ResponsesApiPostprocessor]:
+        all_postprocessors = self._get_all_postprocessors()
+
+        postprocessors: list[Postprocessor | ResponsesApiPostprocessor] = []
+
+        if isinstance(loop_response, ResponsesLanguageModelStreamResponse):
+            """
+            All processore can be executed, since `ResponsesLanguageModelStreamResponse`
+            is a subclass of `LanguageModelStreamResponse`
+            """
+            postprocessors = all_postprocessors
+        else:
+            """
+            Cannot execute Responses API-specific postprocessors
+            """
+            postprocessors = [
+                postprocessor
+                for postprocessor in all_postprocessors
+                if isinstance(postprocessor, Postprocessor)
+            ]
+
+        return postprocessors

unique_toolkit/agentic/reference_manager/reference_manager.py

@@ -0,0 +1,103 @@
+from unique_toolkit.agentic.tools.schemas import ToolCallResponse
+from unique_toolkit.content.schemas import ContentChunk, ContentReference
+
+
+class tool_chunks:
+    def __init__(self, name: str, chunks: list[ContentChunk]) -> None:
+        self.name = name
+        self.chunks = 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] = []
+        self._references: list[list[ContentReference]] = []
+
+    def extract_referenceable_chunks(
+        self, tool_responses: list[ToolCallResponse]
+    ) -> None:
+        for tool_response in tool_responses:
+            if not tool_response.content_chunks:
+                continue
+            self._chunks.extend(tool_response.content_chunks or [])
+            self._tool_chunks[tool_response.id] = tool_chunks(
+                tool_response.name, tool_response.content_chunks
+            )
+
+    def get_chunks(self) -> list[ContentChunk]:
+        return self._chunks
+
+    def get_tool_chunks(self) -> dict[str, tool_chunks]:
+        return self._tool_chunks
+
+    def get_chunks_of_all_tools(self) -> list[list[ContentChunk]]:
+        return [tool_chunks.chunks for tool_chunks in self._tool_chunks.values()]
+
+    def get_chunks_of_tool(self, tool_call_id: str) -> list[ContentChunk]:
+        return self._tool_chunks.get(tool_call_id, tool_chunks("", [])).chunks
+
+    def replace_chunks_of_tool(
+        self, tool_call_id: str, chunks: list[ContentChunk]
+    ) -> None:
+        if tool_call_id in self._tool_chunks:
+            self._tool_chunks[tool_call_id].chunks = chunks
+
+    def replace(self, chunks: list[ContentChunk]):
+        self._chunks = chunks
+
+    def add_references(
+        self,
+        references: list[ContentReference],
+    ):
+        self._references.append(references)
+
+    def get_references(
+        self,
+    ) -> list[list[ContentReference]]:
+        return self._references
+
+    def get_latest_references(
+        self,
+    ) -> list[ContentReference]:
+        if not self._references:
+            return []
+        return self._references[-1]
+
+    def get_latest_referenced_chunks(self) -> list[ContentChunk]:
+        if not self._references:
+            return []
+        return self._get_referenced_chunks_from_references(self._references[-1])
+
+    def _get_referenced_chunks_from_references(
+        self,
+        references: list[ContentReference],
+    ) -> list[ContentChunk]:
+        """
+        Get _referenced_chunks by matching sourceId from _references with merged id and chunk_id from _chunks.
+        """
+        referenced_chunks: list[ContentChunk] = []
+        for ref in references:
+            for chunk in self._chunks:
+                if ref.source_id == f"{chunk.id}_{chunk.chunk_id}":
+                    referenced_chunks.append(chunk)
+        return referenced_chunks

unique_toolkit/agentic/responses_api/__init__.py

@@ -0,0 +1,19 @@
+from unique_toolkit.agentic.responses_api.postprocessors.code_display import (
+    ShowExecutedCodePostprocessor,
+    ShowExecutedCodePostprocessorConfig,
+)
+from unique_toolkit.agentic.responses_api.postprocessors.generated_files import (
+    DisplayCodeInterpreterFilesPostProcessor,
+    DisplayCodeInterpreterFilesPostProcessorConfig,
+)
+from unique_toolkit.agentic.responses_api.stream_handler import (
+    ResponsesStreamingHandler,
+)
+
+__all__ = [
+    "ShowExecutedCodePostprocessor",
+    "ShowExecutedCodePostprocessorConfig",
+    "DisplayCodeInterpreterFilesPostProcessorConfig",
+    "DisplayCodeInterpreterFilesPostProcessor",
+    "ResponsesStreamingHandler",
+]

unique_toolkit/agentic/responses_api/postprocessors/code_display.py

@@ -0,0 +1,71 @@
+import asyncio
+import logging
+import re
+from typing import override
+
+from pydantic import BaseModel, Field
+from pydantic.json_schema import SkipJsonSchema
+
+from unique_toolkit.agentic.postprocessor.postprocessor_manager import (
+    ResponsesApiPostprocessor,
+)
+from unique_toolkit.agentic.tools.config import get_configuration_dict
+from unique_toolkit.language_model.schemas import ResponsesLanguageModelStreamResponse
+
+_TEMPLATE = """
+<details><summary>Code Interpreter Call</summary>    
+
+```python
+{code}
+```  
+
+</details>    
+</br>
+
+""".lstrip()
+
+logger = logging.getLogger(__name__)
+
+
+class ShowExecutedCodePostprocessorConfig(BaseModel):
+    model_config = get_configuration_dict()
+    remove_from_history: SkipJsonSchema[bool] = (
+        Field(  # At the moment, it's not possible to keep executed code in the history
+            default=True,
+            description="If set, the code interpreter call will be removed from the history on subsequent calls to the assistant.",
+        )
+    )
+    sleep_time_before_display: float = Field(
+        default=0.2,
+        description="Time to sleep before displaying the executed code. Please increase this value if you experience rendering issues.",
+    )
+
+
+class ShowExecutedCodePostprocessor(ResponsesApiPostprocessor):
+    def __init__(self, config: ShowExecutedCodePostprocessorConfig):
+        super().__init__(self.__class__.__name__)
+        self._config = config
+
+    @override
+    async def run(self, loop_response: ResponsesLanguageModelStreamResponse) -> None:
+        await asyncio.sleep(self._config.sleep_time_before_display)
+
+    @override
+    def apply_postprocessing_to_response(
+        self, loop_response: ResponsesLanguageModelStreamResponse
+    ) -> bool:
+        prepended_text = ""
+        for output in loop_response.code_interpreter_calls:
+            prepended_text += _TEMPLATE.format(code=output.code)
+
+        loop_response.message.text = prepended_text + loop_response.message.text
+
+        return prepended_text != ""
+
+    @override
+    async def remove_from_text(self, text) -> str:
+        if not self._config.remove_from_history:
+            return text
+        # Remove code interpreter blocks using regex
+        pattern = r"<details><summary>Code Interpreter Call</summary>.*?</details>"
+        return re.sub(pattern, "", text, flags=re.DOTALL)

unique_toolkit/agentic/responses_api/postprocessors/generated_files.py

@@ -0,0 +1,297 @@
+import asyncio
+import logging
+import re
+from mimetypes import guess_type
+from typing import NamedTuple, override
+
+from openai import AsyncOpenAI
+from openai.types.responses.response_output_text import AnnotationContainerFileCitation
+from pydantic import BaseModel, Field, RootModel
+
+from unique_toolkit.agentic.postprocessor.postprocessor_manager import (
+    ResponsesApiPostprocessor,
+)
+from unique_toolkit.agentic.short_term_memory_manager.persistent_short_term_memory_manager import (
+    PersistentShortMemoryManager,
+)
+from unique_toolkit.agentic.tools.config import get_configuration_dict
+from unique_toolkit.agentic.tools.utils import failsafe_async
+from unique_toolkit.content.schemas import ContentReference
+from unique_toolkit.content.service import ContentService
+from unique_toolkit.language_model.schemas import ResponsesLanguageModelStreamResponse
+from unique_toolkit.services.knowledge_base import KnowledgeBaseService
+from unique_toolkit.short_term_memory.service import ShortTermMemoryService
+
+logger = logging.getLogger(__name__)
+
+
+class DisplayCodeInterpreterFilesPostProcessorConfig(BaseModel):
+    model_config = get_configuration_dict()
+    upload_scope_id: str = Field(
+        default="<SCOPE_ID_PLACEHOLDER>",
+        description="The scope ID where the generated files will be uploaded.",
+    )
+
+    file_download_failed_message: str = Field(
+        default="⚠️ File download failed ...",
+        description="The message to display when a file download fails.",
+    )
+    max_concurrent_file_downloads: int = Field(
+        default=10,
+        description="The maximum number of concurrent file downloads.",
+    )
+
+
+class _ContentInfo(NamedTuple):
+    filename: str
+    content_id: str
+
+
+_SHORT_TERM_MEMORY_KEY = "code_interpreter_files"
+_DisplayedFilesShortTermMemorySchema = RootModel[list[_ContentInfo]]
+
+
+def _init_short_term_memory_manager(
+    company_id: str, user_id: str, chat_id: str
+) -> PersistentShortMemoryManager[_DisplayedFilesShortTermMemorySchema]:
+    short_term_memory_service = ShortTermMemoryService(
+        company_id=company_id,
+        user_id=user_id,
+        chat_id=chat_id,
+        message_id=None,
+    )
+    return PersistentShortMemoryManager(
+        short_term_memory_service=short_term_memory_service,
+        short_term_memory_schema=_DisplayedFilesShortTermMemorySchema,
+        short_term_memory_name=_SHORT_TERM_MEMORY_KEY,
+    )
+
+
+class DisplayCodeInterpreterFilesPostProcessor(
+    ResponsesApiPostprocessor,
+):
+    def __init__(
+        self,
+        client: AsyncOpenAI,
+        content_service: ContentService | KnowledgeBaseService,
+        config: DisplayCodeInterpreterFilesPostProcessorConfig,
+        # Short term memory arguments, we prefer to explicitely pass the required auth variables
+        # as it is crucial that we use chat-level short term memory not to leak user files to other chats.
+        # Technically, short term memory can be scoped company-level, we would like to ensure this case is avoided.
+        company_id: str | None = None,
+        user_id: str | None = None,
+        chat_id: str | None = None,
+    ) -> None:
+        super().__init__(self.__class__.__name__)
+        self._content_service = content_service
+        self._config = config
+        self._client = client
+
+        self._short_term_memory_manager = None
+        if chat_id is not None and user_id is not None and company_id is not None:
+            self._short_term_memory_manager = _init_short_term_memory_manager(
+                company_id=company_id,
+                user_id=user_id,
+                chat_id=chat_id,
+            )
+
+    @override
+    async def run(self, loop_response: ResponsesLanguageModelStreamResponse) -> None:
+        logger.info("Fetching and adding code interpreter files to the response")
+
+        container_files = loop_response.container_files
+        logger.info("Found %s container files", len(container_files))
+
+        self._content_map: dict[str, str | None] = await self._load_previous_files()  # type: ignore
+
+        semaphore = asyncio.Semaphore(self._config.max_concurrent_file_downloads)
+        tasks = [
+            self._download_and_upload_container_files_to_knowledge_base(
+                citation, semaphore
+            )
+            for citation in container_files
+        ]
+        results = await asyncio.gather(*tasks)
+
+        for citation, result in zip(container_files, results):
+            # Overwrite if file has been re-generated
+            self._content_map[citation.filename] = (
+                result.content_id if result is not None else None
+            )
+
+        await self._save_generated_files(
+            [
+                _ContentInfo(filename=filename, content_id=content_id)
+                for filename, content_id in self._content_map.items()
+                if content_id is not None
+            ]
+        )
+
+    @override
+    def apply_postprocessing_to_response(
+        self, loop_response: ResponsesLanguageModelStreamResponse
+    ) -> bool:
+        ref_number = _get_next_ref_number(loop_response.message.references)
+        changed = False
+
+        for filename, content_id in self._content_map.items():
+            if content_id is None:
+                loop_response.message.text, replaced = _replace_container_file_error(
+                    text=loop_response.message.text,
+                    filename=filename,
+                    error_message=self._config.file_download_failed_message,
+                )
+                changed |= replaced
+                continue
+
+            is_image = (guess_type(filename)[0] or "").startswith("image/")
+
+            # Images
+            if is_image:
+                loop_response.message.text, replaced = (
+                    _replace_container_image_citation(
+                        text=loop_response.message.text,
+                        filename=filename,
+                        content_id=content_id,
+                    )
+                )
+                changed |= replaced
+
+            # Files
+            else:
+                loop_response.message.text, replaced = _replace_container_file_citation(
+                    text=loop_response.message.text,
+                    filename=filename,
+                    ref_number=ref_number,
+                )
+                changed |= replaced
+
+            if replaced and not is_image:
+                loop_response.message.references.append(
+                    ContentReference(
+                        sequence_number=ref_number,
+                        source_id=content_id,
+                        source="node-ingestion-chunks",
+                        url=f"unique://content/{content_id}",
+                        name=filename,
+                    )
+                )
+                ref_number += 1
+        return changed
+
+    @override
+    async def remove_from_text(self, text) -> str:
+        return text
+
+    @failsafe_async(failure_return_value=None, logger=logger)
+    async def _download_and_upload_container_files_to_knowledge_base(
+        self,
+        container_file: AnnotationContainerFileCitation,
+        semaphore: asyncio.Semaphore,
+    ) -> _ContentInfo | None:
+        async with semaphore:
+            logger.info("Fetching file content for %s", container_file.filename)
+            file_content = await self._client.containers.files.content.retrieve(
+                container_id=container_file.container_id, file_id=container_file.file_id
+            )
+
+            logger.info(
+                "Uploading file content for %s to knowledge base",
+                container_file.filename,
+            )
+            content = await self._content_service.upload_content_from_bytes_async(
+                content=file_content.content,
+                content_name=container_file.filename,
+                skip_ingestion=True,
+                mime_type=guess_type(container_file.filename)[0] or "text/plain",
+                scope_id=self._config.upload_scope_id,
+            )
+
+            return _ContentInfo(filename=container_file.filename, content_id=content.id)
+
+    async def _load_previous_files(self) -> dict[str, str]:
+        if self._short_term_memory_manager is None:
+            return {}
+
+        logger.info(
+            "Loading previously generated code interpreter files from short term memory"
+        )
+        memory = await self._short_term_memory_manager.load_async()
+
+        if memory is None:
+            logger.info(
+                "No previously generated code interpreter files found in short term memory"
+            )
+            return {}
+
+        logger.info(
+            "Found %s previously generated code interpreter files", len(memory.root)
+        )
+
+        return {content.filename: content.content_id for content in memory.root}
+
+    async def _save_generated_files(self, content_infos: list[_ContentInfo]) -> None:
+        if self._short_term_memory_manager is None or len(content_infos) == 0:
+            return
+
+        await self._short_term_memory_manager.save_async(
+            _DisplayedFilesShortTermMemorySchema(root=content_infos)
+        )
+
+
+def _get_next_ref_number(references: list[ContentReference]) -> int:
+    max_ref_number = 0
+    for ref in references:
+        max_ref_number = max(max_ref_number, ref.sequence_number)
+    return max_ref_number + 1
+
+
+def _replace_container_file_error(
+    text: str, filename: str, error_message: str
+) -> tuple[str, bool]:
+    image_markdown = rf"!?\[.*?\]\(sandbox:/mnt/data/{re.escape(filename)}\)"
+
+    if not re.search(image_markdown, text):
+        logger.info("No image markdown found for %s", filename)
+        return text, False
+
+    logger.info("Displaying image %s", filename)
+    return re.sub(
+        image_markdown,
+        error_message,
+        text,
+    ), True
+
+
+def _replace_container_image_citation(
+    text: str, filename: str, content_id: str
+) -> tuple[str, bool]:
+    image_markdown = rf"!?\[.*?\]\(sandbox:/mnt/data/{re.escape(filename)}\)"
+
+    if not re.search(image_markdown, text):
+        logger.info("No image markdown found for %s", filename)
+        return text, False
+
+    logger.info("Displaying image %s", filename)
+    return re.sub(
+        image_markdown,
+        f"![image](unique://content/{content_id})",
+        text,
+    ), True
+
+
+def _replace_container_file_citation(
+    text: str, filename: str, ref_number: int
+) -> tuple[str, bool]:
+    file_markdown = rf"\[.*?\]\(sandbox:/mnt/data/{re.escape(filename)}\)"
+
+    if not re.search(file_markdown, text):
+        logger.info("No file markdown found for %s", filename)
+        return text, False
+
+    logger.info("Displaying file %s", filename)
+    return re.sub(
+        file_markdown,
+        f"<sup>{ref_number}</sup>",
+        text,
+    ), True