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

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

@@ -1,9 +1,9 @@
 import logging
-from typing import override
+from typing import NamedTuple, override
 
 import unique_sdk
 from jinja2 import Template
-from typing_extensions import TypedDict
+from pydantic import BaseModel
 
 from unique_toolkit.agentic.evaluation.evaluation_manager import Evaluation
 from unique_toolkit.agentic.evaluation.schemas import (
@@ -12,15 +12,19 @@ from unique_toolkit.agentic.evaluation.schemas import (
     EvaluationMetricResult,
 )
 from unique_toolkit.agentic.tools.a2a.evaluation._utils import (
-    _get_valid_assessments,
-    _sort_assessments,
-    _worst_label,
+    get_valid_assessments,
+    get_worst_label,
+    sort_assessments,
 )
 from unique_toolkit.agentic.tools.a2a.evaluation.config import (
     SubAgentEvaluationConfig,
     SubAgentEvaluationServiceConfig,
 )
-from unique_toolkit.agentic.tools.a2a.tool import SubAgentTool
+from unique_toolkit.agentic.tools.a2a.response_watcher import (
+    SubAgentResponse,
+    SubAgentResponseWatcher,
+)
+from unique_toolkit.agentic.tools.utils import failsafe
 from unique_toolkit.chat.schemas import (
     ChatMessageAssessmentLabel,
     ChatMessageAssessmentStatus,
@@ -33,12 +37,40 @@ from unique_toolkit.language_model.service import LanguageModelService
 logger = logging.getLogger(__name__)
 
 
-class _SubAgentToolInfo(TypedDict):
-    assessments: dict[int, list[unique_sdk.Space.Assessment]]
+class SubAgentEvaluationSpec(NamedTuple):
     display_name: str
+    assistant_id: str
+    config: SubAgentEvaluationConfig
+
+
+_NO_ASSESSMENTS_FOUND = "NO_ASSESSMENTS_FOUND"
+
+
+class _SingleAssessmentData(BaseModel):
+    name: str
+    explanation: str
+
+
+def _format_single_assessment_found(name: str, explanation: str) -> str:
+    return _SingleAssessmentData(name=name, explanation=explanation).model_dump_json()
+
 
+@failsafe(failure_return_value=None, log_exceptions=False)
+def _parse_single_assesment_found(value: str) -> _SingleAssessmentData | None:
+    return _SingleAssessmentData.model_validate_json(value)
 
-NO_ASSESSMENTS_FOUND = "NO_ASSESSMENTS_FOUND"
+
+def _find_single_assessment(
+    responses: dict[str, list[SubAgentResponse]],
+) -> unique_sdk.Space.Assessment | None:
+    if len(responses) == 1:
+        sub_agent_responses = next(iter(responses.values()))
+        if len(sub_agent_responses) == 1:
+            response = sub_agent_responses[0].message
+            if response["assessment"] is not None and len(response["assessment"]) == 1:
+                return response["assessment"][0]
+
+    return None
 
 
 class SubAgentEvaluationService(Evaluation):
@@ -48,17 +80,67 @@ class SubAgentEvaluationService(Evaluation):
         self,
         config: SubAgentEvaluationServiceConfig,
         language_model_service: LanguageModelService,
-    ):
+        response_watcher: SubAgentResponseWatcher,
+        evaluation_specs: list[SubAgentEvaluationSpec],
+    ) -> None:
         super().__init__(EvaluationMetricName.SUB_AGENT)
         self._config = config
 
-        self._assistant_id_to_tool_info: dict[str, _SubAgentToolInfo] = {}
+        self._response_watcher = response_watcher
         self._language_model_service = language_model_service
 
+        self._evaluation_specs: dict[str, SubAgentEvaluationSpec] = {
+            spec.assistant_id: spec
+            for spec in evaluation_specs
+            if spec.config.include_evaluation
+        }
+
     @override
     def get_assessment_type(self) -> ChatMessageAssessmentType:
         return self._config.assessment_type
 
+    def _get_included_sub_agent_responses(
+        self,
+    ) -> dict[str, list[SubAgentResponse]]:
+        responses = {}
+        for assistant_id, eval_spec in self._evaluation_specs.items():
+            sub_agent_responses = self._response_watcher.get_responses(
+                eval_spec.assistant_id
+            )
+            if len(sub_agent_responses) == 0:
+                logger.debug(
+                    "No responses for sub agent %s (%s)",
+                    eval_spec.display_name,
+                    eval_spec.assistant_id,
+                )
+                continue
+
+            responses_with_assessment = []
+            for response in sub_agent_responses:
+                assessments = response.message["assessment"]
+
+                if assessments is None or len(assessments) == 0:
+                    logger.debug(
+                        "No assessment for sub agent %s (%s) response with sequence number %s",
+                        eval_spec.display_name,
+                        eval_spec.assistant_id,
+                        response.sequence_number,
+                    )
+                    continue
+
+                assessments = get_valid_assessments(
+                    assessments=assessments,
+                    display_name=eval_spec.display_name,
+                    sequence_number=response.sequence_number,
+                )
+
+                if len(assessments) > 0:
+                    responses_with_assessment.append(response)
+
+            responses[assistant_id] = responses_with_assessment
+
+        return responses
+
     @override
     async def run(
         self, loop_response: LanguageModelStreamResponse
@@ -67,47 +149,56 @@ class SubAgentEvaluationService(Evaluation):
 
         sub_agents_display_data = []
 
-        value = ChatMessageAssessmentLabel.GREEN
+        responses = self._get_included_sub_agent_responses()
 
-        for tool_info in self._assistant_id_to_tool_info.values():
-            sub_agent_assessments = tool_info["assessments"] or []
-            display_name = tool_info["display_name"]
+        # No valid assessments found
+        if len(responses) == 0:
+            logger.warning("No valid sub agent assessments found")
 
-            for sequence_number in sorted(sub_agent_assessments):
-                assessments = sub_agent_assessments[sequence_number]
+            return EvaluationMetricResult(
+                name=self.get_name(),
+                # This is a trick to be able to indicate to `evaluation_metric_to_assessment`
+                # that no valid assessments were found
+                value=_NO_ASSESSMENTS_FOUND,
+                reason="No sub agents assessments found",
+            )
 
-                valid_assessments = _get_valid_assessments(
-                    assessments, display_name, sequence_number
-                )
-                if len(valid_assessments) == 0:
-                    logger.info(
-                        "No valid assessment found for assistant %s (sequence number: %s)",
-                        display_name,
-                        sequence_number,
-                    )
-                    continue
+        single_assessment = _find_single_assessment(responses)
+        # Only one valid assessment found, no need to perform summarization
+        if single_assessment is not None:
+            assistant_id = next(iter(responses))
+            explanation = single_assessment["explanation"] or ""
+            name = self._evaluation_specs[assistant_id].display_name
+            label = single_assessment["label"] or ""
 
-                assessments = _sort_assessments(valid_assessments)
-                value = _worst_label(value, assessments[0]["label"])  # type: ignore
+            return EvaluationMetricResult(
+                name=self.get_name(),
+                value=label,
+                # This is a trick to be able to pass the display name to the UI in `evaluation_metric_to_assessment`
+                reason=_format_single_assessment_found(name, explanation),
+                is_positive=label == ChatMessageAssessmentLabel.GREEN,
+            )
+
+        sub_agents_display_data = []
+
+        # Multiple Assessments found
+        value = ChatMessageAssessmentLabel.GREEN
+        for assistant_id, sub_agent_responses in responses.items():
+            display_name = self._evaluation_specs[assistant_id].display_name
+
+            for response in sub_agent_responses:
+                assessments = sort_assessments(response.message["assessment"])  #  type:ignore
+                value = get_worst_label(value, assessments[0]["label"])  # type: ignore
 
                 data = {
-                    "name": tool_info["display_name"],
+                    "name": display_name,
                     "assessments": assessments,
                 }
-                if len(sub_agent_assessments) > 1:
-                    data["name"] += f" {sequence_number}"
+                if len(sub_agent_responses) > 1:
+                    data["name"] += f" {response.sequence_number}"
 
                 sub_agents_display_data.append(data)
 
-        if len(sub_agents_display_data) == 0:
-            logger.warning("No valid sub agent assessments found")
-
-            return EvaluationMetricResult(
-                name=self.get_name(),
-                value=NO_ASSESSMENTS_FOUND,
-                reason="No sub agents assessments found",
-            )
-
         reason = await self._get_reason(sub_agents_display_data)
 
         return EvaluationMetricResult(
@@ -121,7 +212,7 @@ class SubAgentEvaluationService(Evaluation):
     async def evaluation_metric_to_assessment(
         self, evaluation_result: EvaluationMetricResult
     ) -> EvaluationAssessmentMessage:
-        if evaluation_result.value == NO_ASSESSMENTS_FOUND:
+        if evaluation_result.value == _NO_ASSESSMENTS_FOUND:
             return EvaluationAssessmentMessage(
                 status=ChatMessageAssessmentStatus.DONE,
                 explanation="No valid sub agents assessments found to consolidate.",
@@ -130,6 +221,16 @@ class SubAgentEvaluationService(Evaluation):
                 type=self.get_assessment_type(),
             )
 
+        single_assessment_data = _parse_single_assesment_found(evaluation_result.reason)
+        if single_assessment_data is not None:
+            return EvaluationAssessmentMessage(
+                status=ChatMessageAssessmentStatus.DONE,
+                explanation=single_assessment_data.explanation,
+                title=single_assessment_data.name,
+                label=evaluation_result.value,  # type: ignore
+                type=self.get_assessment_type(),
+            )
+
         return EvaluationAssessmentMessage(
             status=ChatMessageAssessmentStatus.DONE,
             explanation=evaluation_result.reason,
@@ -138,56 +239,7 @@ class SubAgentEvaluationService(Evaluation):
             type=self.get_assessment_type(),
         )
 
-    def register_sub_agent_tool(
-        self, tool: SubAgentTool, evaluation_config: SubAgentEvaluationConfig
-    ) -> None:
-        if not evaluation_config.include_evaluation:
-            logger.warning(
-                "Sub agent tool %s has evaluation config `include_evaluation` set to False, responses will be ignored.",
-                tool.config.assistant_id,
-            )
-            return
-
-        if tool.config.assistant_id not in self._assistant_id_to_tool_info:
-            tool.subscribe(self)
-            self._assistant_id_to_tool_info[tool.config.assistant_id] = (
-                _SubAgentToolInfo(
-                    display_name=tool.display_name(),
-                    assessments={},
-                )
-            )
-
-    def notify_sub_agent_response(
-        self,
-        response: unique_sdk.Space.Message,
-        sub_agent_assistant_id: str,
-        sequence_number: int,
-    ) -> None:
-        if sub_agent_assistant_id not in self._assistant_id_to_tool_info:
-            logger.warning(
-                "Unknown assistant id %s received, assessment will be ignored.",
-                sub_agent_assistant_id,
-            )
-            return
-
-        sub_agent_assessments = self._assistant_id_to_tool_info[sub_agent_assistant_id][
-            "assessments"
-        ]
-        sub_agent_assessments[sequence_number] = (
-            response[
-                "assessment"
-            ].copy()  # Shallow copy as we don't modify individual assessments
-            if response["assessment"] is not None
-            else []
-        )
-
     async def _get_reason(self, sub_agents_display_data: list[dict]) -> str:
-        if (
-            len(sub_agents_display_data) == 1
-            and len(sub_agents_display_data[0]["assessments"]) == 1
-        ):
-            return sub_agents_display_data[0]["assessments"][0]["explanation"] or ""
-
         messages = (
             MessagesBuilder()
             .system_message_append(self._config.summarization_system_message)

unique_toolkit/agentic/tools/a2a/manager.py

@@ -1,5 +1,6 @@
 from logging import Logger
 
+from unique_toolkit.agentic.tools.a2a.response_watcher import SubAgentResponseWatcher
 from unique_toolkit.agentic.tools.a2a.tool import SubAgentTool, SubAgentToolConfig
 from unique_toolkit.agentic.tools.config import ToolBuildConfig
 from unique_toolkit.agentic.tools.tool_progress_reporter import ToolProgressReporter
@@ -11,12 +12,16 @@ class A2AManager:
         self,
         logger: Logger,
         tool_progress_reporter: ToolProgressReporter,
+        response_watcher: SubAgentResponseWatcher,
     ):
         self._logger = logger
         self._tool_progress_reporter = tool_progress_reporter
+        self._response_watcher = response_watcher
 
     def get_all_sub_agents(
-        self, tool_configs: list[ToolBuildConfig], event: ChatEvent
+        self,
+        tool_configs: list[ToolBuildConfig],
+        event: ChatEvent,
     ) -> tuple[list[ToolBuildConfig], list[SubAgentTool]]:
         sub_agents = []
 
@@ -39,6 +44,7 @@ class A2AManager:
                     tool_progress_reporter=self._tool_progress_reporter,
                     name=tool_config.name,
                     display_name=tool_config.display_name,
+                    response_watcher=self._response_watcher,
                 )
             )
 

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

@@ -2,12 +2,20 @@ from unique_toolkit.agentic.tools.a2a.postprocessing.config import (
     SubAgentDisplayConfig,
     SubAgentResponseDisplayMode,
 )
-from unique_toolkit.agentic.tools.a2a.postprocessing.postprocessor import (
-    SubAgentResponsesPostprocessor,
+from unique_toolkit.agentic.tools.a2a.postprocessing.display import (
+    SubAgentDisplaySpec,
+    SubAgentResponsesDisplayPostprocessor,
+    SubAgentResponsesPostprocessorConfig,
+)
+from unique_toolkit.agentic.tools.a2a.postprocessing.references import (
+    SubAgentReferencesPostprocessor,
 )
 
 __all__ = [
-    "SubAgentResponsesPostprocessor",
+    "SubAgentResponsesDisplayPostprocessor",
+    "SubAgentResponsesPostprocessorConfig",
+    "SubAgentDisplaySpec",
     "SubAgentResponseDisplayMode",
     "SubAgentDisplayConfig",
+    "SubAgentReferencesPostprocessor",
 ]

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

@@ -0,0 +1,185 @@
+import re
+from typing import Literal
+
+from unique_toolkit.agentic.tools.a2a.postprocessing.config import (
+    SubAgentDisplayConfig,
+    SubAgentResponseDisplayMode,
+)
+
+
+def _wrap_text(text: str, start_text: str, end_text: str) -> str:
+    text = text.strip()
+    start_text = start_text.strip()
+    end_text = end_text.strip()
+
+    if start_text != "":
+        start_text = f"{start_text}\n"
+
+    if end_text != "":
+        end_text = f"\n{end_text}"
+
+    return f"{start_text}{text}{end_text}"
+
+
+def _join_text_blocks(*blocks: str, sep: str = "\n") -> str:
+    return sep.join(block.strip() for block in blocks)
+
+
+def _wrap_with_details_tag(
+    text, mode: Literal["open", "closed"], summary_name: str | None = None
+) -> str:
+    if summary_name is not None:
+        summary_tag = _wrap_text(summary_name, "<summary>", "</summary>")
+        text = _join_text_blocks(summary_tag, text)
+
+    if mode == "open":
+        text = _wrap_text(text, "<details open>", "</details>")
+    else:
+        text = _wrap_text(text, "<details>", "</details>")
+
+    return text
+
+
+_BLOCK_BORDER_STYLE = (
+    "overflow-y: auto; border: 1px solid #ccc; padding: 8px; margin-top: 8px;"
+)
+
+
+def _wrap_with_block_border(text: str) -> str:
+    return _wrap_text(text, f"<div style='{_BLOCK_BORDER_STYLE}'>", "</div>")
+
+
+_QUOTE_BORDER_STYLE = (
+    "margin-left: 20px; border-left: 2px solid #ccc; padding-left: 10px;"
+)
+
+
+def _wrap_with_quote_border(text: str) -> str:
+    return _wrap_text(text, f"<div style='{_QUOTE_BORDER_STYLE}'>", "</div>")
+
+
+def _wrap_strong(text: str) -> str:
+    return _wrap_text(text, "<strong>", "</strong>")
+
+
+def _wrap_hidden_div(text: str) -> str:
+    return _wrap_text(text, '<div style="display: none;">', "</div>")
+
+
+def _add_line_break(text: str, before: bool = True, after: bool = True) -> str:
+    start_tag = ""
+    if before:
+        start_tag = "<br>"
+
+    end_tag = ""
+    if after:
+        end_tag = "<br>"
+
+    return _wrap_text(text, start_tag, end_tag)
+
+
+def _prepare_title_template(
+    display_title_template: str, display_name_placeholder: str
+) -> str:
+    return display_title_template.replace("{}", "{%s}" % display_name_placeholder)
+
+
+def _get_display_template(
+    mode: SubAgentResponseDisplayMode,
+    add_quote_border: bool,
+    add_block_border: bool,
+    display_title_template: str,
+    answer_placeholder: str = "answer",
+    assistant_id_placeholder: str = "assistant_id",
+    display_name_placeholder: str = "display_name",
+) -> str:
+    if mode == SubAgentResponseDisplayMode.HIDDEN:
+        return ""
+
+    assistant_id_placeholder = _wrap_hidden_div("{%s}" % assistant_id_placeholder)
+    title_template = _prepare_title_template(
+        display_title_template, display_name_placeholder
+    )
+    template = _join_text_blocks(
+        assistant_id_placeholder, "{%s}" % answer_placeholder, sep="\n\n"
+    )  # Double line break is needed for markdown formatting
+
+    template = _add_line_break(template, before=True, after=False)
+
+    if add_quote_border:
+        template = _wrap_with_quote_border(template)
+
+    match mode:
+        case SubAgentResponseDisplayMode.DETAILS_OPEN:
+            template = _wrap_with_details_tag(
+                template,
+                "open",
+                title_template,
+            )
+        case SubAgentResponseDisplayMode.DETAILS_CLOSED:
+            template = _wrap_with_details_tag(template, "closed", title_template)
+        case SubAgentResponseDisplayMode.PLAIN:
+            # Add a hidden block border to seperate sub agent answers from the rest of the text.
+            hidden_block_border = _wrap_hidden_div("sub_agent_answer_block")
+            template = _join_text_blocks(title_template, template, hidden_block_border)
+
+    if add_block_border:
+        template = _wrap_with_block_border(template)
+
+    return template
+
+
+def _get_display_removal_re(
+    assistant_id: str,
+    mode: SubAgentResponseDisplayMode,
+    add_quote_border: bool,
+    add_block_border: bool,
+    display_title_template: str,
+) -> re.Pattern[str]:
+    template = _get_display_template(
+        mode=mode,
+        add_quote_border=add_quote_border,
+        add_block_border=add_block_border,
+        display_title_template=display_title_template,
+    )
+
+    pattern = template.format(
+        assistant_id=re.escape(assistant_id), answer=r"(.*?)", display_name=r"(.*?)"
+    )
+
+    return re.compile(pattern, flags=re.DOTALL)
+
+
+def get_sub_agent_answer_display(
+    display_name: str,
+    display_config: SubAgentDisplayConfig,
+    answer: str,
+    assistant_id: str,
+) -> str:
+    template = _get_display_template(
+        mode=display_config.mode,
+        add_quote_border=display_config.add_quote_border,
+        add_block_border=display_config.add_block_border,
+        display_title_template=display_config.display_title_template,
+    )
+    return template.format(
+        display_name=display_name, answer=answer, assistant_id=assistant_id
+    )
+
+
+def remove_sub_agent_answer_from_text(
+    display_config: SubAgentDisplayConfig,
+    text: str,
+    assistant_id: str,
+) -> str:
+    if not display_config.remove_from_history:
+        return text
+
+    pattern = _get_display_removal_re(
+        assistant_id=assistant_id,
+        mode=display_config.mode,
+        add_quote_border=display_config.add_quote_border,
+        add_block_border=display_config.add_block_border,
+        display_title_template=display_config.display_title_template,
+    )
+    return re.sub(pattern, "", text)

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

@@ -0,0 +1,73 @@
+from typing import Callable, Iterable, Mapping, Sequence
+
+from unique_toolkit._common.referencing import get_reference_pattern
+from unique_toolkit._common.string_utilities import replace_in_text
+from unique_toolkit.content import ContentReference
+
+SourceId = str
+SequenceNumber = int
+
+
+def _add_source_ids(
+    existing_refs: Mapping[SourceId, SequenceNumber],
+    new_refs: Iterable[SourceId],
+) -> dict[SourceId, SequenceNumber]:
+    next_seq_num = max(existing_refs.values(), default=0) + 1
+    new_seq_nums: dict[SourceId, SequenceNumber] = {}
+
+    for source_id in new_refs:
+        seq_num = existing_refs.get(source_id, None) or new_seq_nums.get(
+            source_id, None
+        )
+        if seq_num is None:
+            new_seq_nums[source_id] = next_seq_num
+            next_seq_num += 1
+
+    return new_seq_nums
+
+
+def add_content_refs(
+    message_refs: Sequence[ContentReference],
+    new_refs: Sequence[ContentReference],
+) -> list[ContentReference]:
+    message_refs = list(message_refs)
+
+    if len(new_refs) == 0:
+        return message_refs
+
+    existing_refs = {ref.source_id: ref.sequence_number for ref in message_refs}
+    new_refs_by_source_id = {
+        ref.source_id: ref for ref in sorted(new_refs, key=lambda x: x.sequence_number)
+    }
+    new_seq_nums = _add_source_ids(existing_refs, new_refs_by_source_id.keys())
+
+    for source_id, seq_num in new_seq_nums.items():
+        ref = new_refs_by_source_id[source_id]
+        message_refs.append(
+            ref.model_copy(update={"sequence_number": seq_num}, deep=True)
+        )
+
+    return message_refs
+
+
+def add_content_refs_and_replace_in_text(
+    message_text: str,
+    message_refs: Sequence[ContentReference],
+    new_refs: Sequence[ContentReference],
+    ref_pattern_f: Callable[[int], str] = get_reference_pattern,
+    ref_replacement_f: Callable[[int], str] = get_reference_pattern,
+) -> tuple[str, list[ContentReference]]:
+    if len(new_refs) == 0:
+        return message_text, list(message_refs)
+
+    references = add_content_refs(message_refs, new_refs)
+    seq_num_for_source_id = {ref.source_id: ref.sequence_number for ref in references}
+    ref_map = []
+
+    for ref in new_refs:
+        old_seq_num = ref.sequence_number
+        new_seq_num = seq_num_for_source_id[ref.source_id]
+
+        ref_map.append((ref_pattern_f(old_seq_num), ref_replacement_f(new_seq_num)))
+
+    return replace_in_text(message_text, ref_map), references

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

@@ -1,4 +1,5 @@
 from enum import StrEnum
+from typing import Literal
 
 from pydantic import BaseModel, Field
 
@@ -9,6 +10,7 @@ class SubAgentResponseDisplayMode(StrEnum):
     HIDDEN = "hidden"
     DETAILS_OPEN = "details_open"
     DETAILS_CLOSED = "details_closed"
+    PLAIN = "plain"
 
 
 class SubAgentDisplayConfig(BaseModel):
@@ -22,3 +24,22 @@ class SubAgentDisplayConfig(BaseModel):
         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.",
+    )