solace-agent-mesh 1.11.2__py3-none-any.whl

solace_agent_mesh/common/utils/embeds/types.py

@@ -0,0 +1,23 @@
+"""
+Defines types used within the embed processing system, like DataFormat.
+"""
+
+from enum import Enum, auto
+
+
+class ResolutionMode(Enum):
+    """Defines the context in which embed resolution is occurring."""
+
+    A2A_MESSAGE_TO_USER = auto()
+    TOOL_PARAMETER = auto()
+    RECURSIVE_ARTIFACT_CONTENT = auto()
+    ARTIFACT_STREAMING = auto()  # For streaming artifact chunks to browser
+
+
+class DataFormat(Enum):
+    """Represents internal data formats during modifier chain execution."""
+
+    BYTES = auto()
+    STRING = auto()
+    JSON_OBJECT = auto()
+    LIST_OF_DICTS = auto()

solace_agent_mesh/common/utils/in_memory_cache.py

@@ -0,0 +1,108 @@
+"""In Memory Cache utility."""
+
+import threading
+import time
+from typing import Any, Dict, Optional
+
+
+class InMemoryCache:
+    """A thread-safe Singleton class to manage cache data.
+
+    Ensures only one instance of the cache exists across the application.
+    """
+
+    _instance: Optional["InMemoryCache"] = None
+    _lock: threading.Lock = threading.Lock()
+    _initialized: bool = False
+
+    def __new__(cls):
+        """Override __new__ to control instance creation (Singleton pattern).
+
+        Uses a lock to ensure thread safety during the first instantiation.
+
+        Returns:
+            The singleton instance of InMemoryCache.
+        """
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def __init__(self):
+        """Initialize the cache storage.
+
+        Uses a flag (_initialized) to ensure this logic runs only on the very first
+        creation of the singleton instance.
+        """
+        if not self._initialized:
+            with self._lock:
+                if not self._initialized:
+                    self._cache_data: Dict[str, Dict[str, Any]] = {}
+                    self._ttl: Dict[str, float] = {}
+                    self._data_lock: threading.Lock = threading.Lock()
+                    self._initialized = True
+
+    def set(self, key: str, value: Any, ttl: Optional[int] = None) -> None:
+        """Set a key-value pair.
+
+        Args:
+            key: The key for the data.
+            value: The data to store.
+            ttl: Time to live in seconds. If None, data will not expire.
+        """
+        with self._data_lock:
+            self._cache_data[key] = value
+
+            if ttl is not None:
+                self._ttl[key] = time.time() + ttl
+            else:
+                if key in self._ttl:
+                    del self._ttl[key]
+
+    def get(self, key: str, default: Any = None) -> Any:
+        """Get the value associated with a key.
+
+        Args:
+            key: The key for the data within the session.
+            default: The value to return if the session or key is not found.
+
+        Returns:
+            The cached value, or the default value if not found.
+        """
+        with self._data_lock:
+            if key in self._ttl and time.time() > self._ttl[key]:
+                del self._cache_data[key]
+                del self._ttl[key]
+                return default
+            return self._cache_data.get(key, default)
+
+    def delete(self, key: str) -> None:
+        """Delete a specific key-value pair from a cache.
+
+        Args:
+            key: The key to delete.
+
+        Returns:
+            True if the key was found and deleted, False otherwise.
+        """
+
+        with self._data_lock:
+            if key in self._cache_data:
+                del self._cache_data[key]
+                if key in self._ttl:
+                    del self._ttl[key]
+                return True
+            return False
+
+    def clear(self) -> bool:
+        """Remove all data.
+
+        Returns:
+            True if the data was cleared, False otherwise.
+        """
+        with self._data_lock:
+            self._cache_data.clear()
+            self._ttl.clear()
+            return True
+        return False

solace_agent_mesh/common/utils/initializer.py

@@ -0,0 +1,52 @@
+import logging
+import os
+import json
+from solace_ai_connector.main import load_config
+
+log = logging.getLogger(__name__)
+
+_has_initialized_system = False
+
+def initialize():
+    global _has_initialized_system
+    if not _has_initialized_system:
+        _has_initialized_system = True
+    else:
+        return
+    
+    try:
+        from solace_agent_mesh_enterprise.init_enterprise import initialize_enterprise_features
+    except ImportError:
+        # Community edition
+        # Contact Solace support for enterprise features
+        return
+    
+    enterprise_config = os.getenv("SAM_AUTHORIZATION_CONFIG")
+    if enterprise_config and isinstance(enterprise_config, str):
+        if enterprise_config.endswith('.yaml') or enterprise_config.endswith('.yml'):
+            try:
+                enterprise_config = load_config(enterprise_config)
+            except Exception as e:
+                log.error("Failed to load YAML config from SAM_AUTHORIZATION_CONFIG: %s", e, exc_info=True)
+                raise
+        elif enterprise_config.endswith('.json'):
+            try:
+                with open(enterprise_config, 'r', encoding='utf-8') as file:
+                    enterprise_config = json.load(file)
+            except Exception as e:
+                log.error("Failed to load JSON config from SAM_AUTHORIZATION_CONFIG: %s", e, exc_info=True)
+                raise
+        else:
+            try:
+                enterprise_config = json.loads(enterprise_config)
+            except json.JSONDecodeError as e:
+                log.error("Invalid JSON in SAM_AUTHORIZATION_CONFIG: %s", e, exc_info=True)
+                raise
+    else:
+        enterprise_config = {}
+    
+    try:
+        initialize_enterprise_features(enterprise_config)
+    except Exception as e:
+        log.error("Failed to initialize enterprise features: %s", e, exc_info=True)
+        raise

solace_agent_mesh/common/utils/log_formatters.py

@@ -0,0 +1,64 @@
+import logging
+import json
+import os
+import warnings
+from datetime import datetime, timezone
+import traceback
+
+
+class DatadogJsonFormatter(logging.Formatter):
+    """
+    Custom formatter to output logs in Datadog-compatible JSON format.
+    
+    . deprecated::
+        This formatter is deprecated. Please use pythonjsonlogger.json.JsonFormatter instead.
+        For details, consult https://solacelabs.github.io/solace-agent-mesh/docs/documentation/deploying/logging#structured-logging
+
+        Example:
+            formatters:
+              jsonFormatter:
+                "()": pythonjsonlogger.json.JsonFormatter
+                format: "%(timestamp)s %(levelname)s %(threadName)s %(name)s %(message)s"
+    """
+
+    def format(self, record):
+        # Emit deprecation warning once
+        if not hasattr(self.__class__, '_deprecation_warned'):
+            warnings.warn(
+                "DatadogJsonFormatter is deprecated and will be removed in a future version. "
+                "Please use pythonjsonlogger.json.JsonFormatter instead. For details, consult https://solacelabs.github.io/solace-agent-mesh/docs/documentation/deploying/logging#structured-logging",
+                DeprecationWarning,
+                stacklevel=2
+            )
+            self.__class__._deprecation_warned = True
+        log_entry = {
+            "timestamp": datetime.fromtimestamp(
+                record.created, tz=timezone.utc
+            ).isoformat(),
+            "level": record.levelname,
+            "message": record.getMessage(),
+            "logger.name": record.name,
+            "logger.thread_name": record.threadName,
+            "service": os.getenv("SERVICE_NAME", "solace_agent_mesh"),
+            "code.filepath": record.pathname,
+            "code.lineno": record.lineno,
+            "code.module": record.module,
+            "code.funcName": record.funcName,
+        }
+
+        dd_trace_id = getattr(record, "dd.trace_id", None)
+        if dd_trace_id:
+            log_entry["dd.trace_id"] = dd_trace_id
+
+        dd_span_id = getattr(record, "dd.span_id", None)
+        if dd_span_id:
+            log_entry["dd.span_id"] = dd_span_id
+
+        if record.exc_info:
+            log_entry["exception.type"] = record.exc_info[0].__name__
+            log_entry["exception.message"] = str(record.exc_info[1])
+            log_entry["exception.stacktrace"] = "".join(
+                traceback.format_exception(*record.exc_info)
+            )
+
+        return json.dumps(log_entry)

solace_agent_mesh/common/utils/message_utils.py

@@ -0,0 +1,80 @@
+"""Message utility functions for size calculation and validation.
+
+This module provides utilities for calculating and validating message sizes
+to ensure they don't exceed configured limits. The size calculation matches
+the exact serialization format used by the Solace AI Connector (JSON + UTF-8).
+"""
+
+import logging
+import json
+from typing import Any, Dict, Tuple
+
+log = logging.getLogger(__name__)
+
+# Maximum bytes per character in UTF-8 encoding (4 bytes for full Unicode range)
+MAX_UTF8_BYTES_PER_CHARACTER = 4
+
+
+def calculate_message_size(payload: Dict[str, Any]) -> int:
+    """Calculate the exact size of a message payload in bytes.
+
+    Uses JSON serialization followed by UTF-8 encoding to match the exact
+    format used by the Solace AI Connector.
+
+    Args:
+        payload: The message payload dictionary to calculate size for.
+
+    Returns:
+        The size of the payload in bytes.
+
+    Note:
+        If JSON serialization fails, falls back to string representation
+        for size estimation.
+    """
+    try:
+        # Use JSON serialization + UTF-8 encoding to match Solace AI Connector
+        json_str = json.dumps(payload, separators=(",", ":"), ensure_ascii=False)
+        return len(json_str.encode("utf-8"))
+    except (TypeError, ValueError) as e:
+        # Graceful fallback if JSON serialization fails
+        log.warning(
+            f"Failed to serialize payload to JSON for size calculation: {e}. "
+            f"Using string representation fallback."
+        )
+        try:
+            return len(str(payload).encode("utf-8"))
+        except Exception as fallback_error:
+            log.error(
+                f"Fallback size calculation also failed: {fallback_error}. "
+                f"Returning conservative estimate."
+            )
+            # Conservative estimate using maximum UTF-8 bytes per character
+            return len(str(payload)) * MAX_UTF8_BYTES_PER_CHARACTER
+
+
+def validate_message_size(
+    payload: Dict[str, Any], max_size_bytes: int, component_identifier: str = "Unknown"
+) -> Tuple[bool, int]:
+    """Validate that a message payload doesn't exceed the maximum size limit.
+
+    Args:
+        payload: The message payload dictionary to validate.
+        max_size_bytes: The maximum allowed size in bytes.
+        component_identifier: Identifier for the component performing validation
+                             (used in log messages).
+
+    Returns:
+        A tuple containing:
+        - bool: True if the payload is within size limits, False otherwise
+        - int: The actual size of the payload in bytes
+
+    Note:
+        Logs an error if size exceeds the limit.
+    """
+    actual_size = calculate_message_size(payload)
+
+    # Check if size exceeds the limit
+    if actual_size > max_size_bytes:
+        return False, actual_size
+
+    return True, actual_size

solace_agent_mesh/common/utils/mime_helpers.py

@@ -0,0 +1,172 @@
+"""
+Utility functions for handling MIME types.
+"""
+
+from typing import Optional, Set
+
+TEXT_CONTAINER_MIME_TYPES: Set[str] = {
+    "text/plain",
+    "text/markdown",
+    "text/html",
+    "application/json",
+    "application/yaml",
+    "text/yaml",
+    "application/x-yaml",
+    "text/x-yaml",
+    "application/xml",
+    "text/xml",
+    "text/csv",
+}
+
+_TEXT_BASED_PRIMARY_TYPES = {"text"}
+_TEXT_BASED_SUBTYPE_WHOLE = {
+    "json",
+    "xml",
+    "yaml",
+    "x-yaml",
+    "yml",
+    "csv",
+    "javascript",
+    "ecmascript",
+    "xhtml+xml",
+    "svg+xml",
+    "atom+xml",
+    "rss+xml",
+    "sparql-query",
+    "sparql-update",
+    "sql",
+    "graphql",
+    "markdown",
+    "html",
+    "rtf",
+    "sgml",
+}
+_TEXT_BASED_SUBTYPE_SUFFIXES_AFTER_PLUS = {
+    "json",
+    "xml",
+    "yaml",
+    "csv",
+    "svg",
+    "xhtml",
+}
+
+
+def is_text_based_mime_type(mime_type: Optional[str]) -> bool:
+    """
+    Checks if a given MIME type is considered text-based.
+
+    Args:
+        mime_type: The MIME type string (e.g., "text/plain", "application/json").
+
+    Returns:
+        True if the MIME type is text-based, False otherwise.
+    """
+    if not mime_type:
+        return False
+
+    normalized_mime_type = mime_type.lower().strip()
+
+    if normalized_mime_type.startswith("text/"):
+        return True
+
+    if normalized_mime_type in TEXT_CONTAINER_MIME_TYPES:
+        return True
+
+    return False
+
+
+def is_text_based_file(
+    mime_type: Optional[str], content_bytes: Optional[bytes] = None
+) -> bool:
+    """
+    Determines if a file is text-based based on its MIME type and content.
+    Args:
+        mime_type: The MIME type of the file.
+        content_bytes: The content of the file as bytes.
+    Returns:
+        True if the file is text-based, False otherwise.
+    """
+    if not mime_type:
+        return False
+
+    normalized_mime_type = mime_type.lower().strip()
+    primary_type, _, subtype = normalized_mime_type.partition("/")
+
+    if primary_type in _TEXT_BASED_PRIMARY_TYPES:
+        return True
+    elif subtype in _TEXT_BASED_SUBTYPE_WHOLE:
+        return True
+    elif "+" in subtype:
+        specific_format = subtype.split("+")[-1]
+        if specific_format in _TEXT_BASED_SUBTYPE_SUFFIXES_AFTER_PLUS:
+            return True
+    elif (
+        normalized_mime_type == "application/octet-stream" and content_bytes is not None
+    ):
+        try:
+            sample_size = min(1024, len(content_bytes))
+            content_bytes[:sample_size].decode("utf-8")
+            return True
+        except UnicodeDecodeError:
+            return False
+
+    return False
+
+
+def get_extension_for_mime_type(
+    mime_type: Optional[str], default_extension: str = ".dat"
+) -> str:
+    """
+    Returns a file extension for a given MIME type.
+
+    Args:
+        mime_type: The MIME type string (e.g., 'image/png', 'application/json').
+        default_extension: The extension to return if the MIME type is not found.
+
+    Returns:
+        The corresponding file extension (e.g., '.png', '.json').
+    """
+    if not mime_type:
+        return default_extension
+
+    # Comprehensive mapping of common MIME types to file extensions
+    extension_mapping = {
+        # Text formats
+        "text/plain": ".txt",
+        "text/html": ".html",
+        "text/css": ".css",
+        "text/javascript": ".js",
+        "text/csv": ".csv",
+        "text/markdown": ".md",
+        "text/xml": ".xml",
+        # Application formats
+        "application/json": ".json",
+        "application/x-yaml": ".yaml",
+        "application/yaml": ".yaml",
+        "application/pdf": ".pdf",
+        "application/zip": ".zip",
+        "application/octet-stream": ".bin",
+        # Image formats
+        "image/png": ".png",
+        "image/jpeg": ".jpg",
+        "image/jpg": ".jpg",
+        "image/gif": ".gif",
+        "image/bmp": ".bmp",
+        "image/webp": ".webp",
+        "image/svg+xml": ".svg",
+        # Audio formats
+        "audio/wav": ".wav",
+        "audio/mp3": ".mp3",
+        "audio/mpeg": ".mp3",
+        "audio/ogg": ".ogg",
+        "audio/flac": ".flac",
+        "audio/aac": ".aac",
+        "audio/m4a": ".m4a",
+        # Video formats
+        "video/mp4": ".mp4",
+        "video/webm": ".webm",
+        "video/x-msvideo": ".avi",
+        "video/quicktime": ".mov",
+    }
+
+    return extension_mapping.get(mime_type.lower(), default_extension)

solace_agent_mesh/common/utils/push_notification_auth.py

@@ -0,0 +1,135 @@
+import logging
+from jwcrypto import jwk
+import uuid
+from starlette.responses import JSONResponse
+from starlette.requests import Request
+from typing import Any
+
+import jwt
+import time
+import json
+import hashlib
+import httpx
+
+from jwt import PyJWK, PyJWKClient
+
+log = logging.getLogger(__name__)
+
+AUTH_HEADER_PREFIX = "Bearer "
+
+
+class PushNotificationAuth:
+    def _calculate_request_body_sha256(self, data: dict[str, Any]):
+        """Calculates the SHA256 hash of a request body.
+
+        This logic needs to be same for both the agent who signs the payload and the client verifier.
+        """
+        body_str = json.dumps(
+            data,
+            ensure_ascii=False,
+            allow_nan=False,
+            indent=None,
+            separators=(",", ":"),
+        )
+        return hashlib.sha256(body_str.encode()).hexdigest()
+
+
+class PushNotificationSenderAuth(PushNotificationAuth):
+    def __init__(self):
+        self.public_keys = []
+        self.private_key_jwk: PyJWK = None
+
+    @staticmethod
+    async def verify_push_notification_url(url: str) -> bool:
+        async with httpx.AsyncClient(timeout=10) as client:
+            try:
+                validation_token = str(uuid.uuid4())
+                response = await client.get(
+                    url, params={"validationToken": validation_token}
+                )
+                response.raise_for_status()
+                is_verified = response.text == validation_token
+
+                log.info("Verified push-notification URL: %s => %s", url, is_verified)
+                return is_verified
+            except Exception as e:
+                log.warning(
+                    "Error during sending push-notification for URL %s: %s", url, e
+                )
+
+        return False
+
+    def generate_jwk(self):
+        key = jwk.JWK.generate(kty="RSA", size=2048, kid=str(uuid.uuid4()), use="sig")
+        self.public_keys.append(key.export_public(as_dict=True))
+        self.private_key_jwk = PyJWK.from_json(key.export_private())
+
+    def handle_jwks_endpoint(self, _request: Request):
+        """Allow clients to fetch public keys."""
+        return JSONResponse({"keys": self.public_keys})
+
+    def _generate_jwt(self, data: dict[str, Any]):
+        """JWT is generated by signing both the request payload SHA digest and time of token generation.
+
+        Payload is signed with private key and it ensures the integrity of payload for client.
+        Including iat prevents from replay attack.
+        """
+
+        iat = int(time.time())
+
+        return jwt.encode(
+            {
+                "iat": iat,
+                "request_body_sha256": self._calculate_request_body_sha256(data),
+            },
+            key=self.private_key_jwk,
+            headers={"kid": self.private_key_jwk.key_id},
+            algorithm="RS256",
+        )
+
+    async def send_push_notification(self, url: str, data: dict[str, Any]):
+        jwt_token = self._generate_jwt(data)
+        headers = {"Authorization": f"Bearer {jwt_token}"}
+        async with httpx.AsyncClient(timeout=10) as client:
+            try:
+                response = await client.post(url, json=data, headers=headers)
+                response.raise_for_status()
+                log.info("Push-notification sent for URL: %s", url)
+            except Exception as e:
+                log.warning(
+                    "Error during sending push-notification for URL %s: %s", url, e
+                )
+
+
+class PushNotificationReceiverAuth(PushNotificationAuth):
+    def __init__(self):
+        self.public_keys_jwks = []
+        self.jwks_client = None
+
+    async def load_jwks(self, jwks_url: str):
+        self.jwks_client = PyJWKClient(jwks_url)
+
+    async def verify_push_notification(self, request: Request) -> bool:
+        auth_header = request.headers.get("Authorization")
+        if not auth_header or not auth_header.startswith(AUTH_HEADER_PREFIX):
+            log.warning("Invalid authorization header")
+            return False
+
+        token = auth_header[len(AUTH_HEADER_PREFIX) :]
+        signing_key = self.jwks_client.get_signing_key_from_jwt(token)
+
+        decode_token = jwt.decode(
+            token,
+            signing_key,
+            options={"require": ["iat", "request_body_sha256"]},
+            algorithms=["RS256"],
+        )
+
+        actual_body_sha256 = self._calculate_request_body_sha256(await request.json())
+        if actual_body_sha256 != decode_token["request_body_sha256"]:
+            raise ValueError("Invalid request body")
+
+        if time.time() - decode_token["iat"] > 60 * 5:
+            raise ValueError("Token is expired")
+
+        return True