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

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/init_sdk.py

@@ -1,6 +1,11 @@
 import os
+from pathlib import Path
+from typing import overload
 
 import unique_sdk
+from typing_extensions import deprecated
+
+from unique_toolkit.app.unique_settings import UniqueSettings
 
 
 def get_env(var_name, default=None, strict=False):
@@ -24,12 +29,38 @@ def get_env(var_name, default=None, strict=False):
     return val or default
 
 
-def init_sdk(strict_all_vars=False):
+@overload
+def init_unique_sdk(*, env_file: Path | None = None): ...
+
+
+@overload
+def init_unique_sdk(*, unique_settings: UniqueSettings): ...
+
+
+def init_unique_sdk(
+    *, unique_settings: UniqueSettings | None = None, env_file: Path | None = None
+):
+    if unique_settings:
+        unique_sdk.api_key = unique_settings.app.key.get_secret_value()
+        unique_sdk.app_id = unique_settings.app.id.get_secret_value()
+        unique_sdk.api_base = unique_settings.api.sdk_url()
+    elif env_file:
+        unique_settings = UniqueSettings.from_env(env_file=env_file)
+        unique_sdk.api_key = unique_settings.app.key.get_secret_value()
+        unique_sdk.app_id = unique_settings.app.id.get_secret_value()
+        unique_sdk.api_base = unique_settings.api.sdk_url()
+
+
+@deprecated("Use init_unique_sdk instead")
+def init_sdk(
+    strict_all_vars: bool = False,
+):
     """Initialize the SDK.
 
     Args:
         strict_all_vars (bool, optional): This method raises a ValueError if strict and no value is found in the environment. Defaults to False.
     """
+
     unique_sdk.api_key = get_env("API_KEY", default="dummy", strict=strict_all_vars)
     unique_sdk.app_id = get_env("APP_ID", default="dummy", strict=strict_all_vars)
     unique_sdk.api_base = get_env("API_BASE", default=None, strict=strict_all_vars)

unique_toolkit/app/schemas.py

@@ -1,23 +1,43 @@
+import json
 from enum import StrEnum
-from typing import Any, Optional
+from logging import getLogger
+from pathlib import Path
+from typing import Any, Generic, Optional, TypeVar, override
 
 from humps import camelize
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, ConfigDict, Field, JsonValue, field_validator
+from pydantic_settings import BaseSettings
 from typing_extensions import deprecated
 
+from unique_toolkit._common.exception import ConfigurationException
+from unique_toolkit.app.unique_settings import UniqueChatEventFilterOptions
+from unique_toolkit.smart_rules.compile import UniqueQL, parse_uniqueql
+
+FilterOptionsT = TypeVar("FilterOptionsT", bound=BaseSettings)
+
 # set config to convert camelCase to snake_case
 model_config = ConfigDict(
     alias_generator=camelize,
     populate_by_name=True,
     arbitrary_types_allowed=True,
 )
+_logger = getLogger(__name__)
 
 
 class EventName(StrEnum):
     EXTERNAL_MODULE_CHOSEN = "unique.chat.external-module.chosen"
-
-
-class BaseEvent(BaseModel):
+    USER_MESSAGE_CREATED = "unique.chat.user-message.created"
+    INGESTION_CONTENT_UPLOADED = "unique.ingestion.content.uploaded"
+    INGESTION_CONTENT_FINISHED = "unique.ingestion.content.finished"
+    MAGIC_TABLE_IMPORT_COLUMNS = "unique.magic-table.import-columns"
+    MAGIC_TABLE_ADD_META_DATA = "unique.magic-table.add-meta-data"
+    MAGIC_TABLE_ADD_DOCUMENT = "unique.magic-table.add-document"
+    MAGIC_TABLE_DELETE_ROW = "unique.magic-table.delete-row"
+    MAGIC_TABLE_DELETE_COLUMN = "unique.magic-table.delete-column"
+    MAGIC_TABLE_UPDATE_CELL = "unique.magic-table.update-cell"
+
+
+class BaseEvent(BaseModel, Generic[FilterOptionsT]):
     model_config = model_config
 
     id: str
@@ -25,6 +45,72 @@ class BaseEvent(BaseModel):
     user_id: str
     company_id: str
 
+    @classmethod
+    def from_json_file(cls, file_path: Path) -> "BaseEvent":
+        if not file_path.exists():
+            raise FileNotFoundError(f"File not found: {file_path}")
+        with file_path.open("r", encoding="utf-8") as f:
+            data = json.load(f)
+        return cls.model_validate(data)
+
+    def filter_event(self, *, filter_options: FilterOptionsT | None = None) -> bool:
+        """Determine if event should be filtered out and be neglected."""
+        return False
+
+
+###
+# 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.",
+    )
+    tool_format_information: Optional[str] = Field(
+        default=None,
+        description="An optional tool format information. 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
@@ -95,45 +181,134 @@ class ChatEventPayload(BaseModel):
     assistant_id: str
     user_message: ChatEventUserMessage
     assistant_message: ChatEventAssistantMessage
-    text: Optional[str] = None
-    additional_parameters: Optional[ChatEventAdditionalParameters] = None
-    user_metadata: Optional[dict[str, Any]] = None
-    tool_choices: Optional[list[str]] = Field(
-        default=[],
+    text: str | None = None
+    additional_parameters: ChatEventAdditionalParameters | None = None
+    user_metadata: dict[str, Any] | None = Field(
+        default_factory=dict,
+    )
+    tool_choices: list[str] = Field(
+        default_factory=list,
         description="A list containing the tool names the user has chosen to be activated.",
     )
-    tool_parameters: Optional[dict[str, Any]] = None
-    metadata_filter: Optional[dict[str, Any]] = None
+    disabled_tools: list[str] = Field(
+        default_factory=list,
+        description="A list containing the tool names of tools that are disabled at the company level",
+    )
+    tool_parameters: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Parameters extracted from module selection function calling the tool.",
+    )
+    # Default is None as empty dict triggers error in `backend-ingestion`
+    metadata_filter: dict[str, Any] | None = Field(
+        default=None,
+        description="Metadata filter compiled after module selection function calling and scope rules.",
+    )
+    raw_scope_rules: UniqueQL | None = Field(
+        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.",
+    )
+    message_execution_id: str | None = Field(
+        default=None,
+        description="The message execution id for triggering the chat event. Originates from the message execution service.",
+    )
+    session_config: JsonValue | None = Field(
+        default=None,
+        description="The session configuration for the chat session.",
+    )
+
+    @field_validator("raw_scope_rules", mode="before")
+    def validate_scope_rules(cls, value: dict[str, Any] | None) -> UniqueQL | None:
+        if value:
+            return parse_uniqueql(value)
 
 
 @deprecated("""Use `ChatEventPayload` instead.
             This class will be removed in the next major version.""")
 class EventPayload(ChatEventPayload):
-    user_message: EventUserMessage
-    assistant_message: EventAssistantMessage
-    additional_parameters: Optional[EventAdditionalParameters] = None
+    pass
+    # user_message: EventUserMessage
+    # assistant_message: EventAssistantMessage
+    # additional_parameters: Optional[EventAdditionalParameters] = None
 
 
-@deprecated(
-    """Use the more specific `ChatEvent` instead that has the same properties. \
-This class will be removed in the next major version."""
-)
-class Event(BaseModel):
+class ChatEvent(BaseEvent):
     model_config = model_config
 
-    id: str
-    event: EventName
-    user_id: str
-    company_id: str
-    payload: EventPayload
+    payload: ChatEventPayload
     created_at: Optional[int] = None
     version: Optional[str] = None
 
+    @classmethod
+    def from_json_file(cls, file_path: Path) -> "ChatEvent":
+        if not file_path.exists():
+            raise FileNotFoundError(f"File not found: {file_path}")
+        with file_path.open("r", encoding="utf-8") as f:
+            data = json.load(f)
+        return cls.model_validate(data)
+
+    def get_initial_debug_info(self) -> dict[str, Any]:
+        """Get the debug information for the chat event"""
+
+        # TODO: Make sure this coincides with what is shown in the first user message
+        return {
+            "user_metadata": self.payload.user_metadata,
+            "tool_parameters": self.payload.tool_parameters,
+            "chosen_module": self.payload.name,
+            "assistant": {"id": self.payload.assistant_id},
+        }
+
+    @override
+    def filter_event(
+        self, *, filter_options: UniqueChatEventFilterOptions | None = None
+    ) -> bool:
+        # Empty string evals to False
+
+        if filter_options is None:
+            return False  # Don't filter when no options provided
+
+        if not filter_options.assistant_ids and not filter_options.references_in_code:
+            raise ConfigurationException(
+                "No filter options provided, all events will be filtered! \n"
+                "Please define: \n"
+                " - 'UNIQUE_CHAT_EVENT_FILTER_OPTIONS_ASSISTANT_IDS' \n"
+                " - 'UNIQUE_CHAT_EVENT_FILTER_OPTIONS_REFERENCES_IN_CODE' \n"
+                "in your environment variables."
+            )
+
+        # Per reference in code there can be multiple assistants
+        if (
+            filter_options.assistant_ids
+            and self.payload.assistant_id not in filter_options.assistant_ids
+        ):
+            return True
+
+        if (
+            filter_options.references_in_code
+            and self.payload.name not in filter_options.references_in_code
+        ):
+            return True
+
+        return super().filter_event(filter_options=filter_options)
 
-class ChatEvent(BaseEvent):
-    model_config = model_config
 
-    event: EventName
-    payload: ChatEventPayload
-    created_at: Optional[int] = None
-    version: Optional[str] = None
+@deprecated(
+    """Use the more specific `ChatEvent` instead that has the same properties. \
+This class will be removed in the next major version."""
+)
+class Event(ChatEvent):
+    pass
+    # The below should only affect type hints
+    # event: EventName T
+    # payload: EventPayload
+
+    @classmethod
+    def from_json_file(cls, file_path: Path) -> "Event":
+        if not file_path.exists():
+            raise FileNotFoundError(f"File not found: {file_path}")
+        with file_path.open("r", encoding="utf-8") as f:
+            data = json.load(f)
+        return cls.model_validate(data)