solace-agent-mesh 1.1.0__py3-none-any.whl → 1.3.1__py3-none-any.whl

solace_agent_mesh/gateway/http_sse/component.py

@@ -3,34 +3,37 @@ Custom Solace AI Connector Component to host the FastAPI backend for the Web UI.
 """
 
 import asyncio
-import queue
-import uuid
 import json
+import queue
 import re
 import threading
-from typing import Any, Dict, Optional, List, Tuple, Union, Set
+import uuid
 from datetime import datetime, timezone
-from fastapi import UploadFile, Request as FastAPIRequest
+from typing import Any
 
 import uvicorn
-from fastapi import FastAPI
-
+from fastapi import FastAPI, UploadFile
+from fastapi import Request as FastAPIRequest
 from solace_ai_connector.common.log import log
+from solace_ai_connector.components.inputs_outputs.broker_input import BrokerInput
 from solace_ai_connector.flow.app import App as SACApp
-from solace_ai_connector.components.inputs_outputs.broker_input import (
-    BrokerInput,
-)
 
-from ...gateway.http_sse.sse_manager import SSEManager
-
-from .components import VisualizationForwarderComponent
-from ...gateway.http_sse.session_manager import SessionManager
-from ...gateway.base.component import BaseGatewayComponent
 from ...common.agent_registry import AgentRegistry
 from ...core_a2a.service import CoreA2AService
-from google.adk.artifacts import BaseArtifactService
+from ...gateway.base.component import BaseGatewayComponent
+from ...gateway.http_sse.session_manager import SessionManager
+from ...gateway.http_sse.sse_manager import SSEManager
+from .components import VisualizationForwarderComponent
+from .infrastructure.persistence_service import PersistenceService
+
+try:
+    from google.adk.artifacts import BaseArtifactService
+except ImportError:
+
+    class BaseArtifactService:
+        pass
+
 
-from ...common.a2a.types import ContentPart
 from a2a.types import (
     A2ARequest,
     AgentCard,
@@ -40,10 +43,15 @@ from a2a.types import (
     TaskArtifactUpdateEvent,
     TaskStatusUpdateEvent,
 )
+
 from ...common import a2a
-from ...agent.utils.artifact_helpers import save_artifact_with_metadata
+from ...common.a2a.types import ContentPart
 from ...common.middleware.config_resolver import ConfigResolver
-
+from ...common.utils.embeds import (
+    EARLY_EMBED_TYPES,
+    evaluate_embed,
+    resolve_embeds_in_string,
+)
 
 info = {
     "class_name": "WebUIBackendComponent",
@@ -114,27 +122,42 @@ class WebUIBackendComponent(BaseGatewayComponent):
 
         self.sse_manager = SSEManager(max_queue_size=sse_max_queue_size)
 
+        session_config = self._resolve_session_config()
+        if session_config.get("type") == "sql":
+            # SQL type explicitly configured - database_url is required
+            database_url = session_config.get("database_url")
+            if not database_url:
+                raise ValueError(
+                    f"{self.log_identifier} Session service type is 'sql' but no database_url provided. "
+                    "Please provide a database_url in the session_service configuration or use type 'memory'."
+                )
+            self.persistence_service = PersistenceService(database_url)
+        else:
+            # Memory storage or no explicit configuration - no persistence service needed
+            self.persistence_service = None
+
         component_config = self.get_config("component_config", {})
         app_config = component_config.get("app_config", {})
 
         self.session_manager = SessionManager(
             secret_key=self.session_secret_key,
             app_config=app_config,
+            persistence_service=self.persistence_service,
         )
 
-        self.fastapi_app: Optional[FastAPI] = None
-        self.uvicorn_server: Optional[uvicorn.Server] = None
-        self.fastapi_thread: Optional[threading.Thread] = None
-        self.fastapi_event_loop: Optional[asyncio.AbstractEventLoop] = None
+        self.fastapi_app: FastAPI | None = None
+        self.uvicorn_server: uvicorn.Server | None = None
+        self.fastapi_thread: threading.Thread | None = None
+        self.fastapi_event_loop: asyncio.AbstractEventLoop | None = None
 
-        self._visualization_internal_app: Optional[SACApp] = None
-        self._visualization_broker_input: Optional[BrokerInput] = None
+        self._visualization_internal_app: SACApp | None = None
+        self._visualization_broker_input: BrokerInput | None = None
         self._visualization_message_queue: queue.Queue = queue.Queue(maxsize=200)
-        self._active_visualization_streams: Dict[str, Dict[str, Any]] = {}
-        self._visualization_locks: Dict[asyncio.AbstractEventLoop, asyncio.Lock] = {}
+        self._active_visualization_streams: dict[str, dict[str, Any]] = {}
+        self._visualization_locks: dict[asyncio.AbstractEventLoop, asyncio.Lock] = {}
         self._visualization_locks_lock = threading.Lock()
-        self._global_visualization_subscriptions: Dict[str, int] = {}
-        self._visualization_processor_task: Optional[asyncio.Task] = None
+        self._global_visualization_subscriptions: dict[str, int] = {}
+        self._visualization_processor_task: asyncio.Task | None = None
 
         log.info("%s Web UI Backend Component initialized.", self.log_identifier)
 
@@ -297,6 +320,37 @@ class WebUIBackendComponent(BaseGatewayComponent):
             self._visualization_broker_input = None
             raise
 
+    def _resolve_session_config(self) -> dict:
+        """
+        Resolve session service configuration with backward compatibility.
+
+        Priority order:
+        1. Component-specific session_service config (new approach)
+        2. Shared default_session_service config (deprecated, with warning)
+        3. Hardcoded default (SQLite for Web UI)
+        """
+        # Check component-specific session_service config first
+        component_session_config = self.get_config("session_service")
+        if component_session_config:
+            log.debug("Using component-specific session_service configuration")
+            return component_session_config
+
+        # Backward compatibility: check shared config
+        shared_session_config = self.get_config("default_session_service")
+        if shared_session_config:
+            log.warning(
+                "Using session_service from shared config is deprecated. "
+                "Move to component-specific configuration in app_config.session_service"
+            )
+            return shared_session_config
+
+        # Default configuration for Web UI (backward compatibility)
+        default_config = {"type": "memory", "default_behavior": "PERSISTENT"}
+        log.info(
+            "Using default memory session configuration for Web UI (backward compatibility)"
+        )
+        return default_config
+
     async def _visualization_message_processor_loop(self) -> None:
         """
         Asynchronously consumes messages from the _visualization_message_queue,
@@ -564,7 +618,7 @@ class WebUIBackendComponent(BaseGatewayComponent):
                     if not hasattr(
                         self._visualization_broker_input, "add_subscription"
                     ) or not callable(
-                        getattr(self._visualization_broker_input, "add_subscription")
+                        self._visualization_broker_input.add_subscription
                     ):
                         log.error(
                             "%s Visualization BrokerInput does not support dynamic 'add_subscription'. "
@@ -679,9 +733,7 @@ class WebUIBackendComponent(BaseGatewayComponent):
             try:
                 if not hasattr(
                     self._visualization_broker_input, "remove_subscription"
-                ) or not callable(
-                    getattr(self._visualization_broker_input, "remove_subscription")
-                ):
+                ) or not callable(self._visualization_broker_input.remove_subscription):
                     log.error(
                         "%s Visualization BrokerInput does not support dynamic 'remove_subscription'. "
                         "Please upgrade the 'solace-ai-connector' module. Cannot remove subscription '%s'.",
@@ -768,7 +820,7 @@ class WebUIBackendComponent(BaseGatewayComponent):
 
     async def _extract_initial_claims(
         self, external_event_data: Any
-    ) -> Optional[Dict[str, Any]]:
+    ) -> dict[str, Any] | None:
         """
         Extracts initial identity claims from the incoming external event.
         For the WebUI, this means inspecting the FastAPIRequest.
@@ -824,16 +876,12 @@ class WebUIBackendComponent(BaseGatewayComponent):
             return
 
         try:
-            from ...gateway.http_sse.main import (
-                app as fastapi_app_instance,
-            )
-            from ...gateway.http_sse.main import (
-                setup_dependencies,
-            )
+            from ...gateway.http_sse.main import app as fastapi_app_instance
+            from ...gateway.http_sse.main import setup_dependencies
 
             self.fastapi_app = fastapi_app_instance
 
-            setup_dependencies(self)
+            setup_dependencies(self, self.persistence_service)
 
             port = (
                 self.fastapi_https_port
@@ -936,7 +984,7 @@ class WebUIBackendComponent(BaseGatewayComponent):
             raise
 
     def publish_a2a(
-        self, topic: str, payload: Dict, user_properties: Optional[Dict] = None
+        self, topic: str, payload: dict, user_properties: dict | None = None
     ):
         """
         Publishes an A2A message using the SAC App's send_message method.
@@ -994,8 +1042,8 @@ class WebUIBackendComponent(BaseGatewayComponent):
         log.info("%s Visualization resources cleaned up.", self.log_identifier)
 
     def _infer_visualization_event_details(
-        self, topic: str, payload: Dict[str, Any]
-    ) -> Dict[str, Any]:
+        self, topic: str, payload: dict[str, Any]
+    ) -> dict[str, Any]:
         """
         Infers details for the visualization SSE payload from the Solace topic and A2A message.
         This version is updated to parse the official A2A SDK message formats.
@@ -1135,19 +1183,19 @@ class WebUIBackendComponent(BaseGatewayComponent):
                 (summary_str[:100] + "...") if len(summary_str) > 100 else summary_str
             )
         except Exception:
-            details["payload_summary"][
-                "params_preview"
-            ] = "[Could not serialize payload]"
+            details["payload_summary"]["params_preview"] = (
+                "[Could not serialize payload]"
+            )
 
         return details
 
     def _extract_involved_agents_for_viz(
-        self, topic: str, payload_dict: Dict[str, Any]
-    ) -> Set[str]:
+        self, topic: str, payload_dict: dict[str, Any]
+    ) -> set[str]:
         """
         Extracts agent names involved in a message from its topic and payload.
         """
-        agents: Set[str] = set()
+        agents: set[str] = set()
         log_id_prefix = f"{self.log_identifier}[ExtractAgentsViz]"
 
         topic_agent_match = re.match(
@@ -1274,13 +1322,13 @@ class WebUIBackendComponent(BaseGatewayComponent):
         """Returns the unique identifier for this gateway instance."""
         return self.gateway_id
 
-    def get_cors_origins(self) -> List[str]:
+    def get_cors_origins(self) -> list[str]:
         return self.cors_allowed_origins
 
-    def get_shared_artifact_service(self) -> Optional[BaseArtifactService]:
+    def get_shared_artifact_service(self) -> BaseArtifactService | None:
         return self.shared_artifact_service
 
-    def get_embed_config(self) -> Dict[str, Any]:
+    def get_embed_config(self) -> dict[str, Any]:
         """Returns embed-related configuration needed by dependencies."""
         return {
             "enable_embed_resolution": self.enable_embed_resolution,
@@ -1296,6 +1344,52 @@ class WebUIBackendComponent(BaseGatewayComponent):
         """Returns the instance of the ConfigResolver."""
         return self._config_resolver
 
+    async def _resolve_embeds_for_persistence(
+        self, message_content: str, session_id: str, user_id: str, log_identifier: str
+    ) -> str:
+        """
+        Resolves embeds in a message for database storage.
+        Returns the resolved text.
+
+        Args:
+            message_content: The message text that may contain embeds
+            session_id: The A2A session ID
+            user_id: The user ID
+            log_identifier: Logging identifier
+
+        Returns:
+            The message with embeds resolved (or original if resolution fails)
+        """
+        try:
+            embed_context = {
+                "artifact_service": self.shared_artifact_service,
+                "session_context": {
+                    "app_name": self.gateway_id,
+                    "user_id": user_id,
+                    "session_id": session_id,
+                },
+                "config": self.get_embed_config(),
+            }
+
+            resolved_text, _, _ = await resolve_embeds_in_string(
+                text=message_content,
+                context=embed_context,
+                resolver_func=evaluate_embed,
+                types_to_resolve=EARLY_EMBED_TYPES,
+                log_identifier=log_identifier,
+                config=embed_context["config"],
+            )
+
+            return resolved_text
+
+        except Exception as e:
+            log.warning(
+                "%s Error resolving embeds for storage: %s. Using original message.",
+                log_identifier,
+                e,
+            )
+            return message_content
+
     def _start_listener(self) -> None:
         """
         GDK Hook: Starts the FastAPI/Uvicorn server.
@@ -1317,8 +1411,8 @@ class WebUIBackendComponent(BaseGatewayComponent):
         pass
 
     async def _translate_external_input(
-        self, external_event_data: Dict[str, Any]
-    ) -> Tuple[str, List[ContentPart], Dict[str, Any]]:
+        self, external_event_data: dict[str, Any]
+    ) -> tuple[str, list[ContentPart], dict[str, Any]]:
         """
         Translates raw HTTP request data (from FastAPI form) into A2A task parameters.
 
@@ -1342,10 +1436,9 @@ class WebUIBackendComponent(BaseGatewayComponent):
 
         target_agent_name: str = external_event_data.get("agent_name")
         user_message: str = external_event_data.get("message", "")
-        files: Optional[List[UploadFile]] = external_event_data.get("files")
+        files: list[UploadFile] | None = external_event_data.get("files")
         client_id: str = external_event_data.get("client_id")
         a2a_session_id: str = external_event_data.get("a2a_session_id")
-
         if not target_agent_name:
             raise ValueError("Target agent name is missing in external_event_data.")
         if not client_id or not a2a_session_id:
@@ -1353,7 +1446,7 @@ class WebUIBackendComponent(BaseGatewayComponent):
                 "Client ID or A2A Session ID is missing in external_event_data."
             )
 
-        a2a_parts: List[ContentPart] = []
+        a2a_parts: list[ContentPart] = []
 
         if files:
             for upload_file in files:
@@ -1413,18 +1506,24 @@ class WebUIBackendComponent(BaseGatewayComponent):
 
     async def _send_update_to_external(
         self,
-        external_request_context: Dict[str, Any],
-        event_data: Union[TaskStatusUpdateEvent, TaskArtifactUpdateEvent],
+        external_request_context: dict[str, Any],
+        event_data: TaskStatusUpdateEvent | TaskArtifactUpdateEvent,
         is_final_chunk_of_update: bool,
     ) -> None:
         """
         Sends an intermediate update (TaskStatusUpdateEvent or TaskArtifactUpdateEvent)
-        to the external platform (Web UI via SSE).
+        to the external platform (Web UI via SSE) and stores agent messages in the database.
         """
         log_id_prefix = f"{self.log_identifier}[SendUpdate]"
         sse_task_id = external_request_context.get("a2a_task_id_for_event")
         a2a_task_id = event_data.task_id
 
+        log.debug(
+            "%s _send_update_to_external called with event_type: %s",
+            log_id_prefix,
+            type(event_data).__name__,
+        )
+
         if not sse_task_id:
             log.error(
                 "%s Cannot send update: 'a2a_task_id_for_event' missing from external_request_context.",
@@ -1459,6 +1558,10 @@ class WebUIBackendComponent(BaseGatewayComponent):
                 sse_event_type,
                 a2a_task_id,
             )
+
+            # Note: Agent message storage is handled in _send_final_response_to_external
+            # to avoid duplicate storage of intermediate status updates
+
         except Exception as e:
             log.exception(
                 "%s Failed to send %s via SSE for A2A Task ID %s: %s",
@@ -1469,7 +1572,7 @@ class WebUIBackendComponent(BaseGatewayComponent):
             )
 
     async def _send_final_response_to_external(
-        self, external_request_context: Dict[str, Any], task_data: Task
+        self, external_request_context: dict[str, Any], task_data: Task
     ) -> None:
         """
         Sends the final A2A Task result to the external platform (Web UI via SSE).
@@ -1478,6 +1581,8 @@ class WebUIBackendComponent(BaseGatewayComponent):
         sse_task_id = external_request_context.get("a2a_task_id_for_event")
         a2a_task_id = task_data.id
 
+        log.debug("%s _send_final_response_to_external called", log_id_prefix)
+
         if not sse_task_id:
             log.error(
                 "%s Cannot send final response: 'a2a_task_id_for_event' missing from external_request_context.",
@@ -1506,6 +1611,63 @@ class WebUIBackendComponent(BaseGatewayComponent):
                 log_id_prefix,
                 a2a_task_id,
             )
+
+            # Store final agent response in persistence layer if available
+            if hasattr(self, "persistence_service") and self.persistence_service:
+                try:
+                    session_id = external_request_context.get("a2a_session_id")
+                    user_id = external_request_context.get("user_id_for_a2a")
+                    agent_name = external_request_context.get(
+                        "target_agent_name", "agent"
+                    )
+
+                    # Extract message content from the task status
+                    message_text = ""
+                    if task_data.status and task_data.status.message:
+                        parts = a2a.get_parts_from_message(task_data.status.message)
+                        for part in parts:
+                            if hasattr(part, "text") and part.text:
+                                if message_text:
+                                    message_text += "\n"
+                                message_text += part.text
+
+                    log.info(
+                        "%s Final agent response storage debug - session_id: %s, user_id: %s, message_text: '%s', parts_count: %s",
+                        log_id_prefix,
+                        session_id,
+                        user_id,
+                        message_text[:100] if message_text else None,
+                        len(a2a.get_parts_from_message(task_data.status.message))
+                        if task_data.status and task_data.status.message
+                        else 0,
+                    )
+
+                    if message_text and session_id and user_id:
+                        from .dependencies import get_session_service
+                        from .shared.enums import SenderType
+
+                        session_service = get_session_service(self)
+                        session_service.add_message_to_session(
+                            session_id=session_id,
+                            user_id=user_id,
+                            message=message_text,
+                            sender_type=SenderType.AGENT,
+                            sender_name=agent_name,
+                            agent_id=agent_name,
+                        )
+                        log.info(
+                            "%s Final agent response stored in session %s",
+                            log_id_prefix,
+                            session_id,
+                        )
+                except Exception as storage_error:
+                    log.warning(
+                        "%s Failed to store final agent response: %s",
+                        log_id_prefix,
+                        storage_error,
+                    )
+                    # Don't fail the SSE send if storage fails
+
         except Exception as e:
             log.exception(
                 "%s Failed to send final_response via SSE for A2A Task ID %s: %s",
@@ -1522,7 +1684,7 @@ class WebUIBackendComponent(BaseGatewayComponent):
             )
 
     async def _send_error_to_external(
-        self, external_request_context: Dict[str, Any], error_data: JSONRPCError
+        self, external_request_context: dict[str, Any], error_data: JSONRPCError
     ) -> None:
         """
         Sends an error notification to the external platform (Web UI via SSE).