unique_toolkit 0.8.4__py3-none-any.whl → 0.8.5__py3-none-any.whl

unique_toolkit/app/schemas.py

@@ -51,6 +51,7 @@ class BaseEvent(BaseModel):
 # MCP schemas
 ###
 
+
 class McpTool(BaseModel):
     model_config = model_config
 
@@ -79,6 +80,7 @@ class McpTool(BaseModel):
         description="Whether the tool is connected to the MCP server. This is a Unique specific field.",
     )
 
+
 class McpServer(BaseModel):
     model_config = model_config
 
@@ -94,6 +96,7 @@ class McpServer(BaseModel):
     )
     tools: list[McpTool] = []
 
+
 ###
 # ChatEvent schemas
 ###

unique_toolkit/language_model/schemas.py

@@ -86,6 +86,24 @@ class LanguageModelFunction(BaseModel):
         return seralization
 
 
+    def __eq__(self, other:Self) -> bool:
+        """
+        Compare two tool calls based on name and arguments.
+        """
+        if not isinstance(other, LanguageModelFunction):
+            return False
+
+        if self.id != other.id:
+            return False
+
+        if self.name != other.name:
+            return False
+
+        if self.arguments != other.arguments:
+            return False
+        
+        return True
+
 # This is tailored to the unique backend
 class LanguageModelStreamResponse(BaseModel):
     model_config = model_config

unique_toolkit/reference_manager/reference_manager.py

@@ -0,0 +1,72 @@
+from unique_toolkit.content.schemas import ContentChunk, ContentReference
+from unique_toolkit.tools.schemas import ToolCallResponse
+
+
+class tool_chunks:
+    def __init__(self, name: str, chunks: list) -> None:
+        self.name = name
+        self.chunks = chunks
+
+
+class ReferenceManager:
+    def __init__(self):
+        self._tool_chunks: dict[str, tool_chunks] = {}
+        self._chunks: list[ContentChunk] = []
+        self._references: list[list[ContentReference]] = []
+
+    def extract_referenceable_chunks(
+        self, tool_responses: list[ToolCallResponse]
+    ) -> None:
+        for tool_response in tool_responses:
+            if not tool_response.content_chunks:
+                continue
+            self._chunks.extend(tool_response.content_chunks or [])
+            self._tool_chunks[tool_response.id] = tool_chunks(
+                tool_response.name, tool_response.content_chunks
+            )
+
+    def get_chunks(self) -> list[ContentChunk]:
+        return self._chunks
+
+    def get_tool_chunks(self) -> dict:
+        return self._tool_chunks
+
+    def replace(self, chunks: list[ContentChunk]):
+        self._chunks = chunks
+
+    def add_references(
+        self,
+        references: list[ContentReference],
+    ):
+        self._references.append(references)
+
+    def get_references(
+        self,
+    ) -> list[list[ContentReference]]:
+        return self._references
+
+    def get_latest_references(
+        self,
+    ) -> list[ContentReference]:
+        if not self._references:
+            return []
+        return self._references[-1]
+
+    def get_latest_referenced_chunks(self) -> list[ContentChunk]:
+        if not self._references:
+            return []
+        return self._get_referenced_chunks_from_references(self._references[-1])
+
+    def _get_referenced_chunks_from_references(
+        self,
+        references: list[ContentReference],
+    ) -> list[ContentChunk]:
+        """
+        Get _referenced_chunks by matching sourceId from _references with merged id and chunk_id from _chunks.
+        """
+        referenced_chunks: list[ContentChunk] = []
+        for ref in references:
+            for chunk in self._chunks:
+                if ref.source_id == f"{chunk.id}-{chunk.chunk_id}":
+                    referenced_chunks.append(chunk)
+        return referenced_chunks

unique_toolkit/tools/agent_chunks_handler.py

@@ -0,0 +1,62 @@
+from unique_toolkit.content.schemas import ContentChunk, ContentReference
+
+
+class AgentChunksHandler:
+    def __init__(self):
+        self._tool_chunks = {}
+        self._chunks: list[ContentChunk] = []
+        self._references: list[list[ContentReference]] = []
+
+    @property
+    def chunks(self) -> list[ContentChunk]:
+        return self._chunks
+
+    @property
+    def tool_chunks(self) -> dict:
+        return self._tool_chunks
+
+    def extend(self, chunks: list[ContentChunk]):
+        self._chunks.extend(chunks)
+
+    def replace(self, chunks: list[ContentChunk]):
+        self._chunks = chunks
+
+    def add_references(
+        self,
+        references: list[ContentReference],
+    ):
+        self._references.append(references)
+
+    @property
+    def all_references(
+        self,
+    ) -> list[list[ContentReference]]:
+        return self._references
+
+    @property
+    def latest_references(
+        self,
+    ) -> list[ContentReference]:
+        if not self._references:
+            return []
+        return self._references[-1]
+
+    @property
+    def latest_referenced_chunks(self) -> list[ContentChunk]:
+        if not self._references:
+            return []
+        return self._get_referenced_chunks_from_references(self._references[-1])
+
+    def _get_referenced_chunks_from_references(
+        self,
+        references: list[ContentReference],
+    ) -> list[ContentChunk]:
+        """
+        Get _referenced_chunks by matching sourceId from _references with merged id and chunk_id from _chunks.
+        """
+        referenced_chunks: list[ContentChunk] = []
+        for ref in references:
+            for chunk in self._chunks:
+                if ref.source_id == str(chunk.id) + "_" + str(chunk.chunk_id):
+                    referenced_chunks.append(chunk)
+        return referenced_chunks

unique_toolkit/tools/config.py

@@ -0,0 +1,108 @@
+from enum import StrEnum
+import humps
+from typing import Any
+from pydantic.fields import ComputedFieldInfo, FieldInfo
+from pydantic.alias_generators import to_camel
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    ValidationInfo,
+    model_validator,
+)
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from unique_toolkit.tools.schemas import BaseToolConfig
+
+
+def field_title_generator(
+    title: str,
+    info: FieldInfo | ComputedFieldInfo,
+) -> str:
+    return humps.decamelize(title).replace("_", " ").title()
+
+
+def model_title_generator(model: type) -> str:
+    return humps.decamelize(model.__name__).replace("_", " ").title()
+
+
+def get_configuration_dict(**kwargs) -> ConfigDict:
+    return ConfigDict(
+        alias_generator=to_camel,
+        field_title_generator=field_title_generator,
+        model_title_generator=model_title_generator,
+        populate_by_name=True,
+        protected_namespaces=(),
+        **kwargs,
+    )
+
+
+class ToolIcon(StrEnum):
+    ANALYTICS = "IconAnalytics"
+    BOOK = "IconBook"
+    FOLDERDATA = "IconFolderData"
+    INTEGRATION = "IconIntegration"
+    TEXT_COMPARE = "IconTextCompare"
+    WORLD = "IconWorld"
+    QUICK_REPLY = "IconQuickReply"
+    CHAT_PLUS = "IconChatPlus"
+
+
+class ToolSelectionPolicy(StrEnum):
+    """Determine the usage policy of tools."""
+
+    FORCED_BY_DEFAULT = "ForcedByDefault"
+    ON_BY_DEFAULT = "OnByDefault"
+    BY_USER = "ByUser"
+
+
+class ToolBuildConfig(BaseModel):
+    model_config = get_configuration_dict()
+    """Main tool configuration"""
+
+    name: str
+    configuration: "BaseToolConfig"
+    display_name: str = ""
+    icon: ToolIcon = ToolIcon.BOOK
+    selection_policy: ToolSelectionPolicy = Field(
+        default=ToolSelectionPolicy.BY_USER,
+    )
+    is_exclusive: bool = Field(
+        default=False,
+        description="This tool must be chosen by the user and no other tools are used for this iteration.",
+    )
+
+    is_enabled: bool = Field(default=True)
+
+    @model_validator(mode="before")
+    def initialize_config_based_on_tool_name(
+        cls,
+        value: Any,
+        info: ValidationInfo,
+    ) -> Any:
+        """Check the given values for."""
+        if not isinstance(value, dict):
+            return value
+
+        configuration = value.get("configuration", {})
+        if isinstance(configuration, dict):
+            # Local import to avoid circular import at module import time
+            from unique_toolkit.tools.factory import ToolFactory
+
+            config = ToolFactory.build_tool_config(
+                value["name"],
+                **configuration,
+            )
+        else:
+            # Check that the type of config matches the tool name
+            from unique_toolkit.tools.factory import ToolFactory
+
+            assert isinstance(
+                configuration,
+                ToolFactory.tool_config_map[value["name"]],  # type: ignore
+            )
+            config = configuration
+        value["configuration"] = config
+        return value

unique_toolkit/tools/{tool_factory.py → factory.py}

@@ -1,7 +1,11 @@
 from typing import Callable
 
-from unique_toolkit.unique_toolkit.tools.tool_definitions import BaseToolConfig, Tool
+from typing import TYPE_CHECKING
+from unique_toolkit.tools.schemas import BaseToolConfig
+from unique_toolkit.tools.tool import Tool
 
+if TYPE_CHECKING:
+    from unique_toolkit.tools.config import ToolBuildConfig
 
 
 class ToolFactory:
@@ -18,14 +22,20 @@ class ToolFactory:
         cls.tool_config_map[tool.name] = tool_config
 
     @classmethod
-    def build_tool(cls, tool_name: str, *args, **kwargs) -> Tool:
+    def build_tool(cls, tool_name: str, *args, **kwargs) -> Tool[BaseToolConfig]:
         tool = cls.tool_map[tool_name](*args, **kwargs)
         return tool
 
     @classmethod
-    def build_tool_config(
-        cls, tool_name: str, **kwargs
-    ) -> BaseToolConfig:
+    def build_tool_with_settings(
+        cls, tool_name: str, settings: "ToolBuildConfig", *args, **kwargs
+    ) -> Tool[BaseToolConfig]:
+        tool = cls.tool_map[tool_name](*args, **kwargs)
+        tool.settings = settings
+        return tool
+
+    @classmethod
+    def build_tool_config(cls, tool_name: str, **kwargs) -> BaseToolConfig:
         if tool_name not in cls.tool_config_map:
             raise ValueError(f"Tool {tool_name} not found")
         return cls.tool_config_map[tool_name](**kwargs)

unique_toolkit/tools/schemas.py

@@ -0,0 +1,138 @@
+import base64
+import gzip
+import re
+from typing import Any, Optional
+
+from pydantic import BaseModel, ConfigDict, Field, field_serializer, field_validator
+from unique_toolkit.content.schemas import ContentChunk
+
+from unique_toolkit.tools.config import get_configuration_dict
+from unique_toolkit.tools.utils.source_handling.schema import SourceFormatConfig
+
+
+# TODO: this needs to be more general as the tools can potentially return anything maybe make a base class and then derive per "type" of tool
+class ToolCallResponse(BaseModel):
+    id: str
+    name: str
+    debug_info: Optional[dict] = None  # TODO: Make the default {}
+    content_chunks: Optional[list[ContentChunk]] = None  # TODO: Make the default []
+    reasoning_result: Optional[dict] = None  # TODO: Make the default {}
+    error_message: str = ""
+
+    @property
+    def successful(self) -> bool:
+        return self.error_message == ""
+
+
+class BaseToolConfig(BaseModel):
+    model_config = get_configuration_dict()
+    # TODO: add a check for the parameters to all be consistent within the tool config
+    pass
+
+
+class Source(BaseModel):
+    """Represents the sources in the tool call response that the llm will see
+
+    Args:
+        source_number: The number of the source
+        content: The content of the source
+    """
+
+    model_config = ConfigDict(
+        validate_by_alias=True, serialize_by_alias=True, validate_by_name=True
+    )
+
+    source_number: int | None = Field(
+        default=None,
+        serialization_alias="[source_number] - Used for citations!",
+        validation_alias="[source_number] - Used for citations!",
+    )
+    content: str = Field(
+        serialization_alias="[content] - Content of source",
+        validation_alias="[content] - Content of source",
+    )
+    order: int = Field(
+        serialization_alias="[order] - Index in the document!",
+        validation_alias="[order] - Index in the document!",
+    )
+    chunk_id: str | None = Field(
+        default=None,
+        serialization_alias="[chunk_id] - IGNORE",
+        validation_alias="[chunk_id] - IGNORE",
+    )
+    id: str = Field(
+        serialization_alias="[id] - IGNORE",
+        validation_alias="[id] - IGNORE",
+    )
+    key: str | None = Field(
+        default=None,
+        serialization_alias="[key] - IGNORE",
+        validation_alias="[key] - IGNORE",
+    )
+    metadata: dict[str, str] | str | None = Field(
+        default=None,
+        serialization_alias="[metadata] - Formatted metadata",
+        validation_alias="[metadata] - Formatted metadata",
+    )
+    url: str | None = Field(
+        default=None,
+        serialization_alias="[url] - IGNORE",
+        validation_alias="[url] - IGNORE",
+    )
+
+    @field_validator("metadata", mode="before")
+    def _metadata_str_to_dict(
+        cls, v: str | dict[str, str] | None
+    ) -> dict[str, str] | None:
+        """
+        Accept   • dict   → keep as-is
+                 • str    → parse tag-string back to dict
+        """
+        if v is None or isinstance(v, dict):
+            return v
+
+        # v is the rendered string.  Build a dict by matching the
+        # patterns defined in SourceFormatConfig.sections.
+        cfg = SourceFormatConfig()  # or inject your app-wide config
+        out: dict[str, str] = {}
+        for key, tmpl in cfg.sections.items():
+            pattern = cfg.template_to_pattern(tmpl)
+            m = re.search(pattern, v, flags=re.S)
+            if m:
+                out[key] = m.group(1).strip()
+
+        return out if out else v  # type: ignore
+
+    # Compression + Base64 for url to hide it from the LLM
+    @field_serializer("url")
+    def serialize_url(self, value: str | None) -> str | None:
+        if value is None:
+            return None
+        # Compress then base64 encode
+        compressed = gzip.compress(value.encode())
+        return base64.b64encode(compressed).decode()
+
+    @field_validator("url", mode="before")
+    @classmethod
+    def validate_url(cls, value: Any) -> str | None:
+        if value is None or isinstance(value, str) and not value:
+            return None
+        if isinstance(value, str):
+            try:
+                # Try to decode base64 then decompress
+                decoded_bytes = base64.b64decode(value.encode())
+                decompressed = gzip.decompress(decoded_bytes).decode()
+                return decompressed
+            except Exception:
+                # If decoding/decompression fails, assume it's plain text
+                return value
+        return str(value)
+
+
+class ToolPrompts(BaseModel):
+    name: str
+    display_name: str
+    tool_description: str
+    tool_format_information_for_system_prompt: str
+    tool_format_information_for_system_prompt: str
+    input_model: dict[str, Any]

unique_toolkit/tools/test/test_tool_progress_reporter.py

@@ -0,0 +1,204 @@
+from unittest.mock import AsyncMock
+
+import pytest
+from unique_toolkit.chat.service import ChatService
+from unique_toolkit.content.schemas import ContentReference
+from unique_toolkit.language_model.schemas import LanguageModelFunction
+from unique_toolkit.tools.tool_progress_reporter import (
+    DUMMY_REFERENCE_PLACEHOLDER,
+    ProgressState,
+    ToolExecutionStatus,
+    ToolProgressReporter,
+    ToolWithToolProgressReporter,
+    track_tool_progress,
+)
+
+
+@pytest.fixture
+def chat_service():
+    return AsyncMock(spec=ChatService)
+
+
+@pytest.fixture
+def tool_progress_reporter(chat_service):
+    return ToolProgressReporter(chat_service)
+
+
+@pytest.fixture
+def tool_call():
+    return LanguageModelFunction(id="test_id", name="test_tool")
+
+
+class TestToolProgressReporter:
+    @pytest.mark.asyncio
+    async def test_notify_from_tool_call(self, tool_progress_reporter, tool_call):
+        # Arrange
+        name = "Test Tool"
+        message = "Processing..."
+        state = ProgressState.RUNNING
+        references = [
+            ContentReference(
+                sequence_number=1,
+                id="1",
+                message_id="1",
+                name="1",
+                source="1",
+                source_id="1",
+                url="1",
+            )
+        ]
+
+        # Act
+        await tool_progress_reporter.notify_from_tool_call(
+            tool_call=tool_call,
+            name=name,
+            message=message,
+            state=state,
+            references=references,
+            requires_new_assistant_message=True,
+        )
+
+        # Assert
+        assert tool_call.id in tool_progress_reporter.tool_statuses
+        status = tool_progress_reporter.tool_statuses[tool_call.id]
+        assert status.name == name
+        assert status.message == message
+        assert status.state == state
+        assert status.references == references
+        assert tool_progress_reporter.requires_new_assistant_message is True
+
+    def test_replace_placeholders(self, tool_progress_reporter):
+        # Arrange
+        message = (
+            f"Test{DUMMY_REFERENCE_PLACEHOLDER}message{DUMMY_REFERENCE_PLACEHOLDER}"
+        )
+
+        # Act
+        result = tool_progress_reporter._replace_placeholders(message, start_number=1)
+
+        # Assert
+        assert result == "Test<sup>1</sup>message<sup>2</sup>"
+
+    def test_correct_reference_sequence(self, tool_progress_reporter):
+        # Arrange
+        references = [
+            ContentReference(
+                sequence_number=0,
+                id="1",
+                message_id="1",
+                name="1",
+                source="1",
+                source_id="1",
+                url="1",
+            ),
+            ContentReference(
+                sequence_number=0,
+                id="2",
+                message_id="2",
+                name="2",
+                source="2",
+                source_id="2",
+                url="2",
+            ),
+        ]
+
+        # Act
+        result = tool_progress_reporter._correct_reference_sequence(
+            references, start_number=1
+        )
+
+        # Assert
+        assert len(result) == 2
+        assert result[0].sequence_number == 1
+        assert result[1].sequence_number == 2
+
+    @pytest.mark.asyncio
+    async def test_publish_updates_chat_service(
+        self, tool_progress_reporter, tool_call
+    ):
+        # Arrange
+        status = ToolExecutionStatus(
+            name="Test Tool",
+            message="Test message",
+            state=ProgressState.FINISHED,
+            references=[
+                ContentReference(
+                    sequence_number=1,
+                    id="1",
+                    message_id="1",
+                    name="1",
+                    source="1",
+                    source_id="1",
+                    url="1",
+                )
+            ],
+        )
+        tool_progress_reporter.tool_statuses[tool_call.id] = status
+
+        # Act
+        await tool_progress_reporter.publish()
+
+        # Assert
+        tool_progress_reporter.chat_service.modify_assistant_message_async.assert_called_once()
+
+
+class TestToolProgressDecorator:
+    class DummyTool(ToolWithToolProgressReporter):
+        def __init__(self, tool_progress_reporter):
+            self.tool_progress_reporter = tool_progress_reporter
+
+        @track_tool_progress(
+            message="Processing",
+            on_start_state=ProgressState.STARTED,
+            on_success_state=ProgressState.FINISHED,
+            on_success_message="Completed",
+            requires_new_assistant_message=True,
+        )
+        async def execute(self, tool_call, notification_tool_name):
+            return {
+                "references": [
+                    ContentReference(
+                        sequence_number=1,
+                        id="1",
+                        message_id="1",
+                        name="1",
+                        source="1",
+                        source_id="1",
+                        url="1",
+                    )
+                ]
+            }
+
+    @pytest.mark.asyncio
+    async def test_decorator_success_flow(self, tool_progress_reporter, tool_call):
+        # Arrange
+        tool = self.DummyTool(tool_progress_reporter)
+
+        # Act
+        await tool.execute(tool_call, "Test Tool")
+
+        # Assert
+        assert len(tool_progress_reporter.tool_statuses) == 1
+        status = tool_progress_reporter.tool_statuses[tool_call.id]
+        assert status.state == ProgressState.FINISHED
+        assert status.message == "Completed"
+
+    @pytest.mark.asyncio
+    async def test_decorator_error_flow(self, tool_progress_reporter, tool_call):
+        # Arrange
+        class ErrorTool(ToolWithToolProgressReporter):
+            def __init__(self, tool_progress_reporter):
+                self.tool_progress_reporter = tool_progress_reporter
+
+            @track_tool_progress(message="Processing")
+            async def execute(self, tool_call, notification_tool_name):
+                raise ValueError("Test error")
+
+        tool = ErrorTool(tool_progress_reporter)
+
+        # Act & Assert
+        with pytest.raises(ValueError):
+            await tool.execute(tool_call, "Test Tool")
+
+        status = tool_progress_reporter.tool_statuses[tool_call.id]
+        assert status.state == ProgressState.FAILED