unique_toolkit 2026.36.0.dev0__py3-none-any.whl → 2026.36.0.dev1__py3-none-any.whl

unique_toolkit/agentic/tools/mcp/tool_wrapper.py

@@ -4,9 +4,10 @@ import json
 import logging
 import mimetypes
 import uuid
-from typing import Any, Dict
+from typing import Any
 
 import unique_sdk
+from pydantic import ValidationError
 
 from unique_toolkit.agentic.evaluation.schemas import EvaluationMetricName
 from unique_toolkit.agentic.feature_flags import FeatureFlagNames
@@ -21,6 +22,7 @@ from unique_toolkit.app.schemas import ChatEvent, McpServer, McpTool
 from unique_toolkit.chat.schemas import MessageLog, MessageLogStatus
 from unique_toolkit.chat.service import ChatService
 from unique_toolkit.content.functions import upload_content_from_bytes
+from unique_toolkit.content.schemas import ContentChunk, ContentReference
 from unique_toolkit.experimental.resources.feature_flags import is_flag_enabled
 from unique_toolkit.language_model import LanguageModelService
 from unique_toolkit.language_model.schemas import (
@@ -30,6 +32,8 @@ from unique_toolkit.language_model.schemas import (
 
 logger = logging.getLogger(__name__)
 
+REFERENCE_META_KEY = "unique.app/reference"
+
 
 class MCPToolWrapper(Tool[MCPToolConfig]):
     """Wrapper class for MCP tools that implements the Tool interface"""
@@ -72,7 +76,7 @@ class MCPToolWrapper(Tool[MCPToolConfig]):
             parameters=self._mcp_tool.input_schema,
         )
 
-    def _json_schema_to_python_type(self, schema: Dict[str, Any]) -> type:
+    def _json_schema_to_python_type(self, schema: dict[str, Any]) -> type:
         """Convert JSON schema type to Python type"""
         json_type = schema.get("type", "string")
 
@@ -152,9 +156,14 @@ class MCPToolWrapper(Tool[MCPToolConfig]):
             # Use SDK to call the public API
             result = await self._call_mcp_tool_via_sdk(arguments)
 
+            content_chunks, references = self._extract_structured_references(result)
             content_str, image_data_urls = await asyncio.to_thread(
                 self._process_mcp_result, result
             )
+            if content_chunks:
+                # Structured references already carry the text; keep any
+                # extracted images but drop the redundant flat text.
+                content_str = ""
 
             # TODO: Why result here not applied directly to the body of the tool_response? like so how does it know the results in the history?
             tool_response = ToolCallResponse(
@@ -167,6 +176,7 @@ class MCPToolWrapper(Tool[MCPToolConfig]):
                 },
                 error_message="",
                 content=content_str,
+                content_chunks=content_chunks or None,
                 image_data_urls=image_data_urls,
             )
 
@@ -187,6 +197,7 @@ class MCPToolWrapper(Tool[MCPToolConfig]):
                 progress_message="_Completed MCP tool_",
                 status=MessageLogStatus.COMPLETED,
                 active_message_log=active_message_log,
+                references=references,
             )
 
             return tool_response
@@ -230,17 +241,19 @@ class MCPToolWrapper(Tool[MCPToolConfig]):
         progress_message: str | None = None,
         status: MessageLogStatus = MessageLogStatus.RUNNING,
         active_message_log: MessageLog | None = None,
+        references: list[ContentReference] | None = None,
     ) -> MessageLog | None:
         return self._message_step_logger.create_or_update_message_log(
             active_message_log=active_message_log,
             header=self.display_name(),
             progress_message=progress_message,
             status=status,
+            references=references or [],
         )
 
     def _extract_and_validate_arguments(
         self, tool_call: LanguageModelFunction
-    ) -> Dict[str, Any]:
+    ) -> dict[str, Any]:
         """
         Extract and validate arguments from tool call, handling various formats robustly.
 
@@ -287,7 +300,39 @@ class MCPToolWrapper(Tool[MCPToolConfig]):
             f"Unexpected arguments type for MCP tool {self.name}: {type(raw_arguments)}"
         )
 
-    def _process_mcp_result(self, result: Dict[str, Any]) -> tuple[str, list[str]]:
+    def _extract_structured_references(
+        self, result: dict[str, Any]
+    ) -> tuple[list[ContentChunk], list[ContentReference]]:
+        """Opt-in: content items with valid ``unique.app/reference`` + chunk ``_meta``."""
+        chunks: list[ContentChunk] = []
+        references: list[ContentReference] = []
+        content_items = self._extract_content_items(result)
+        for item in content_items:
+            if not isinstance(item, dict) or item.get("type") != "text":
+                continue
+            meta = item.get("_meta")
+            if not isinstance(meta, dict):
+                continue
+            ref_payload = meta.get(REFERENCE_META_KEY)
+            if not isinstance(ref_payload, dict):
+                continue
+            try:
+                reference = ContentReference.model_validate(ref_payload)
+                chunk = ContentChunk.model_validate(
+                    {k: v for k, v in meta.items() if k != REFERENCE_META_KEY}
+                )
+            except ValidationError:
+                self.logger.warning(
+                    "MCP tool %s: ignoring invalid %s metadata",
+                    self.name,
+                    REFERENCE_META_KEY,
+                )
+                continue
+            references.append(reference)
+            chunks.append(chunk)
+        return chunks, references
+
+    def _process_mcp_result(self, result: dict[str, Any]) -> tuple[str, list[str]]:
         """Parse MCP result content array, upload images, return content string and data URLs."""
         content_items = self._extract_content_items(result)
         if not content_items:
@@ -327,7 +372,7 @@ class MCPToolWrapper(Tool[MCPToolConfig]):
             content_str = json.dumps(result)
         return content_str, image_data_urls
 
-    def _extract_content_items(self, result: Dict[str, Any]) -> list[Any]:
+    def _extract_content_items(self, result: dict[str, Any]) -> list[Any]:
         """Extract content array from MCP result (handles various response shapes)."""
         if "content" in result and isinstance(result["content"], list):
             return result["content"]
@@ -378,7 +423,7 @@ class MCPToolWrapper(Tool[MCPToolConfig]):
 
         return data_url, content_id
 
-    async def _call_mcp_tool_via_sdk(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
+    async def _call_mcp_tool_via_sdk(self, arguments: dict[str, Any]) -> dict[str, Any]:
         """Call MCP tool via SDK to public API"""
         try:
             self.logger.info(

{unique_toolkit-2026.36.0.dev0.dist-info → unique_toolkit-2026.36.0.dev1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: unique-toolkit
-Version: 2026.36.0.dev0
+Version: 2026.36.0.dev1
 Summary: 
 Author: Cedric Klinkert, Martin Fadler, Sadique Sheik, Fabian Schläpfer, Pascal Hauri, Ejiroghene Oki-Peter, Rami Azouz, Ricardo Visini, Sanzio Monti
 Author-email: Cedric Klinkert <cedric.klinkert@unique.ch>, Martin Fadler <martin.fadler@unique.ch>, Sadique Sheik <sadique@unique.ch>, Fabian Schläpfer <fabian@unique.ch>, Pascal Hauri <pascal@unique.ch>, Ejiroghene Oki-Peter <oki-peter.ejiroghene@unique.ai>, Rami Azouz <rami.ext@unique.ch>, Ricardo Visini <ricardo.visini@unique.ch>, Sanzio Monti <sanzio.monti@unique.ai>

{unique_toolkit-2026.36.0.dev0.dist-info → unique_toolkit-2026.36.0.dev1.dist-info}/RECORD

@@ -201,7 +201,7 @@ unique_toolkit/agentic/tools/factory.py,sha256=YK9uh2EqotQFb9stsXs-A8lsNaxklTCd_
 unique_toolkit/agentic/tools/mcp/__init__.py,sha256=RLF_p-LDRC7GhiB3fdCi4u3bh6V9PY_w26fg61BLyco,122
 unique_toolkit/agentic/tools/mcp/manager.py,sha256=EfWbfNhB8Lnx2EqBgUXhK7klFLaBv2eezvYq8J0TXyo,3808
 unique_toolkit/agentic/tools/mcp/models.py,sha256=DaBVzwJ4dqtgzm8Mv0-NOGBU04zyVe-RgWntBhkpmnE,742
-unique_toolkit/agentic/tools/mcp/tool_wrapper.py,sha256=fE5nVALJ-wVZkkbbJ_ySRLBhAAR6sCAAu1lUxQQVWC4,15596
+unique_toolkit/agentic/tools/mcp/tool_wrapper.py,sha256=-WZaPjhbxozrPWgwZDD48hfYIp7s1cGAKU3IXqLXFVA,17575
 unique_toolkit/agentic/tools/names.py,sha256=MDF0ONLrxeC-iz--EtyrHAox_RN6zasxo5MSbXQYGyw,412
 unique_toolkit/agentic/tools/openai_builtin/__init__.py,sha256=Q6d9mhsLRVF3erT2yMxLld62D3FTZ8nbovI4vLv4Nso,701
 unique_toolkit/agentic/tools/openai_builtin/base.py,sha256=QWF-GL_fkJqzGPpxvtg69-TB_IjsrxokeJVQESzH-18,1415
@@ -420,6 +420,6 @@ unique_toolkit/short_term_memory/service.py,sha256=2JMBEBn6GoyK-a0xvdcWpPft-3C74
 unique_toolkit/smart_rules/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 unique_toolkit/smart_rules/compile.py,sha256=vRe0aMJWuVIK-VsqomgxQImzke7MDmwgSxxJfYCa7zw,1386
 unique_toolkit/test_utilities/events.py,sha256=FLrhMh0A_mQutAUjoD0FdMcVQ4Gq31r35OVGzRjvLro,6760
-unique_toolkit-2026.36.0.dev0.dist-info/WHEEL,sha256=eh7sammvW2TypMMMGKgsM83HyA_3qQ5Lgg3ynoecH3M,79
-unique_toolkit-2026.36.0.dev0.dist-info/METADATA,sha256=FDHFO1DJ6xFdOS6tl_Nise6nYmhJmGppXAW5CkEmlL4,7821
-unique_toolkit-2026.36.0.dev0.dist-info/RECORD,,
+unique_toolkit-2026.36.0.dev1.dist-info/WHEEL,sha256=eh7sammvW2TypMMMGKgsM83HyA_3qQ5Lgg3ynoecH3M,79
+unique_toolkit-2026.36.0.dev1.dist-info/METADATA,sha256=TRfcAs3vlbrpma_tG96J6ZZq6eAu9L9NWGIrq2Y2eVY,7821
+unique_toolkit-2026.36.0.dev1.dist-info/RECORD,,