solace-agent-mesh 1.4.12__py3-none-any.whl → 1.5.1__py3-none-any.whl

solace_agent_mesh/gateway/http_sse/services/session_service.py

@@ -1,21 +1,22 @@
+import logging
 import uuid
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Optional, List, Dict, Any
 
-from solace_ai_connector.common.log import log
 from sqlalchemy.orm import Session as DbSession
 
 from ..repository import (
-    IMessageRepository,
     ISessionRepository,
-    Message,
     Session,
-    SessionHistory,
 )
-from ..shared.enums import MessageType, SenderType
-from ..shared.types import PaginationInfo, SessionId, UserId
+from ..repository.chat_task_repository import ChatTaskRepository
+from ..repository.entities import ChatTask
+from ..shared.enums import SenderType
+from ..shared.types import SessionId, UserId
 from ..shared import now_epoch_ms
 from ..shared.pagination import PaginationParams, PaginatedResponse, get_pagination_or_default
 
+log = logging.getLogger(__name__)
+
 if TYPE_CHECKING:
     from ..component import WebUIBackendComponent
 
@@ -28,11 +29,10 @@ class SessionService:
         self.component = component
 
     def _get_repositories(self, db: DbSession):
-        """Create repositories for the given database session."""
-        from ..repository import SessionRepository, MessageRepository
+        """Create session repository for the given database session."""
+        from ..repository import SessionRepository
         session_repository = SessionRepository(db)
-        message_repository = MessageRepository(db)
-        return session_repository, message_repository
+        return session_repository
 
     def is_persistence_enabled(self) -> bool:
         """Checks if the service is configured with a persistent backend."""
@@ -54,18 +54,10 @@ class SessionService:
             raise ValueError("User ID cannot be empty")
 
         pagination = get_pagination_or_default(pagination)
-        session_repository, _ = self._get_repositories(db)
-
-        pagination_info = PaginationInfo(
-            page=pagination.page_number,
-            page_size=pagination.page_size,
-            total_items=0,
-            total_pages=0,
-            has_next=False,
-            has_previous=False,
-        )
+        session_repository = self._get_repositories(db)
 
-        sessions = session_repository.find_by_user(user_id, pagination_info)
+        # Pass pagination params directly - repository will handle offset calculation
+        sessions = session_repository.find_by_user(user_id, pagination)
         total_count = session_repository.count_by_user(user_id)
 
         return PaginatedResponse.create(sessions, total_count, pagination)
@@ -76,33 +68,9 @@ class SessionService:
         if not self._is_valid_session_id(session_id):
             return None
 
-        session_repository, _ = self._get_repositories(db)
+        session_repository = self._get_repositories(db)
         return session_repository.find_user_session(session_id, user_id)
 
-    def get_session_history(
-        self,
-        db: DbSession,
-        session_id: SessionId,
-        user_id: UserId,
-        pagination: PaginationInfo | None = None,
-    ) -> SessionHistory | None:
-        if not self._is_valid_session_id(session_id):
-            return None
-
-        session_repository, _ = self._get_repositories(db)
-        result = session_repository.find_user_session_with_messages(
-            session_id, user_id, pagination
-        )
-        if not result:
-            return None
-
-        session, messages = result
-        return SessionHistory(
-            session=session,
-            messages=messages,
-            total_message_count=len(messages),
-        )
-
     def create_session(
         self,
         db: DbSession,
@@ -131,7 +99,7 @@ class SessionService:
             updated_time=now_ms,
         )
 
-        session_repository, _ = self._get_repositories(db)
+        session_repository = self._get_repositories(db)
         created_session = session_repository.save(session)
         log.info("Created new session %s for user %s", created_session.id, user_id)
 
@@ -152,7 +120,7 @@ class SessionService:
         if len(name.strip()) > 255:
             raise ValueError("Session name cannot exceed 255 characters")
 
-        session_repository, _ = self._get_repositories(db)
+        session_repository = self._get_repositories(db)
         session = session_repository.find_user_session(session_id, user_id)
         if not session:
             return None
@@ -169,7 +137,7 @@ class SessionService:
         if not self._is_valid_session_id(session_id):
             raise ValueError("Invalid session ID")
 
-        session_repository, _ = self._get_repositories(db)
+        session_repository = self._get_repositories(db)
         session = session_repository.find_user_session(session_id, user_id)
         if not session:
             log.warning(
@@ -198,50 +166,151 @@ class SessionService:
 
         return True
 
-    def add_message_to_session(
+    def save_task(
         self,
         db: DbSession,
-        session_id: SessionId,
-        user_id: UserId,
-        message: str,
-        sender_type: SenderType,
-        sender_name: str,
-        agent_id: str | None = None,
-        message_type: MessageType = MessageType.TEXT,
-    ) -> Message:
-        if not self._is_valid_session_id(session_id):
-            raise ValueError("Invalid session ID")
-
-        if not message or message.strip() == "":
-            raise ValueError("Message cannot be empty")
-
-        session_repository, message_repository = self._get_repositories(db)
+        task_id: str,
+        session_id: str,
+        user_id: str,
+        user_message: Optional[str],
+        message_bubbles: str,  # JSON string (opaque)
+        task_metadata: Optional[str] = None  # JSON string (opaque)
+    ) -> ChatTask:
+        """
+        Save a complete task interaction.
+        
+        Args:
+            db: Database session
+            task_id: A2A task ID
+            session_id: Session ID
+            user_id: User ID
+            user_message: Original user input text
+            message_bubbles: Array of all message bubbles displayed during this task
+            task_metadata: Task-level metadata (status, feedback, agent name, etc.)
+            
+        Returns:
+            Saved ChatTask entity
+            
+        Raises:
+            ValueError: If session not found or validation fails
+        """
+        # Validate session exists and belongs to user
+        session_repository = self._get_repositories(db)
         session = session_repository.find_user_session(session_id, user_id)
         if not session:
-            session = self.create_session(
-                db=db,
-                user_id=user_id,
-                agent_id=agent_id,
-                session_id=session_id,
-            )
-
-        message_entity = Message(
-            id=str(uuid.uuid4()),
+            raise ValueError(f"Session {session_id} not found for user {user_id}")
+        
+        # Create task entity - pass strings directly
+        task = ChatTask(
+            id=task_id,
             session_id=session_id,
-            message=message.strip(),
-            sender_type=sender_type,
-            sender_name=sender_name,
-            message_type=message_type,
+            user_id=user_id,
+            user_message=user_message,
+            message_bubbles=message_bubbles,  # Already a string
+            task_metadata=task_metadata,      # Already a string
             created_time=now_epoch_ms(),
+            updated_time=None
         )
-
-        saved_message = message_repository.save(message_entity)
-
+        
+        # Save via repository
+        task_repo = ChatTaskRepository(db)
+        saved_task = task_repo.save(task)
+        
+        # Update session activity
         session.mark_activity()
         session_repository.save(session)
+        
+        log.info(f"Saved task {task_id} for session {session_id}")
+        return saved_task
 
-        log.info("Added message to session %s from %s", session_id, sender_name)
-        return saved_message
+    def get_session_tasks(
+        self,
+        db: DbSession,
+        session_id: str,
+        user_id: str
+    ) -> List[ChatTask]:
+        """
+        Get all tasks for a session.
+        
+        Args:
+            db: Database session
+            session_id: Session ID
+            user_id: User ID
+            
+        Returns:
+            List of ChatTask entities in chronological order
+            
+        Raises:
+            ValueError: If session not found
+        """
+        # Validate session exists and belongs to user
+        session_repository = self._get_repositories(db)
+        session = session_repository.find_user_session(session_id, user_id)
+        if not session:
+            raise ValueError(f"Session {session_id} not found for user {user_id}")
+        
+        # Load tasks
+        task_repo = ChatTaskRepository(db)
+        return task_repo.find_by_session(session_id, user_id)
+
+    def get_session_messages_from_tasks(
+        self,
+        db: DbSession,
+        session_id: str,
+        user_id: str
+    ) -> List[Dict[str, Any]]:
+        """
+        Get session messages by flattening task message_bubbles.
+        This provides backward compatibility with the old message-based API.
+        
+        Args:
+            db: Database session
+            session_id: Session ID
+            user_id: User ID
+            
+        Returns:
+            List of message dictionaries flattened from tasks
+            
+        Raises:
+            ValueError: If session not found
+        """
+        # Load tasks
+        tasks = self.get_session_tasks(db, session_id, user_id)
+        
+        # Flatten message_bubbles from all tasks
+        messages = []
+        for task in tasks:
+            import json
+            message_bubbles = json.loads(task.message_bubbles) if isinstance(task.message_bubbles, str) else task.message_bubbles
+            
+            for bubble in message_bubbles:
+                # Determine sender type from bubble type
+                bubble_type = bubble.get("type", "agent")
+                sender_type = "user" if bubble_type == "user" else "agent"
+                
+                # Get sender name
+                if bubble_type == "user":
+                    sender_name = user_id
+                else:
+                    # Try to get agent name from task metadata, fallback to "agent"
+                    sender_name = "agent"
+                    if task.task_metadata:
+                        task_metadata = json.loads(task.task_metadata) if isinstance(task.task_metadata, str) else task.task_metadata
+                        sender_name = task_metadata.get("agent_name", "agent")
+                
+                # Create message dictionary
+                message = {
+                    "id": bubble.get("id", str(uuid.uuid4())),
+                    "session_id": session_id,
+                    "message": bubble.get("text", ""),
+                    "sender_type": sender_type,
+                    "sender_name": sender_name,
+                    "message_type": "text",
+                    "created_time": task.created_time
+                }
+                messages.append(message)
+        
+        return messages
 
     def _is_valid_session_id(self, session_id: SessionId) -> bool:
         return (
@@ -289,4 +358,4 @@ class SessionService:
                 "Failed to publish session deletion event to agent %s: %s",
                 agent_id,
                 e,
-            )
+            )

solace_agent_mesh/gateway/http_sse/services/task_logger_service.py

@@ -0,0 +1,318 @@
+"""
+Service for logging A2A tasks and events to the database.
+"""
+
+import copy
+import logging
+import uuid
+from typing import Any, Callable, Dict, Union
+
+from a2a.types import (
+    A2ARequest,
+    JSONRPCError,
+    JSONRPCResponse,
+    Task as A2ATask,
+    TaskArtifactUpdateEvent,
+    TaskStatusUpdateEvent,
+)
+from sqlalchemy.orm import Session as DBSession
+
+from ....common import a2a
+from ..repository.entities import Task, TaskEvent
+from ..repository.task_repository import TaskRepository
+from ..shared import now_epoch_ms
+
+log = logging.getLogger(__name__)
+
+class TaskLoggerService:
+    """Service for logging A2A tasks and events to the database."""
+
+    def __init__(
+        self, session_factory: Callable[[], DBSession] | None, config: Dict[str, Any]
+    ):
+        self.session_factory = session_factory
+        self.config = config
+        self.log_identifier = "[TaskLoggerService]"
+        log.info(f"{self.log_identifier} Initialized.")
+
+    def log_event(self, event_data: Dict[str, Any]):
+        """
+        Parses a raw A2A message and logs it as a task event.
+        Creates or updates the master task record as needed.
+        """
+        if not self.config.get("enabled", False):
+            return
+
+        if not self.session_factory:
+            log.warning(
+                f"{self.log_identifier} Task logging is enabled but no database is configured. Skipping event."
+            )
+            return
+
+        topic = event_data.get("topic")
+        payload = event_data.get("payload")
+        user_properties = event_data.get("user_properties", {})
+
+        if not topic or not payload:
+            log.warning(
+                f"{self.log_identifier} Received event with missing topic or payload."
+            )
+            return
+
+        if "discovery" in topic:
+            # Ignore discovery messages
+            return
+
+        # Parse the event into a Pydantic model first.
+        parsed_event = self._parse_a2a_event(topic, payload)
+        if parsed_event is None:
+            # Parsing failed or event should be ignored.
+            return
+
+        db = self.session_factory()
+        try:
+            repo = TaskRepository(db)
+
+            # Infer details from the parsed event
+            direction, task_id, user_id = self._infer_event_details(
+                topic, parsed_event, user_properties
+            )
+
+            if not task_id:
+                log.debug(
+                    f"{self.log_identifier} Could not determine task_id for event on topic {topic}. Skipping."
+                )
+                return
+
+            # Check if we should log this event type
+            if not self._should_log_event(topic, parsed_event):
+                log.debug(
+                    f"{self.log_identifier} Event on topic {topic} is configured to be skipped."
+                )
+                return
+
+            # Sanitize the original raw payload before storing
+            sanitized_payload = self._sanitize_payload(payload)
+
+            # Check for existing task or create a new one
+            task = repo.find_by_id(task_id)
+            if not task:
+                if direction == "request":
+                    initial_text = self._extract_initial_text(parsed_event)
+                    new_task = Task(
+                        id=task_id,
+                        user_id=user_id or "unknown",
+                        start_time=now_epoch_ms(),
+                        initial_request_text=(
+                            initial_text[:1024] if initial_text else None
+                        ),  # Truncate
+                    )
+                    repo.save_task(new_task)
+                    log.info(
+                        f"{self.log_identifier} Created new task record for ID: {task_id}"
+                    )
+                else:
+                    # We received an event for a task we haven't seen the start of.
+                    # This can happen if the logger starts mid-conversation. Create a placeholder.
+                    placeholder_task = Task(
+                        id=task_id,
+                        user_id=user_id or "unknown",
+                        start_time=now_epoch_ms(),
+                        initial_request_text="[Task started before logger was active]",
+                    )
+                    repo.save_task(placeholder_task)
+                    log.info(
+                        f"{self.log_identifier} Created placeholder task record for ID: {task_id}"
+                    )
+
+            # Create and save the event using the sanitized raw payload
+            task_event = TaskEvent(
+                id=str(uuid.uuid4()),
+                task_id=task_id,
+                user_id=user_id,
+                created_time=now_epoch_ms(),
+                topic=topic,
+                direction=direction,
+                payload=sanitized_payload,
+            )
+            repo.save_event(task_event)
+
+            # If it's a final event, update the master task record
+            final_status = self._get_final_status(parsed_event)
+            if final_status:
+                task_to_update = repo.find_by_id(task_id)
+                if task_to_update:
+                    task_to_update.end_time = now_epoch_ms()
+                    task_to_update.status = final_status
+                    
+                    # Extract and store token usage if present
+                    if isinstance(parsed_event, A2ATask) and parsed_event.metadata:
+                        token_usage = parsed_event.metadata.get("token_usage")
+                        if token_usage and isinstance(token_usage, dict):
+                            task_to_update.total_input_tokens = token_usage.get("total_input_tokens")
+                            task_to_update.total_output_tokens = token_usage.get("total_output_tokens")
+                            task_to_update.total_cached_input_tokens = token_usage.get("total_cached_input_tokens")
+                            task_to_update.token_usage_details = token_usage
+                            log.info(
+                                f"{self.log_identifier} Stored token usage for task {task_id}: "
+                                f"input={token_usage.get('total_input_tokens')}, "
+                                f"output={token_usage.get('total_output_tokens')}, "
+                                f"cached={token_usage.get('total_cached_input_tokens')}"
+                            )
+                    
+                    repo.save_task(task_to_update)
+                    log.info(
+                        f"{self.log_identifier} Finalized task record for ID: {task_id} with status: {final_status}"
+                    )
+            db.commit()
+        except Exception as e:
+            log.exception(
+                f"{self.log_identifier} Error logging event on topic {topic}: {e}"
+            )
+            db.rollback()
+        finally:
+            db.close()
+
+    def _parse_a2a_event(self, topic: str, payload: dict) -> Union[
+        A2ARequest,
+        A2ATask,
+        TaskStatusUpdateEvent,
+        TaskArtifactUpdateEvent,
+        JSONRPCError,
+        None,
+    ]:
+        """
+        Safely parses a raw A2A message payload into a Pydantic model.
+        Returns the parsed model or None if parsing fails or is not applicable.
+        """
+        # Ignore discovery messages
+        if "/discovery/agentcards" in topic:
+            return None
+
+        try:
+            # Check if it's a response (has 'result' or 'error')
+            if "result" in payload or "error" in payload:
+                rpc_response = JSONRPCResponse.model_validate(payload)
+                error = a2a.get_response_error(rpc_response)
+                if error:
+                    return error
+                result = a2a.get_response_result(rpc_response)
+                if result:
+                    # The result is already a parsed Pydantic model
+                    return result
+            # Check if it's a request
+            elif "method" in payload:
+                return A2ARequest.model_validate(payload)
+
+            log.warning(
+                f"{self.log_identifier} Payload for topic '{topic}' is not a recognizable JSON-RPC request or response. Payload: {payload}"
+            )
+            return None
+
+        except Exception as e:
+            log.error(
+                f"{self.log_identifier} Failed to parse A2A event for topic '{topic}': {e}. Payload: {payload}"
+            )
+            return None
+
+    def _infer_event_details(
+        self, topic: str, parsed_event: Any, user_props: Dict | None
+    ) -> tuple[str, str | None, str | None]:
+        """Infers direction, task_id, and user_id from a parsed A2A event."""
+        direction = "unknown"
+        task_id = None
+        # Ensure user_props is a dict, not None
+        user_props = user_props or {}
+        user_id = user_props.get("userId")
+
+        if isinstance(parsed_event, A2ARequest):
+            direction = "request"
+            task_id = a2a.get_request_id(parsed_event)
+        elif isinstance(
+            parsed_event, (A2ATask, TaskStatusUpdateEvent, TaskArtifactUpdateEvent)
+        ):
+            direction = "response" if isinstance(parsed_event, A2ATask) else "status"
+            task_id = getattr(parsed_event, "task_id", None) or getattr(
+                parsed_event, "id", None
+            )
+        elif isinstance(parsed_event, JSONRPCError):
+            direction = "error"
+            if isinstance(parsed_event.data, dict):
+                task_id = parsed_event.data.get("taskId")
+
+        if not user_id:
+            user_config = user_props.get("a2aUserConfig") or user_props.get("a2a_user_config")
+            if isinstance(user_config, dict):
+                user_profile = user_config.get("user_profile", {})
+                if isinstance(user_profile, dict):
+                    user_id = user_profile.get("id")
+
+        return direction, str(task_id) if task_id else None, user_id
+
+    def _extract_initial_text(self, parsed_event: Any) -> str | None:
+        """Extracts the initial text from a send message request."""
+        try:
+            if isinstance(parsed_event, A2ARequest):
+                message = a2a.get_message_from_send_request(parsed_event)
+                if message:
+                    return a2a.get_text_from_message(message)
+        except Exception:
+            return None
+        return None
+
+    def _get_final_status(self, parsed_event: Any) -> str | None:
+        """Checks if a parsed event represents a final task status and returns the state."""
+        if isinstance(parsed_event, A2ATask):
+            return parsed_event.status.state.value
+        elif isinstance(parsed_event, JSONRPCError):
+            return "failed"
+        return None
+
+    def _should_log_event(self, topic: str, parsed_event: Any) -> bool:
+        """Determines if an event should be logged based on configuration."""
+        if not self.config.get("log_status_updates", True):
+            if "status" in topic:
+                return False
+        if not self.config.get("log_artifact_events", True):
+            if isinstance(parsed_event, TaskArtifactUpdateEvent):
+                return False
+        return True
+
+    def _sanitize_payload(self, payload: Dict) -> Dict:
+        """Strips or truncates file content from payload based on configuration."""
+        new_payload = copy.deepcopy(payload)
+
+        def walk_and_sanitize(node):
+            if isinstance(node, dict):
+                for key, value in list(node.items()):
+                    if key == "parts" and isinstance(value, list):
+                        new_parts = []
+                        for part in value:
+                            if isinstance(part, dict) and "file" in part:
+                                if not self.config.get("log_file_parts", True):
+                                    continue  # Skip this part entirely
+
+                                file_dict = part.get("file")
+                                if isinstance(file_dict, dict) and "bytes" in file_dict:
+                                    max_bytes = self.config.get(
+                                        "max_file_part_size_bytes", 102400
+                                    )
+                                    file_bytes_b64 = file_dict.get("bytes")
+                                    if isinstance(file_bytes_b64, str):
+                                        if (len(file_bytes_b64) * 3 / 4) > max_bytes:
+                                            file_dict["bytes"] = (
+                                                f"[Content stripped, size > {max_bytes} bytes]"
+                                            )
+                                new_parts.append(part)
+                            else:
+                                walk_and_sanitize(part)
+                                new_parts.append(part)
+                        node["parts"] = new_parts
+                    else:
+                        walk_and_sanitize(value)
+            elif isinstance(node, list):
+                for item in node:
+                    walk_and_sanitize(item)
+
+        walk_and_sanitize(new_payload)
+        return new_payload

solace_agent_mesh/gateway/http_sse/services/task_service.py

@@ -3,15 +3,16 @@ Service layer for handling A2A task submissions and cancellations.
 Uses CoreA2AService for logic and a provided function for publishing.
 """
 
+import logging
 import threading
 from typing import Callable, Dict, Optional
 
-from solace_ai_connector.common.log import log
-
 from ....common import a2a
 from ....gateway.http_sse.sse_manager import SSEManager
 from ....core_a2a.service import CoreA2AService
 
+log = logging.getLogger(__name__)
+
 PublishFunc = Callable[[str, Dict, Optional[Dict]], None]
 
 

solace_agent_mesh/gateway/http_sse/session_manager.py

@@ -2,13 +2,14 @@
 Manages web user sessions and mapping to A2A Client IDs.
 """
 
+import logging
 import uuid
 from collections.abc import Callable
 from typing import Any
 
-from solace_ai_connector.common.log import log
 from starlette.requests import Request
 
+log = logging.getLogger(__name__)
 
 SESSION_KEY_CLIENT_ID = "a2a_client_id"
 SESSION_KEY_SESSION_ID = "a2a_session_id"