unique_toolkit 1.28.8__py3-none-any.whl → 1.33.3__py3-none-any.whl

unique_toolkit/agentic/tools/a2a/tool/config.py

@@ -1,10 +1,74 @@
-from typing import Literal
+import re
+from enum import StrEnum
+from typing import Annotated, Generic, Literal, TypeVar
 
 from pydantic import Field
+from pydantic.main import BaseModel
 
 from unique_toolkit._common.pydantic_helpers import get_configuration_dict
 from unique_toolkit.agentic.tools.schemas import BaseToolConfig
 
+
+class SubAgentSystemReminderType(StrEnum):
+    FIXED = "fixed"
+    REGEXP = "regexp"
+    REFERENCE = "reference"
+
+
+T = TypeVar("T", bound=SubAgentSystemReminderType)
+
+
+class SystemReminderConfig(BaseModel, Generic[T]):
+    model_config = get_configuration_dict()
+
+    type: T
+
+
+_SYSTEM_REMINDER_FIELD_DESCRIPTION = """
+The reminder to add to the tool response. The reminder can be a Jinja template and can contain the following placeholders:
+- {{ display_name }}: The display name of the sub agent.
+- {{ tool_name }}: The tool name. 
+""".strip()
+
+
+class ReferenceSystemReminderConfig(SystemReminderConfig):
+    type: Literal[SubAgentSystemReminderType.REFERENCE] = (
+        SubAgentSystemReminderType.REFERENCE
+    )
+    reminder: str = Field(
+        default="Rememeber to properly reference EACH fact from sub agent {{ display_name }}'s response with the correct format INLINE.",
+        description=_SYSTEM_REMINDER_FIELD_DESCRIPTION,
+    )
+
+
+class FixedSystemReminderConfig(SystemReminderConfig):
+    type: Literal[SubAgentSystemReminderType.FIXED] = SubAgentSystemReminderType.FIXED
+    reminder: str = Field(
+        description=_SYSTEM_REMINDER_FIELD_DESCRIPTION,
+    )
+
+
+_REGEXP_DETECTED_REMINDER_FIELD_DESCRIPTION = """
+The reminder to add to the tool response. The reminder can be a Jinja template and can contain the following placeholders:
+- {{ display_name }}: The display name of the sub agent.
+- {{ tool_name }}: The tool name. 
+- {{ text_matches }}: Will be replaced with the portions of the text that triggered the reminder.
+""".strip()
+
+
+class RegExpDetectedSystemReminderConfig(SystemReminderConfig):
+    """A system reminder that is only added if the sub agent response matches a regular expression."""
+
+    type: Literal[SubAgentSystemReminderType.REGEXP] = SubAgentSystemReminderType.REGEXP
+
+    regexp: re.Pattern[str] = Field(
+        description="The regular expression to use to detect whether the system reminder should be added.",
+    )
+    reminder: str = Field(
+        description=_REGEXP_DETECTED_REMINDER_FIELD_DESCRIPTION,
+    )
+
+
 DEFAULT_PARAM_DESCRIPTION_SUB_AGENT_USER_MESSAGE = """
 This is the message that will be sent to the sub-agent.
 """.strip()
@@ -80,3 +144,15 @@ class SubAgentToolConfig(BaseToolConfig):
         default=False,
         description="If set, the sub-agent response will be interpreted as a list of content chunks.",
     )
+
+    system_reminders_config: list[
+        Annotated[
+            FixedSystemReminderConfig
+            | RegExpDetectedSystemReminderConfig
+            | ReferenceSystemReminderConfig,
+            Field(discriminator="type"),
+        ]
+    ] = Field(
+        default=[],
+        description="Configuration for the system reminders to add to the tool response.",
+    )

unique_toolkit/agentic/tools/a2a/tool/service.py

@@ -4,7 +4,7 @@ import json
 import logging
 import re
 from datetime import datetime
-from typing import override
+from typing import cast, override
 
 import unique_sdk
 from pydantic import Field, TypeAdapter, create_model
@@ -15,6 +15,7 @@ from unique_toolkit._common.referencing import (
     remove_all_refs,
     replace_ref_number,
 )
+from unique_toolkit._common.utils.jinja.render import render_template
 from unique_toolkit.agentic.evaluation.schemas import EvaluationMetricName
 from unique_toolkit.agentic.tools.a2a.response_watcher import SubAgentResponseWatcher
 from unique_toolkit.agentic.tools.a2a.tool._memory import (
@@ -25,6 +26,8 @@ from unique_toolkit.agentic.tools.a2a.tool._schema import (
     SubAgentToolInput,
 )
 from unique_toolkit.agentic.tools.a2a.tool.config import (
+    RegExpDetectedSystemReminderConfig,
+    SubAgentSystemReminderType,
     SubAgentToolConfig,
 )
 from unique_toolkit.agentic.tools.factory import ToolFactory
@@ -215,12 +218,63 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
             )
 
             return ToolCallResponse(
-                id=tool_call.id,  # type: ignore
+                id=tool_call.id,
                 name=tool_call.name,
-                content=content,
+                content=_format_response(
+                    tool_name=self.name,
+                    text=content,
+                    system_reminders=self._get_system_reminders(response),
+                ),
                 content_chunks=content_chunks,
             )
 
+    def _get_system_reminders(self, message: unique_sdk.Space.Message) -> list[str]:
+        reminders = []
+        for reminder_config in self.config.system_reminders_config:
+            if reminder_config.type == SubAgentSystemReminderType.FIXED:
+                reminders.append(
+                    render_template(
+                        reminder_config.reminder,
+                        display_name=self.display_name(),
+                        tool_name=self.name,
+                    )
+                )
+            elif (
+                reminder_config.type == SubAgentSystemReminderType.REFERENCE
+                and self.config.use_sub_agent_references
+                and message["references"] is not None
+                and len(message["references"]) > 0
+            ):
+                reminders.append(
+                    render_template(
+                        reminder_config.reminder,
+                        display_name=self.display_name(),
+                        tool_name=self.name,
+                    )
+                )
+            elif (
+                reminder_config.type == SubAgentSystemReminderType.REGEXP
+                and message["text"] is not None
+            ):
+                reminder_config = cast(
+                    RegExpDetectedSystemReminderConfig, reminder_config
+                )
+                text_matches = [
+                    match.group(0)
+                    for match in reminder_config.regexp.finditer(message["text"])
+                ]
+                if len(text_matches) > 0:
+                    reminders.append(
+                        render_template(
+                            reminder_config.reminder,
+                            display_name=self.display_name(),
+                            tool_name=self.name,
+                            text_matches=text_matches,
+                        )
+                    )
+
+        return reminders
+
     async def _get_chat_id(self) -> str | None:
         if not self.config.reuse_chat:
             return None
@@ -326,4 +380,14 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
             ) from e
 
 
+def _format_response(tool_name: str, text: str, system_reminders: list[str]) -> str:
+    if len(system_reminders) == 0:
+        return text
+
+    reponse_key = f"{tool_name} response"
+    response = {reponse_key: text, "SYSTEM_REMINDERS": system_reminders}
+
+    return json.dumps(response, indent=2)
+
+
 ToolFactory.register_tool(SubAgentTool, SubAgentToolConfig)

unique_toolkit/agentic/tools/config.py

@@ -1,12 +1,12 @@
-import json
 from enum import StrEnum
-from typing import Annotated, Any, Dict
+from typing import Annotated, Any
 
 from pydantic import (
     BaseModel,
     BeforeValidator,
     Field,
     ValidationInfo,
+    field_serializer,
     model_validator,
 )
 
@@ -123,46 +123,6 @@ class ToolBuildConfig(BaseModel):
         value["configuration"] = config
         return value
 
-    def model_dump(self) -> Dict[str, Any]:
-        """
-        Returns a dict representation of the tool config that preserves
-        subclass fields from `configuration` by delegating to its own
-        model_dump. This prevents `{}` when `configuration` is typed
-        as `BaseToolConfig` but holds a subclass instance.
-        """
-        data: Dict[str, Any] = {
-            "name": self.name,
-            "configuration": self.configuration.model_dump()
-            if self.configuration
-            else None,
-            "display_name": self.display_name,
-            "icon": self.icon,
-            "selection_policy": self.selection_policy,
-            "is_exclusive": self.is_exclusive,
-            "is_sub_agent": self.is_sub_agent,
-            "is_enabled": self.is_enabled,
-        }
-        return data
-
-    def model_dump_json(self) -> str:
-        """
-        Returns a JSON string representation of the tool config.
-        Ensures `configuration` is fully serialized by using the
-        subclass's `model_dump_json()` when available.
-        """
-        config_json = (
-            self.configuration.model_dump_json() if self.configuration else None
-        )
-        config = json.loads(config_json) if config_json else None
-
-        data: Dict[str, Any] = {
-            "name": self.name,
-            "configuration": config,
-            "display_name": self.display_name,
-            "icon": self.icon,
-            "selection_policy": self.selection_policy,
-            "is_exclusive": self.is_exclusive,
-            "is_sub_agent": self.is_sub_agent,
-            "is_enabled": self.is_enabled,
-        }
-        return json.dumps(data)
+    @field_serializer("configuration")
+    def serialize_config(self, value: BaseToolConfig) -> dict[str, Any]:
+        return value.__class__.model_dump(value)

unique_toolkit/agentic/tools/openai_builtin/base.py

@@ -40,3 +40,7 @@ class OpenAIBuiltInTool(ABC, Generic[ToolType]):
     @abstractmethod
     def takes_control(self) -> bool:
         raise NotImplementedError()
+
+    @abstractmethod
+    def display_name(self) -> str:
+        raise NotImplementedError()

unique_toolkit/agentic/tools/openai_builtin/code_interpreter/service.py

@@ -244,3 +244,7 @@ class OpenAICodeInterpreterTool(OpenAIBuiltInTool[CodeInterpreter]):
             tool_format_information_for_user_prompt=self._config.tool_format_information_for_user_prompt,
             input_model={},
         )
+
+    @override
+    def display_name(self) -> str:
+        return self.DISPLAY_NAME

unique_toolkit/agentic/tools/tool_manager.py

@@ -231,25 +231,6 @@ class _ToolManager(Generic[_ApiMode]):
         self,
         tool_calls: list[LanguageModelFunction],
     ) -> list[ToolCallResponse]:
-        tool_calls = tool_calls
-
-        tool_calls = self.filter_duplicate_tool_calls(
-            tool_calls=tool_calls,
-        )
-        num_tool_calls = len(tool_calls)
-
-        if num_tool_calls > self._config.max_tool_calls:
-            self._logger.warning(
-                (
-                    "Number of tool calls %s exceeds the allowed maximum of %s."
-                    "The tool calls will be reduced to the first %s."
-                ),
-                num_tool_calls,
-                self._config.max_tool_calls,
-                self._config.max_tool_calls,
-            )
-            tool_calls = tool_calls[: self._config.max_tool_calls]
-
         tool_call_responses = await self._execute_parallelized(tool_calls)
         return tool_call_responses
 
@@ -358,6 +339,22 @@ class _ToolManager(Generic[_ApiMode]):
             )
         return unique_tool_calls
 
+    def filter_tool_calls_by_max_tool_calls_allowed(
+        self, tool_calls: list[LanguageModelFunction]
+    ) -> list[LanguageModelFunction]:
+        if len(tool_calls) > self._config.max_tool_calls:
+            self._logger.warning(
+                (
+                    "Number of tool calls %s exceeds the allowed maximum of %s."
+                    "The tool calls will be reduced to the first %s."
+                ),
+                len(tool_calls),
+                self._config.max_tool_calls,
+                self._config.max_tool_calls,
+            )
+            return tool_calls[: self._config.max_tool_calls]
+        return tool_calls
+
     @overload
     def get_tool_by_name(
         self: "_ToolManager[Literal['completions']]", name: str

unique_toolkit/app/__init__.py

@@ -47,5 +47,8 @@ from .schemas import (
 from .verification import (
     verify_signature_and_construct_event as verify_signature_and_construct_event,
 )
+from .webhook import (
+    is_webhook_signature_valid as is_webhook_signature_valid,
+)
 
 DOMAIN_NAME = "app"

unique_toolkit/app/fast_api_factory.py

@@ -0,0 +1,131 @@
+import json
+from logging import getLogger
+from typing import TYPE_CHECKING, Any, Callable, TypeVar
+
+from pydantic import ValidationError
+
+if TYPE_CHECKING:
+    from fastapi import BackgroundTasks, FastAPI, Request, status
+    from fastapi.responses import JSONResponse
+else:
+    try:
+        from fastapi import BackgroundTasks, FastAPI, Request, status
+        from fastapi.responses import JSONResponse
+    except ImportError:
+        FastAPI = None  # type: ignore[assignment, misc]
+        Request = None  # type: ignore[assignment, misc]
+        status = None  # type: ignore[assignment, misc]
+        JSONResponse = None  # type: ignore[assignment, misc]
+        BackgroundTasks = None  # type: ignore[assignment, misc]
+
+from unique_toolkit.app.schemas import BaseEvent, ChatEvent, EventName
+from unique_toolkit.app.unique_settings import UniqueSettings
+
+logger = getLogger(__name__)
+
+
+def default_event_handler(event: Any) -> int:
+    logger.info("Event received at event handler")
+    if status is not None:
+        return status.HTTP_200_OK
+    else:
+        # No fastapi installed
+        return 200
+
+
+T = TypeVar("T", bound=BaseEvent)
+
+
+def build_unique_custom_app(
+    *,
+    title: str = "Unique Chat App",
+    webhook_path: str = "/webhook",
+    settings: UniqueSettings,
+    event_handler: Callable[[T], int] = default_event_handler,
+    event_constructor: Callable[..., T] = ChatEvent,
+    subscribed_event_names: list[str] | None = None,
+) -> "FastAPI":
+    """Factory class for creating FastAPI apps with Unique webhook handling."""
+    if FastAPI is None:
+        raise ImportError(
+            "FastAPI is not installed. Install it with: poetry install --with fastapi"
+        )
+
+    app = FastAPI(title=title)
+
+    if subscribed_event_names is None:
+        subscribed_event_names = [EventName.EXTERNAL_MODULE_CHOSEN]
+
+    @app.get(path="/")
+    async def health_check() -> JSONResponse:
+        """Health check endpoint."""
+        return JSONResponse(content={"status": "healthy", "service": title})
+
+    @app.post(path=webhook_path)
+    async def webhook_handler(
+        request: Request, background_tasks: BackgroundTasks
+    ) -> JSONResponse:
+        """
+        Webhook endpoint for receiving events from Unique platform.
+
+        This endpoint:
+        1. Verifies the webhook signature
+        2. Constructs an event from the payload
+        3. Calls the configured event handler
+        """
+        # Get raw body and headers
+        body = await request.body()
+        headers = dict(request.headers)
+
+        from unique_toolkit.app.webhook import is_webhook_signature_valid
+
+        if not is_webhook_signature_valid(
+            headers=headers,
+            payload=body,
+            endpoint_secret=settings.app.endpoint_secret.get_secret_value(),
+        ):
+            return JSONResponse(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                content={"error": "Invalid webhook signature"},
+            )
+
+        try:
+            event_data = json.loads(body.decode(encoding="utf-8"))
+        except json.JSONDecodeError as e:
+            logger.error(f"Error parsing event: {e}", exc_info=True)
+            return JSONResponse(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                content={"error": f"Invalid event format: {str(e)}"},
+            )
+
+        if event_data["event"] not in subscribed_event_names:
+            return JSONResponse(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                content={"error": "Not subscribed event"},
+            )
+
+        try:
+            event = event_constructor(**event_data)
+            if event.filter_event(filter_options=settings.chat_event_filter_options):
+                return JSONResponse(
+                    status_code=status.HTTP_200_OK,
+                    content={"error": "Event filtered out"},
+                )
+        except ValidationError as e:
+            # pydantic errors https://docs.pydantic.dev/2.10/errors/errors/
+            logger.error(f"Validation error with model: {e.json()}", exc_info=True)
+            raise e
+        except ValueError as e:
+            logger.error(f"Error deserializing event: {e}", exc_info=True)
+            return JSONResponse(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                content={"error": "Invalid event"},
+            )
+
+        # Run the task in background so that we don't block for long running tasks
+        background_tasks.add_task(event_handler, event)
+        return JSONResponse(
+            status_code=status.HTTP_200_OK, content={"message": "Event received"}
+        )
+
+    return app

unique_toolkit/app/webhook.py

@@ -0,0 +1,77 @@
+"""
+Webhook signature verification for Unique platform.
+
+Extracted from unique_sdk to provide standalone verification without event construction.
+"""
+
+import hashlib
+import hmac
+import logging
+import time
+
+_LOGGER = logging.getLogger(__name__)
+
+
+def is_webhook_signature_valid(
+    headers: dict[str, str],
+    payload: bytes,
+    endpoint_secret: str,
+    tolerance: int = 300,
+) -> bool:
+    """
+    Verify webhook signature from Unique platform.
+
+    Args:
+        headers: Request headers with X-Unique-Signature and X-Unique-Created-At
+        payload: Raw request body bytes
+        endpoint_secret: App endpoint secret from Unique platform
+        tolerance: Max seconds between timestamp and now (default: 300)
+
+    Returns:
+        True if signature is valid, False otherwise
+    """
+    # Extract headers
+    signature = headers.get("X-Unique-Signature") or headers.get("x-unique-signature")
+    timestamp_str = headers.get("X-Unique-Created-At") or headers.get(
+        "x-unique-created-at"
+    )
+
+    if not signature:
+        _LOGGER.error("Missing X-Unique-Signature header")
+        return False
+
+    if not timestamp_str:
+        _LOGGER.error("Missing X-Unique-Created-At header")
+        return False
+
+    # Convert timestamp to int
+    try:
+        timestamp = int(timestamp_str)
+    except ValueError:
+        _LOGGER.error(f"Invalid timestamp: {timestamp_str}")
+        return False
+
+    # Decode payload if bytes
+    message = payload.decode("utf-8") if isinstance(payload, bytes) else payload
+
+    # Compute expected signature: HMAC-SHA256(message, secret)
+    expected_signature = hmac.new(
+        endpoint_secret.encode("utf-8"),
+        msg=message.encode("utf-8"),
+        digestmod=hashlib.sha256,
+    ).hexdigest()
+
+    # Compare signatures (constant-time to prevent timing attacks)
+    if not hmac.compare_digest(expected_signature, signature):
+        _LOGGER.error("Signature mismatch. Ensure you're using the raw request body.")
+        return False
+
+    # Check timestamp tolerance (prevent replay attacks)
+    if tolerance and timestamp < time.time() - tolerance:
+        _LOGGER.error(
+            f"Timestamp outside tolerance ({tolerance}s). Possible replay attack."
+        )
+        return False
+
+    _LOGGER.debug("✅ Webhook signature verified successfully")
+    return True

unique_toolkit/chat/functions.py

@@ -263,7 +263,7 @@ def create_message(
     references: list[ContentReference] | None = None,
     debug_info: dict | None = None,
     set_completed_at: bool | None = False,
-):
+) -> ChatMessage:
     """Creates a message in the chat session synchronously.
 
     Args:

unique_toolkit/content/functions.py

@@ -144,7 +144,7 @@ def search_contents(
     chat_id: str,
     where: dict,
     include_failed_content: bool = False,
-):
+) -> list[Content]:
     """
     Performs an asynchronous search for content files in the knowledge base by filter.
 
@@ -297,7 +297,7 @@ def upload_content_from_bytes(
     skip_ingestion: bool = False,
     ingestion_config: unique_sdk.Content.IngestionConfig | None = None,
     metadata: dict[str, Any] | None = None,
-):
+) -> Content:
     """
     Uploads content to the knowledge base.
 
@@ -347,7 +347,7 @@ def upload_content(
     skip_excel_ingestion: bool = False,
     ingestion_config: unique_sdk.Content.IngestionConfig | None = None,
     metadata: dict[str, Any] | None = None,
-):
+) -> Content:
     """
     Uploads content to the knowledge base.
 
@@ -399,7 +399,7 @@ def _trigger_upload_content(
     skip_excel_ingestion: bool = False,
     ingestion_config: unique_sdk.Content.IngestionConfig | None = None,
     metadata: dict[str, Any] | None = None,
-):
+) -> Content:
     """
     Uploads content to the knowledge base.
 

unique_toolkit/content/service.py

@@ -528,7 +528,7 @@ class ContentService:
         skip_excel_ingestion: bool = False,
         ingestion_config: unique_sdk.Content.IngestionConfig | None = None,
         metadata: dict[str, Any] | None = None,
-    ):
+    ) -> Content:
         """
         Uploads content to the knowledge base.