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

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

@@ -0,0 +1,78 @@
+import re
+from enum import StrEnum
+from typing import Literal
+
+from pydantic import BaseModel, Field
+
+from unique_toolkit._common.pydantic_helpers import get_configuration_dict
+
+
+class SubAgentResponseDisplayMode(StrEnum):
+    HIDDEN = "hidden"
+    DETAILS_OPEN = "details_open"
+    DETAILS_CLOSED = "details_closed"
+    PLAIN = "plain"
+
+
+class SubAgentAnswerSubstringConfig(BaseModel):
+    model_config = get_configuration_dict()
+
+    regexp: re.Pattern[str] = Field(
+        description="The regular expression to use to extract the substring. The first capture group will always be used.",
+    )
+    display_template: str = Field(
+        default="{}",
+        description="The template to use to display the substring. It should contain exactly one empty placeholder '{}' for the substring.",
+    )
+
+
+_ANSWER_SUBSTRINGS_JINJA_TEMPLATE = """
+{% for substring in substrings %}
+{{ substring }}
+{% endfor %}
+""".strip()
+
+
+class SubAgentDisplayConfig(BaseModel):
+    model_config = get_configuration_dict()
+
+    mode: SubAgentResponseDisplayMode = Field(
+        default=SubAgentResponseDisplayMode.HIDDEN,
+        description="Controls how to display the sub agent response.",
+    )
+    remove_from_history: bool = Field(
+        default=True,
+        description="If set, sub agent responses will be removed from the history on subsequent calls to the assistant.",
+    )
+    add_quote_border: bool = Field(
+        default=True,
+        description="If set, a quote border is added to the left of the sub agent response.",
+    )
+    add_block_border: bool = Field(
+        default=False,
+        description="If set, a block border is added around the sub agent response.",
+    )
+    display_title_template: str = Field(
+        default="Answer from <strong>{}</strong>",
+        description=(
+            "The template to use for the display title of the sub agent response."
+            "If a placeholder '{}' is present, it will be replaced with the display name of the sub agent."
+        ),
+    )
+    position: Literal["before", "after"] = Field(
+        default="before",
+        description="The position of the sub agent response in the main agent response.",
+    )
+    force_include_references: bool = Field(
+        default=False,
+        description="If set, the sub agent references will be added to the main agent response references even in not mentioned in the main agent response text.",
+    )
+
+    answer_substrings_config: list[SubAgentAnswerSubstringConfig] = Field(
+        default=[],
+        description="If set, only parts of the answer matching the provided regular expressions will be displayed.",
+    )
+    answer_substrings_jinja_template: str = Field(
+        default=_ANSWER_SUBSTRINGS_JINJA_TEMPLATE,
+        description="The template to use in order to format the different answer substrings, if any.",
+    )

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

@@ -0,0 +1,264 @@
+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._common.utils.jinja.render import render_template
+from unique_toolkit.agentic.postprocessor.postprocessor_manager import Postprocessor
+from unique_toolkit.agentic.tools.a2a.postprocessing._display_utils import (
+    SubAgentAnswerPart,
+    get_sub_agent_answer_display,
+    get_sub_agent_answer_from_parts,
+    get_sub_agent_answer_parts,
+    remove_sub_agent_answer_from_text,
+)
+from unique_toolkit.agentic.tools.a2a.postprocessing._ref_utils import (
+    add_content_refs_and_replace_in_text,
+    remove_unused_refs,
+)
+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
+
+
+_ANSWERS_JINJA_TEMPLATE = """
+{% for answer in answers %}
+{{ answer }}
+{% endfor %}
+""".strip()
+
+
+class SubAgentResponsesPostprocessorConfig(BaseModel):
+    model_config = get_configuration_dict()
+
+    sleep_time_before_update: float = Field(
+        default=0, description="Time to sleep before updating the main agent message."
+    )
+    answers_jinja_template: str = Field(
+        default=_ANSWERS_JINJA_TEMPLATE,
+        description="The template to use to display the sub agent answers.",
+    )
+    filter_duplicate_answers: bool = Field(
+        default=True,
+        description="If set, duplicate answers will be filtered out.",
+    )
+
+
+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
+            # We should keep track of these messages even if they are hidden
+            or display_spec.display_config.force_include_references
+        }
+
+    @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 = []
+        all_displayed_answers = set()
+
+        for assistant_id, responses in displayed_sub_agent_responses.items():
+            tool_info = self._display_specs[assistant_id]
+            tool_name = tool_info.display_name
+
+            for response in responses:
+                message = response.message
+
+                if tool_info.display_config.mode == SubAgentResponseDisplayMode.HIDDEN:
+                    # Add references and continue
+                    _add_response_references_to_message_in_place(
+                        loop_response=loop_response,
+                        response=message,
+                        remove_unused_references=False,
+                    )
+                    continue
+
+                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,
+                    )
+                    continue
+
+                answer_parts = get_sub_agent_answer_parts(
+                    answer=message["text"],
+                    display_config=tool_info.display_config,
+                )
+
+                if self._config.filter_duplicate_answers:
+                    answer_parts, all_displayed_answers = (
+                        _filter_and_update_duplicate_answers(
+                            answers=answer_parts,
+                            existing_answers=all_displayed_answers,
+                        )
+                    )
+
+                answer = get_sub_agent_answer_from_parts(
+                    answer_parts=answer_parts,
+                    config=tool_info.display_config,
+                )
+                message["text"] = answer
+
+                _add_response_references_to_message_in_place(
+                    loop_response=loop_response,
+                    response=message,
+                    remove_unused_references=not tool_info.display_config.force_include_references,
+                )
+
+                if len(answer_parts) == 0:
+                    continue
+
+                display_name = tool_name
+                if len(responses) > 1:
+                    display_name = tool_name + f" {response.sequence_number}"
+
+                answer = get_sub_agent_answer_display(
+                    display_name=display_name,
+                    display_config=tool_info.display_config,
+                    answer=answer,
+                    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,
+            template=self._config.answers_jinja_template,
+        )
+
+        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,
+    remove_unused_references: bool = True,
+) -> 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]
+
+    if remove_unused_references:
+        content_refs = remove_unused_refs(
+            references=content_refs,
+            text=text,
+        )
+
+    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],
+    template: str = _ANSWERS_JINJA_TEMPLATE,
+) -> str:
+    if len(answers_before) > 0:
+        text = render_template(template, {"answers": answers_before}) + text
+
+    if len(answers_after) > 0:
+        text = text + render_template(template, {"answers": answers_after})
+
+    return text.strip()
+
+
+def _filter_and_update_duplicate_answers(
+    answers: list[SubAgentAnswerPart],
+    existing_answers: set[str],
+) -> tuple[list[SubAgentAnswerPart], set[str]]:
+    new_answers = []
+
+    for answer in answers:
+        if answer.matching_text in existing_answers:
+            continue
+        existing_answers.add(answer.matching_text)
+        new_answers.append(answer)
+
+    return new_answers, existing_answers

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)