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

unique_toolkit/agentic/tools/a2a/postprocessing/display.py

@@ -0,0 +1,180 @@
+import asyncio
+import logging
+from collections import defaultdict
+from typing import NamedTuple, override
+
+import unique_sdk
+from pydantic import BaseModel, Field
+
+from unique_toolkit._common.pydantic_helpers import get_configuration_dict
+from unique_toolkit.agentic.postprocessor.postprocessor_manager import Postprocessor
+from unique_toolkit.agentic.tools.a2a.postprocessing._display_utils import (
+    get_sub_agent_answer_display,
+    remove_sub_agent_answer_from_text,
+)
+from unique_toolkit.agentic.tools.a2a.postprocessing._ref_utils import (
+    add_content_refs_and_replace_in_text,
+)
+from unique_toolkit.agentic.tools.a2a.postprocessing.config import (
+    SubAgentDisplayConfig,
+    SubAgentResponseDisplayMode,
+)
+from unique_toolkit.agentic.tools.a2a.response_watcher import (
+    SubAgentResponse,
+    SubAgentResponseWatcher,
+)
+from unique_toolkit.content import ContentReference
+from unique_toolkit.language_model.schemas import LanguageModelStreamResponse
+
+logger = logging.getLogger(__name__)
+
+SpaceMessage = unique_sdk.Space.Message
+
+
+class SubAgentDisplaySpec(NamedTuple):
+    assistant_id: str
+    display_name: str
+    display_config: SubAgentDisplayConfig
+
+
+class SubAgentResponsesPostprocessorConfig(BaseModel):
+    model_config = get_configuration_dict()
+
+    sleep_time_before_update: float = Field(
+        default=1, description="Time to sleep before updating the main agent message."
+    )
+
+
+class SubAgentResponsesDisplayPostprocessor(Postprocessor):
+    def __init__(
+        self,
+        config: SubAgentResponsesPostprocessorConfig,
+        response_watcher: SubAgentResponseWatcher,
+        display_specs: list[SubAgentDisplaySpec],
+    ) -> None:
+        super().__init__(name=self.__class__.__name__)
+
+        self._config = config
+        self._response_watcher = response_watcher
+        self._display_specs: dict[str, SubAgentDisplaySpec] = {
+            display_spec.assistant_id: display_spec
+            for display_spec in display_specs
+            if display_spec.display_config.mode != SubAgentResponseDisplayMode.HIDDEN
+        }
+
+    @override
+    async def run(self, loop_response: LanguageModelStreamResponse) -> None:
+        await asyncio.sleep(
+            self._config.sleep_time_before_update
+        )  # Frontend rendering issues
+
+    def _get_displayed_sub_agent_responses(
+        self,
+    ) -> dict[str, list[SubAgentResponse]]:
+        responses = defaultdict(list)
+        all_responses = self._response_watcher.get_all_responses()
+        for response in all_responses:
+            assistant_id = response.assistant_id
+            if assistant_id in self._display_specs:
+                responses[assistant_id].append(response)
+        return responses
+
+    @override
+    def apply_postprocessing_to_response(
+        self, loop_response: LanguageModelStreamResponse
+    ) -> bool:
+        displayed_sub_agent_responses = self._get_displayed_sub_agent_responses()
+
+        if len(displayed_sub_agent_responses) == 0:
+            logger.info("No sub agent responses to prepend")
+            return False
+
+        logger.info("Prepending sub agent responses to the main agent response")
+
+        answers_displayed_before = []
+        answers_displayed_after = []
+
+        for assistant_id, responses in displayed_sub_agent_responses.items():
+            for response in responses:
+                message = response.message
+                tool_info = self._display_specs[assistant_id]
+
+                _add_response_references_to_message_in_place(
+                    loop_response=loop_response, response=message
+                )
+
+                display_name = tool_info.display_name
+                if len(responses) > 1:
+                    display_name += f" {response.sequence_number}"
+
+                if message["text"] is None:
+                    logger.warning(
+                        "Sub agent response for assistant %s with sequence number %s does not contain any text",
+                        assistant_id,
+                        response.sequence_number,
+                    )
+
+                answer = get_sub_agent_answer_display(
+                    display_name=display_name,
+                    display_config=tool_info.display_config,
+                    answer=message["text"] or "",
+                    assistant_id=assistant_id,
+                )
+
+                if tool_info.display_config.position == "before":
+                    answers_displayed_before.append(answer)
+                else:
+                    answers_displayed_after.append(answer)
+
+        loop_response.message.text = _get_final_answer_display(
+            text=loop_response.message.text,
+            answers_before=answers_displayed_before,
+            answers_after=answers_displayed_after,
+        )
+
+        return True
+
+    @override
+    async def remove_from_text(self, text) -> str:
+        for display_info in self._display_specs.values():
+            text = remove_sub_agent_answer_from_text(
+                display_config=display_info.display_config,
+                text=text,
+                assistant_id=display_info.assistant_id,
+            )
+        return text
+
+
+def _add_response_references_to_message_in_place(
+    loop_response: LanguageModelStreamResponse, response: unique_sdk.Space.Message
+) -> None:
+    references = response["references"]
+    text = response["text"]
+
+    if references is None or len(references) == 0 or text is None:
+        return
+
+    content_refs = [ContentReference.from_sdk_reference(ref) for ref in references]
+
+    text, refs = add_content_refs_and_replace_in_text(
+        message_text=text,
+        message_refs=loop_response.message.references,
+        new_refs=content_refs,
+    )
+
+    response["text"] = text  # Diplayed at a later stage
+    loop_response.message.references = refs
+
+
+def _get_final_answer_display(
+    text: str,
+    answers_before: list[str],
+    answers_after: list[str],
+    sep: str = "<br>\n\n",
+) -> str:
+    if len(answers_before) > 0:
+        text = sep.join(answers_before) + sep + text
+
+    if len(answers_after) > 0:
+        text = text + sep + sep.join(answers_after)
+    return text

unique_toolkit/agentic/tools/a2a/postprocessing/references.py

@@ -0,0 +1,101 @@
+import logging
+import re
+from typing import override
+
+from unique_toolkit._common.referencing import (
+    get_reference_pattern,
+    remove_consecutive_ref_space,
+)
+from unique_toolkit.agentic.postprocessor.postprocessor_manager import Postprocessor
+from unique_toolkit.agentic.tools.a2a.postprocessing._ref_utils import (
+    add_content_refs_and_replace_in_text,
+)
+from unique_toolkit.agentic.tools.a2a.response_watcher import (
+    SubAgentResponse,
+    SubAgentResponseWatcher,
+)
+from unique_toolkit.agentic.tools.a2a.tool import SubAgentTool
+from unique_toolkit.content import ContentReference
+from unique_toolkit.language_model.schemas import LanguageModelStreamResponse
+
+logger = logging.getLogger(__name__)
+
+
+class SubAgentReferencesPostprocessor(Postprocessor):
+    def __init__(self, response_watcher: SubAgentResponseWatcher) -> None:
+        super().__init__(name=self.__class__.__name__)
+        self._response_watcher = response_watcher
+
+    @override
+    async def run(self, loop_response: LanguageModelStreamResponse) -> None:
+        return
+
+    @override
+    def apply_postprocessing_to_response(
+        self, loop_response: LanguageModelStreamResponse
+    ) -> bool:
+        logger.info("Adding sub agent references to the main agent response")
+
+        num_sources = len(loop_response.message.references)
+
+        # At the moment, the `PostprocessorManager` expects modifications to happen in place
+        _add_sub_agent_references_in_place(
+            loop_response=loop_response,
+            responses=self._response_watcher.get_all_responses(),
+        )
+
+        return num_sources != len(
+            loop_response.message.references
+        )  # We only add references
+
+    @override
+    async def remove_from_text(self, text: str) -> str:
+        """
+        It is not possible to **only** remove sub agent references from the text,
+        as they are identical to normal references.
+        """
+        return text
+
+
+def _add_sub_agent_references_in_place(
+    loop_response: LanguageModelStreamResponse,
+    responses: list[SubAgentResponse],
+) -> None:
+    text = loop_response.message.text
+    refs = []
+
+    for response in responses:
+        sub_agent_refs = []
+        references = response.message["references"]
+
+        if references is None or len(references) == 0:
+            continue
+
+        for reference in sorted(references, key=lambda r: r["sequenceNumber"]):
+            reference_re = SubAgentTool.get_sub_agent_reference_re(
+                name=response.name,
+                sequence_number=response.sequence_number,
+                reference_number=reference["sequenceNumber"],
+            )
+
+            if re.search(reference_re, text) is None:
+                # Reference not used
+                continue
+
+            sub_agent_refs.append(ContentReference.from_sdk_reference(reference))
+
+        text, refs = add_content_refs_and_replace_in_text(
+            message_text=text,
+            message_refs=refs,
+            new_refs=sub_agent_refs,
+            ref_pattern_f=lambda x: r"\s*"  # Normalize spaces
+            + SubAgentTool.get_sub_agent_reference_re(
+                name=response.name,
+                sequence_number=response.sequence_number,
+                reference_number=x,
+            ),
+            ref_replacement_f=lambda x: " " + get_reference_pattern(x),
+        )
+
+    loop_response.message.references = refs
+    loop_response.message.text = remove_consecutive_ref_space(text)