unique_toolkit 0.8.30__py3-none-any.whl → 0.8.32__py3-none-any.whl

unique_toolkit/_common/chunk_relevancy_sorter/tests/test_service.py

@@ -0,0 +1,273 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from unique_toolkit._common.chunk_relevancy_sorter.config import (
+    ChunkRelevancySortConfig,
+)
+from unique_toolkit._common.chunk_relevancy_sorter.exception import (
+    ChunkRelevancySorterException,
+)
+from unique_toolkit._common.chunk_relevancy_sorter.schemas import (
+    ChunkRelevancy,
+    ChunkRelevancySorterResult,
+)
+from unique_toolkit._common.chunk_relevancy_sorter.service import ChunkRelevancySorter
+from unique_toolkit._common.default_language_model import (
+    DEFAULT_GPT_35_TURBO,
+    DEFAULT_GPT_4o,
+)
+from unique_toolkit.app.schemas import ChatEvent
+from unique_toolkit.content.schemas import ContentChunk
+from unique_toolkit.evals.context_relevancy.schema import StructuredOutputConfig
+from unique_toolkit.evals.schemas import EvaluationMetricName, EvaluationMetricResult
+from unique_toolkit.language_model.infos import LanguageModelInfo
+
+
+@pytest.fixture
+def event():
+    event = MagicMock(spec=ChatEvent)
+    event.payload = MagicMock()
+    event.payload.user_message = MagicMock()
+    event.payload.user_message.text = "Test query"
+    event.user_id = "user_0"
+    event.company_id = "company_0"
+    return event
+
+
+@pytest.fixture
+def mock_chunks():
+    return [
+        ContentChunk(
+            id=f"chunk_{i}",
+            order=i,
+            chunk_id=f"chunk_{i}",
+            text=f"Test content {i}",
+        )
+        for i in range(3)
+    ]
+
+
+@pytest.fixture
+def config():
+    return ChunkRelevancySortConfig(
+        enabled=True,
+        relevancy_levels_to_consider=["high", "medium", "low"],
+        relevancy_level_order={"high": 0, "medium": 1, "low": 2},
+        language_model=LanguageModelInfo.from_name(DEFAULT_GPT_4o),
+        fallback_language_model=LanguageModelInfo.from_name(DEFAULT_GPT_35_TURBO),
+        structured_output_config=StructuredOutputConfig(
+            enabled=False,
+            extract_fact_list=False,
+        ),
+    )
+
+
+@pytest.fixture
+def chunk_relevancy_sorter(event):
+    return ChunkRelevancySorter(event)
+
+
+@pytest.mark.asyncio
+async def test_run_disabled_config(chunk_relevancy_sorter, mock_chunks, config):
+    config.enabled = False
+    result = await chunk_relevancy_sorter.run("test input", mock_chunks, config)
+
+    assert isinstance(result, ChunkRelevancySorterResult)
+    assert result.content_chunks == mock_chunks
+    assert len(result.content_chunks) == len(mock_chunks)
+
+
+@pytest.mark.asyncio
+async def test_run_enabled_config(chunk_relevancy_sorter, mock_chunks, config):
+    with patch.object(chunk_relevancy_sorter, "_run_chunk_relevancy_sort") as mock_sort:
+        mock_sort.return_value = ChunkRelevancySorterResult.from_chunks(mock_chunks)
+
+        result = await chunk_relevancy_sorter.run("test input", mock_chunks, config)
+
+        assert isinstance(result, ChunkRelevancySorterResult)
+        assert result.content_chunks == mock_chunks
+        mock_sort.assert_called_once_with("test input", mock_chunks, config)
+
+
+@pytest.mark.asyncio
+async def test_evaluate_chunks_relevancy(chunk_relevancy_sorter, mock_chunks, config):
+    mock_relevancy = EvaluationMetricResult(
+        value="high",
+        name=EvaluationMetricName.CONTEXT_RELEVANCY,
+        reason="Test reason",
+    )
+
+    with patch.object(
+        chunk_relevancy_sorter, "_process_relevancy_evaluation"
+    ) as mock_process:
+        mock_process.return_value = ChunkRelevancy(
+            chunk=mock_chunks[0], relevancy=mock_relevancy
+        )
+
+        result = await chunk_relevancy_sorter._evaluate_chunks_relevancy(
+            "test input", mock_chunks, config
+        )
+
+        assert len(result) == len(mock_chunks)
+        assert all(isinstance(r, ChunkRelevancy) for r in result)
+        assert mock_process.call_count == len(mock_chunks)
+
+
+@pytest.mark.asyncio
+async def test_evaluate_chunk_relevancy(chunk_relevancy_sorter, mock_chunks, config):
+    with patch(
+        "unique_toolkit._common.chunk_relevancy_sorter.service.ContextRelevancyEvaluator.analyze"
+    ) as mock_analyze:
+        mock_analyze.return_value = EvaluationMetricResult(
+            value="high",
+            name=EvaluationMetricName.CONTEXT_RELEVANCY,
+            reason="Test reason",
+        )
+
+        result = await chunk_relevancy_sorter._evaluate_chunk_relevancy(
+            input_text="test input",
+            chunk=mock_chunks[0],
+            langugage_model=config.language_model,
+            structured_output_config=config.structured_output_config,
+            additional_llm_options=config.additional_llm_options,
+        )
+
+        assert isinstance(result, EvaluationMetricResult)
+        assert result.value == "high"
+        mock_analyze.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_process_relevancy_evaluation_success(
+    chunk_relevancy_sorter, mock_chunks, config
+):
+    with patch.object(
+        chunk_relevancy_sorter, "_evaluate_chunk_relevancy"
+    ) as mock_evaluate:
+        mock_evaluate.return_value = EvaluationMetricResult(
+            value="high",
+            name=EvaluationMetricName.CONTEXT_RELEVANCY,
+            reason="Test reason",
+        )
+
+        result = await chunk_relevancy_sorter._process_relevancy_evaluation(
+            "test input", mock_chunks[0], config
+        )
+
+        assert isinstance(result, ChunkRelevancy)
+        assert result.chunk == mock_chunks[0]
+        assert result.relevancy is not None
+        assert result.relevancy.value == "high"
+
+
+@pytest.mark.asyncio
+async def test_process_relevancy_evaluation_fallback(
+    chunk_relevancy_sorter, mock_chunks, config
+):
+    with patch.object(
+        chunk_relevancy_sorter, "_evaluate_chunk_relevancy"
+    ) as mock_evaluate:
+        # First call raises exception, second call succeeds
+        mock_evaluate.side_effect = [
+            Exception("Test error"),
+            EvaluationMetricResult(
+                value="medium",
+                name=EvaluationMetricName.CONTEXT_RELEVANCY,
+                reason="Test reason",
+            ),
+        ]
+
+        with pytest.raises(ChunkRelevancySorterException):
+            await chunk_relevancy_sorter._process_relevancy_evaluation(
+                "test input", mock_chunks[0], config
+            )
+
+
+@pytest.mark.asyncio
+async def test_validate_and_sort_relevant_chunks(
+    chunk_relevancy_sorter, mock_chunks, config
+):
+    chunk_relevancies = [
+        ChunkRelevancy(
+            chunk=mock_chunks[0],
+            relevancy=EvaluationMetricResult(
+                value="low",
+                name=EvaluationMetricName.CONTEXT_RELEVANCY,
+                reason="Test reason",
+            ),
+        )
+    ]
+    chunk_relevancies.append(
+        ChunkRelevancy(
+            chunk=mock_chunks[1],
+            relevancy=EvaluationMetricResult(
+                value="medium",
+                name=EvaluationMetricName.CONTEXT_RELEVANCY,
+                reason="Test reason",
+            ),
+        )
+    )
+    chunk_relevancies.append(
+        ChunkRelevancy(
+            chunk=mock_chunks[2],
+            relevancy=EvaluationMetricResult(
+                value="high",
+                name=EvaluationMetricName.CONTEXT_RELEVANCY,
+                reason="Test reason",
+            ),
+        )
+    )
+
+    result = await chunk_relevancy_sorter._validate_and_sort_relevant_chunks(
+        config, chunk_relevancies
+    )
+
+    assert isinstance(result, list)
+    assert len(result) == len(mock_chunks)
+    assert all(isinstance(relevancy.chunk, ContentChunk) for relevancy in result)
+    assert result[0].chunk == mock_chunks[2]
+    assert result[1].chunk == mock_chunks[1]
+    assert result[2].chunk == mock_chunks[0]
+
+
+@pytest.mark.asyncio
+async def test_validate_chunk_relevancies_invalid(chunk_relevancy_sorter):
+    invalid_relevancies = [
+        ChunkRelevancy(
+            chunk=ContentChunk(chunk_id="test", text="test", id="test", order=0),
+            relevancy=None,
+        )
+    ]
+
+    with pytest.raises(ChunkRelevancySorterException):
+        await chunk_relevancy_sorter._validate_chunk_relevancies(invalid_relevancies)
+
+
+def test_count_distinct_values(chunk_relevancy_sorter, mock_chunks):
+    chunk_relevancies = [
+        ChunkRelevancy(
+            chunk=chunk,
+            relevancy=EvaluationMetricResult(
+                value="high",
+                name=EvaluationMetricName.CONTEXT_RELEVANCY,
+                reason="Test reason",
+            ),
+        )
+        for chunk in mock_chunks[:2]
+    ]
+    chunk_relevancies.append(
+        ChunkRelevancy(
+            chunk=mock_chunks[2],
+            relevancy=EvaluationMetricResult(
+                value="medium",
+                name=EvaluationMetricName.CONTEXT_RELEVANCY,
+                reason="Test reason",
+            ),
+        )
+    )
+
+    value_counts = chunk_relevancy_sorter._count_distinct_values(chunk_relevancies)
+
+    assert value_counts["high"] == 2
+    assert value_counts["medium"] == 1

unique_toolkit/_common/feature_flags/schema.py

@@ -0,0 +1,13 @@
+from pydantic import BaseModel, Field
+
+from unique_toolkit.tools.config import get_configuration_dict
+
+
+class FeatureExtendedSourceSerialization(BaseModel):
+    """Mixin for experimental feature in Source serialization"""
+
+    model_config = get_configuration_dict()
+    full_sources_serialize_dump: bool = Field(
+        default=False,
+        description="Whether to include the full source object in the tool response. If True, includes the full Source object. If False, uses the old format with only source_number and content.",
+    )

unique_toolkit/_common/utils/structured_output/schema.py

@@ -0,0 +1,5 @@
+from pydantic import BaseModel, ConfigDict
+
+
+class StructuredOutputModel(BaseModel):
+    model_config = ConfigDict(extra="forbid")

unique_toolkit/evals/context_relevancy/schema.py

@@ -1,13 +1,10 @@
-from pydantic import BaseModel, ConfigDict, Field, create_model
+from pydantic import BaseModel, Field, create_model
 from pydantic.json_schema import SkipJsonSchema
 
+from unique_toolkit._common.utils.structured_output.schema import StructuredOutputModel
 from unique_toolkit.tools.config import get_configuration_dict
 
 
-class StructuredOutputModel(BaseModel):
-    model_config = ConfigDict(extra="forbid")
-
-
 class StructuredOutputConfig(BaseModel):
     model_config = get_configuration_dict()
 

unique_toolkit/evals/context_relevancy/service.py

@@ -1,9 +1,14 @@
 import logging
+from typing import overload
 
 from pydantic import BaseModel, ValidationError
+from typing_extensions import deprecated
 
-from unique_toolkit.app.schemas import ChatEvent
-from unique_toolkit.chat.service import ChatService
+from unique_toolkit._common.default_language_model import DEFAULT_GPT_35_TURBO
+from unique_toolkit._common.validate_required_values import (
+    validate_required_values,
+)
+from unique_toolkit.app.schemas import BaseEvent, ChatEvent
 from unique_toolkit.evals.config import EvaluationMetricConfig
 from unique_toolkit.evals.context_relevancy.schema import (
     EvaluationSchemaStructuredOutput,
@@ -21,7 +26,6 @@ from unique_toolkit.evals.schemas import (
 )
 from unique_toolkit.language_model.infos import (
     LanguageModelInfo,
-    LanguageModelName,
     ModelCapabilities,
 )
 from unique_toolkit.language_model.prompt import Prompt
@@ -45,9 +49,7 @@ USER_MSG_KEY = "userPrompt"
 default_config = EvaluationMetricConfig(
     enabled=False,
     name=EvaluationMetricName.CONTEXT_RELEVANCY,
-    language_model=LanguageModelInfo.from_name(
-        LanguageModelName.AZURE_GPT_4o_2024_1120
-    ),
+    language_model=LanguageModelInfo.from_name(DEFAULT_GPT_35_TURBO),
     custom_prompts={
         SYSTEM_MSG_KEY: CONTEXT_RELEVANCY_METRIC_SYSTEM_MSG,
         USER_MSG_KEY: CONTEXT_RELEVANCY_METRIC_USER_MSG,
@@ -61,13 +63,42 @@ relevancy_required_input_fields = [
 
 
 class ContextRelevancyEvaluator:
+    @deprecated(
+        "Use __init__ with company_id and user_id instead or use the classmethod `from_event`"
+    )
+    @overload
+    def __init__(self, event: ChatEvent | BaseEvent):
+        """
+        Initialize the ContextRelevancyEvaluator with an event (deprecated)
+        """
+
+    @overload
+    def __init__(self, *, company_id: str, user_id: str):
+        """
+        Initialize the ContextRelevancyEvaluator with a company_id and user_id
+        """
+
     def __init__(
         self,
-        event: ChatEvent,
+        event: ChatEvent | BaseEvent | None = None,
+        company_id: str | None = None,
+        user_id: str | None = None,
     ):
-        self.chat_service = ChatService(event)
-        self.language_model_service = LanguageModelService(event)
-        self.logger = logging.getLogger(f"ContextRelevancyEvaluator.{__name__}")
+        if isinstance(event, (ChatEvent, BaseEvent)):
+            self.language_model_service = LanguageModelService.from_event(event)
+        else:
+            [company_id, user_id] = validate_required_values([company_id, user_id])
+            self.language_model_service = LanguageModelService(
+                company_id=company_id, user_id=user_id
+            )
+
+        # Setup the logger
+        module_name = "ContextRelevancyEvaluator"
+        self.logger = logging.getLogger(f"{module_name}.{__name__}")
+
+    @classmethod
+    def from_event(cls, event: ChatEvent | BaseEvent):
+        return cls(company_id=event.company_id, user_id=event.user_id)
 
     async def analyze(
         self,
@@ -95,6 +126,7 @@ class ContextRelevancyEvaluator:
 
         input.validate_required_fields(relevancy_required_input_fields)
 
+        # TODO: Was already there in monorepo
         if len(input.context_texts) == 0:  # type: ignore
             error_message = "No context texts provided."
             raise EvaluatorException(

unique_toolkit/language_model/infos.py

@@ -48,6 +48,14 @@ class LanguageModelName(StrEnum):
     LITELLM_OPENAI_GPT_5_MINI = "litellm:openai-gpt-5-mini"
     LITELLM_OPENAI_GPT_5_NANO = "litellm:openai-gpt-5-nano"
     LITELLM_OPENAI_GPT_5_CHAT = "litellm:openai-gpt-5-chat"
+    LITELLM_OPENAI_O1 = "litellm:openai-o1"
+    LITELLM_OPENAI_O3 = "litellm:openai-o3"
+    LITELLM_OPENAI_O3_DEEP_RESEARCH = "litellm:openai-o3-deep-research"
+    LITELLM_OPENAI_O3_PRO = "litellm:openai-o3-pro"
+    LITELLM_OPENAI_O4_MINI = "litellm:openai-o4-mini"
+    LITELLM_OPENAI_O4_MINI_DEEP_RESEARCH = "litellm:openai-o4-mini-deep-research"
+    LITELLM_OPENAI_GPT_4_1_MINI = "litellm:openai-gpt-4-1-mini"
+    LITELLM_OPENAI_GPT_4_1_NANO = "litellm:openai-gpt-4-1-nano"
     LITELLM_DEEPSEEK_R1 = "litellm:deepseek-r1"
     LITELLM_DEEPSEEK_V3 = "litellm:deepseek-v3-1"
     LITELLM_QWEN_3 = "litellm:qwen-3-235B-A22B"
@@ -83,6 +91,14 @@ def get_encoder_name(model_name: LanguageModelName) -> EncoderName:
             | LMN.LITELLM_OPENAI_GPT_5_MINI
             | LMN.LITELLM_OPENAI_GPT_5_NANO
             | LMN.LITELLM_OPENAI_GPT_5_CHAT
+            | LMN.LITELLM_OPENAI_O1
+            | LMN.LITELLM_OPENAI_O3
+            | LMN.LITELLM_OPENAI_O3_DEEP_RESEARCH
+            | LMN.LITELLM_OPENAI_O4_MINI
+            | LMN.LITELLM_OPENAI_O4_MINI_DEEP_RESEARCH
+            | LMN.LITELLM_OPENAI_GPT_4_1_MINI
+            | LMN.LITELLM_OPENAI_GPT_4_1_NANO
+            | LMN.LITELLM_OPENAI_O3_PRO
         ):
             return EncoderName.O200K_BASE
         case _:
@@ -879,6 +895,146 @@ class LanguageModelInfo(BaseModel):
                     deprecated_at=date(2026, 8, 7),
                     retirement_at=date(2026, 8, 7),
                 )
+            case LanguageModelName.LITELLM_OPENAI_O1:
+                return cls(
+                    name=model_name,
+                    provider=LanguageModelProvider.LITELLM,
+                    version="2024-12-17",
+                    encoder_name=EncoderName.O200K_BASE,
+                    capabilities=[
+                        ModelCapabilities.STRUCTURED_OUTPUT,
+                        ModelCapabilities.FUNCTION_CALLING,
+                        ModelCapabilities.STREAMING,
+                        ModelCapabilities.VISION,
+                        ModelCapabilities.REASONING,
+                    ],
+                    token_limits=LanguageModelTokenLimits(
+                        token_limit_input=200_000, token_limit_output=100_000
+                    ),
+                    info_cutoff_at=date(2023, 10, 1),
+                    published_at=date(2024, 12, 17),
+                    temperature_bounds=TemperatureBounds(
+                        min_temperature=1.0, max_temperature=1.0
+                    ),
+                )
+            case LanguageModelName.LITELLM_OPENAI_O3:
+                return cls(
+                    name=model_name,
+                    provider=LanguageModelProvider.LITELLM,
+                    version="2025-04-16",
+                    encoder_name=EncoderName.O200K_BASE,
+                    capabilities=[
+                        ModelCapabilities.FUNCTION_CALLING,
+                        ModelCapabilities.STRUCTURED_OUTPUT,
+                        ModelCapabilities.STREAMING,
+                        ModelCapabilities.REASONING,
+                    ],
+                    token_limits=LanguageModelTokenLimits(
+                        token_limit_input=200_000, token_limit_output=100_000
+                    ),
+                    temperature_bounds=TemperatureBounds(
+                        min_temperature=1.0, max_temperature=1.0
+                    ),
+                    published_at=date(2025, 4, 16),
+                    info_cutoff_at=date(2024, 6, 1),
+                )
+            case LanguageModelName.LITELLM_OPENAI_O3_DEEP_RESEARCH:
+                return cls(
+                    name=model_name,
+                    provider=LanguageModelProvider.LITELLM,
+                    version="2025-06-26",
+                    encoder_name=EncoderName.O200K_BASE,
+                    token_limits=LanguageModelTokenLimits(
+                        token_limit_input=200_000, token_limit_output=100_000
+                    ),
+                    published_at=date(2025, 4, 16),
+                    capabilities=[ModelCapabilities.STREAMING],
+                    info_cutoff_at=date(2024, 6, 1),
+                )
+            case LanguageModelName.LITELLM_OPENAI_O3_PRO:
+                return cls(
+                    name=model_name,
+                    provider=LanguageModelProvider.LITELLM,
+                    version="2025-06-10",
+                    encoder_name=EncoderName.O200K_BASE,
+                    capabilities=[
+                        ModelCapabilities.FUNCTION_CALLING,
+                        ModelCapabilities.REASONING,
+                        ModelCapabilities.STRUCTURED_OUTPUT,
+                    ],
+                    token_limits=LanguageModelTokenLimits(
+                        token_limit_input=200_000, token_limit_output=100_000
+                    ),
+                    published_at=date(2025, 6, 10),
+                    info_cutoff_at=date(2024, 6, 1),
+                )
+            case LanguageModelName.LITELLM_OPENAI_O4_MINI:
+                return cls(
+                    name=model_name,
+                    provider=LanguageModelProvider.LITELLM,
+                    version="2025-04-16",
+                    encoder_name=EncoderName.O200K_BASE,
+                    capabilities=[
+                        ModelCapabilities.FUNCTION_CALLING,
+                        ModelCapabilities.STREAMING,
+                        ModelCapabilities.STRUCTURED_OUTPUT,
+                    ],
+                    token_limits=LanguageModelTokenLimits(
+                        token_limit_input=200_000, token_limit_output=100_000
+                    ),
+                    published_at=date(2025, 4, 16),
+                    info_cutoff_at=date(2024, 6, 1),
+                    temperature_bounds=TemperatureBounds(
+                        min_temperature=1.0, max_temperature=1.0
+                    ),
+                )
+            case LanguageModelName.LITELLM_OPENAI_O4_MINI_DEEP_RESEARCH:
+                return cls(
+                    name=model_name,
+                    provider=LanguageModelProvider.LITELLM,
+                    version="2025-06-26",
+                    encoder_name=EncoderName.O200K_BASE,
+                    token_limits=LanguageModelTokenLimits(
+                        token_limit_input=200_000, token_limit_output=100_000
+                    ),
+                    published_at=date(2025, 4, 16),
+                    capabilities=[ModelCapabilities.STREAMING],
+                    info_cutoff_at=date(2024, 6, 1),
+                )
+            case LanguageModelName.LITELLM_OPENAI_GPT_4_1_MINI:
+                return cls(
+                    name=model_name,
+                    provider=LanguageModelProvider.LITELLM,
+                    version="2025-04-14",
+                    encoder_name=EncoderName.O200K_BASE,
+                    published_at=date(2025, 4, 14),
+                    info_cutoff_at=date(2024, 6, 1),
+                    token_limits=LanguageModelTokenLimits(
+                        token_limit_input=1_047_576, token_limit_output=32_768
+                    ),
+                    capabilities=[
+                        ModelCapabilities.STREAMING,
+                        ModelCapabilities.FUNCTION_CALLING,
+                        ModelCapabilities.STRUCTURED_OUTPUT,
+                    ],
+                )
+            case LanguageModelName.LITELLM_OPENAI_GPT_4_1_NANO:
+                return cls(
+                    name=model_name,
+                    provider=LanguageModelProvider.LITELLM,
+                    version="2025-04-14",
+                    encoder_name=EncoderName.O200K_BASE,
+                    published_at=date(2025, 4, 14),
+                    info_cutoff_at=date(2024, 6, 1),
+                    token_limits=LanguageModelTokenLimits(
+                        token_limit_input=1_047_576, token_limit_output=32_768
+                    ),
+                    capabilities=[
+                        ModelCapabilities.STREAMING,
+                        ModelCapabilities.FUNCTION_CALLING,
+                        ModelCapabilities.STRUCTURED_OUTPUT,
+                    ],
+                )
             case LanguageModelName.LITELLM_DEEPSEEK_R1:
                 return cls(
                     name=model_name,

unique_toolkit/tools/a2a/config.py

@@ -1,7 +1,5 @@
-from unique_toolkit.tools.schemas import BaseToolConfig
-
 from unique_toolkit.tools.config import get_configuration_dict
-
+from unique_toolkit.tools.schemas import BaseToolConfig
 
 DEFAULT_PARAM_DESCRIPTION_SUB_AGENT_USER_MESSAGE = """
 This is the message that will be sent to the sub-agent.

unique_toolkit/tools/a2a/manager.py

@@ -1,9 +1,8 @@
 from logging import Logger
 
 from unique_toolkit.app.schemas import ChatEvent
-from unique_toolkit.tools.a2a.service import SubAgentTool, ToolProgressReporter
-
 from unique_toolkit.tools.a2a.config import SubAgentToolConfig
+from unique_toolkit.tools.a2a.service import SubAgentTool, ToolProgressReporter
 from unique_toolkit.tools.config import ToolBuildConfig
 from unique_toolkit.tools.schemas import BaseToolConfig
 from unique_toolkit.tools.tool import Tool

unique_toolkit/tools/a2a/memory.py

@@ -1,5 +1,4 @@
 from unique_toolkit import ShortTermMemoryService
-
 from unique_toolkit.short_term_memory.persistent_short_term_memory_manager import (
     PersistentShortMemoryManager,
 )

unique_toolkit/tools/a2a/service.py

@@ -1,25 +1,23 @@
 from pydantic import Field, create_model
 from unique_sdk.utils.chat_in_space import send_message_and_wait_for_completion
+
 from unique_toolkit.app import ChatEvent
+from unique_toolkit.evaluators.schemas import EvaluationMetricName
 from unique_toolkit.language_model import (
     LanguageModelFunction,
     LanguageModelMessage,
     LanguageModelToolDescription,
 )
-
-from unique_toolkit.evaluators.schemas import EvaluationMetricName
-from unique_toolkit.tools.schemas import ToolCallResponse
 from unique_toolkit.tools.a2a.config import SubAgentToolConfig
 from unique_toolkit.tools.a2a.memory import (
     get_sub_agent_short_term_memory_manager,
 )
 from unique_toolkit.tools.a2a.schema import (
     SubAgentShortTermMemorySchema,
-    SubAgentToolCallResponse,
     SubAgentToolInput,
 )
+from unique_toolkit.tools.schemas import ToolCallResponse
 from unique_toolkit.tools.tool import Tool
-
 from unique_toolkit.tools.tool_progress_reporter import (
     ProgressState,
     ToolProgressReporter,
@@ -148,7 +146,7 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
         tool_response: ToolCallResponse,
     ) -> LanguageModelMessage:
         return ToolCallResponse(
-            id=tool_call.id,  # type: ignore
-            name=tool_call.name,
-            content=response["text"],
+            id=tool_response.id,
+            name=tool_response.name,
+            content=tool_response["content"],
         )