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

solace_agent_mesh/gateway/http_sse/repository/session_repository.py

@@ -6,22 +6,17 @@ from sqlalchemy.orm import Session as DBSession
 
 from ..shared.base_repository import PaginatedRepository
 from ..shared.pagination import PaginationParams
-from ..shared.types import PaginationInfo, SessionId, UserId
+from ..shared.types import SessionId, UserId
 from .entities import Session
 from .interfaces import ISessionRepository
-from .models import (
-    SessionModel,
-    CreateSessionModel,
-    UpdateSessionModel,
-)
+from .models import CreateSessionModel, SessionModel, UpdateSessionModel
 
 
 class SessionRepository(PaginatedRepository[SessionModel, Session], ISessionRepository):
     """SQLAlchemy implementation of session repository using BaseRepository."""
 
-    def __init__(self, db: DBSession):
+    def __init__(self):
         super().__init__(SessionModel, Session)
-        self.db = db
 
     @property
     def entity_name(self) -> str:
@@ -29,29 +24,30 @@ class SessionRepository(PaginatedRepository[SessionModel, Session], ISessionRepo
         return "session"
 
     def find_by_user(
-        self, user_id: UserId, pagination: PaginationParams | None = None
+        self, session: DBSession, user_id: UserId, pagination: PaginationParams | None = None
     ) -> list[Session]:
         """Find all sessions for a specific user."""
-        query = self.db.query(SessionModel).filter(SessionModel.user_id == user_id)
+        query = session.query(SessionModel).filter(SessionModel.user_id == user_id)
         query = query.order_by(SessionModel.updated_time.desc())
 
         if pagination:
-            offset = (pagination.page_number - 1) * pagination.page_size
-            query = query.offset(offset).limit(pagination.page_size)
+            query = query.offset(pagination.offset).limit(pagination.page_size)
 
         models = query.all()
         return [Session.model_validate(model) for model in models]
 
-    def count_by_user(self, user_id: UserId) -> int:
+    def count_by_user(self, session: DBSession, user_id: UserId) -> int:
         """Count total sessions for a specific user."""
-        return self.db.query(SessionModel).filter(SessionModel.user_id == user_id).count()
+        return (
+            session.query(SessionModel).filter(SessionModel.user_id == user_id).count()
+        )
 
     def find_user_session(
-        self, session_id: SessionId, user_id: UserId
+        self, session: DBSession, session_id: SessionId, user_id: UserId
     ) -> Session | None:
         """Find a specific session belonging to a user."""
         model = (
-            self.db.query(SessionModel)
+            session.query(SessionModel)
             .filter(
                 SessionModel.id == session_id,
                 SessionModel.user_id == user_id,
@@ -60,9 +56,11 @@ class SessionRepository(PaginatedRepository[SessionModel, Session], ISessionRepo
         )
         return Session.model_validate(model) if model else None
 
-    def save(self, session: Session) -> Session:
+    def save(self, db_session: DBSession, session: Session) -> Session:
         """Save or update a session."""
-        existing_model = self.db.query(SessionModel).filter(SessionModel.id == session.id).first()
+        existing_model = (
+            db_session.query(SessionModel).filter(SessionModel.id == session.id).first()
+        )
 
         if existing_model:
             update_model = UpdateSessionModel(
@@ -70,7 +68,9 @@ class SessionRepository(PaginatedRepository[SessionModel, Session], ISessionRepo
                 agent_id=session.agent_id,
                 updated_time=session.updated_time,
             )
-            return self.update(self.db, session.id, update_model.model_dump(exclude_none=True))
+            return self.update(
+                db_session, session.id, update_model.model_dump(exclude_none=True)
+            )
         else:
             create_model = CreateSessionModel(
                 id=session.id,
@@ -80,19 +80,23 @@ class SessionRepository(PaginatedRepository[SessionModel, Session], ISessionRepo
                 created_time=session.created_time,
                 updated_time=session.updated_time,
             )
-            return self.create(self.db, create_model.model_dump())
+            return self.create(db_session, create_model.model_dump())
 
-    def delete(self, session_id: SessionId, user_id: UserId) -> bool:
+    def delete(self, db_session: DBSession, session_id: SessionId, user_id: UserId) -> bool:
         """Delete a session belonging to a user."""
         # Check if session belongs to user first
-        session_model = self.db.query(SessionModel).filter(
-            SessionModel.id == session_id,
-            SessionModel.user_id == user_id,
-        ).first()
+        session_model = (
+            db_session.query(SessionModel)
+            .filter(
+                SessionModel.id == session_id,
+                SessionModel.user_id == user_id,
+            )
+            .first()
+        )
 
         if not session_model:
             return False
 
         # Use BaseRepository delete method
-        super().delete(self.db, session_id)
+        super().delete(db_session, session_id)
         return True

solace_agent_mesh/gateway/http_sse/repository/task_repository.py

@@ -4,7 +4,8 @@ Task repository implementation using SQLAlchemy.
 
 from sqlalchemy.orm import Session as DBSession
 
-from ..shared.types import PaginationInfo, PaginationParams, UserId
+from ..shared.pagination import PaginationParams
+from ..shared.types import UserId
 from .entities import Task, TaskEvent
 from .interfaces import ITaskRepository
 from .models import TaskEventModel, TaskModel
@@ -13,15 +14,11 @@ from .models import TaskEventModel, TaskModel
 class TaskRepository(ITaskRepository):
     """SQLAlchemy implementation of task repository."""
 
-    def __init__(self, db: DBSession):
-        self.db = db
-
-    def save_task(self, task: Task) -> Task:
+    def save_task(self, session: DBSession, task: Task) -> Task:
         """Create or update a task."""
-        model = self.db.query(TaskModel).filter(TaskModel.id == task.id).first()
+        model = session.query(TaskModel).filter(TaskModel.id == task.id).first()
 
         if model:
-            # Update existing
             model.end_time = task.end_time
             model.status = task.status
             model.total_input_tokens = task.total_input_tokens
@@ -29,7 +26,6 @@ class TaskRepository(ITaskRepository):
             model.total_cached_input_tokens = task.total_cached_input_tokens
             model.token_usage_details = task.token_usage_details
         else:
-            # Create new
             model = TaskModel(
                 id=task.id,
                 user_id=task.user_id,
@@ -42,13 +38,13 @@ class TaskRepository(ITaskRepository):
                 total_cached_input_tokens=task.total_cached_input_tokens,
                 token_usage_details=task.token_usage_details,
             )
-            self.db.add(model)
+            session.add(model)
 
-        self.db.commit()
-        self.db.refresh(model)
+        session.flush()
+        session.refresh(model)
         return self._task_model_to_entity(model)
 
-    def save_event(self, event: TaskEvent) -> TaskEvent:
+    def save_event(self, session: DBSession, event: TaskEvent) -> TaskEvent:
         """Save a task event."""
         model = TaskEventModel(
             id=event.id,
@@ -59,26 +55,26 @@ class TaskRepository(ITaskRepository):
             direction=event.direction,
             payload=event.payload,
         )
-        self.db.add(model)
-        self.db.commit()
-        self.db.refresh(model)
+        session.add(model)
+        session.flush()
+        session.refresh(model)
         return self._event_model_to_entity(model)
 
-    def find_by_id(self, task_id: str) -> Task | None:
+    def find_by_id(self, session: DBSession, task_id: str) -> Task | None:
         """Find a task by its ID."""
-        model = self.db.query(TaskModel).filter(TaskModel.id == task_id).first()
+        model = session.query(TaskModel).filter(TaskModel.id == task_id).first()
         return self._task_model_to_entity(model) if model else None
 
     def find_by_id_with_events(
-        self, task_id: str
+        self, session: DBSession, task_id: str
     ) -> tuple[Task, list[TaskEvent]] | None:
         """Find a task with all its events."""
-        task_model = self.db.query(TaskModel).filter(TaskModel.id == task_id).first()
+        task_model = session.query(TaskModel).filter(TaskModel.id == task_id).first()
         if not task_model:
             return None
 
         event_models = (
-            self.db.query(TaskEventModel)
+            session.query(TaskEventModel)
             .filter(TaskEventModel.task_id == task_id)
             .order_by(TaskEventModel.created_time.asc())
             .all()
@@ -90,36 +86,31 @@ class TaskRepository(ITaskRepository):
 
     def search(
         self,
+        session: DBSession,
         user_id: UserId,
         start_date: int | None = None,
         end_date: int | None = None,
-        search_query: str | None = None,
         pagination: PaginationParams | None = None,
     ) -> list[Task]:
         """Search for tasks with filters."""
-        query = self.db.query(TaskModel)
-        if user_id != "*":  # Allow wildcard for admin/system searches
+        query = session.query(TaskModel)
+        if user_id != "*":
             query = query.filter(TaskModel.user_id == user_id)
 
         if start_date:
             query = query.filter(TaskModel.start_time >= start_date)
         if end_date:
             query = query.filter(TaskModel.start_time <= end_date)
-        if search_query:
-            query = query.filter(
-                TaskModel.initial_request_text.ilike(f"%{search_query}%")
-            )
 
         query = query.order_by(TaskModel.start_time.desc())
 
         if pagination:
-            offset = (pagination.page - 1) * pagination.page_size
-            query = query.offset(offset).limit(pagination.page_size)
+            query = query.offset(pagination.offset).limit(pagination.page_size)
 
         models = query.all()
         return [self._task_model_to_entity(model) for model in models]
 
-    def delete_tasks_older_than(self, cutoff_time_ms: int, batch_size: int) -> int:
+    def delete_tasks_older_than(self, session: DBSession, cutoff_time_ms: int, batch_size: int) -> int:
         """
         Delete tasks (and their events via cascade) older than the cutoff time.
         Uses batch deletion to avoid long-running transactions.
@@ -130,38 +121,38 @@ class TaskRepository(ITaskRepository):
 
         Returns:
             Total number of tasks deleted
+
+        Note:
+            This method commits each batch internally due to the nature of batch processing.
+            Each batch is an atomic operation to prevent long-running transactions.
         """
         total_deleted = 0
-        
+
         while True:
-            # Find a batch of task IDs to delete
             task_ids_to_delete = (
-                self.db.query(TaskModel.id)
+                session.query(TaskModel.id)
                 .filter(TaskModel.start_time < cutoff_time_ms)
                 .limit(batch_size)
                 .all()
             )
-            
+
             if not task_ids_to_delete:
                 break
-            
-            # Extract IDs from the result tuples
+
             ids = [task_id[0] for task_id in task_ids_to_delete]
-            
-            # Delete this batch
+
             deleted_count = (
-                self.db.query(TaskModel)
+                session.query(TaskModel)
                 .filter(TaskModel.id.in_(ids))
                 .delete(synchronize_session=False)
             )
-            
-            self.db.commit()
+
+            session.commit()
             total_deleted += deleted_count
-            
-            # If we deleted fewer than batch_size, we're done
+
             if deleted_count < batch_size:
                 break
-        
+
         return total_deleted
 
     def _task_model_to_entity(self, model: TaskModel) -> Task:

solace_agent_mesh/gateway/http_sse/routers/agent_cards.py

@@ -1,14 +1,15 @@
 """
-API Router for agent discovery.
+API Router for agent discovery and management.
 """
 
 import logging
 from fastapi import APIRouter, Depends, HTTPException, status
-from typing import List
+from typing import Dict, List
 
 from ....common.agent_registry import AgentRegistry
 from a2a.types import AgentCard
-from ..dependencies import get_agent_registry
+from ..dependencies import get_agent_registry, get_sac_component
+from ..component import WebUIBackendComponent
 
 log = logging.getLogger(__name__)