remdb 0.3.181__py3-none-any.whl → 0.3.223__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/postgres/repository.py

@@ -31,17 +31,27 @@ if TYPE_CHECKING:
     from .service import PostgresService
 
 
+# Singleton instance for connection pool reuse
+_postgres_instance: "PostgresService | None" = None
+
+
 def get_postgres_service() -> "PostgresService | None":
     """
-    Get PostgresService instance with connection string from settings.
+    Get PostgresService singleton instance.
 
     Returns None if Postgres is disabled.
+    Uses singleton pattern to prevent connection pool exhaustion.
     """
+    global _postgres_instance
+
     if not settings.postgres.enabled:
         return None
-    
-    from .service import PostgresService
-    return PostgresService()
+
+    if _postgres_instance is None:
+        from .service import PostgresService
+        _postgres_instance = PostgresService()
+
+    return _postgres_instance
 
 T = TypeVar("T", bound=BaseModel)
 

rem/services/session/__init__.py

@@ -1,6 +1,13 @@
 """Session management services for conversation persistence and compression."""
 
 from .compression import MessageCompressor, SessionMessageStore
+from .pydantic_messages import audit_session_history, session_to_pydantic_messages
 from .reload import reload_session
 
-__all__ = ["MessageCompressor", "SessionMessageStore", "reload_session"]
+__all__ = [
+    "MessageCompressor",
+    "SessionMessageStore",
+    "audit_session_history",
+    "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,39 @@ 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 UUID from X-Session-Id header
+            user_id: Optional user identifier
+        """
+        try:
+            # Check if session already exists by UUID
+            existing = await self._session_repo.get_by_id(session_id)
+            if existing:
+                return  # Session already exists
+
+            # Create new session with the provided UUID as id
+            session = Session(
+                id=session_id,  # Use the provided UUID as session id
+                name=session_id,  # Default name to UUID, can be updated later
+                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 +316,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 UUID
             messages: List of messages to store
             user_id: Optional user identifier
             compress: Whether to compress messages (default: True)
@@ -296,6 +331,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,276 @@
+"""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
+
+
+def audit_session_history(
+    session_id: str,
+    agent_name: str,
+    prompt: str,
+    raw_session_history: list[dict[str, Any]],
+    pydantic_messages_count: int,
+) -> None:
+    """
+    Dump session history to a YAML file for debugging.
+
+    Only runs when DEBUG__AUDIT_SESSION=true. Writes to DEBUG__AUDIT_DIR (default /tmp).
+    Appends to the same file for a session, so all agent invocations are in one place.
+
+    Args:
+        session_id: The session identifier
+        agent_name: Name of the agent being invoked
+        prompt: The prompt being sent to the agent
+        raw_session_history: The raw session messages from the database
+        pydantic_messages_count: Count of converted pydantic-ai messages
+    """
+    from ...settings import settings
+
+    if not settings.debug.audit_session:
+        return
+
+    try:
+        import yaml
+        from pathlib import Path
+        from ...utils.date_utils import utc_now, to_iso
+
+        audit_dir = Path(settings.debug.audit_dir)
+        audit_dir.mkdir(parents=True, exist_ok=True)
+        audit_file = audit_dir / f"{session_id}.yaml"
+
+        # Create entry for this agent invocation
+        entry = {
+            "timestamp": to_iso(utc_now()),
+            "agent_name": agent_name,
+            "prompt": prompt,
+            "raw_history_count": len(raw_session_history),
+            "pydantic_messages_count": pydantic_messages_count,
+            "raw_session_history": raw_session_history,
+        }
+
+        # Load existing data or create new
+        existing_data: dict[str, Any] = {"session_id": session_id, "invocations": []}
+        if audit_file.exists():
+            with open(audit_file) as f:
+                loaded = yaml.safe_load(f)
+                if loaded:
+                    # Ensure session_id is always present (backfill if missing)
+                    existing_data = {
+                        "session_id": loaded.get("session_id", session_id),
+                        "invocations": loaded.get("invocations", []),
+                    }
+
+        # Append this invocation
+        existing_data["invocations"].append(entry)
+
+        with open(audit_file, "w") as f:
+            yaml.dump(existing_data, f, default_flow_style=False, allow_unicode=True)
+        logger.info(f"DEBUG: Session audit updated: {audit_file}")
+    except Exception as e:
+        logger.warning(f"DEBUG: Failed to dump session audit: {e}")

rem/settings.py

@@ -1651,6 +1651,33 @@ class EmailSettings(BaseSettings):
         return kwargs
 
 
+class DebugSettings(BaseSettings):
+    """
+    Debug settings for development and troubleshooting.
+
+    Environment variables:
+        DEBUG__AUDIT_SESSION - Dump session history to /tmp/{session_id}.yaml
+        DEBUG__AUDIT_DIR - Directory for session audit files (default: /tmp)
+    """
+
+    model_config = SettingsConfigDict(
+        env_prefix="DEBUG__",
+        env_file=".env",
+        env_file_encoding="utf-8",
+        extra="ignore",
+    )
+
+    audit_session: bool = Field(
+        default=False,
+        description="When true, dump full session history to audit files for debugging",
+    )
+
+    audit_dir: str = Field(
+        default="/tmp",
+        description="Directory for session audit files",
+    )
+
+
 class TestSettings(BaseSettings):
     """
     Test environment settings.
@@ -1767,6 +1794,7 @@ class Settings(BaseSettings):
     schema_search: SchemaSettings = Field(default_factory=SchemaSettings)
     email: EmailSettings = Field(default_factory=EmailSettings)
     test: TestSettings = Field(default_factory=TestSettings)
+    debug: DebugSettings = Field(default_factory=DebugSettings)
 
 
 # Auto-load .env file from current directory if it exists