remdb 0.3.181__py3-none-any.whl → 0.3.202__py3-none-any.whl

rem/services/postgres/register_type.py

@@ -94,14 +94,14 @@ def generate_table_schema(
     # Always add id as primary key
     columns.append("id UUID PRIMARY KEY DEFAULT uuid_generate_v4()")
 
-    # Add tenant_id if tenant scoped
+    # Add tenant_id if tenant scoped (nullable - NULL means public/shared)
     if tenant_scoped:
-        columns.append("tenant_id VARCHAR(100) NOT NULL")
-        indexes.append(f"CREATE INDEX idx_{table_name}_tenant ON {table_name} (tenant_id);")
+        columns.append("tenant_id VARCHAR(100)")
+        indexes.append(f"CREATE INDEX IF NOT EXISTS idx_{table_name}_tenant ON {table_name} (tenant_id);")
 
     # Add user_id (owner field)
     columns.append("user_id VARCHAR(256)")
-    indexes.append(f"CREATE INDEX idx_{table_name}_user ON {table_name} (user_id);")
+    indexes.append(f"CREATE INDEX IF NOT EXISTS idx_{table_name}_user ON {table_name} (user_id);")
 
     # Process Pydantic fields (skip system fields)
     for field_name, field_info in model.model_fields.items():
@@ -125,19 +125,19 @@ def generate_table_schema(
     # Add graph_edges JSONB field
     columns.append("graph_edges JSONB DEFAULT '[]'::jsonb")
     indexes.append(
-        f"CREATE INDEX idx_{table_name}_graph_edges ON {table_name} USING GIN (graph_edges);"
+        f"CREATE INDEX IF NOT EXISTS idx_{table_name}_graph_edges ON {table_name} USING GIN (graph_edges);"
     )
 
     # Add metadata JSONB field
     columns.append("metadata JSONB DEFAULT '{}'::jsonb")
     indexes.append(
-        f"CREATE INDEX idx_{table_name}_metadata ON {table_name} USING GIN (metadata);"
+        f"CREATE INDEX IF NOT EXISTS idx_{table_name}_metadata ON {table_name} USING GIN (metadata);"
     )
 
     # Add tags field (TEXT[] for list[str])
     columns.append("tags TEXT[] DEFAULT ARRAY[]::TEXT[]")
     indexes.append(
-        f"CREATE INDEX idx_{table_name}_tags ON {table_name} USING GIN (tags);"
+        f"CREATE INDEX IF NOT EXISTS idx_{table_name}_tags ON {table_name} USING GIN (tags);"
     )
 
     # Generate CREATE TABLE statement
@@ -202,10 +202,10 @@ CREATE TABLE IF NOT EXISTS {embeddings_table} (
 );
 
 -- Index for entity lookup (get all embeddings for entity)
-CREATE INDEX idx_{embeddings_table}_entity ON {embeddings_table} (entity_id);
+CREATE INDEX IF NOT EXISTS idx_{embeddings_table}_entity ON {embeddings_table} (entity_id);
 
 -- Index for field + provider lookup
-CREATE INDEX idx_{embeddings_table}_field_provider ON {embeddings_table} (field_name, provider);
+CREATE INDEX IF NOT EXISTS idx_{embeddings_table}_field_provider ON {embeddings_table} (field_name, provider);
 
 -- HNSW index for vector similarity search (created in background)
 -- Note: This will be created by background thread after data load
@@ -258,6 +258,7 @@ BEGIN
         RETURN OLD;
     ELSIF (TG_OP = 'INSERT' OR TG_OP = 'UPDATE') THEN
         -- Upsert to KV_STORE (O(1) lookup by entity_key)
+        -- tenant_id can be NULL (meaning public/shared data)
         INSERT INTO kv_store (
             entity_key,
             entity_type,
@@ -277,7 +278,7 @@ BEGIN
             COALESCE(NEW.graph_edges, '[]'::jsonb),
             CURRENT_TIMESTAMP
         )
-        ON CONFLICT (tenant_id, entity_key)
+        ON CONFLICT (COALESCE(tenant_id, ''), entity_key)
         DO UPDATE SET
             entity_id = EXCLUDED.entity_id,
             user_id = EXCLUDED.user_id,

rem/services/session/__init__.py

@@ -1,6 +1,12 @@
 """Session management services for conversation persistence and compression."""
 
 from .compression import MessageCompressor, SessionMessageStore
+from .pydantic_messages import session_to_pydantic_messages
 from .reload import reload_session
 
-__all__ = ["MessageCompressor", "SessionMessageStore", "reload_session"]
+__all__ = [
+    "MessageCompressor",
+    "SessionMessageStore",
+    "reload_session",
+    "session_to_pydantic_messages",
+]

rem/services/session/compression.py

@@ -65,7 +65,7 @@ def truncate_key(key: str, max_length: int = MAX_ENTITY_KEY_LENGTH) -> str:
     logger.warning(f"Truncated key from {len(key)} to {len(truncated)} chars: {key[:50]}...")
     return truncated
 
-from rem.models.entities import Message
+from rem.models.entities import Message, Session
 from rem.services.postgres import PostgresService, Repository
 from rem.settings import settings
 
@@ -177,6 +177,41 @@ class SessionMessageStore:
         self.user_id = user_id
         self.compressor = compressor or MessageCompressor()
         self.repo = Repository(Message)
+        self._session_repo = Repository(Session, table_name="sessions")
+
+    async def _ensure_session_exists(
+        self,
+        session_id: str,
+        user_id: str | None = None,
+    ) -> None:
+        """
+        Ensure session exists, creating it if necessary.
+
+        Args:
+            session_id: Session identifier (maps to Session.name)
+            user_id: Optional user identifier
+        """
+        try:
+            # Check if session already exists by name
+            existing = await self._session_repo.find(
+                filters={"name": session_id},
+                limit=1,
+            )
+            if existing:
+                return  # Session already exists
+
+            # Create new session
+            session = Session(
+                name=session_id,
+                user_id=user_id or self.user_id,
+                tenant_id=self.user_id,  # tenant_id set to user_id for scoping
+            )
+            await self._session_repo.upsert(session)
+            logger.info(f"Created session {session_id} for user {user_id or self.user_id}")
+
+        except Exception as e:
+            # Log but don't fail - session creation is best-effort
+            logger.warning(f"Failed to ensure session exists: {e}")
 
     async def store_message(
         self,
@@ -283,8 +318,10 @@ class SessionMessageStore:
         """
         Store all session messages and return compressed versions.
 
+        Ensures session exists before storing messages.
+
         Args:
-            session_id: Session identifier
+            session_id: Session identifier (maps to Session.name)
             messages: List of messages to store
             user_id: Optional user identifier
             compress: Whether to compress messages (default: True)
@@ -296,6 +333,9 @@ class SessionMessageStore:
             logger.debug("Postgres disabled, returning messages uncompressed")
             return messages
 
+        # Ensure session exists before storing messages
+        await self._ensure_session_exists(session_id, user_id)
+
         compressed_messages = []
 
         for idx, message in enumerate(messages):

rem/services/session/pydantic_messages.py

@@ -0,0 +1,210 @@
+"""Convert stored session messages to pydantic-ai native message format.
+
+This module enables proper conversation history replay by converting our simplified
+storage format into pydantic-ai's native ModelRequest/ModelResponse types.
+
+Key insight: When we store tool results, we only store the result (ToolReturnPart).
+But LLM APIs require matching ToolCallPart for each ToolReturnPart. So we synthesize
+the ToolCallPart from stored metadata (tool_name, tool_call_id, tool_arguments).
+
+Storage format (our simplified format):
+    {"role": "user", "content": "..."}
+    {"role": "assistant", "content": "..."}
+    {"role": "tool", "content": "{...}", "tool_name": "...", "tool_call_id": "...", "tool_arguments": {...}}
+
+Pydantic-ai format (what the LLM expects):
+    ModelRequest(parts=[UserPromptPart(content="...")])
+    ModelResponse(parts=[TextPart(content="..."), ToolCallPart(...)])  # Call
+    ModelRequest(parts=[ToolReturnPart(...)])  # Result
+
+Example usage:
+    from rem.services.session.pydantic_messages import session_to_pydantic_messages
+
+    # Load session history
+    session_history = await store.load_session_messages(session_id)
+
+    # Convert to pydantic-ai format
+    message_history = session_to_pydantic_messages(session_history)
+
+    # Use with agent.run()
+    result = await agent.run(user_prompt, message_history=message_history)
+"""
+
+import json
+from typing import Any
+
+from loguru import logger
+from pydantic_ai.messages import (
+    ModelMessage,
+    ModelRequest,
+    ModelResponse,
+    SystemPromptPart,
+    TextPart,
+    ToolCallPart,
+    ToolReturnPart,
+    UserPromptPart,
+)
+
+
+def session_to_pydantic_messages(
+    session_history: list[dict[str, Any]],
+    system_prompt: str | None = None,
+) -> list[ModelMessage]:
+    """Convert stored session messages to pydantic-ai ModelMessage format.
+
+    Handles the conversion of our simplified storage format to pydantic-ai's
+    native message types, including synthesizing ToolCallPart for tool results.
+
+    IMPORTANT: pydantic-ai only auto-adds system prompts when message_history is empty.
+    When passing message_history to agent.run(), you MUST include the system prompt
+    via the system_prompt parameter here.
+
+    Args:
+        session_history: List of message dicts from SessionMessageStore.load_session_messages()
+            Each dict has: role, content, and optionally tool_name, tool_call_id, tool_arguments
+        system_prompt: The agent's system prompt (from schema description). This is REQUIRED
+            for proper agent behavior on subsequent turns, as pydantic-ai won't add it
+            automatically when message_history is provided.
+
+    Returns:
+        List of ModelMessage (ModelRequest | ModelResponse) ready for agent.run(message_history=...)
+
+    Note:
+        - System prompts ARE included as SystemPromptPart when system_prompt is provided
+        - Tool results require synthesized ToolCallPart to satisfy LLM API requirements
+        - The first message in session_history should be "user" role (from context builder)
+    """
+    messages: list[ModelMessage] = []
+
+    # CRITICAL: Prepend agent's system prompt if provided
+    # This ensures the agent's instructions are present on every turn
+    # pydantic-ai only auto-adds system prompts when message_history is empty
+    if system_prompt:
+        messages.append(ModelRequest(parts=[SystemPromptPart(content=system_prompt)]))
+        logger.debug(f"Prepended agent system prompt ({len(system_prompt)} chars) to message history")
+
+    # Track pending tool results to batch them with assistant responses
+    # When we see a tool message, we need to:
+    # 1. Add a ModelResponse with ToolCallPart (synthesized)
+    # 2. Add a ModelRequest with ToolReturnPart (actual result)
+
+    i = 0
+    while i < len(session_history):
+        msg = session_history[i]
+        role = msg.get("role", "")
+        content = msg.get("content", "")
+
+        if role == "user":
+            # User messages become ModelRequest with UserPromptPart
+            messages.append(ModelRequest(parts=[UserPromptPart(content=content)]))
+
+        elif role == "assistant":
+            # Assistant text becomes ModelResponse with TextPart
+            # Check if there are following tool messages that should be grouped
+            tool_calls = []
+            tool_returns = []
+
+            # Look ahead for tool messages that follow this assistant message
+            j = i + 1
+            while j < len(session_history) and session_history[j].get("role") == "tool":
+                tool_msg = session_history[j]
+                tool_name = tool_msg.get("tool_name", "unknown_tool")
+                tool_call_id = tool_msg.get("tool_call_id", f"call_{j}")
+                tool_arguments = tool_msg.get("tool_arguments", {})
+                tool_content = tool_msg.get("content", "{}")
+
+                # Parse tool content if it's a JSON string
+                if isinstance(tool_content, str):
+                    try:
+                        tool_result = json.loads(tool_content)
+                    except json.JSONDecodeError:
+                        tool_result = {"raw": tool_content}
+                else:
+                    tool_result = tool_content
+
+                # Synthesize ToolCallPart (what the model "called")
+                tool_calls.append(ToolCallPart(
+                    tool_name=tool_name,
+                    args=tool_arguments if tool_arguments else {},
+                    tool_call_id=tool_call_id,
+                ))
+
+                # Create ToolReturnPart (the actual result)
+                tool_returns.append(ToolReturnPart(
+                    tool_name=tool_name,
+                    content=tool_result,
+                    tool_call_id=tool_call_id,
+                ))
+
+                j += 1
+
+            # Build the assistant's ModelResponse
+            response_parts = []
+
+            # Add tool calls first (if any)
+            response_parts.extend(tool_calls)
+
+            # Add text content (if any)
+            if content:
+                response_parts.append(TextPart(content=content))
+
+            # Only add ModelResponse if we have parts
+            if response_parts:
+                messages.append(ModelResponse(
+                    parts=response_parts,
+                    model_name="recovered",  # We don't store model name
+                ))
+
+            # Add tool returns as ModelRequest (required by LLM API)
+            if tool_returns:
+                messages.append(ModelRequest(parts=tool_returns))
+
+            # Skip the tool messages we just processed
+            i = j - 1
+
+        elif role == "tool":
+            # Orphan tool message (no preceding assistant) - synthesize both parts
+            tool_name = msg.get("tool_name", "unknown_tool")
+            tool_call_id = msg.get("tool_call_id", f"call_{i}")
+            tool_arguments = msg.get("tool_arguments", {})
+            tool_content = msg.get("content", "{}")
+
+            # Parse tool content
+            if isinstance(tool_content, str):
+                try:
+                    tool_result = json.loads(tool_content)
+                except json.JSONDecodeError:
+                    tool_result = {"raw": tool_content}
+            else:
+                tool_result = tool_content
+
+            # Synthesize the tool call (ModelResponse with ToolCallPart)
+            messages.append(ModelResponse(
+                parts=[ToolCallPart(
+                    tool_name=tool_name,
+                    args=tool_arguments if tool_arguments else {},
+                    tool_call_id=tool_call_id,
+                )],
+                model_name="recovered",
+            ))
+
+            # Add the tool return (ModelRequest with ToolReturnPart)
+            messages.append(ModelRequest(
+                parts=[ToolReturnPart(
+                    tool_name=tool_name,
+                    content=tool_result,
+                    tool_call_id=tool_call_id,
+                )]
+            ))
+
+        elif role == "system":
+            # Skip system messages - pydantic-ai handles these via Agent.system_prompt
+            logger.debug("Skipping system message in session history (handled by Agent)")
+
+        else:
+            logger.warning(f"Unknown message role in session history: {role}")
+
+        i += 1
+
+    logger.debug(f"Converted {len(session_history)} stored messages to {len(messages)} pydantic-ai messages")
+    return messages

rem/sql/migrations/001_install.sql

@@ -121,18 +121,18 @@ CREATE UNLOGGED TABLE IF NOT EXISTS kv_store (
     entity_key VARCHAR(255) NOT NULL,
     entity_type VARCHAR(100) NOT NULL,
     entity_id UUID NOT NULL,
-    tenant_id VARCHAR(100) NOT NULL,
+    tenant_id VARCHAR(100),  -- NULL = public/shared data
     user_id VARCHAR(100),
     content_summary TEXT,
     metadata JSONB DEFAULT '{}',
     graph_edges JSONB DEFAULT '[]'::jsonb,  -- Cached edges for fast graph traversal
     created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
-    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
-
-    -- Composite primary key: entity_key unique per tenant
-    PRIMARY KEY (tenant_id, entity_key)
+    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
 );
 
+-- Unique constraint on (tenant_id, entity_key) using COALESCE to handle NULL tenant_id
+CREATE UNIQUE INDEX IF NOT EXISTS idx_kv_store_tenant_key ON kv_store (COALESCE(tenant_id, ''), entity_key);
+
 -- Index for user-scoped lookups (when user_id IS NOT NULL)
 CREATE INDEX IF NOT EXISTS idx_kv_store_user ON kv_store (tenant_id, user_id)
 WHERE user_id IS NOT NULL;
@@ -173,7 +173,7 @@ COMMENT ON COLUMN kv_store.entity_id IS
 'UUID from primary table for reverse lookup';
 
 COMMENT ON COLUMN kv_store.tenant_id IS
-'Tenant identifier for multi-tenancy isolation';
+'Tenant identifier for multi-tenancy isolation. NULL = public/shared data visible to all.';
 
 COMMENT ON COLUMN kv_store.user_id IS
 'Optional user scoping. NULL = system-level entity, visible to all users in tenant';
@@ -271,8 +271,12 @@ BEGIN
     AND kv.entity_key = normalize_key(p_entity_key)
     LIMIT 1;
 
-    -- If not found, return empty
+    -- If not found, check if cache is empty and maybe trigger rebuild
     IF entity_table IS NULL THEN
+        -- SELF-HEALING: Check if this is because cache is empty
+        IF rem_kv_store_empty(effective_user_id) THEN
+            PERFORM maybe_trigger_kv_rebuild(effective_user_id, 'rem_lookup');
+        END IF;
         RETURN;
     END IF;
 
@@ -357,6 +361,7 @@ DECLARE
     entities_by_table JSONB := '{}'::jsonb;
     table_keys JSONB;
     effective_user_id VARCHAR(100);
+    v_found_any BOOLEAN := FALSE;
 BEGIN
     effective_user_id := COALESCE(p_user_id, p_tenant_id);
 
@@ -373,6 +378,7 @@ BEGIN
         ORDER BY sim_score DESC
         LIMIT p_limit
     LOOP
+        v_found_any := TRUE;
         -- Build JSONB mapping {table: [keys]}
         IF entities_by_table ? kv_matches.entity_type THEN
             table_keys := entities_by_table->kv_matches.entity_type;
@@ -390,6 +396,11 @@ BEGIN
         END IF;
     END LOOP;
 
+    -- SELF-HEALING: If no matches and cache is empty, trigger rebuild
+    IF NOT v_found_any AND rem_kv_store_empty(effective_user_id) THEN
+        PERFORM maybe_trigger_kv_rebuild(effective_user_id, 'rem_fuzzy');
+    END IF;
+
     -- Fetch full records using rem_fetch (which now supports NULL user_id)
     RETURN QUERY
     SELECT
@@ -436,9 +447,25 @@ DECLARE
     entities_by_table JSONB := '{}'::jsonb;
     table_keys JSONB;
     effective_user_id VARCHAR(100);
+    v_found_start BOOLEAN := FALSE;
 BEGIN
     effective_user_id := COALESCE(p_user_id, p_tenant_id);
 
+    -- Check if start entity exists in kv_store
+    SELECT TRUE INTO v_found_start
+    FROM kv_store kv
+    WHERE (kv.user_id = effective_user_id OR kv.user_id IS NULL)
+    AND kv.entity_key = normalize_key(p_entity_key)
+    LIMIT 1;
+
+    -- SELF-HEALING: If start not found and cache is empty, trigger rebuild
+    IF NOT COALESCE(v_found_start, FALSE) THEN
+        IF rem_kv_store_empty(effective_user_id) THEN
+            PERFORM maybe_trigger_kv_rebuild(effective_user_id, 'rem_traverse');
+        END IF;
+        RETURN;
+    END IF;
+
     FOR graph_keys IN
         WITH RECURSIVE graph_traversal AS (
             -- Base case: Find starting entity (user-owned OR public)
@@ -789,6 +816,97 @@ $$ LANGUAGE plpgsql STABLE;
 COMMENT ON FUNCTION fn_get_shared_messages IS
 'Get messages from sessions shared by a specific user with the recipient.';
 
+-- ============================================================================
+-- SESSIONS WITH USER INFO
+-- ============================================================================
+-- Function to list sessions with user details (name, email) for admin views
+
+-- List sessions with user info, CTE pagination
+-- Note: messages.session_id stores the session name (not UUID), so we join on sessions.name
+CREATE OR REPLACE FUNCTION fn_list_sessions_with_user(
+    p_user_id VARCHAR(256) DEFAULT NULL,  -- Filter by user_id (NULL = all users, admin only)
+    p_user_name VARCHAR(256) DEFAULT NULL,  -- Filter by user name (partial match, admin only)
+    p_user_email VARCHAR(256) DEFAULT NULL,  -- Filter by user email (partial match, admin only)
+    p_mode VARCHAR(50) DEFAULT NULL,  -- Filter by session mode
+    p_page INTEGER DEFAULT 1,
+    p_page_size INTEGER DEFAULT 50
+)
+RETURNS TABLE(
+    id UUID,
+    name VARCHAR(256),
+    mode TEXT,
+    description TEXT,
+    user_id VARCHAR(256),
+    user_name VARCHAR(256),
+    user_email VARCHAR(256),
+    message_count INTEGER,
+    total_tokens INTEGER,
+    created_at TIMESTAMP,
+    updated_at TIMESTAMP,
+    metadata JSONB,
+    total_count BIGINT
+) AS $$
+BEGIN
+    RETURN QUERY
+    WITH session_msg_counts AS (
+        -- Count messages per session (joining on session name since messages.session_id = sessions.name)
+        SELECT
+            m.session_id as session_name,
+            COUNT(*)::INTEGER as actual_message_count
+        FROM messages m
+        GROUP BY m.session_id
+    ),
+    filtered_sessions AS (
+        SELECT
+            s.id,
+            s.name,
+            s.mode,
+            s.description,
+            s.user_id,
+            COALESCE(u.name, s.user_id)::VARCHAR(256) AS user_name,
+            u.email::VARCHAR(256) AS user_email,
+            COALESCE(mc.actual_message_count, 0) AS message_count,
+            s.total_tokens,
+            s.created_at,
+            s.updated_at,
+            s.metadata
+        FROM sessions s
+        LEFT JOIN users u ON u.id::text = s.user_id
+        LEFT JOIN session_msg_counts mc ON mc.session_name = s.name
+        WHERE s.deleted_at IS NULL
+          AND (p_user_id IS NULL OR s.user_id = p_user_id)
+          AND (p_user_name IS NULL OR u.name ILIKE '%' || p_user_name || '%')
+          AND (p_user_email IS NULL OR u.email ILIKE '%' || p_user_email || '%')
+          AND (p_mode IS NULL OR s.mode = p_mode)
+    ),
+    counted AS (
+        SELECT *, COUNT(*) OVER () AS total_count
+        FROM filtered_sessions
+    )
+    SELECT
+        c.id,
+        c.name,
+        c.mode,
+        c.description,
+        c.user_id,
+        c.user_name,
+        c.user_email,
+        c.message_count,
+        c.total_tokens,
+        c.created_at,
+        c.updated_at,
+        c.metadata,
+        c.total_count
+    FROM counted c
+    ORDER BY c.created_at DESC
+    LIMIT p_page_size
+    OFFSET (p_page - 1) * p_page_size;
+END;
+$$ LANGUAGE plpgsql STABLE;
+
+COMMENT ON FUNCTION fn_list_sessions_with_user IS
+'List sessions with user details and computed message counts. Joins messages on session name.';
+
 -- ============================================================================
 -- RECORD INSTALLATION
 -- ============================================================================