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

unique_toolkit/__init__.py

@@ -1,3 +0,0 @@
-from .observability.log_config import load_logging_config  # noqa: F401
-from .sdk.init import init_sdk  # noqa: F401
-from .security.verify import verify_signature_and_construct_event  # noqa: F401

unique_toolkit/{observability/log_config.py → app/init_logging.py}

@@ -7,7 +7,7 @@ class UTCFormatter(Formatter):
     converter = gmtime
 
 
-log_config = {
+unique_log_config = {
     "version": 1,
     "root": {"level": "DEBUG", "handlers": ["console"]},
     "handlers": {
@@ -15,17 +15,17 @@ log_config = {
             "class": "logging.StreamHandler",
             "level": "DEBUG",
             "formatter": "utc",
-        },
-        "formatters": {
-            "utc": {
-                "()": UTCFormatter,
-                "format": "%(asctime)s: %(message)s",
-                "datefmt": "%Y-%m-%d %H:%M:%S",
-            },
+        }
+    },
+    "formatters": {
+        "utc": {
+            "()": UTCFormatter,
+            "format": "%(asctime)s: %(message)s",
+            "datefmt": "%Y-%m-%d %H:%M:%S",
         },
     },
 }
 
 
-def load_logging_config():
-    return dictConfig(log_config)
+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/{performance → app/performance}/async_executor.py

@@ -1,19 +1,46 @@
 import asyncio
+import contextlib
+import logging
 import threading
 import time
 from concurrent.futures import ThreadPoolExecutor
 from math import ceil
-from typing import Awaitable, Sequence, TypeVar, Union
-
-from quart import Quart
+from typing import (
+    AsyncContextManager,
+    Awaitable,
+    Callable,
+    Optional,
+    Sequence,
+    TypeVar,
+    Union,
+)
 
 T = TypeVar("T")
 Result = Union[T, BaseException]
 
 
 class AsyncExecutor:
-    def __init__(self, app_instance: Quart) -> None:
-        self.app = app_instance
+    """
+    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,
@@ -21,7 +48,7 @@ class AsyncExecutor:
         max_tasks: int,
     ) -> list[Result]:
         """
-        Executes the given async tasks within one thread non-blocking and returns the results.
+        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.
@@ -35,23 +62,23 @@ class AsyncExecutor:
             thread = threading.current_thread()
             start_time = time.time()
 
-            self.app.logger.info(
+            self.logger.info(
                 f"Thread {thread.name} (ID: {thread.ident}) starting task {task_id}"
             )
 
             try:
-                async with self.app.app_context():
+                async with self.context_manager():
                     result = await task
                 return result
             except Exception as e:
-                self.app.logger.error(
+                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.app.logger.debug(
+                self.logger.debug(
                     f"Thread {thread.name} (ID: {thread.ident}) - Task {task_id} finished in {duration:.2f} seconds"
                 )
 
@@ -78,7 +105,7 @@ class AsyncExecutor:
         max_tasks: int,
     ) -> list[Result[T]]:
         """
-        Executes the given async tasks in parallel threads and returns the results.
+        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.
@@ -92,14 +119,14 @@ class AsyncExecutor:
         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.app.app_context():
+            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.app.logger.info(
+            self.logger.info(
                 f"Thread {thread.name} (ID: {thread.ident}) starting chunk {chunk_id} with {len(task_chunk)} tasks"
             )
 
@@ -111,12 +138,12 @@ class AsyncExecutor:
                 results = loop.run_until_complete(run_in_thread(task_chunk))
                 end_time = time.time()
                 duration = end_time - start_time
-                self.app.logger.info(
+                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.app.logger.error(
+                self.logger.error(
                     f"Thread {thread.name} (ID: {thread.ident}) encountered an error in chunk {chunk_id}: {str(e)}"
                 )
                 raise
@@ -133,7 +160,7 @@ class AsyncExecutor:
             for i in range(0, len(tasks), tasks_per_thread)
         ]
 
-        self.app.logger.info(
+        self.logger.info(
             f"Splitting {len(tasks)} tasks into {len(task_chunks)} chunks across {max_threads} threads"
         )
 
@@ -152,7 +179,7 @@ class AsyncExecutor:
         results: list[Result] = [item for sublist in future_results for item in sublist]
         end_time = time.time()
         duration = end_time - start_time
-        self.app.logger.info(
+        self.logger.info(
             f"All threads completed. Total results: {len(results)}. Duration: {duration:.2f} seconds"
         )
 

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/{event/schema.py → app/schemas.py}

@@ -1,13 +1,20 @@
+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)
+model_config = ConfigDict(
+    alias_generator=camelize, populate_by_name=True, arbitrary_types_allowed=True
+)
 
 
-class UserMessage(BaseModel):
+class EventName(StrEnum):
+    EXTERNAL_MODULE_CHOSEN = "unique.chat.external-module.chosen"
+
+
+class EventUserMessage(BaseModel):
     model_config = model_config
 
     id: str
@@ -15,23 +22,23 @@ class UserMessage(BaseModel):
     created_at: str
 
 
-class AssistantMessage(BaseModel):
+class EventAssistantMessage(BaseModel):
     model_config = model_config
 
     id: str
     created_at: str
 
 
-class Payload(BaseModel):
+class EventPayload(BaseModel):
     model_config = model_config
 
-    name: str
+    name: EventName
     description: str
     configuration: dict[str, Any]
     chat_id: str
     assistant_id: str
-    user_message: UserMessage
-    assistant_message: AssistantMessage
+    user_message: EventUserMessage
+    assistant_message: EventAssistantMessage
     text: str | None = None
 
 
@@ -39,9 +46,9 @@ class Event(BaseModel):
     model_config = model_config
 
     id: str
-    version: str
     event: str
-    created_at: int
     user_id: str
     company_id: str
-    payload: Payload
+    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()