unique_toolkit 0.0.2__py3-none-any.whl → 0.5.0__py3-none-any.whl

unique_toolkit/app/init_logging.py

@@ -0,0 +1,31 @@
+from logging import Formatter
+from logging.config import dictConfig
+from time import gmtime
+
+
+class UTCFormatter(Formatter):
+    converter = gmtime
+
+
+unique_log_config = {
+    "version": 1,
+    "root": {"level": "DEBUG", "handlers": ["console"]},
+    "handlers": {
+        "console": {
+            "class": "logging.StreamHandler",
+            "level": "DEBUG",
+            "formatter": "utc",
+        }
+    },
+    "formatters": {
+        "utc": {
+            "()": UTCFormatter,
+            "format": "%(asctime)s: %(message)s",
+            "datefmt": "%Y-%m-%d %H:%M:%S",
+        },
+    },
+}
+
+
+def init_logging(config: dict = unique_log_config):
+    return dictConfig(config)

unique_toolkit/app/init_sdk.py

@@ -0,0 +1,41 @@
+import os
+
+import unique_sdk
+
+
+def get_env(var_name, default=None, strict=False):
+    """Get the environment variable.
+
+    Args:
+        var_name (str): Name of the environment variable.
+        default (str, optional): Default value. Defaults to None.
+        strict (bool, optional): This method raises a ValueError, if strict, and no value is found in the environment. Defaults to False.
+
+    Raises:
+        ValueError: _description_
+
+    Returns:
+        _type_: _description_
+    """
+    val = os.environ.get(var_name)
+    if not val:
+        if strict:
+            raise ValueError(f"{var_name} is not set")
+    return val or default
+
+
+def init_sdk(strict_all_vars=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)
+
+
+def get_endpoint_secret():
+    """Fetch endpoint secret from the environment."""
+    endpoint_secret = os.getenv("ENDPOINT_SECRET")
+    return endpoint_secret

unique_toolkit/app/performance/async_executor.py

@@ -0,0 +1,186 @@
+import asyncio
+import contextlib
+import logging
+import threading
+import time
+from concurrent.futures import ThreadPoolExecutor
+from math import ceil
+from typing import (
+    AsyncContextManager,
+    Awaitable,
+    Callable,
+    Optional,
+    Sequence,
+    TypeVar,
+    Union,
+)
+
+T = TypeVar("T")
+Result = Union[T, BaseException]
+
+
+class AsyncExecutor:
+    """
+    A class for executing asynchronous tasks concurrently, with optional threading support.
+
+    This class provides methods to run multiple asynchronous tasks in parallel, with
+    the ability to limit the number of concurrent tasks and distribute work across
+    multiple threads if needed.
+
+    Attributes:
+        logger (logging.Logger): Logger instance for recording execution information.
+        context_manager (Callable[[], AsyncContextManager]): A factory function that returns
+            an async context manager to be used for each task execution, e.g., quart.current_app.app_context().
+
+    """
+
+    def __init__(
+        self,
+        logger: Optional[logging.Logger] = None,
+        context_manager: Optional[Callable[[], AsyncContextManager]] = None,
+    ) -> None:
+        self.logger = logger or logging.getLogger(__name__)
+        self.context_manager = context_manager or (lambda: contextlib.nullcontext())
+
+    async def run_async_tasks(
+        self,
+        tasks: Sequence[Awaitable[T]],
+        max_tasks: int,
+    ) -> list[Result]:
+        """
+        Executes the a set of given async tasks and returns the results.
+
+        Args:
+        tasks (list[Awaitable[T]]): list of async callables to execute in parallel.
+        max_tasks (int): Maximum number of tasks for the asyncio Semaphore.
+
+        Returns:
+        list[Result]: list of results from the executed tasks.
+        """
+
+        async def logging_wrapper(task: Awaitable[T], task_id: int) -> Result:
+            thread = threading.current_thread()
+            start_time = time.time()
+
+            self.logger.info(
+                f"Thread {thread.name} (ID: {thread.ident}) starting task {task_id}"
+            )
+
+            try:
+                async with self.context_manager():
+                    result = await task
+                return result
+            except Exception as e:
+                self.logger.error(
+                    f"Thread {thread.name} (ID: {thread.ident}) - Task {task_id} failed with error: {e}"
+                )
+                return e
+            finally:
+                end_time = time.time()
+                duration = end_time - start_time
+                self.logger.debug(
+                    f"Thread {thread.name} (ID: {thread.ident}) - Task {task_id} finished in {duration:.2f} seconds"
+                )
+
+        sem = asyncio.Semaphore(max_tasks)
+
+        async def sem_task(task: Awaitable[T], task_id: int) -> Result:
+            async with sem:
+                return await logging_wrapper(task, task_id)
+
+        wrapped_tasks: list[Awaitable[Result]] = [
+            sem_task(task, i) for i, task in enumerate(tasks)
+        ]
+
+        results: list[Result] = await asyncio.gather(
+            *wrapped_tasks, return_exceptions=True
+        )
+
+        return results
+
+    async def run_async_tasks_in_threads(
+        self,
+        tasks: Sequence[Awaitable[T]],
+        max_threads: int,
+        max_tasks: int,
+    ) -> list[Result[T]]:
+        """
+        Executes the given async tasks in multiple threads and returns the results.
+
+        Args:
+        tasks (list[Awaitable[T]]): list of async callables to execute in parallel.
+        max_threads (int): Maximum number of threads.
+        max_tasks (int): Maximum number of tasks per thread run in parallel.
+
+        Returns:
+        list[Result]: list of results from the executed tasks.
+        """
+
+        async def run_in_thread(task_chunk: list[Awaitable[T]]) -> list[Result]:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            async with self.context_manager():
+                return await self.run_async_tasks(task_chunk, max_tasks)
+
+        def thread_worker(
+            task_chunk: list[Awaitable[T]], chunk_id: int
+        ) -> list[Result]:
+            thread = threading.current_thread()
+            self.logger.info(
+                f"Thread {thread.name} (ID: {thread.ident}) starting chunk {chunk_id} with {len(task_chunk)} tasks"
+            )
+
+            start_time = time.time()
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+            try:
+                results = loop.run_until_complete(run_in_thread(task_chunk))
+                end_time = time.time()
+                duration = end_time - start_time
+                self.logger.info(
+                    f"Thread {thread.name} (ID: {thread.ident}) finished chunk {chunk_id} in {duration:.2f} seconds"
+                )
+                return results
+            except Exception as e:
+                self.logger.error(
+                    f"Thread {thread.name} (ID: {thread.ident}) encountered an error in chunk {chunk_id}: {str(e)}"
+                )
+                raise
+            finally:
+                loop.close()
+
+        start_time = time.time()
+        # Calculate the number of tasks per thread
+        tasks_per_thread: int = ceil(len(tasks) / max_threads)
+
+        # Split tasks into chunks
+        task_chunks: list[Sequence[Awaitable[T]]] = [
+            tasks[i : i + tasks_per_thread]
+            for i in range(0, len(tasks), tasks_per_thread)
+        ]
+
+        self.logger.info(
+            f"Splitting {len(tasks)} tasks into {len(task_chunks)} chunks across {max_threads} threads"
+        )
+
+        # Use ThreadPoolExecutor to manage threads
+        with ThreadPoolExecutor(max_workers=max_threads) as executor:
+            # Submit each chunk of tasks to a thread
+            future_results: list[list[Result]] = list(
+                executor.map(
+                    thread_worker,
+                    task_chunks,
+                    range(len(task_chunks)),  # chunk_id
+                )
+            )
+
+        # Flatten the results from all threads
+        results: list[Result] = [item for sublist in future_results for item in sublist]
+        end_time = time.time()
+        duration = end_time - start_time
+        self.logger.info(
+            f"All threads completed. Total results: {len(results)}. Duration: {duration:.2f} seconds"
+        )
+
+        return results

unique_toolkit/app/performance/async_wrapper.py

@@ -0,0 +1,28 @@
+import asyncio
+import warnings
+from functools import wraps
+from typing import Any, Callable, Coroutine, TypeVar
+
+T = TypeVar("T")
+
+
+def to_async(func: Callable[..., T]) -> Callable[..., Coroutine[Any, Any, T]]:
+    @wraps(func)
+    async def wrapper(*args, **kwargs) -> T:
+        return await asyncio.to_thread(func, *args, **kwargs)
+
+    return wrapper
+
+
+def async_warning(func):
+    @wraps(func)
+    async def wrapper(*args, **kwargs):
+        warnings.warn(
+            f"The function '{func.__name__}' is not purely async. It uses a thread pool executor underneath, "
+            "which may impact performance for CPU-bound operations.",
+            RuntimeWarning,
+            stacklevel=2,
+        )
+        return await func(*args, **kwargs)
+
+    return wrapper

unique_toolkit/app/schemas.py

@@ -0,0 +1,54 @@
+from enum import StrEnum
+from typing import Any
+
+from humps import camelize
+from pydantic import BaseModel, ConfigDict
+
+# set config to convert camelCase to snake_case
+model_config = ConfigDict(
+    alias_generator=camelize, populate_by_name=True, arbitrary_types_allowed=True
+)
+
+
+class EventName(StrEnum):
+    EXTERNAL_MODULE_CHOSEN = "unique.chat.external-module.chosen"
+
+
+class EventUserMessage(BaseModel):
+    model_config = model_config
+
+    id: str
+    text: str
+    created_at: str
+
+
+class EventAssistantMessage(BaseModel):
+    model_config = model_config
+
+    id: str
+    created_at: str
+
+
+class EventPayload(BaseModel):
+    model_config = model_config
+
+    name: EventName
+    description: str
+    configuration: dict[str, Any]
+    chat_id: str
+    assistant_id: str
+    user_message: EventUserMessage
+    assistant_message: EventAssistantMessage
+    text: str | None = None
+
+
+class Event(BaseModel):
+    model_config = model_config
+
+    id: str
+    event: str
+    user_id: str
+    company_id: str
+    payload: EventPayload
+    created_at: int | None = None
+    version: str | None = None

unique_toolkit/app/verification.py

@@ -0,0 +1,58 @@
+import logging
+
+import unique_sdk
+
+from unique_toolkit.app.schemas import Event
+
+
+class WebhookVerificationError(Exception):
+    """Custom exception for webhook verification errors."""
+
+    pass
+
+
+def verify_signature_and_construct_event(
+    headers: dict[str, str],
+    payload: bytes,
+    endpoint_secret: str,
+    logger: logging.Logger = logging.getLogger(__name__),
+):
+    """
+    Verify the signature of a webhook and construct an event object.
+
+    Args:
+        headers (Dict[str, str]): The headers of the webhook request.
+        payload (bytes): The raw payload of the webhook request.
+        endpoint_secret (str): The secret used to verify the webhook signature.
+        logger (logging.Logger): A logger instance for logging messages.
+
+    Returns:
+        Union[Event, Tuple[Dict[str, bool], int]]:
+            If successful, returns an Event object.
+            If unsuccessful, returns a tuple with an error response and HTTP status code.
+
+    Raises:
+        WebhookVerificationError: If there's an error during verification or event construction.
+    """
+
+    # Only verify the event if there is an endpoint secret defined
+    # Otherwise use the basic event deserialized with json
+    sig_header = headers.get("X-Unique-Signature")
+    timestamp = headers.get("X-Unique-Created-At")
+
+    if not sig_header or not timestamp:
+        logger.error("⚠️  Webhook signature or timestamp headers missing.")
+        raise WebhookVerificationError("Signature or timestamp headers missing")
+
+    try:
+        event = unique_sdk.Webhook.construct_event(
+            payload,
+            sig_header,
+            timestamp,
+            endpoint_secret,
+        )
+        logger.info("✅  Webhook signature verification successful.")
+        return Event(**event)
+    except unique_sdk.SignatureVerificationError as e:
+        logger.error("⚠️  Webhook signature verification failed. " + str(e))
+        raise WebhookVerificationError(f"Signature verification failed: {str(e)}")

unique_toolkit/chat/schemas.py

@@ -0,0 +1,30 @@
+from enum import Enum
+
+from humps import camelize
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+
+# set config to convert camelCase to snake_case
+model_config = ConfigDict(
+    alias_generator=camelize, populate_by_name=True, arbitrary_types_allowed=True
+)
+
+
+class ChatMessageRole(str, Enum):
+    USER = "user"
+    ASSISTANT = "assistant"
+
+
+class ChatMessage(BaseModel):
+    model_config = model_config
+
+    id: str | None = None
+    object: str | None = None
+    content: str = Field(alias="text")
+    role: ChatMessageRole
+    debug_info: dict = {}
+
+    # TODO make sdk return role consistently in lowercase
+    # Currently needed as sdk returns role in uppercase
+    @field_validator("role", mode="before")
+    def set_role(cls, value: str):
+        return value.lower()