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

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

@@ -0,0 +1,260 @@
+import logging
+from typing import NamedTuple, override
+
+import unique_sdk
+from jinja2 import Template
+from pydantic import BaseModel
+
+from unique_toolkit.agentic.evaluation.evaluation_manager import Evaluation
+from unique_toolkit.agentic.evaluation.schemas import (
+    EvaluationAssessmentMessage,
+    EvaluationMetricName,
+    EvaluationMetricResult,
+)
+from unique_toolkit.agentic.tools.a2a.evaluation._utils import (
+    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.response_watcher import (
+    SubAgentResponse,
+    SubAgentResponseWatcher,
+)
+from unique_toolkit.agentic.tools.utils import failsafe
+from unique_toolkit.chat.schemas import (
+    ChatMessageAssessmentLabel,
+    ChatMessageAssessmentStatus,
+    ChatMessageAssessmentType,
+)
+from unique_toolkit.language_model.builder import MessagesBuilder
+from unique_toolkit.language_model.schemas import LanguageModelStreamResponse
+from unique_toolkit.language_model.service import LanguageModelService
+
+logger = logging.getLogger(__name__)
+
+
+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)
+
+
+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):
+    DISPLAY_NAME = "Sub Agents"
+
+    def __init__(
+        self,
+        config: SubAgentEvaluationServiceConfig,
+        language_model_service: LanguageModelService,
+        response_watcher: SubAgentResponseWatcher,
+        evaluation_specs: list[SubAgentEvaluationSpec],
+    ) -> None:
+        super().__init__(EvaluationMetricName.SUB_AGENT)
+        self._config = config
+
+        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
+    ) -> EvaluationMetricResult:
+        logger.info("Running sub agents evaluation")
+
+        sub_agents_display_data = []
+
+        responses = self._get_included_sub_agent_responses()
+
+        # No valid assessments found
+        if len(responses) == 0:
+            logger.warning("No valid sub agent assessments found")
+
+            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",
+            )
+
+        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 ""
+
+            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": display_name,
+                    "assessments": assessments,
+                }
+                if len(sub_agent_responses) > 1:
+                    data["name"] += f" {response.sequence_number}"
+
+                sub_agents_display_data.append(data)
+
+        reason = await self._get_reason(sub_agents_display_data)
+
+        return EvaluationMetricResult(
+            name=self.get_name(),
+            value=value,
+            reason=reason,
+            is_positive=value == ChatMessageAssessmentLabel.GREEN,
+        )
+
+    @override
+    async def evaluation_metric_to_assessment(
+        self, evaluation_result: EvaluationMetricResult
+    ) -> EvaluationAssessmentMessage:
+        if evaluation_result.value == _NO_ASSESSMENTS_FOUND:
+            return EvaluationAssessmentMessage(
+                status=ChatMessageAssessmentStatus.DONE,
+                explanation="No valid sub agents assessments found to consolidate.",
+                title=self.DISPLAY_NAME,
+                label=ChatMessageAssessmentLabel.GREEN,
+                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,
+            title=self.DISPLAY_NAME,
+            label=evaluation_result.value,  # type: ignore
+            type=self.get_assessment_type(),
+        )
+
+    async def _get_reason(self, sub_agents_display_data: list[dict]) -> str:
+        messages = (
+            MessagesBuilder()
+            .system_message_append(self._config.summarization_system_message)
+            .user_message_append(
+                Template(self._config.summarization_user_message_template).render(
+                    sub_agents=sub_agents_display_data,
+                )
+            )
+            .build()
+        )
+
+        reason = await self._language_model_service.complete_async(
+            messages=messages,
+            model_name=self._config.summarization_model.name,
+            temperature=0.0,
+        )
+
+        return str(reason.choices[0].message.content)

unique_toolkit/agentic/tools/a2a/evaluation/summarization_user_message.j2

@@ -0,0 +1,9 @@
+Here are the sub_agent(s) assessments:
+
+{% for sub_agent in sub_agents %}
+- Agent name: {{ sub_agent.name }} 
+{% for assessment in sub_agent.assessments %}
+    - {{ assessment.title }}: {{ assessment.label }}
+        {{ assessment.explanation }}
+{% endfor %}
+{% endfor %}

unique_toolkit/agentic/tools/a2a/manager.py

@@ -0,0 +1,55 @@
+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
+from unique_toolkit.app.schemas import ChatEvent
+
+
+class A2AManager:
+    def __init__(
+        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,
+    ) -> tuple[list[ToolBuildConfig], list[SubAgentTool]]:
+        sub_agents = []
+
+        for tool_config in tool_configs:
+            if not tool_config.is_sub_agent:
+                continue
+
+            if not isinstance(tool_config.configuration, SubAgentToolConfig):
+                self._logger.error(
+                    "tool_config.configuration must be of type SubAgentToolConfig"
+                )
+                continue
+
+            sub_agent_tool_config = tool_config.configuration
+
+            sub_agents.append(
+                SubAgentTool(
+                    configuration=sub_agent_tool_config,
+                    event=event,
+                    tool_progress_reporter=self._tool_progress_reporter,
+                    name=tool_config.name,
+                    display_name=tool_config.display_name,
+                    response_watcher=self._response_watcher,
+                )
+            )
+
+        filtered_tool_config = [
+            tool_config for tool_config in tool_configs if not tool_config.is_sub_agent
+        ]
+
+        return filtered_tool_config, sub_agents

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

@@ -0,0 +1,21 @@
+from unique_toolkit.agentic.tools.a2a.postprocessing.config import (
+    SubAgentDisplayConfig,
+    SubAgentResponseDisplayMode,
+)
+from unique_toolkit.agentic.tools.a2a.postprocessing.display import (
+    SubAgentDisplaySpec,
+    SubAgentResponsesDisplayPostprocessor,
+    SubAgentResponsesPostprocessorConfig,
+)
+from unique_toolkit.agentic.tools.a2a.postprocessing.references import (
+    SubAgentReferencesPostprocessor,
+)
+
+__all__ = [
+    "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

@@ -0,0 +1,45 @@
+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 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.",
+    )