remdb 0.3.14__py3-none-any.whl → 0.3.157__py3-none-any.whl

rem/services/session/compression.py

@@ -1,19 +1,70 @@
 """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
+This module implements message storage and compression to keep conversation history
+within context windows while preserving full content via REM LOOKUP.
+
+Message Types and Storage Strategy
+===================================
+
+All messages are stored UNCOMPRESSED in the database for full audit/analysis.
+Compression happens only on RELOAD when reconstructing context for the LLM.
+
+Message Types:
+- `user`: User messages - stored and reloaded as-is
+- `tool`: Tool call messages (e.g., register_metadata) - stored and reloaded as-is
+         NEVER compressed - contains important structured metadata
+- `assistant`: Assistant text responses - stored uncompressed, but MAY BE
+              compressed on reload if long (>400 chars) with REM LOOKUP hints
+
+Example Session Flow:
+```
+Turn 1 (stored uncompressed):
+  - user: "I have a headache"
+  - tool: register_metadata({confidence: 0.3, collected_fields: {...}})
+  - assistant: "I'm sorry to hear that. How long has this been going on?"
+
+Turn 2 (stored uncompressed):
+  - user: "About 3 days, really bad"
+  - tool: register_metadata({confidence: 0.6, collected_fields: {...}})
+  - assistant: "Got it - 3 days. On a scale of 1-10..."
+
+On reload (for LLM context):
+  - user messages: returned as-is
+  - tool messages: returned as-is (never compressed)
+  - assistant messages: compressed if long, with REM LOOKUP hint for full retrieval
+```
+
+REM LOOKUP Pattern:
+- Long assistant messages get truncated with hint: "... [REM LOOKUP session-{id}-msg-{idx}] ..."
+- Agent can retrieve full content on-demand using the LOOKUP key
+- Keeps context window efficient while preserving data integrity
+
+Key Design Decisions:
+1. Store everything uncompressed - full audit trail in database
+2. Compress only on reload - optimize for LLM context window
+3. Never compress tool messages - structured metadata must stay intact
+4. REM LOOKUP enables on-demand retrieval of full assistant responses
 """
 
 from typing import Any
 
 from loguru import logger
 
+# Max length for entity keys (kv_store.entity_key is varchar(255))
+MAX_ENTITY_KEY_LENGTH = 255
+
+
+def truncate_key(key: str, max_length: int = MAX_ENTITY_KEY_LENGTH) -> str:
+    """Truncate a key to max length, preserving useful suffix if possible."""
+    if len(key) <= max_length:
+        return key
+    # Keep first part and add hash suffix for uniqueness
+    import hashlib
+    hash_suffix = hashlib.md5(key.encode()).hexdigest()[:8]
+    truncated = key[:max_length - 9] + "-" + hash_suffix
+    logger.warning(f"Truncated key from {len(key)} to {len(truncated)} chars: {key[:50]}...")
+    return truncated
+
 from rem.models.entities import Message
 from rem.services.postgres import PostgresService, Repository
 from rem.settings import settings
@@ -151,15 +202,20 @@ class SessionMessageStore:
             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}"
+        # Truncate to avoid exceeding kv_store.entity_key varchar(255) limit
+        entity_key = truncate_key(f"session-{session_id}-msg-{message_index}")
 
         # Create Message entity for assistant response
+        # Use pre-generated id from message dict if available (for frontend feedback)
         msg = Message(
+            id=message.get("id"),  # Use pre-generated ID if provided
             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,
+            trace_id=message.get("trace_id"),
+            span_id=message.get("span_id"),
             metadata={
                 "message_index": message_index,
                 "entity_key": entity_key,  # Store entity key for LOOKUP
@@ -265,18 +321,33 @@ class SessionMessageStore:
                     msg_copy["_entity_key"] = entity_key
                     compressed_messages.append(msg_copy)
             else:
-                # Short assistant messages, user messages, and system messages stored as-is
+                # Short assistant messages, user messages, tool messages, and system messages stored as-is
                 # Store ALL messages in database for full audit trail
+                # Build metadata dict with standard fields
+                msg_metadata = {
+                    "message_index": idx,
+                    "timestamp": message.get("timestamp"),
+                }
+
+                # For tool messages, include tool call details in metadata
+                if message.get("role") == "tool":
+                    if message.get("tool_call_id"):
+                        msg_metadata["tool_call_id"] = message.get("tool_call_id")
+                    if message.get("tool_name"):
+                        msg_metadata["tool_name"] = message.get("tool_name")
+                    if message.get("tool_arguments"):
+                        msg_metadata["tool_arguments"] = message.get("tool_arguments")
+
                 msg = Message(
+                    id=message.get("id"),  # Use pre-generated ID if provided
                     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"),
-                    },
+                    trace_id=message.get("trace_id"),
+                    span_id=message.get("span_id"),
+                    metadata=msg_metadata,
                 )
                 await self.repo.upsert(msg)
                 compressed_messages.append(message.copy())
@@ -284,18 +355,24 @@ class SessionMessageStore:
         return compressed_messages
 
     async def load_session_messages(
-        self, session_id: str, user_id: str | None = None, decompress: bool = False
+        self, session_id: str, user_id: str | None = None, compress_on_load: bool = True
     ) -> list[dict[str, Any]]:
         """
-        Load session messages from database.
+        Load session messages from database, optionally compressing long assistant messages.
+
+        Compression on Load:
+        - Tool messages (role: "tool") are NEVER compressed - they contain structured metadata
+        - User messages are returned as-is
+        - Assistant messages MAY be compressed if long (>400 chars) with REM LOOKUP hints
 
         Args:
             session_id: Session identifier
             user_id: Optional user identifier for filtering
-            decompress: Whether to decompress messages (default: False)
+            compress_on_load: Whether to compress long assistant messages (default: True)
 
         Returns:
-            List of session messages in chronological order
+            List of session messages in chronological order, with long assistant
+            messages optionally compressed with REM LOOKUP hints
         """
         if not settings.postgres.enabled:
             logger.debug("Postgres disabled, returning empty message list")
@@ -312,49 +389,58 @@ class SessionMessageStore:
 
             # Convert Message entities to dict format
             message_dicts = []
-            for msg in messages:
+            for idx, msg in enumerate(messages):
+                role = msg.message_type or "assistant"
                 msg_dict = {
-                    "role": msg.message_type or "assistant",
+                    "role": role,
                     "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)
+                # For tool messages, reconstruct tool call metadata
+                if role == "tool" and msg.metadata:
+                    if msg.metadata.get("tool_call_id"):
+                        msg_dict["tool_call_id"] = msg.metadata["tool_call_id"]
+                    if msg.metadata.get("tool_name"):
+                        msg_dict["tool_name"] = msg.metadata["tool_name"]
+                    if msg.metadata.get("tool_arguments"):
+                        msg_dict["tool_arguments"] = msg.metadata["tool_arguments"]
+
+                # Compress long ASSISTANT messages on load (never tool messages)
+                if (
+                    compress_on_load
+                    and role == "assistant"
+                    and len(msg.content) > self.compressor.min_length_for_compression
+                ):
+                    # Generate entity key for REM LOOKUP
+                    entity_key = truncate_key(f"session-{session_id}-msg-{idx}")
+                    msg_dict = self.compressor.compress_message(msg_dict, entity_key)
 
                 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
-
+            logger.debug(
+                f"Loaded {len(message_dicts)} messages for session {session_id} "
+                f"(compress_on_load={compress_on_load})"
+            )
             return message_dicts
 
         except Exception as e:
             logger.error(f"Failed to load session messages: {e}")
             return []
+
+    async def retrieve_full_message(self, session_id: str, message_index: int) -> str | None:
+        """
+        Retrieve full message content by session and message index (for REM LOOKUP).
+
+        This is used when an agent needs to recover full content from a compressed
+        message that has a REM LOOKUP hint.
+
+        Args:
+            session_id: Session identifier
+            message_index: Index of message in session (from REM LOOKUP key)
+
+        Returns:
+            Full message content or None if not found
+        """
+        entity_key = truncate_key(f"session-{session_id}-msg-{message_index}")
+        return await self.retrieve_message(entity_key)

rem/services/session/reload.py

@@ -6,8 +6,14 @@ 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
+- Long assistant messages compressed on load with REM LOOKUP hints
+- Tool messages (register_metadata, etc.) are NEVER compressed
 - Gracefully handles missing database (returns empty history)
+
+Message Types on Reload:
+- user: Returned as-is
+- tool: Returned as-is with metadata (tool_call_id, tool_name, tool_arguments)
+- assistant: Compressed on load if long (>400 chars), with REM LOOKUP for recovery
 """
 
 from loguru import logger
@@ -19,7 +25,7 @@ from rem.settings import settings
 async def reload_session(
     session_id: str,
     user_id: str,
-    decompress_messages: bool = False,
+    compress_on_load: bool = True,
 ) -> list[dict]:
     """
     Reload all messages for a session from the database.
@@ -27,7 +33,8 @@ async def reload_session(
     Args:
         session_id: Session/conversation identifier
         user_id: User identifier for data isolation
-        decompress_messages: Whether to decompress long messages via REM LOOKUP
+        compress_on_load: Whether to compress long assistant messages (default: True)
+                         Tool messages are NEVER compressed.
 
     Returns:
         List of message dicts in chronological order (oldest first)
@@ -41,7 +48,7 @@ async def reload_session(
         history = await reload_session(
             session_id=context.session_id,
             user_id=context.user_id,
-            decompress_messages=False,  # Use compressed versions for efficiency
+            compress_on_load=True,  # Compress long assistant messages
         )
 
         # Combine with new user message
@@ -60,14 +67,14 @@ async def reload_session(
         # Create message store for this session
         store = SessionMessageStore(user_id=user_id)
 
-        # Load messages (optionally decompressed)
+        # Load messages (assistant messages compressed on load, tool messages never compressed)
         messages = await store.load_session_messages(
-            session_id=session_id, user_id=user_id, decompress=decompress_messages
+            session_id=session_id, user_id=user_id, compress_on_load=compress_on_load
         )
 
-        logger.info(
+        logger.debug(
             f"Reloaded {len(messages)} messages for session {session_id} "
-            f"(decompressed={decompress_messages})"
+            f"(compress_on_load={compress_on_load})"
         )
 
         return messages