solace-agent-mesh 1.4.6__py3-none-any.whl → 1.4.8__py3-none-any.whl

solace_agent_mesh/gateway/http_sse/repository/message_repository.py

@@ -4,19 +4,26 @@ Message repository implementation using SQLAlchemy.
 
 from sqlalchemy.orm import Session as DBSession
 
+from ..shared.base_repository import PaginatedRepository
 from ..shared.enums import MessageType, SenderType
 from ..shared.types import PaginationInfo, SessionId
 from .entities import Message
 from .interfaces import IMessageRepository
-from .models import MessageModel
+from .models import MessageModel, CreateMessageModel, UpdateMessageModel
 
 
-class MessageRepository(IMessageRepository):
-    """SQLAlchemy implementation of message repository."""
+class MessageRepository(PaginatedRepository[MessageModel, Message], IMessageRepository):
+    """SQLAlchemy implementation of message repository using BaseRepository."""
 
     def __init__(self, db: DBSession):
+        super().__init__(MessageModel, Message)
         self.db = db
 
+    @property
+    def entity_name(self) -> str:
+        """Return the entity name for error messages."""
+        return "message"
+
     def find_by_session(
         self, session_id: SessionId, pagination: PaginationInfo | None = None
     ) -> list[Message]:
@@ -24,28 +31,28 @@ class MessageRepository(IMessageRepository):
         query = self.db.query(MessageModel).filter(
             MessageModel.session_id == session_id
         )
+        query = query.order_by(MessageModel.created_time.asc())
 
         if pagination:
             offset = (pagination.page - 1) * pagination.page_size
             query = query.offset(offset).limit(pagination.page_size)
 
-        models = query.order_by(MessageModel.created_time.asc()).all()
-        return [self._model_to_entity(model) for model in models]
+        models = query.all()
+        return [self._convert_model_to_entity(model) for model in models]
 
     def save(self, message: Message) -> Message:
         """Save or update a message."""
-        model = (
-            self.db.query(MessageModel).filter(MessageModel.id == message.id).first()
-        )
+        existing_model = self.db.query(MessageModel).filter(MessageModel.id == message.id).first()
 
-        if model:
-            # Update existing
-            model.message = message.message
-            model.sender_type = message.sender_type.value
-            model.sender_name = message.sender_name
+        if existing_model:
+            update_model = UpdateMessageModel(
+                message=message.message,
+                sender_type=message.sender_type.value,
+                sender_name=message.sender_name,
+            )
+            return self.update(self.db, message.id, update_model.model_dump())
         else:
-            # Create new
-            model = MessageModel(
+            create_model = CreateMessageModel(
                 id=message.id,
                 session_id=message.session_id,
                 message=message.message,
@@ -53,11 +60,7 @@ class MessageRepository(IMessageRepository):
                 sender_name=message.sender_name,
                 created_time=message.created_time,
             )
-            self.db.add(model)
-
-        self.db.commit()
-        self.db.refresh(model)
-        return self._model_to_entity(model)
+            return self.create(self.db, create_model.model_dump())
 
     def delete_by_session(self, session_id: SessionId) -> bool:
         """Delete all messages in a session."""
@@ -66,11 +69,10 @@ class MessageRepository(IMessageRepository):
             .filter(MessageModel.session_id == session_id)
             .delete()
         )
-        self.db.commit()
         return result > 0
 
-    def _model_to_entity(self, model: MessageModel) -> Message:
-        """Convert SQLAlchemy model to domain entity."""
+    def _convert_model_to_entity(self, model: MessageModel) -> Message:
+        """Convert SQLAlchemy model to domain entity with enum handling."""
         return Message(
             id=model.id,
             session_id=model.session_id,

solace_agent_mesh/gateway/http_sse/repository/models/__init__.py

@@ -1,9 +1,17 @@
 """
-SQLAlchemy models for database persistence.
+SQLAlchemy models and Pydantic models for database persistence.
 """
 
 from .base import Base
-from .message_model import MessageModel
-from .session_model import SessionModel
+from .message_model import MessageModel, CreateMessageModel, UpdateMessageModel
+from .session_model import SessionModel, CreateSessionModel, UpdateSessionModel
 
-__all__ = ["Base", "MessageModel", "SessionModel"]
+__all__ = [
+    "Base",
+    "MessageModel",
+    "SessionModel",
+    "CreateMessageModel",
+    "UpdateMessageModel",
+    "CreateSessionModel",
+    "UpdateSessionModel",
+]

solace_agent_mesh/gateway/http_sse/repository/models/message_model.py

@@ -1,7 +1,8 @@
 """
-Message SQLAlchemy model.
+Message SQLAlchemy model and Pydantic models for strongly-typed operations.
 """
 
+from pydantic import BaseModel
 from sqlalchemy import BigInteger, Column, ForeignKey, String, Text
 from sqlalchemy.orm import relationship
 
@@ -25,3 +26,20 @@ class MessageModel(Base):
 
     # Relationship to session
     session = relationship("SessionModel", back_populates="messages")
+
+
+class CreateMessageModel(BaseModel):
+    """Pydantic model for creating a message."""
+    id: str
+    session_id: str
+    message: str
+    sender_type: str
+    sender_name: str
+    created_time: int
+
+
+class UpdateMessageModel(BaseModel):
+    """Pydantic model for updating a message."""
+    message: str
+    sender_type: str
+    sender_name: str

solace_agent_mesh/gateway/http_sse/repository/models/session_model.py

@@ -1,7 +1,8 @@
 """
-Session SQLAlchemy model.
+Session SQLAlchemy model and Pydantic models for strongly-typed operations.
 """
 
+from pydantic import BaseModel
 from sqlalchemy import BigInteger, Column, String
 from sqlalchemy.orm import relationship
 
@@ -27,3 +28,20 @@ class SessionModel(Base):
     messages = relationship(
         "MessageModel", back_populates="session", cascade="all, delete-orphan"
     )
+
+
+class CreateSessionModel(BaseModel):
+    """Pydantic model for creating a session."""
+    id: str
+    name: str | None
+    user_id: str
+    agent_id: str | None
+    created_time: int
+    updated_time: int
+
+
+class UpdateSessionModel(BaseModel):
+    """Pydantic model for updating a session."""
+    name: str | None = None
+    agent_id: str | None = None
+    updated_time: int

solace_agent_mesh/gateway/http_sse/repository/session_repository.py

@@ -4,30 +4,47 @@ Session repository implementation using SQLAlchemy.
 
 from sqlalchemy.orm import Session as DBSession
 
+from ..shared.base_repository import PaginatedRepository
 from ..shared.types import PaginationInfo, SessionId, UserId
 from .entities import Message, Session
 from .interfaces import ISessionRepository
-from .models import MessageModel, SessionModel
+from .models import (
+    MessageModel,
+    SessionModel,
+    CreateSessionModel,
+    UpdateSessionModel,
+)
 
 
-class SessionRepository(ISessionRepository):
-    """SQLAlchemy implementation of session repository."""
+class SessionRepository(PaginatedRepository[SessionModel, Session], ISessionRepository):
+    """SQLAlchemy implementation of session repository using BaseRepository."""
 
     def __init__(self, db: DBSession):
+        super().__init__(SessionModel, Session)
         self.db = db
 
+    @property
+    def entity_name(self) -> str:
+        """Return the entity name for error messages."""
+        return "session"
+
     def find_by_user(
         self, user_id: UserId, pagination: PaginationInfo | None = None
     ) -> list[Session]:
         """Find all sessions for a specific user."""
         query = self.db.query(SessionModel).filter(SessionModel.user_id == user_id)
+        query = query.order_by(SessionModel.updated_time.desc())
 
         if pagination:
             offset = (pagination.page - 1) * pagination.page_size
             query = query.offset(offset).limit(pagination.page_size)
 
-        models = query.order_by(SessionModel.updated_time.desc()).all()
-        return [self._model_to_entity(model) for model in models]
+        models = query.all()
+        return [Session.model_validate(model) for model in models]
+
+    def count_by_user(self, user_id: UserId) -> int:
+        """Count total sessions for a specific user."""
+        return self.db.query(SessionModel).filter(SessionModel.user_id == user_id).count()
 
     def find_user_session(
         self, session_id: SessionId, user_id: UserId
@@ -41,22 +58,21 @@ class SessionRepository(ISessionRepository):
             )
             .first()
         )
-        return self._model_to_entity(model) if model else None
+        return Session.model_validate(model) if model else None
 
     def save(self, session: Session) -> Session:
         """Save or update a session."""
-        model = (
-            self.db.query(SessionModel).filter(SessionModel.id == session.id).first()
-        )
+        existing_model = self.db.query(SessionModel).filter(SessionModel.id == session.id).first()
 
-        if model:
-            # Update existing
-            model.name = session.name
-            model.agent_id = session.agent_id
-            model.updated_time = session.updated_time
+        if existing_model:
+            update_model = UpdateSessionModel(
+                name=session.name,
+                agent_id=session.agent_id,
+                updated_time=session.updated_time,
+            )
+            return self.update(self.db, session.id, update_model.model_dump(exclude_none=True))
         else:
-            # Create new
-            model = SessionModel(
+            create_model = CreateSessionModel(
                 id=session.id,
                 name=session.name,
                 user_id=session.user_id,
@@ -64,24 +80,22 @@ class SessionRepository(ISessionRepository):
                 created_time=session.created_time,
                 updated_time=session.updated_time,
             )
-            self.db.add(model)
-
-        self.db.commit()
-        self.db.refresh(model)
-        return self._model_to_entity(model)
+            return self.create(self.db, create_model.model_dump())
 
     def delete(self, session_id: SessionId, user_id: UserId) -> bool:
         """Delete a session belonging to a user."""
-        result = (
-            self.db.query(SessionModel)
-            .filter(
-                SessionModel.id == session_id,
-                SessionModel.user_id == user_id,
-            )
-            .delete()
-        )
-        self.db.commit()
-        return result > 0
+        # Check if session belongs to user first
+        session_model = self.db.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)
+        return True
 
     def find_user_session_with_messages(
         self,
@@ -112,21 +126,11 @@ class SessionRepository(ISessionRepository):
 
         message_models = message_query.order_by(MessageModel.created_time.asc()).all()
 
-        session = self._model_to_entity(session_model)
+        session = Session.model_validate(session_model)
         messages = [self._message_model_to_entity(model) for model in message_models]
 
         return session, messages
 
-    def _model_to_entity(self, model: SessionModel) -> Session:
-        """Convert SQLAlchemy model to domain entity."""
-        return Session(
-            id=model.id,
-            user_id=model.user_id,
-            name=model.name,
-            agent_id=model.agent_id,
-            created_time=model.created_time,
-            updated_time=model.updated_time,
-        )
 
     def _message_model_to_entity(self, model: MessageModel) -> Message:
         """Convert SQLAlchemy message model to domain entity."""