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

unique_toolkit/agentic/postprocessor/postprocessor_manager.py

@@ -6,6 +6,7 @@ from unique_toolkit.agentic.tools.utils.execution.execution import SafeTaskExecu
 from unique_toolkit.chat.service import ChatService
 from unique_toolkit.language_model.schemas import (
     LanguageModelStreamResponse,
+    ResponsesLanguageModelStreamResponse,
 )
 
 
@@ -26,7 +27,30 @@ class Postprocessor(ABC):
             "Subclasses must implement this method to apply post-processing to the response."
         )
 
-    async def remove_from_text(self, text) -> str:
+    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."
         )
@@ -59,13 +83,41 @@ class PostprocessorManager:
     ):
         self._logger = logger
         self._chat_service = chat_service
-        self._postprocessors: list[Postprocessor] = []
+        self._postprocessors: list[Postprocessor | ResponsesApiPostprocessor] = []
 
-    def add_postprocessor(self, postprocessor: Postprocessor):
+        # 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 get_postprocessors(self, name: str) -> list[Postprocessor]:
-        return self._postprocessors
+    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,
@@ -75,25 +127,26 @@ class PostprocessorManager:
             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 self._postprocessors
+            for postprocessor in 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."
-                )
+        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)
-            for postprocessor in self._postprocessors
+            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)
@@ -108,14 +161,52 @@ class PostprocessorManager:
     async def execute_postprocessors(
         self,
         loop_response: LanguageModelStreamResponse,
-        postprocessor_instance: Postprocessor,
+        postprocessor_instance: Postprocessor | ResponsesApiPostprocessor,
     ) -> None:
-        await postprocessor_instance.run(loop_response)
+        await postprocessor_instance.run(loop_response)  # type: ignore
 
     async def remove_from_text(
         self,
         text: str,
     ) -> str:
-        for postprocessor in self._postprocessors:
+        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

@@ -98,6 +98,6 @@ class ReferenceManager:
         referenced_chunks: list[ContentChunk] = []
         for ref in references:
             for chunk in self._chunks:
-                if ref.source_id == f"{chunk.id}-{chunk.chunk_id}":
+                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,63 @@
+import logging
+import re
+from typing import override
+
+from pydantic import BaseModel, Field
+
+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: bool = Field(
+        default=False,
+        description="If set, the code interpreter call will be removed from the history on subsequent calls to the assistant.",
+    )
+
+
+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:
+        return None
+
+    @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,145 @@
+import logging
+import re
+from mimetypes import guess_type
+from typing import override
+
+from openai import AsyncOpenAI
+from pydantic import BaseModel
+from unique_sdk import Content
+
+from unique_toolkit.agentic.postprocessor.postprocessor_manager import (
+    ResponsesApiPostprocessor,
+)
+from unique_toolkit.agentic.tools.config import get_configuration_dict
+from unique_toolkit.content.schemas import ContentReference
+from unique_toolkit.content.service import ContentService
+from unique_toolkit.language_model.schemas import ResponsesLanguageModelStreamResponse
+
+logger = logging.getLogger(__name__)
+
+
+class DisplayCodeInterpreterFilesPostProcessorConfig(BaseModel):
+    model_config = get_configuration_dict()
+    upload_scope_id: str
+
+
+class DisplayCodeInterpreterFilesPostProcessor(
+    ResponsesApiPostprocessor,
+):
+    def __init__(
+        self,
+        client: AsyncOpenAI,
+        content_service: ContentService,
+        config: DisplayCodeInterpreterFilesPostProcessorConfig,
+    ) -> None:
+        super().__init__(self.__class__.__name__)
+        self._content_service = content_service
+        self._config = config
+        self._client = client
+        self._content_map = {}
+
+    @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 = {}
+        for container_file in container_files:
+            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 = self._content_service.upload_content_from_bytes(
+                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,
+            )
+            self._content_map[container_file.filename] = content
+
+    @override
+    def apply_postprocessing_to_response(
+        self, loop_response: ResponsesLanguageModelStreamResponse
+    ) -> bool:
+        ref_number = _get_next_ref_number(loop_response.message.references)
+        changed = False
+        # images
+        for filename, content in self._content_map.items():
+            # Images
+            loop_response.message.text, replaced = _replace_container_image_citation(
+                text=loop_response.message.text, filename=filename, content=content
+            )
+            changed |= replaced
+
+            # Files
+            loop_response.message.text, replaced = _replace_container_file_citation(
+                text=loop_response.message.text,
+                filename=filename,
+                ref_number=ref_number,
+            )
+            changed |= replaced
+            if replaced:
+                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
+
+
+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_image_citation(
+    text: str, filename: str, content: Content
+) -> 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

unique_toolkit/agentic/responses_api/stream_handler.py

@@ -0,0 +1,15 @@
+from unique_toolkit.protocols.support import ResponsesSupportCompleteWithReferences
+from unique_toolkit.services.chat_service import ChatService
+
+
+class ResponsesStreamingHandler(ResponsesSupportCompleteWithReferences):
+    def __init__(self, chat_service: ChatService):
+        self._chat_service = chat_service
+
+    def complete_with_references(self, *args, **kwargs):
+        return self._chat_service.complete_responses_with_references(*args, **kwargs)
+
+    async def complete_with_references_async(self, *args, **kwargs):
+        return await self._chat_service.complete_responses_with_references_async(
+            *args, **kwargs
+        )

unique_toolkit/agentic/tools/a2a/__init__.py

@@ -2,19 +2,35 @@ 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 (
-    SubAgentResponsesPostprocessor,
+    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",
-    "SubAgentResponsesPostprocessor",
+    "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/evaluation/__init__.py

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

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

@@ -16,7 +16,7 @@ _ASSESSMENT_LABEL_COMPARISON_DICT: dict[str, int] = {
 }
 
 
-def _sort_assessments(
+def sort_assessments(
     assessments: list[unique_sdk.Space.Assessment],
 ) -> list[unique_sdk.Space.Assessment]:
     return sorted(
@@ -25,7 +25,7 @@ def _sort_assessments(
     )
 
 
-def _worst_label(
+def get_worst_label(
     *labels: str,
 ) -> str:
     return min(
@@ -34,7 +34,7 @@ def _worst_label(
     )
 
 
-def _get_valid_assessments(
+def get_valid_assessments(
     assessments: list[unique_sdk.Space.Assessment],
     display_name: str,
     sequence_number: int,

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

@@ -2,12 +2,12 @@ from pathlib import Path
 
 from pydantic import AliasChoices, BaseModel, Field
 
-from unique_toolkit._common.default_language_model import DEFAULT_GPT_4o
 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.