unique_toolkit 0.7.41__tar.gz → 0.8.0__tar.gz

{unique_toolkit-0.7.41 → unique_toolkit-0.8.0}/CHANGELOG.md

@@ -5,6 +5,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), 
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.8.0] - 2025-08-04
+- Add MCP support
+
+## [0.7.42] - 2025-08-01
+- Added tool definitions
+
 ## [0.7.41] - 2025-07-31
 - Add new chat event attribute indicating tools disabled on a company level
 

{unique_toolkit-0.7.41 → unique_toolkit-0.8.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: unique_toolkit
-Version: 0.7.41
+Version: 0.8.0
 Summary: 
 License: Proprietary
 Author: Martin Fadler
@@ -19,7 +19,7 @@ Requires-Dist: regex (>=2024.5.15,<2025.0.0)
 Requires-Dist: sseclient (>=0.0.27,<0.0.28)
 Requires-Dist: tiktoken (>=0.7.0,<0.8.0)
 Requires-Dist: typing-extensions (>=4.9.0,<5.0.0)
-Requires-Dist: unique-sdk (>=0.9.40,<0.10.0)
+Requires-Dist: unique-sdk (>=0.10.0,<0.11.0)
 Description-Content-Type: text/markdown
 
 # Unique Toolkit
@@ -113,6 +113,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), 
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.8.0] - 2025-08-04
+- Add MCP support
+
+## [0.7.42] - 2025-08-01
+- Added tool definitions
+
 ## [0.7.41] - 2025-07-31
 - Add new chat event attribute indicating tools disabled on a company level
 

{unique_toolkit-0.7.41 → unique_toolkit-0.8.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "unique_toolkit"
-version = "0.7.41"
+version = "0.8.0"
 description = ""
 authors = [
     "Martin Fadler <martin.fadler@unique.ch>",
@@ -20,7 +20,7 @@ numpy = "^1.26.4"
 python-dotenv = "^1.0.1"
 regex = "^2024.5.15"
 tiktoken = "^0.7.0"
-unique-sdk = "^0.9.40"
+unique-sdk = "^0.10.0"
 pydantic-settings = "^2.10.1"
 sseclient = "^0.0.27"
 
@@ -47,7 +47,7 @@ pytest-cov = "^4.1.0"
 pre-commit = "^3.7.1"
 pytest-asyncio = "^0.23.8"
 pytest-mock = "^3.14.0"
-unique-sdk = {path = "../unique_sdk"}
+unique-sdk = { path = "../unique_sdk" }
 ipykernel = "^6.29.5"
 
 [build-system]

{unique_toolkit-0.7.41 → unique_toolkit-0.8.0}/unique_toolkit/app/__init__.py

@@ -41,5 +41,11 @@ from .schemas import (
 from .verification import (
     verify_signature_and_construct_event as verify_signature_and_construct_event,
 )
+from .schemas import (
+    McpServer as McpServer,
+)
+from .schemas import (
+    McpTool as McpTool,
+)
 
 DOMAIN_NAME = "app"

{unique_toolkit-0.7.41 → unique_toolkit-0.8.0}/unique_toolkit/app/schemas.py

@@ -38,6 +38,53 @@ class BaseEvent(BaseModel):
         return cls.model_validate(data)
 
 
+###
+# MCP schemas
+###
+
+class McpTool(BaseModel):
+    model_config = model_config
+
+    name: str
+    description: Optional[str] = None
+    input_schema: dict[str, Any]
+    output_schema: Optional[dict[str, Any]] = None
+    annotations: Optional[dict[str, Any]] = None
+    title: Optional[str] = Field(
+        default=None,
+        description="The display title for a tool. This is a Unique specific field.",
+    )
+    icon: Optional[str] = Field(
+        default=None,
+        description="An icon name from the Lucide icon set for the tool. This is a Unique specific field.",
+    )
+    system_prompt: Optional[str] = Field(
+        default=None,
+        description="An optional system prompt for the tool. This is a Unique specific field.",
+    )
+    user_prompt: Optional[str] = Field(
+        default=None,
+        description="An optional user prompt for the tool. This is a Unique specific field.",
+    )
+    is_connected: bool = Field(
+        description="Whether the tool is connected to the MCP server. This is a Unique specific field.",
+    )
+
+class McpServer(BaseModel):
+    model_config = model_config
+
+    id: str
+    name: str
+    system_prompt: Optional[str] = Field(
+        default=None,
+        description="An optional system prompt for the MCP server.",
+    )
+    user_prompt: Optional[str] = Field(
+        default=None,
+        description="An optional user prompt for the MCP server.",
+    )
+    tools: list[McpTool] = []
+
 ###
 # ChatEvent schemas
 ###
@@ -133,6 +180,10 @@ class ChatEventPayload(BaseModel):
         default=None,
         description="Raw UniqueQL rule that can be compiled to a metadata filter.",
     )
+    mcp_servers: list[McpServer] = Field(
+        default_factory=list,
+        description="A list of MCP servers with tools available for the chat session.",
+    )
 
     @field_validator("raw_scope_rules", mode="before")
     def validate_scope_rules(cls, value: dict[str, Any] | None) -> UniqueQL | None:

unique_toolkit-0.8.0/unique_toolkit/tools/tool_definitions.py

@@ -0,0 +1,145 @@
+from abc import ABC, abstractmethod
+from enum import StrEnum
+from typing import Any, Generic, List, Self, TypeVar
+from unique_toolkit.language_model import LanguageModelToolDescription
+from typing_extensions import deprecated
+# import baseModel from pedantic
+from unique_toolkit.language_model import LanguageModelFunction
+from pydantic import BaseModel, Field, model_validator, root_validator
+
+from unique_toolkit.unique_toolkit.app.schemas import ChatEvent
+from unique_toolkit.unique_toolkit.tools.tool_progress_reporter import ToolProgressReporter
+
+class ToolSelectionPolicy(StrEnum):
+    """Determine the usage policy of tools."""
+
+    FORCED_BY_DEFAULT = "ForcedByDefault"
+    ON_BY_DEFAULT = "OnByDefault"
+    BY_USER = "ByUser"
+
+class UqToolName(StrEnum):
+    WEB_SEARCH = "WebSearch"
+    INTERNAL_SEARCH = "InternalSearch"
+    DOCUMENT_SUMMARIZER = "DocumentSummarizer"
+    CHART_GENERATOR = "ChartGenerator"
+    DOCUMENT_GENERATOR = "DocumentGenerator"
+    DOCUMENT_PARSER = "DocumentParser"
+    IMAGE_CONTENT = "ImageContent"
+    TABLE_SEARCH = "TableSearch"
+    BAR_CHART = "BarChart"
+    LINE_CHART = "LineChart"
+    PIE_CHART = "PieChart"
+    BASE_TOOL = "BaseTool"
+
+
+class BaseToolConfig(BaseModel):
+    pass
+    
+
+ConfigType = TypeVar("ConfigType", bound=BaseToolConfig)
+
+
+class ToolSettings(Generic[ConfigType]):
+    configuration: ConfigType
+    display_name: str
+    icon: str
+    selection_policy: ToolSelectionPolicy = Field(
+        default=ToolSelectionPolicy.BY_USER,
+    )
+    is_exclusive: bool = Field(default=False)
+    is_enabled: bool = Field(default=True)
+
+    @classmethod
+    def from_service_dict(cls, service_dict: dict[str, Any]) -> Self | None:
+        try:
+            return cls(**service_dict)
+        except (ValueError, TypeError) as e:
+            print(e)
+            return None
+
+
+class ToolCallResponse(BaseModel):
+    id: str
+    name: str
+    debug_info: dict = {}
+
+
+class ToolPromptInstructions(BaseModel):
+    system_prompt: str = Field(
+        default="",
+        description=("Helps the LLM understand how to use the tool. "
+                     "This is injected into the system prompt."
+                     "This might not be needed for every tool but some of the work better with user prompt "
+                     "instructions while others work better with system prompt instructions."),
+    )
+
+    user_prompt: str = Field(
+        default="",
+        description=("Helps the LLM understand how to use the tool. "
+                     "This is injected into the user prompt. " 
+                     "This might not be needed for every tool but some of the work better with user prompt "
+                     "instructions while others work better with system prompt instructions.")
+    )
+
+    system_prompt_tool_chosen: str = Field(
+        default="",
+        description=("Once the tool is chosen, this is injected into the system prompt"
+                     " to help the LLM understand how work with the tools results."),
+    )
+
+    user_prompt_tool_chosen: str = Field(
+        default="",
+        description=("Once the tool is chosen, this is injected into the user prompt " 
+                     "to help the LLM understand how to work with the tools results."),
+    )
+
+
+class Tool(ABC, Generic[ConfigType]):
+    name: str
+
+    def tool_description(self) -> LanguageModelToolDescription:
+        raise NotImplementedError
+    
+    
+    def get_prompt_instructions(self) -> ToolPromptInstructions:
+        return ToolPromptInstructions(
+            system_prompt="",
+            user_prompt="",
+            system_prompt_tool_chosen="",
+            user_prompt_tool_chosen="",
+        )
+
+
+    def is_exclusive(self) -> bool:
+        return self.settings.is_exclusive
+    
+    def is_enabled(self) -> bool:
+        return self.settings.is_enabled
+    
+    def display_name(self) -> str:
+        return self.settings.display_name
+    
+    def icon(self) -> str:
+        return self.settings.icon
+    
+    def tool_selection_policy(self) -> ToolSelectionPolicy:
+        return self.settings.selection_policy
+
+
+    @abstractmethod
+    async def run(self, tool_call: LanguageModelFunction) -> ToolCallResponse:
+        raise NotImplementedError
+
+   
+  
+    def __init__(
+        self,
+        settings: ToolSettings[ConfigType],
+        event: ChatEvent,
+        tool_progress_reporter: ToolProgressReporter
+    ):
+        self.settings = settings
+        self.tool_progress_reporter = tool_progress_reporter
+        self.event = event
+
+

unique_toolkit-0.8.0/unique_toolkit/tools/tool_definitionsV2.py

@@ -0,0 +1,137 @@
+from abc import ABC, abstractmethod
+from enum import StrEnum
+from typing import Any, Generic, List, Self, TypeVar
+from unique_toolkit.language_model import LanguageModelToolDescription
+from typing_extensions import deprecated
+# import baseModel from pedantic
+from unique_toolkit.language_model import LanguageModelFunction
+from pydantic import BaseModel, Field, model_validator, root_validator
+
+from unique_toolkit.unique_toolkit.app.schemas import ChatEvent
+from unique_toolkit.unique_toolkit.tools.tool_definitions import BaseToolConfig, Tool, ToolCallResponse, ToolPromptInstructions, ToolSettings
+from unique_toolkit.unique_toolkit.tools.tool_progress_reporter import ToolProgressReporter
+
+
+class BaseToolConfigV2(BaseToolConfig):
+    class ToolCallConfig(BaseModel):
+        description: str = Field(
+            default="Base",  
+            description="The tool description must be set by subclasses",
+        )
+
+        parameters: type[BaseModel] = Field(
+            default=BaseModel,
+            description="The tool parameters configuration must be set by subclasses",
+        )
+
+    class PromptInstructionsConfig(BaseModel):
+        system_prompt: str = Field(
+        default="",
+        description=("Helps the LLM understand how to use the tool. "
+                     "This is injected into the system prompt."
+                     "This might not be needed for every tool but some of the work better with user prompt "
+                     "instructions while others work better with system prompt instructions."),
+        )
+
+        user_prompt: str = Field(
+            default="",
+            description=("Helps the LLM understand how to use the tool. "
+                        "This is injected into the user prompt. " 
+                        "This might not be needed for every tool but some of the work better with user prompt "
+                        "instructions while others work better with system prompt instructions.")
+        )
+
+        system_prompt_tool_chosen: str = Field(
+            default="",
+            description=("Once the tool is chosen, this is injected into the system prompt"
+                        " to help the LLM understand how work with the tools results."),
+        )
+
+        user_prompt_tool_chosen: str = Field(
+            default="",
+            description=("Once the tool is chosen, this is injected into the user prompt " 
+                        "to help the LLM understand how to work with the tools results."),
+        )
+
+    tool_call: ToolCallConfig = Field(
+        default_factory=ToolCallConfig,
+        description="Configuration for the tool, including description and parameters",
+    )
+
+    prompts: PromptInstructionsConfig = Field(
+        default_factory=PromptInstructionsConfig,
+        description="Configuration for prompts related to the tool",
+    )
+
+    # This makes sure that the settings are all present in all subclasses and that they define a default value. 
+    @model_validator(mode="after")
+    def validate_tool_description(cls):
+        if cls.__class__ is BaseToolConfig:
+            return cls  # Skip validation for the base class
+        if cls.tool_call.description == "Base":
+            raise ValueError(
+                f"Subclass {cls.__class__.__name__} must define a default value for 'tool_description'."
+            )
+        if cls.tool_call.parameters == BaseModel:
+            raise ValueError(
+                f"Subclass {cls.__class__.__name__} must define a default value for 'tool_parameters_config'."
+            )
+        if cls.prompts.system_prompt == "":
+            raise ValueError(
+                f"Subclass {cls.__class__.__name__} must define a default value for 'system_prompt_base_instructions'."
+            )
+        if cls.prompts.user_prompt == "":
+            raise ValueError(
+                f"Subclass {cls.__class__.__name__} must define a default value for 'user_prompt_base_instructions'."
+            )
+        if cls.prompts.system_prompt_tool_chosen == "":
+            raise ValueError(
+                f"Subclass {cls.__class__.__name__} must define a default value for 'system_prompt_tool_chosen_instructions'."
+            )
+        if cls.prompts.user_prompt_tool_chosen == "":
+            raise ValueError(
+                f"Subclass {cls.__class__.__name__} must define a default value for 'user_prompt_tool_chosen_instructions'."
+            )
+        return cls
+    
+
+ConfigTypeV2 = TypeVar("ConfigTypeV2", bound=BaseToolConfigV2)
+
+
+
+class ToolV2(Tool[ConfigTypeV2]):
+    name: str
+
+    def tool_description(self) -> LanguageModelToolDescription:
+         return LanguageModelToolDescription(
+            name=self.name,
+            description=self.settings.configuration.tool_call.description,
+            parameters=self.settings.configuration.tool_call.parameters,
+        )
+
+    def get_prompt_instructions(self) -> ToolPromptInstructions:
+        return ToolPromptInstructions(
+            system_prompt=self.settings.configuration.prompts.system_prompt,
+            user_prompt=self.settings.configuration.prompts.user_prompt,
+            system_prompt_tool_chosen=self.settings.configuration.prompts.system_prompt_tool_chosen,
+            user_prompt_tool_chosen=self.settings.configuration.prompts.user_prompt_tool_chosen
+        )
+
+
+    @abstractmethod
+    async def run(self, tool_call: LanguageModelFunction) -> ToolCallResponse:
+        raise NotImplementedError
+
+   
+  
+    def __init__(
+        self,
+        settings: ToolSettings[ConfigTypeV2],
+        event: ChatEvent,
+        tool_progress_reporter: ToolProgressReporter
+    ):
+        self.settings = settings
+        self.tool_progress_reporter = tool_progress_reporter
+        self.event = event
+
+

unique_toolkit-0.8.0/unique_toolkit/tools/tool_factory.py

@@ -0,0 +1,31 @@
+from typing import Callable
+
+from unique_toolkit.unique_toolkit.tools.tool_definitions import BaseToolConfig, Tool
+
+
+
+class ToolFactory:
+    tool_map: dict[str, type[Tool]] = {}
+    tool_config_map: dict[str, Callable] = {}
+
+    @classmethod
+    def register_tool(
+        cls,
+        tool: type[Tool],
+        tool_config: type[BaseToolConfig],
+    ):
+        cls.tool_map[tool.name] = tool
+        cls.tool_config_map[tool.name] = tool_config
+
+    @classmethod
+    def build_tool(cls, tool_name: str, *args, **kwargs) -> Tool:
+        tool = cls.tool_map[tool_name](*args, **kwargs)
+        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-0.8.0/unique_toolkit/tools/tool_progress_reporter.py

@@ -0,0 +1,225 @@
+import re
+from datetime import datetime
+from enum import Enum
+from functools import wraps
+from typing import Protocol
+
+from pydantic import BaseModel
+from unique_toolkit.chat.service import ChatService
+from unique_toolkit.content.schemas import ContentReference
+from unique_toolkit.language_model.schemas import (
+    LanguageModelFunction,
+    LanguageModelStreamResponse,
+)
+
+ARROW = "→ "
+DUMMY_REFERENCE_PLACEHOLDER = "<sup></sup>"
+
+
+class ProgressState(Enum):
+    STARTED = "⚪"
+    RUNNING = "🟡"
+    FAILED = "🔴"
+    FINISHED = "🟢"
+
+
+class ToolExecutionStatus(BaseModel):
+    name: str
+    message: str
+    state: ProgressState
+    references: list[ContentReference] = []
+    timestamp: datetime = datetime.now()
+
+
+class ToolProgressReporter:
+    def __init__(self, chat_service: ChatService):
+        self.chat_service = chat_service
+        self.tool_statuses: dict[str, ToolExecutionStatus] = {}
+        self._progress_start_text = ""
+        self._requires_new_assistant_message = False
+
+    @property
+    def requires_new_assistant_message(self):
+        return self._requires_new_assistant_message
+
+    @requires_new_assistant_message.setter
+    def requires_new_assistant_message(self, value: bool):
+        self._requires_new_assistant_message = value
+
+    @property
+    def tool_statuses_is_empty(self):
+        return len(self.tool_statuses) == 0
+
+    def empty_tool_statuses_if_stream_has_text(
+        self, stream_response: LanguageModelStreamResponse
+    ):
+        if stream_response.message.text:
+            self.tool_statuses = {}
+
+    async def notify_from_tool_call(
+        self,
+        tool_call: LanguageModelFunction,
+        name: str,
+        message: str,
+        state: ProgressState,
+        references: list[ContentReference] = [],
+        requires_new_assistant_message: bool = False,
+    ):
+        """
+        Notifies about a tool call execution status and updates the assistant message.
+
+        Args:
+            tool_call (LanguageModelFunction): The tool call being executed
+            name (str): Name of the tool being executed
+            message (str): Status message to display
+            state (ProgressState): Current execution state of the tool
+            references (list[ContentReference], optional): List of content references. Defaults to [].
+            requires_new_assistant_message (bool, optional): Whether a new assistant message is needed when tool call is finished.
+            Defaults to False. If yes, the agentic steps will remain in chat history and will be overwritten by the stream response.
+
+        Raises:
+            AssertionError: If tool_call.id is None
+        """
+        assert tool_call.id is not None
+        self.tool_statuses[tool_call.id] = ToolExecutionStatus(
+            name=name,
+            message=message,
+            state=state,
+            references=references,
+        )
+        self.requires_new_assistant_message = (
+            self.requires_new_assistant_message
+            or requires_new_assistant_message
+        )
+        await self.publish()
+
+    async def publish(self):
+        messages = []
+        all_references = []
+        for item in sorted(
+            self.tool_statuses.values(), key=lambda x: x.timestamp
+        ):
+            references = item.references
+            start_number = len(all_references) + 1
+            message = self._replace_placeholders(item.message, start_number)
+            references = self._correct_reference_sequence(
+                references, start_number
+            )
+            all_references.extend(references)
+
+            messages.append(
+                f"{ARROW}**{item.name} {item.state.value}**: {message}"
+            )
+
+        await self.chat_service.modify_assistant_message_async(
+            content=self._progress_start_text + "\n\n" + "\n\n".join(messages),
+            references=all_references,
+        )
+
+    @staticmethod
+    def _replace_placeholders(message: str, start_number: int = 1) -> str:
+        counter = start_number
+
+        def replace_match(match):
+            nonlocal counter
+            result = f"<sup>{counter}</sup>"
+            counter += 1
+            return result
+
+        return re.sub(r"<sup></sup>", replace_match, message)
+
+    @staticmethod
+    def _correct_reference_sequence(
+        references: list[ContentReference], start_number: int = 1
+    ) -> list[ContentReference]:
+        for i, reference in enumerate(references, start_number):
+            reference.sequence_number = i
+        return references
+
+
+class ToolWithToolProgressReporter(Protocol):
+    tool_progress_reporter: ToolProgressReporter
+
+
+def track_tool_progress(
+    message: str,
+    on_start_state: ProgressState = ProgressState.RUNNING,
+    on_success_state: ProgressState = ProgressState.RUNNING,
+    on_success_message: str | None = None,
+    on_error_message: str = "Unexpected error occurred",
+    requires_new_assistant_message: bool = False,
+):
+    """
+    Decorator to add progress reporting and status tracking steps to tool functions. Can be used with async and sync functions.
+
+    Args:
+        name (str): Display name for the tool progress status
+        message (str): Message to show during tool execution
+        on_error_message (str, optional): Message to show if tool execution fails. Defaults to empty string.
+        on_success_state (ProgressState, optional): State to set after successful execution. Defaults to RUNNING.
+        requires_new_assistant_message (bool, optional): Whether to create a new assistant message. Defaults to False.
+
+    The decorator will:
+    1. Show a RUNNING status when the tool starts executing
+    2. Update the status to on_success_state if execution succeeds
+    3. Update the status to FAILED if execution fails
+    4. Include any references from the tool result in the status update if the result has a 'references' attribute or item.
+    5. Create a new assistant message if requires_new_assistant_message is True
+
+    The decorated function must be a method of a class that implements ToolWithToolProgressReporter.
+    """
+
+    def decorator(func):
+        @wraps(func)  # Preserve the original function's metadata
+        async def async_wrapper(
+            self: ToolWithToolProgressReporter,
+            tool_call: LanguageModelFunction,
+            notification_tool_name: str,
+            *args,
+            **kwargs,
+        ):
+            try:
+                # Start status
+                await self.tool_progress_reporter.notify_from_tool_call(
+                    tool_call=tool_call,
+                    name=notification_tool_name,
+                    message=message,
+                    state=on_start_state,
+                )
+
+                # Execute the tool function
+                result = await func(
+                    self, tool_call, notification_tool_name, *args, **kwargs
+                )
+
+                # Success status
+                await self.tool_progress_reporter.notify_from_tool_call(
+                    tool_call=tool_call,
+                    name=notification_tool_name,
+                    message=on_success_message or message,
+                    state=on_success_state,
+                    references=_get_references_from_results(result),
+                    requires_new_assistant_message=requires_new_assistant_message,
+                )
+                return result
+
+            except Exception as e:
+                # Failure status
+                await self.tool_progress_reporter.notify_from_tool_call(
+                    tool_call=tool_call,
+                    name=notification_tool_name,
+                    message=on_error_message,
+                    state=ProgressState.FAILED,
+                    requires_new_assistant_message=requires_new_assistant_message,
+                )
+                raise e
+
+        return async_wrapper
+
+    return decorator
+
+
+def _get_references_from_results(result):
+    if isinstance(result, dict):
+        return result.get("references", [])
+    return getattr(result, "references", [])