remdb 0.3.171__py3-none-any.whl → 0.3.230__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,
@@ -268,7 +269,7 @@ BEGIN
             graph_edges,
             updated_at
         ) VALUES (
-            NEW.{entity_key_field}::VARCHAR,
+            normalize_key(NEW.{entity_key_field}::VARCHAR),
             '{table_name}',
             NEW.id,
             NEW.tenant_id,
@@ -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

@@ -25,7 +25,6 @@ from .sql_builder import (
     build_select,
     build_upsert,
 )
-from ...settings import settings
 
 if TYPE_CHECKING:
     from .service import PostgresService
@@ -33,15 +32,15 @@ if TYPE_CHECKING:
 
 def get_postgres_service() -> "PostgresService | None":
     """
-    Get PostgresService instance with connection string from settings.
+    Get PostgresService singleton from parent module.
 
-    Returns None if Postgres is disabled.
+    Uses late import to avoid circular import issues.
+    Previously had a separate _postgres_instance here which caused
+    "pool not connected" errors due to duplicate connection pools.
     """
-    if not settings.postgres.enabled:
-        return None
-    
-    from .service import PostgresService
-    return PostgresService()
+    # Late import to avoid circular import (repository.py imported by __init__.py)
+    from rem.services.postgres import get_postgres_service as _get_singleton
+    return _get_singleton()
 
 T = TypeVar("T", bound=BaseModel)
 
@@ -74,7 +73,7 @@ class Repository(Generic[T]):
         self,
         records: T | list[T],
         embeddable_fields: list[str] | None = None,
-        generate_embeddings: bool = False,
+        generate_embeddings: bool = True,
     ) -> T | list[T]:
         """
         Upsert single record or list of records (create or update on ID conflict).
@@ -84,8 +83,9 @@ class Repository(Generic[T]):
 
         Args:
             records: Single model instance or list of model instances
-            embeddable_fields: Optional list of fields to generate embeddings for
-            generate_embeddings: Whether to queue embedding generation tasks
+            embeddable_fields: Optional list of fields to generate embeddings for.
+                              If None, auto-detects 'content' field if present.
+            generate_embeddings: Whether to queue embedding generation tasks (default: True)
 
         Returns:
             Single record or list of records with generated IDs (matches input type)
@@ -118,25 +118,35 @@ class Repository(Generic[T]):
                     record.id = row["id"]  # type: ignore[attr-defined]
 
         # Queue embedding generation if requested and worker is available
-        if generate_embeddings and embeddable_fields and self.db.embedding_worker:
+        if generate_embeddings and self.db.embedding_worker:
             from rem.services.embeddings import EmbeddingTask
-
-            for record in records_list:
-                for field_name in embeddable_fields:
-                    content = getattr(record, field_name, None)
-                    if content and isinstance(content, str):
-                        task = EmbeddingTask(
-                            task_id=f"{record.id}-{field_name}",  # type: ignore[attr-defined]
-                            entity_id=str(record.id),  # type: ignore[attr-defined]
-                            table_name=self.table_name,
-                            field_name=field_name,
-                            content=content,
-                            provider="openai",  # Default provider
-                            model="text-embedding-3-small",  # Default model
-                        )
-                        await self.db.embedding_worker.queue_task(task)
-
-            logger.debug(f"Queued {len(records_list) * len(embeddable_fields)} embedding tasks")
+            from .register_type import should_embed_field
+
+            # Auto-detect embeddable fields if not specified
+            if embeddable_fields is None:
+                embeddable_fields = [
+                    field_name
+                    for field_name, field_info in self.model_class.model_fields.items()
+                    if should_embed_field(field_name, field_info)
+                ]
+
+            if embeddable_fields:
+                for record in records_list:
+                    for field_name in embeddable_fields:
+                        content = getattr(record, field_name, None)
+                        if content and isinstance(content, str):
+                            task = EmbeddingTask(
+                                task_id=f"{record.id}-{field_name}",  # type: ignore[attr-defined]
+                                entity_id=str(record.id),  # type: ignore[attr-defined]
+                                table_name=self.table_name,
+                                field_name=field_name,
+                                content=content,
+                                provider="openai",  # Default provider
+                                model="text-embedding-3-small",  # Default model
+                            )
+                            await self.db.embedding_worker.queue_task(task)
+
+                logger.debug(f"Queued {len(records_list) * len(embeddable_fields)} embedding tasks")
 
         # Return single item or list to match input type
         return records_list[0] if is_single else records_list

rem/services/postgres/schema_generator.py

@@ -351,10 +351,10 @@ class SchemaGenerator:
 
         Priority:
         1. Field with json_schema_extra={\"entity_key\": True}
-        2. Field named \"name\"
+        2. Field named \"name\" (human-readable identifier)
         3. Field named \"key\"
-        4. Field named \"label\"
-        5. First string field
+        4. Field named \"uri\"
+        5. Field named \"id\" (fallback)
 
         Args:
             model: Pydantic model class
@@ -369,9 +369,9 @@ class SchemaGenerator:
                 if json_extra.get("entity_key"):
                     return field_name
 
-        # Check for key fields in priority order: id -> uri -> key -> name
+        # Check for key fields in priority order: name -> key -> uri -> id
         # (matching sql_builder.get_entity_key convention)
-        for candidate in ["id", "uri", "key", "name"]:
+        for candidate in ["name", "key", "uri", "id"]:
             if candidate in model.model_fields:
                 return candidate
 

rem/services/postgres/sql_builder.py

@@ -35,10 +35,11 @@ def get_natural_key(model: BaseModel) -> str | None:
 
 def get_entity_key(model: BaseModel) -> str:
     """
-    Get entity key for KV store following precedence: id -> uri -> key -> name.
+    Get entity key for KV store following precedence: name -> key -> uri -> id.
 
-    For KV store lookups, we prefer globally unique identifiers first (id),
-    then natural keys (uri/key/name). Always returns a value (id as fallback).
+    For KV store lookups, we prefer human-readable identifiers first (name/key),
+    then URIs, with id as the fallback. This allows users to lookup entities
+    by their natural names like "panic-disorder" instead of UUIDs.
 
     Args:
         model: Pydantic model instance
@@ -46,13 +47,13 @@ def get_entity_key(model: BaseModel) -> str:
     Returns:
         Entity key string (guaranteed to exist)
     """
-    for field in ["id", "uri", "key", "name"]:
+    for field in ["name", "key", "uri", "id"]:
         if hasattr(model, field):
             value = getattr(model, field)
             if value:
                 return str(value)
     # Should never reach here since id always exists in CoreModel
-    raise ValueError(f"Model {type(model)} has no id, uri, key, or name field")
+    raise ValueError(f"Model {type(model)} has no name, key, uri, or id field")
 
 
 def generate_deterministic_id(user_id: str | None, entity_key: str) -> uuid.UUID:

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,292 @@
+"""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
+import re
+from typing import Any
+
+from loguru import logger
+from pydantic_ai.messages import (
+    ModelMessage,
+    ModelRequest,
+    ModelResponse,
+    SystemPromptPart,
+    TextPart,
+    ToolCallPart,
+    ToolReturnPart,
+    UserPromptPart,
+)
+
+
+def _sanitize_tool_name(tool_name: str) -> str:
+    """Sanitize tool name for OpenAI API compatibility.
+
+    OpenAI requires tool names to match pattern: ^[a-zA-Z0-9_-]+$
+    This replaces invalid characters (like colons) with underscores.
+    """
+    return re.sub(r'[^a-zA-Z0-9_-]', '_', tool_name)
+
+
+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
+
+                # Sanitize tool name for OpenAI API compatibility
+                safe_tool_name = _sanitize_tool_name(tool_name)
+
+                # Synthesize ToolCallPart (what the model "called")
+                tool_calls.append(ToolCallPart(
+                    tool_name=safe_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=safe_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
+
+            # Sanitize tool name for OpenAI API compatibility
+            safe_tool_name = _sanitize_tool_name(tool_name)
+
+            # Synthesize the tool call (ModelResponse with ToolCallPart)
+            messages.append(ModelResponse(
+                parts=[ToolCallPart(
+                    tool_name=safe_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=safe_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}")