remdb 0.3.0__py3-none-any.whl

rem/services/session/compression.py

@@ -0,0 +1,360 @@
+"""Session message compression and rehydration for efficient context loading.
+
+This module implements message compression to keep conversation history within
+context windows while preserving full content via REM LOOKUP.
+
+Design Pattern:
+- Long assistant messages (>400 chars) are stored as separate Message entities
+- In-memory conversation uses truncated versions with REM lookup hints
+- Full content retrieved on-demand via LOOKUP queries
+- Compression disabled when Postgres is disabled
+"""
+
+from typing import Any
+
+from loguru import logger
+
+from rem.models.entities import Message
+from rem.services.postgres import PostgresService, Repository
+from rem.settings import settings
+
+
+class MessageCompressor:
+    """Compress and decompress session messages with REM lookup keys."""
+
+    def __init__(self, truncate_length: int = 200):
+        """
+        Initialize message compressor.
+
+        Args:
+            truncate_length: Number of characters to keep from start/end (default: 200)
+        """
+        self.truncate_length = truncate_length
+        self.min_length_for_compression = truncate_length * 2
+
+    def compress_message(
+        self, message: dict[str, Any], entity_key: str | None = None
+    ) -> dict[str, Any]:
+        """
+        Compress a message by truncating long content and adding REM lookup key.
+
+        Args:
+            message: Message dict with role and content
+            entity_key: Optional REM lookup key for full message recovery
+
+        Returns:
+            Compressed message dict
+        """
+        content = message.get("content", "")
+
+        # Don't compress short messages or system messages
+        if (
+            len(content) <= self.min_length_for_compression
+            or message.get("role") == "system"
+        ):
+            return message.copy()
+
+        # Compress long messages
+        n = self.truncate_length
+        start = content[:n]
+        end = content[-n:]
+
+        # Create compressed content with REM lookup hint
+        if entity_key:
+            compressed_content = f"{start}\n\n... [Message truncated - REM LOOKUP {entity_key} to recover full content] ...\n\n{end}"
+        else:
+            compressed_content = f"{start}\n\n... [Message truncated - {len(content) - 2*n} characters omitted] ...\n\n{end}"
+
+        compressed_message = message.copy()
+        compressed_message["content"] = compressed_content
+        compressed_message["_compressed"] = True
+        compressed_message["_original_length"] = len(content)
+        if entity_key:
+            compressed_message["_entity_key"] = entity_key
+
+        logger.debug(
+            f"Compressed message from {len(content)} to {len(compressed_content)} chars (key={entity_key})"
+        )
+
+        return compressed_message
+
+    def decompress_message(
+        self, message: dict[str, Any], full_content: str
+    ) -> dict[str, Any]:
+        """
+        Decompress a message by restoring full content.
+
+        Args:
+            message: Compressed message dict
+            full_content: Full content to restore
+
+        Returns:
+            Decompressed message dict
+        """
+        decompressed = message.copy()
+        decompressed["content"] = full_content
+        decompressed.pop("_compressed", None)
+        decompressed.pop("_original_length", None)
+        decompressed.pop("_entity_key", None)
+
+        return decompressed
+
+    def is_compressed(self, message: dict[str, Any]) -> bool:
+        """Check if a message is compressed."""
+        return message.get("_compressed", False)
+
+    def get_entity_key(self, message: dict[str, Any]) -> str | None:
+        """Get REM lookup key from compressed message."""
+        return message.get("_entity_key")
+
+
+class SessionMessageStore:
+    """Store and retrieve session messages with compression."""
+
+    def __init__(
+        self,
+        user_id: str,
+        compressor: MessageCompressor | None = None,
+    ):
+        """
+        Initialize session message store.
+
+        Args:
+            user_id: User identifier for data isolation
+            compressor: Optional message compressor (creates default if None)
+        """
+        self.user_id = user_id
+        self.compressor = compressor or MessageCompressor()
+        self.repo = Repository(Message)
+
+    async def store_message(
+        self,
+        session_id: str,
+        message: dict[str, Any],
+        message_index: int,
+        user_id: str | None = None,
+    ) -> str:
+        """
+        Store a long assistant message as a Message entity for REM lookup.
+
+        Args:
+            session_id: Parent session identifier
+            message: Message dict to store
+            message_index: Index of message in conversation
+            user_id: Optional user identifier
+
+        Returns:
+            Entity key for REM lookup (message ID)
+        """
+        if not settings.postgres.enabled:
+            logger.debug("Postgres disabled, skipping message storage")
+            return f"msg-{message_index}"
+
+        # Create entity key for REM LOOKUP: session-{session_id}-msg-{index}
+        entity_key = f"session-{session_id}-msg-{message_index}"
+
+        # Create Message entity for assistant response
+        msg = Message(
+            content=message.get("content", ""),
+            message_type=message.get("role", "assistant"),
+            session_id=session_id,
+            tenant_id=self.user_id,  # Set tenant_id to user_id (application scoped to user)
+            user_id=user_id or self.user_id,
+            metadata={
+                "message_index": message_index,
+                "entity_key": entity_key,  # Store entity key for LOOKUP
+                "timestamp": message.get("timestamp"),
+            },
+        )
+
+        # Store in database
+        await self.repo.upsert(msg)
+
+        logger.debug(f"Stored assistant response: {entity_key} (id={msg.id})")
+        return entity_key
+
+    async def retrieve_message(self, entity_key: str) -> str | None:
+        """
+        Retrieve full message content by REM lookup key.
+
+        Uses LOOKUP query pattern: finds message by entity_key in metadata.
+
+        Args:
+            entity_key: REM lookup key (session-{id}-msg-{index})
+
+        Returns:
+            Full message content or None if not found
+        """
+        if not settings.postgres.enabled:
+            logger.debug("Postgres disabled, cannot retrieve message")
+            return None
+
+        try:
+            # LOOKUP pattern: find message by entity_key in metadata
+            query = """
+                SELECT * FROM messages
+                WHERE metadata->>'entity_key' = $1
+                  AND user_id = $2
+                  AND deleted_at IS NULL
+                LIMIT 1
+            """
+
+            if not self.repo.db:
+                logger.warning("Database not available for message lookup")
+                return None
+
+            row = await self.repo.db.fetchrow(query, entity_key, self.user_id)
+
+            if row:
+                msg = Message.model_validate(dict(row))
+                logger.debug(f"Retrieved message via LOOKUP: {entity_key}")
+                return msg.content
+
+            logger.warning(f"Message not found via LOOKUP: {entity_key}")
+            return None
+
+        except Exception as e:
+            logger.error(f"Failed to retrieve message {entity_key}: {e}")
+            return None
+
+    async def store_session_messages(
+        self,
+        session_id: str,
+        messages: list[dict[str, Any]],
+        user_id: str | None = None,
+        compress: bool = True,
+    ) -> list[dict[str, Any]]:
+        """
+        Store all session messages and return compressed versions.
+
+        Args:
+            session_id: Session identifier
+            messages: List of messages to store
+            user_id: Optional user identifier
+            compress: Whether to compress messages (default: True)
+
+        Returns:
+            List of compressed messages with REM lookup keys
+        """
+        if not settings.postgres.enabled:
+            logger.debug("Postgres disabled, returning messages uncompressed")
+            return messages
+
+        compressed_messages = []
+
+        for idx, message in enumerate(messages):
+            content = message.get("content", "")
+
+            # Only store and compress long assistant responses
+            if (
+                message.get("role") == "assistant"
+                and len(content) > self.compressor.min_length_for_compression
+            ):
+                # Store full message as separate Message entity
+                entity_key = await self.store_message(
+                    session_id, message, idx, user_id
+                )
+
+                if compress:
+                    compressed_msg = self.compressor.compress_message(
+                        message, entity_key
+                    )
+                    compressed_messages.append(compressed_msg)
+                else:
+                    msg_copy = message.copy()
+                    msg_copy["_entity_key"] = entity_key
+                    compressed_messages.append(msg_copy)
+            else:
+                # Short assistant messages, user messages, and system messages stored as-is
+                # Store ALL messages in database for full audit trail
+                msg = Message(
+                    content=content,
+                    message_type=message.get("role", "user"),
+                    session_id=session_id,
+                    tenant_id=self.user_id,  # Set tenant_id to user_id (application scoped to user)
+                    user_id=user_id or self.user_id,
+                    metadata={
+                        "message_index": idx,
+                        "timestamp": message.get("timestamp"),
+                    },
+                )
+                await self.repo.upsert(msg)
+                compressed_messages.append(message.copy())
+
+        return compressed_messages
+
+    async def load_session_messages(
+        self, session_id: str, user_id: str | None = None, decompress: bool = False
+    ) -> list[dict[str, Any]]:
+        """
+        Load session messages from database.
+
+        Args:
+            session_id: Session identifier
+            user_id: Optional user identifier for filtering
+            decompress: Whether to decompress messages (default: False)
+
+        Returns:
+            List of session messages in chronological order
+        """
+        if not settings.postgres.enabled:
+            logger.debug("Postgres disabled, returning empty message list")
+            return []
+
+        try:
+            # Load messages from repository
+            # Note: tenant_id column in messages table maps to user_id (user-scoped partitioning)
+            filters = {"session_id": session_id, "tenant_id": self.user_id}
+            if user_id:
+                filters["user_id"] = user_id
+
+            messages = await self.repo.find(filters, order_by="created_at ASC")
+
+            # Convert Message entities to dict format
+            message_dicts = []
+            for msg in messages:
+                msg_dict = {
+                    "role": msg.message_type or "assistant",
+                    "content": msg.content,
+                    "timestamp": msg.created_at.isoformat() if msg.created_at else None,
+                }
+
+                # Check if message was compressed
+                entity_key: str | None = msg.metadata.get("entity_key") if msg.metadata else None
+                if entity_key and len(msg.content) <= self.compressor.min_length_for_compression:
+                    # This is a compressed reference, mark it
+                    msg_dict["_compressed"] = True
+                    msg_dict["_entity_key"] = entity_key
+                    msg_dict["_original_length"] = msg.metadata.get("original_length", 0)
+
+                message_dicts.append(msg_dict)
+
+            # Decompress if requested
+            if decompress:
+                decompressed_messages = []
+                for message in message_dicts:
+                    if self.compressor.is_compressed(message):
+                        entity_key = self.compressor.get_entity_key(message)
+                        if entity_key:
+                            full_content = await self.retrieve_message(entity_key)
+                            if full_content:
+                                decompressed_messages.append(
+                                    self.compressor.decompress_message(
+                                        message, full_content
+                                    )
+                                )
+                            else:
+                                # Fallback to compressed version if retrieval fails
+                                decompressed_messages.append(message)
+                        else:
+                            decompressed_messages.append(message)
+                    else:
+                        decompressed_messages.append(message)
+
+                return decompressed_messages
+
+            return message_dicts
+
+        except Exception as e:
+            logger.error(f"Failed to load session messages: {e}")
+            return []

rem/services/session/reload.py

@@ -0,0 +1,77 @@
+"""Session reloading logic for conversation history restoration.
+
+This module implements session history loading from the database,
+allowing conversations to be resumed across multiple API calls.
+
+Design Pattern:
+- Session identified by session_id from X-Session-Id header
+- All messages for session loaded in chronological order
+- Optional decompression of long assistant messages via REM LOOKUP
+- Gracefully handles missing database (returns empty history)
+"""
+
+from loguru import logger
+
+from rem.services.session.compression import SessionMessageStore
+from rem.settings import settings
+
+
+async def reload_session(
+    session_id: str,
+    user_id: str,
+    decompress_messages: bool = False,
+) -> list[dict]:
+    """
+    Reload all messages for a session from the database.
+
+    Args:
+        session_id: Session/conversation identifier
+        user_id: User identifier for data isolation
+        decompress_messages: Whether to decompress long messages via REM LOOKUP
+
+    Returns:
+        List of message dicts in chronological order (oldest first)
+
+    Example:
+        ```python
+        # In completions endpoint
+        context = AgentContext.from_headers(dict(request.headers))
+
+        # Reload previous conversation history
+        history = await reload_session(
+            session_id=context.session_id,
+            user_id=context.user_id,
+            decompress_messages=False,  # Use compressed versions for efficiency
+        )
+
+        # Combine with new user message
+        messages = history + [{"role": "user", "content": prompt}]
+        ```
+    """
+    if not settings.postgres.enabled:
+        logger.debug("Postgres disabled, returning empty session history")
+        return []
+
+    if not session_id:
+        logger.debug("No session_id provided, returning empty history")
+        return []
+
+    try:
+        # Create message store for this session
+        store = SessionMessageStore(user_id=user_id)
+
+        # Load messages (optionally decompressed)
+        messages = await store.load_session_messages(
+            session_id=session_id, user_id=user_id, decompress=decompress_messages
+        )
+
+        logger.info(
+            f"Reloaded {len(messages)} messages for session {session_id} "
+            f"(decompressed={decompress_messages})"
+        )
+
+        return messages
+
+    except Exception as e:
+        logger.error(f"Failed to reload session {session_id}: {e}")
+        return []