remdb 0.3.114__py3-none-any.whl → 0.3.172__py3-none-any.whl

rem/services/session/compression.py

@@ -1,13 +1,49 @@
 """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
@@ -170,12 +206,16 @@ class SessionMessageStore:
         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
@@ -281,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())
@@ -300,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")
@@ -328,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.debug(
             f"Reloaded {len(messages)} messages for session {session_id} "
-            f"(decompressed={decompress_messages})"
+            f"(compress_on_load={compress_on_load})"
         )
 
         return messages

rem/services/user_service.py

@@ -4,7 +4,8 @@ User Service - User account management.
 Handles user creation, profile updates, and session linking.
 """
 
-from datetime import datetime
+from rem.utils.date_utils import utc_now
+from rem.utils.user_id import email_to_user_id
 from typing import Optional
 
 from loguru import logger
@@ -51,28 +52,59 @@ class UserService:
                     updated = True
             
             if updated:
-                user.updated_at = datetime.utcnow()
+                user.updated_at = utc_now()
                 await self.repo.upsert(user)
             
             return user
         
         # Create new user
+        # id and user_id = UUID5 hash of email (deterministic bijection)
+        # name = email (entity_key for LOOKUP by email in KV store)
+        hashed_id = email_to_user_id(email)
         user = User(
+            id=hashed_id,  # Database id = hash of email
             tenant_id=tenant_id,
-            user_id=email, # Use email as user_id for now? Or UUID?
-            # The User model has 'user_id' field but also 'id' UUID.
-            # Usually user_id is the external ID or email.
-            name=name,
+            user_id=hashed_id,  # user_id = hash of email (same as id)
+            name=email,  # Email as entity_key for REM LOOKUP
             email=email,
             tier=UserTier.FREE,
-            created_at=datetime.utcnow(),
-            updated_at=datetime.utcnow(),
+            created_at=utc_now(),
+            updated_at=utc_now(),
             metadata={"avatar_url": avatar_url} if avatar_url else {},
         )
         await self.repo.upsert(user)
         logger.info(f"Created new user: {email}")
         return user
 
+    async def get_user_by_id(self, user_id: str) -> Optional[User]:
+        """
+        Get a user by their UUID.
+
+        Args:
+            user_id: The user's UUID
+
+        Returns:
+            User if found, None otherwise
+        """
+        try:
+            return await self.repo.get_by_id(user_id)
+        except Exception as e:
+            logger.warning(f"Could not find user by id {user_id}: {e}")
+            return None
+
+    async def get_user_by_email(self, email: str) -> Optional[User]:
+        """
+        Get a user by their email address.
+
+        Args:
+            email: The user's email
+
+        Returns:
+            User if found, None otherwise
+        """
+        users = await self.repo.find(filters={"email": email}, limit=1)
+        return users[0] if users else None
+
     async def link_anonymous_session(self, user: User, anon_id: str) -> None:
         """
         Link an anonymous session ID to a user account.
@@ -88,7 +120,7 @@ class UserService:
 
         # Add to list
         user.anonymous_ids.append(anon_id)
-        user.updated_at = datetime.utcnow()
+        user.updated_at = utc_now()
         
         # Save
         await self.repo.upsert(user)