solace-agent-mesh 0.2.4__py3-none-any.whl → 1.0.2__py3-none-any.whl

solace_agent_mesh/common/server/task_manager.py

@@ -0,0 +1,291 @@
+from abc import ABC, abstractmethod
+from typing import Union, AsyncIterable, List
+import asyncio
+from solace_ai_connector.common.log import log
+from ...common.types import (
+    Task,
+    JSONRPCResponse,
+    TaskIdParams,
+    TaskQueryParams,
+    GetTaskRequest,
+    TaskNotFoundError,
+    SendTaskRequest,
+    CancelTaskRequest,
+    TaskNotCancelableError,
+    SetTaskPushNotificationRequest,
+    GetTaskPushNotificationRequest,
+    GetTaskResponse,
+    CancelTaskResponse,
+    SendTaskResponse,
+    SetTaskPushNotificationResponse,
+    GetTaskPushNotificationResponse,
+    TaskSendParams,
+    TaskStatus,
+    TaskState,
+    TaskResubscriptionRequest,
+    SendTaskStreamingRequest,
+    SendTaskStreamingResponse,
+    Artifact,
+    PushNotificationConfig,
+    TaskStatusUpdateEvent,
+    JSONRPCError,
+    TaskPushNotificationConfig,
+    InternalError,
+)
+from ...common.server.utils import (
+    new_not_implemented_error,
+)
+
+
+class TaskManager(ABC):
+    @abstractmethod
+    async def on_get_task(self, request: GetTaskRequest) -> GetTaskResponse:
+        pass
+
+    @abstractmethod
+    async def on_cancel_task(self, request: CancelTaskRequest) -> CancelTaskResponse:
+        pass
+
+    @abstractmethod
+    async def on_send_task(self, request: SendTaskRequest) -> SendTaskResponse:
+        pass
+
+    @abstractmethod
+    async def on_send_task_subscribe(
+        self, request: SendTaskStreamingRequest
+    ) -> Union[AsyncIterable[SendTaskStreamingResponse], JSONRPCResponse]:
+        pass
+
+    @abstractmethod
+    async def on_set_task_push_notification(
+        self, request: SetTaskPushNotificationRequest
+    ) -> SetTaskPushNotificationResponse:
+        pass
+
+    @abstractmethod
+    async def on_get_task_push_notification(
+        self, request: GetTaskPushNotificationRequest
+    ) -> GetTaskPushNotificationResponse:
+        pass
+
+    @abstractmethod
+    async def on_resubscribe_to_task(
+        self, request: TaskResubscriptionRequest
+    ) -> Union[AsyncIterable[SendTaskResponse], JSONRPCResponse]:
+        pass
+
+
+class InMemoryTaskManager(TaskManager):
+    def __init__(self):
+        self.tasks: dict[str, Task] = {}
+        self.push_notification_infos: dict[str, PushNotificationConfig] = {}
+        self.lock = asyncio.Lock()
+        self.task_sse_subscribers: dict[str, List[asyncio.Queue]] = {}
+        self.subscriber_lock = asyncio.Lock()
+
+    async def on_get_task(self, request: GetTaskRequest) -> GetTaskResponse:
+        log.info("Getting task %s", request.params.id)
+        task_query_params: TaskQueryParams = request.params
+
+        async with self.lock:
+            task = self.tasks.get(task_query_params.id)
+            if task is None:
+                return GetTaskResponse(id=request.id, error=TaskNotFoundError())
+
+            task_result = self.append_task_history(
+                task, task_query_params.historyLength
+            )
+
+        return GetTaskResponse(id=request.id, result=task_result)
+
+    async def on_cancel_task(self, request: CancelTaskRequest) -> CancelTaskResponse:
+        log.info("Cancelling task %s", request.params.id)
+        task_id_params: TaskIdParams = request.params
+
+        async with self.lock:
+            task = self.tasks.get(task_id_params.id)
+            if task is None:
+                return CancelTaskResponse(id=request.id, error=TaskNotFoundError())
+
+        return CancelTaskResponse(id=request.id, error=TaskNotCancelableError())
+
+    @abstractmethod
+    async def on_send_task(self, request: SendTaskRequest) -> SendTaskResponse:
+        pass
+
+    @abstractmethod
+    async def on_send_task_subscribe(
+        self, request: SendTaskStreamingRequest
+    ) -> Union[AsyncIterable[SendTaskStreamingResponse], JSONRPCResponse]:
+        pass
+
+    async def set_push_notification_info(
+        self, task_id: str, notification_config: PushNotificationConfig
+    ):
+        async with self.lock:
+            task = self.tasks.get(task_id)
+            if task is None:
+                raise ValueError("Task not found for %s" % task_id)
+
+            self.push_notification_infos[task_id] = notification_config
+
+        return
+
+    async def get_push_notification_info(self, task_id: str) -> PushNotificationConfig:
+        async with self.lock:
+            task = self.tasks.get(task_id)
+            if task is None:
+                raise ValueError("Task not found for %s" % task_id)
+
+            return self.push_notification_infos[task_id]
+
+        return
+
+    async def has_push_notification_info(self, task_id: str) -> bool:
+        async with self.lock:
+            return task_id in self.push_notification_infos
+
+    async def on_set_task_push_notification(
+        self, request: SetTaskPushNotificationRequest
+    ) -> SetTaskPushNotificationResponse:
+        log.info("Setting task push notification %s", request.params.id)
+        task_notification_params: TaskPushNotificationConfig = request.params
+
+        try:
+            await self.set_push_notification_info(
+                task_notification_params.id,
+                task_notification_params.pushNotificationConfig,
+            )
+        except Exception as e:
+            log.error(
+                "Error while setting push notification info: %s", e, exc_info=True
+            )
+            return JSONRPCResponse(
+                id=request.id,
+                error=InternalError(
+                    message="An error occurred while setting push notification info"
+                ),
+            )
+
+        return SetTaskPushNotificationResponse(
+            id=request.id, result=task_notification_params
+        )
+
+    async def on_get_task_push_notification(
+        self, request: GetTaskPushNotificationRequest
+    ) -> GetTaskPushNotificationResponse:
+        log.info("Getting task push notification %s", request.params.id)
+        task_params: TaskIdParams = request.params
+
+        try:
+            notification_info = await self.get_push_notification_info(task_params.id)
+        except Exception as e:
+            log.error(
+                "Error while getting push notification info: %s", e, exc_info=True
+            )
+            return GetTaskPushNotificationResponse(
+                id=request.id,
+                error=InternalError(
+                    message="An error occurred while getting push notification info"
+                ),
+            )
+
+        return GetTaskPushNotificationResponse(
+            id=request.id,
+            result=TaskPushNotificationConfig(
+                id=task_params.id, pushNotificationConfig=notification_info
+            ),
+        )
+
+    async def upsert_task(self, task_send_params: TaskSendParams) -> Task:
+        log.info("Upserting task %s", task_send_params.id)
+        async with self.lock:
+            task = self.tasks.get(task_send_params.id)
+            if task is None:
+                task = Task(
+                    id=task_send_params.id,
+                    sessionId=task_send_params.sessionId,
+                    messages=[task_send_params.message],
+                    status=TaskStatus(state=TaskState.SUBMITTED),
+                    history=[task_send_params.message],
+                )
+                self.tasks[task_send_params.id] = task
+            else:
+                task.history.append(task_send_params.message)
+
+            return task
+
+    async def on_resubscribe_to_task(
+        self, request: TaskResubscriptionRequest
+    ) -> Union[AsyncIterable[SendTaskStreamingResponse], JSONRPCResponse]:
+        return new_not_implemented_error(request.id)
+
+    async def update_store(
+        self, task_id: str, status: TaskStatus, artifacts: list[Artifact]
+    ) -> Task:
+        async with self.lock:
+            try:
+                task = self.tasks[task_id]
+            except KeyError:
+                log.error("Task %s not found for updating the task", task_id)
+                raise ValueError("Task %s not found" % task_id)
+
+            task.status = status
+
+            if status.message is not None:
+                task.history.append(status.message)
+
+            if artifacts is not None:
+                if task.artifacts is None:
+                    task.artifacts = []
+                task.artifacts.extend(artifacts)
+
+            return task
+
+    def append_task_history(self, task: Task, historyLength: int | None):
+        new_task = task.model_copy()
+        if historyLength is not None and historyLength > 0:
+            new_task.history = new_task.history[-historyLength:]
+        else:
+            new_task.history = []
+
+        return new_task
+
+    async def setup_sse_consumer(self, task_id: str, is_resubscribe: bool = False):
+        async with self.subscriber_lock:
+            if task_id not in self.task_sse_subscribers:
+                if is_resubscribe:
+                    raise ValueError("Task not found for resubscription")
+                else:
+                    self.task_sse_subscribers[task_id] = []
+
+            sse_event_queue = asyncio.Queue(maxsize=0)  # <=0 is unlimited
+            self.task_sse_subscribers[task_id].append(sse_event_queue)
+            return sse_event_queue
+
+    async def enqueue_events_for_sse(self, task_id, task_update_event):
+        async with self.subscriber_lock:
+            if task_id not in self.task_sse_subscribers:
+                return
+
+            current_subscribers = self.task_sse_subscribers[task_id]
+            for subscriber in current_subscribers:
+                await subscriber.put(task_update_event)
+
+    async def dequeue_events_for_sse(
+        self, request_id, task_id, sse_event_queue: asyncio.Queue
+    ) -> AsyncIterable[SendTaskStreamingResponse] | JSONRPCResponse:
+        try:
+            while True:
+                event = await sse_event_queue.get()
+                if isinstance(event, JSONRPCError):
+                    yield SendTaskStreamingResponse(id=request_id, error=event)
+                    break
+
+                yield SendTaskStreamingResponse(id=request_id, result=event)
+                if isinstance(event, TaskStatusUpdateEvent) and event.final:
+                    break
+        finally:
+            async with self.subscriber_lock:
+                if task_id in self.task_sse_subscribers:
+                    self.task_sse_subscribers[task_id].remove(sse_event_queue)

solace_agent_mesh/common/server/utils.py

@@ -0,0 +1,28 @@
+from ...common.types import (
+    JSONRPCResponse,
+    ContentTypeNotSupportedError,
+    UnsupportedOperationError,
+)
+from typing import List
+
+
+def are_modalities_compatible(
+    server_output_modes: List[str], client_output_modes: List[str]
+):
+    """Modalities are compatible if they are both non-empty
+    and there is at least one common element."""
+    if client_output_modes is None or len(client_output_modes) == 0:
+        return True
+
+    if server_output_modes is None or len(server_output_modes) == 0:
+        return True
+
+    return any(x in server_output_modes for x in client_output_modes)
+
+
+def new_incompatible_types_error(request_id):
+    return JSONRPCResponse(id=request_id, error=ContentTypeNotSupportedError())
+
+
+def new_not_implemented_error(request_id):
+    return JSONRPCResponse(id=request_id, error=UnsupportedOperationError())

solace_agent_mesh/common/services/__init__.py

@@ -0,0 +1,4 @@
+"""
+This package contains shared, reusable services for the Solace AI Connector,
+such as the IdentityService.
+"""

solace_agent_mesh/common/services/employee_service.py

@@ -0,0 +1,162 @@
+"""
+Defines the abstract base class and factory for creating Employee Service providers.
+"""
+
+import importlib
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional
+import importlib.metadata as metadata
+
+from solace_ai_connector.common.log import log
+from ..utils.in_memory_cache import InMemoryCache
+import pandas as pd
+
+
+class BaseEmployeeService(ABC):
+    """
+    Abstract base class for Employee Service providers.
+
+    This class defines a "thin provider" contract. Implementations should focus
+    solely on fetching data from a source system (like an HR platform) and
+    mapping it to the canonical format. All complex business logic, such as
+    building organizational charts or calculating availability schedules, should
+    be handled by the tools that consume this service.
+
+    Canonical Employee Schema
+    This schema defines the standard fields that all EmployeeService providers should aim to return.
+    - id (string): A unique, stable, and lowercase identifier for the employee (e.g., email, GUID).
+    - displayName (string): The employee's full name for display purposes.
+    - workEmail (string): The employee's primary work email address.
+    - jobTitle (string): The employee's official job title.
+    - department (string): The department the employee belongs to.
+    - location (string): The physical or regional location of the employee.
+    - supervisorId (string): The unique id of the employee's direct manager.
+    - hireDate (string): The date the employee was hired (ISO 8601: YYYY-MM-DD).
+    - mobilePhone (string): The employee's mobile phone number (optional).
+
+    """
+
+    def __init__(self, config: Dict[str, Any]):
+        """
+        Initializes the service with its specific configuration block.
+
+        Args:
+            config: The dictionary of configuration parameters for this provider.
+        """
+        self.config = config
+        self.log_identifier = f"[{self.__class__.__name__}]"
+        self.cache_ttl = config.get("cache_ttl_seconds", 3600)
+        self.cache = InMemoryCache() if self.cache_ttl > 0 else None
+        log.info(
+            "%s Initialized. Cache TTL: %d seconds.",
+            self.log_identifier,
+            self.cache_ttl,
+        )
+
+    @abstractmethod
+    async def get_employee_dataframe(self) -> pd.DataFrame:
+        """Returns the entire employee directory as a pandas DataFrame."""
+        pass
+
+    @abstractmethod
+    async def get_employee_profile(self, employee_id: str) -> Optional[Dict[str, Any]]:
+        """Fetches the profile for a single employee."""
+        pass
+
+    @abstractmethod
+    async def get_time_off_data(self, employee_id: str) -> List[Dict[str, Any]]:
+        """
+        Retrieves a list of raw time-off entries for an employee.
+
+        The tool consuming this data is responsible for interpreting it.
+        Each dictionary in the list MUST contain the following keys:
+        - 'start' (str): The start date of the leave (YYYY-MM-DD).
+        - 'end' (str): The end date of the leave (YYYY-MM-DD).
+        - 'type' (str): The category of leave (e.g., 'Vacation', 'Sick', 'Holiday').
+        - 'amount' (str): The amount of time taken, must be one of 'full_day' or 'half_day'.
+
+        Example:
+        [
+            {'start': '2025-07-04', 'end': '2025-07-04', 'type': 'Holiday', 'amount': 'full_day'},
+            {'start': '2025-08-15', 'end': '2025-08-15', 'type': 'Vacation', 'amount': 'full_day'}
+        ]
+
+        Args:
+            employee_id: The unique identifier for the employee.
+
+        Returns:
+            A list of dictionaries, each representing a time-off entry.
+        """
+        pass
+
+    @abstractmethod
+    async def get_employee_profile_picture(self, employee_id: str) -> Optional[str]:
+        """
+        Fetches an employee's profile picture and returns it as a data URI.
+
+        Args:
+            employee_id: The unique identifier for the employee.
+
+        Returns:
+            A string containing the data URI (e.g., 'data:image/jpeg;base64,...')
+            or None if not available.
+        """
+        pass
+
+
+def create_employee_service(
+    config: Optional[Dict[str, Any]],
+) -> Optional[BaseEmployeeService]:
+    """
+    Factory function to create an instance of an Employee Service provider
+    based on the provided configuration.
+    """
+    if not config:
+        log.info(
+            "[EmployeeFactory] No 'employee_service' configuration found. Skipping creation."
+        )
+        return None
+
+    provider_type = config.get("type")
+    if not provider_type:
+        raise ValueError("Employee service config must contain a 'type' key.")
+
+    log.info(
+        f"[EmployeeFactory] Attempting to create employee service of type: {provider_type}"
+    )
+
+    try:
+        entry_points = metadata.entry_points(group="solace_agent_mesh.plugins")
+        provider_info_entry = next(
+            (ep for ep in entry_points if ep.name == provider_type), None
+        )
+
+        if not provider_info_entry:
+            raise ValueError(
+                f"No plugin provider found for type '{provider_type}' under the 'solace_agent_mesh.plugins' entry point."
+            )
+
+        provider_info = provider_info_entry.load()
+        class_path = provider_info.get("class_path")
+        if not class_path:
+            raise ValueError(
+                f"Plugin '{provider_type}' is missing 'class_path' in its info dictionary."
+            )
+
+        module_path, class_name = class_path.rsplit(".", 1)
+        module = importlib.import_module(module_path)
+        provider_class = getattr(module, class_name)
+
+        if not issubclass(provider_class, BaseEmployeeService):
+            raise TypeError(
+                f"Provider class '{class_path}' does not inherit from BaseEmployeeService."
+            )
+
+        log.info(f"Successfully loaded employee provider plugin: {provider_type}")
+        return provider_class(config)
+    except (ImportError, AttributeError, TypeError, ValueError) as e:
+        log.exception(
+            f"[EmployeeFactory] Failed to load employee provider plugin '{provider_type}'. "
+            "Ensure the plugin is installed and the entry point is correct."
+        )
+        raise ValueError(f"Could not load employee provider plugin: {e}") from e

solace_agent_mesh/common/services/identity_service.py

@@ -0,0 +1,129 @@
+"""
+Defines the abstract base class and factory for creating Identity Service providers.
+"""
+
+import importlib
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional
+import importlib.metadata as metadata
+
+from solace_ai_connector.common.log import log
+from ...common.utils.in_memory_cache import InMemoryCache
+
+
+class BaseIdentityService(ABC):
+    """Abstract base class for all Identity Service providers."""
+
+    def __init__(self, config: Dict[str, Any]):
+        """
+        Initializes the service with its specific configuration block.
+
+        Args:
+            config: The dictionary of configuration parameters for this provider.
+        """
+        self.config = config
+        self.log_identifier = f"[{self.__class__.__name__}]"
+        self.cache_ttl = config.get("cache_ttl_seconds", 3600)
+        self.cache = InMemoryCache() if self.cache_ttl > 0 else None
+        log.info(
+            "%s Initialized. Cache TTL: %d seconds.",
+            self.log_identifier,
+            self.cache_ttl,
+        )
+
+    @abstractmethod
+    async def get_user_profile(
+        self, auth_claims: Dict[str, Any]
+    ) -> Optional[Dict[str, Any]]:
+        """
+        Fetches additional profile details for an already authenticated user.
+
+        Args:
+            auth_claims: A dictionary of claims from the primary authentication
+                         system (e.g., decoded JWT, session data). It's guaranteed
+                         to contain at least a primary user identifier.
+
+        Returns:
+            A dictionary containing additional user details (e.g., title, manager)
+            or None if the user is not found in this identity system.
+        """
+        pass
+
+    @abstractmethod
+    async def search_users(self, query: str, limit: int = 10) -> List[Dict[str, Any]]:
+        """
+        Searches for users based on a query string for autocomplete.
+
+        Args:
+            query: The partial name or email to search for.
+            limit: The maximum number of results to return.
+
+        Returns:
+            A list of dictionaries, each containing basic user info
+            (e.g., id, name, email, title).
+        """
+        pass
+
+
+def create_identity_service(
+    config: Optional[Dict[str, Any]],
+) -> Optional[BaseIdentityService]:
+    """
+    Factory function to create an instance of an Identity Service provider
+    based on the provided configuration.
+    """
+    if not config:
+        log.info(
+            "[IdentityFactory] No 'identity_service' configuration found. Skipping creation."
+        )
+        return None
+
+    provider_type = config.get("type")
+    if not provider_type:
+        raise ValueError("Identity service config must contain a 'type' key.")
+
+    log.info(
+        f"[IdentityFactory] Attempting to create identity service of type: {provider_type}"
+    )
+
+    if provider_type == "local_file":
+        from .providers.local_file_identity_service import LocalFileIdentityService
+
+        return LocalFileIdentityService(config)
+
+    else:
+        try:
+            entry_points = metadata.entry_points(group="solace_agent_mesh.plugins")
+            provider_info_entry = next(
+                (ep for ep in entry_points if ep.name == provider_type), None
+            )
+
+            if not provider_info_entry:
+                raise ValueError(
+                    f"No plugin provider found for type '{provider_type}' under the 'solace_agent_mesh.plugins' entry point."
+                )
+
+            provider_info = provider_info_entry.load()
+            class_path = provider_info.get("class_path")
+            if not class_path:
+                raise ValueError(
+                    f"Plugin '{provider_type}' is missing 'class_path' in its info dictionary."
+                )
+
+            module_path, class_name = class_path.rsplit(".", 1)
+            module = importlib.import_module(module_path)
+            provider_class = getattr(module, class_name)
+
+            if not issubclass(provider_class, BaseIdentityService):
+                raise TypeError(
+                    f"Provider class '{class_path}' does not inherit from BaseIdentityService."
+                )
+
+            log.info(f"Successfully loaded identity provider plugin: {provider_type}")
+            return provider_class(config)
+        except (ImportError, AttributeError, TypeError, ValueError) as e:
+            log.exception(
+                f"[IdentityFactory] Failed to load identity provider plugin '{provider_type}'. "
+                "Ensure the plugin is installed and the entry point is correct."
+            )
+            raise ValueError(f"Could not load identity provider plugin: {e}") from e

solace_agent_mesh/common/services/providers/__init__.py

@@ -0,0 +1,4 @@
+"""
+This package contains concrete implementations (providers) for the services
+defined in the parent 'services' package.
+"""