remdb 0.3.7__py3-none-any.whl

rem/services/postgres/register_type.py

@@ -0,0 +1,352 @@
+"""
+Dynamic table and embeddings schema generator from Pydantic models.
+
+Generates:
+1. Primary table for entity storage
+2. embeddings_<table> for vector embeddings (one row per field per provider)
+3. Registers entity in KV_STORE cache
+4. Background index creation for performance
+
+Design Patterns:
+- Fields marked with json_schema_extra={\"embed\": True} get embeddings
+- Content fields (TextField, description, etc.) embed by default
+- Multiple embedding providers supported (OpenAI, Cohere, etc.)
+- UNLOGGED KV_STORE for O(1) lookups
+- Background index creation to avoid blocking writes
+"""
+
+from typing import Any, Type
+
+from loguru import logger
+from pydantic import BaseModel
+
+from ...utils.sql_types import get_column_definition
+
+# Embedding configuration
+DEFAULT_EMBEDDING_PROVIDER = "openai"
+DEFAULT_EMBEDDING_MODEL = "text-embedding-3-small"
+DEFAULT_EMBEDDING_DIMENSIONS = 1536
+
+# Fields that embed by default (if not explicitly disabled)
+DEFAULT_EMBED_FIELD_NAMES = {
+    "content",
+    "description",
+    "summary",
+    "text",
+    "body",
+    "message",
+    "notes",
+}
+
+
+def should_embed_field(field_name: str, field_info: Any) -> bool:
+    """
+    Determine if a field should have embeddings generated.
+
+    Rules:
+    1. If json_schema_extra.embed = True, always embed
+    2. If json_schema_extra.embed = False, never embed
+    3. If field name in DEFAULT_EMBED_FIELD_NAMES, embed by default
+    4. Otherwise, don't embed
+
+    Args:
+        field_name: Field name from Pydantic model
+        field_info: Field metadata from model.model_fields
+
+    Returns:
+        True if field should have embeddings
+    """
+    # Check json_schema_extra for explicit embed configuration
+    json_extra = getattr(field_info, "json_schema_extra", None)
+    if json_extra:
+        if isinstance(json_extra, dict):
+            embed = json_extra.get("embed")
+            if embed is not None:
+                return bool(embed)
+
+    # Default: embed if field name matches common content fields
+    return field_name.lower() in DEFAULT_EMBED_FIELD_NAMES
+
+
+def generate_table_schema(
+    model: Type[BaseModel], table_name: str, tenant_scoped: bool = True
+) -> str:
+    """
+    Generate CREATE TABLE SQL for Pydantic model.
+
+    Args:
+        model: Pydantic model class
+        table_name: Table name (e.g., "resources", "moments")
+        tenant_scoped: If True, add tenant_id column and indexes
+
+    Returns:
+        SQL CREATE TABLE statement
+    """
+    columns = []
+    indexes = []
+
+    # System fields that we add separately (skip if in model)
+    SYSTEM_FIELDS = {
+        "id", "created_at", "updated_at", "deleted_at",
+        "tenant_id", "user_id", "graph_edges", "metadata", "tags", "column"
+    }
+
+    # Always add id as primary key
+    columns.append("id UUID PRIMARY KEY DEFAULT uuid_generate_v4()")
+
+    # Add tenant_id if tenant scoped
+    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);")
+
+    # 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);")
+
+    # Process Pydantic fields (skip system fields)
+    for field_name, field_info in model.model_fields.items():
+        if field_name in SYSTEM_FIELDS:
+            continue  # Skip system fields - we add them separately
+
+        # Use sql_types utility for consistent type mapping
+        column_def = get_column_definition(
+            field_info,
+            field_name,
+            nullable=not field_info.is_required(),
+            primary_key=False
+        )
+        columns.append(column_def)
+
+    # Add system fields (timestamps)
+    columns.append("created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP")
+    columns.append("updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP")
+    columns.append("deleted_at TIMESTAMP")
+
+    # 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);"
+    )
+
+    # 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);"
+    )
+
+    # 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);"
+    )
+
+    # Generate CREATE TABLE statement
+    create_table = f"""
+CREATE TABLE IF NOT EXISTS {table_name} (
+    {',\n    '.join(columns)}
+);
+""".strip()
+
+    # Generate indexes
+    index_sql = "\n".join(indexes)
+
+    return f"{create_table}\n\n{index_sql}"
+
+
+def generate_embeddings_schema(
+    model: Type[BaseModel], table_name: str, embedding_provider: str = DEFAULT_EMBEDDING_PROVIDER
+) -> tuple[str, list[str]]:
+    """
+    Generate embeddings table schema for a model.
+
+    Creates embeddings_<table_name> with:
+    - One row per entity per field per provider
+    - Unique constraint on (entity_id, field_name, provider)
+    - Vector column with pgvector
+    - HNSW index for fast similarity search
+
+    Args:
+        model: Pydantic model class
+        table_name: Base table name
+        embedding_provider: Default provider (e.g., "openai", "cohere")
+
+    Returns:
+        Tuple of (CREATE TABLE sql, list of embeddable field names)
+    """
+    embeddings_table = f"embeddings_{table_name}"
+    embeddable_fields = []
+
+    # Find fields that should have embeddings
+    for field_name, field_info in model.model_fields.items():
+        if should_embed_field(field_name, field_info):
+            embeddable_fields.append(field_name)
+
+    if not embeddable_fields:
+        logger.warning(f"No embeddable fields found for {table_name}")
+        return "", []
+
+    # Generate embeddings table
+    create_sql = f"""
+CREATE TABLE IF NOT EXISTS {embeddings_table} (
+    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
+    entity_id UUID NOT NULL REFERENCES {table_name}(id) ON DELETE CASCADE,
+    field_name VARCHAR(100) NOT NULL,
+    provider VARCHAR(50) NOT NULL DEFAULT '{embedding_provider}',
+    model VARCHAR(100) NOT NULL DEFAULT '{DEFAULT_EMBEDDING_MODEL}',
+    embedding vector({DEFAULT_EMBEDDING_DIMENSIONS}) NOT NULL,
+    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+
+    -- Unique: one embedding per entity per field per provider
+    UNIQUE (entity_id, field_name, provider)
+);
+
+-- Index for entity lookup (get all embeddings for entity)
+CREATE INDEX 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);
+
+-- HNSW index for vector similarity search (created in background)
+-- Note: This will be created by background thread after data load
+-- CREATE INDEX idx_{embeddings_table}_vector_hnsw ON {embeddings_table}
+-- USING hnsw (embedding vector_cosine_ops);
+""".strip()
+
+    logger.info(
+        f"Generated embeddings schema for {table_name} with fields: {embeddable_fields}"
+    )
+
+    return create_sql, embeddable_fields
+
+
+# NOTE: _map_pydantic_to_postgres_type is now replaced by utils.sql_types.get_sql_type
+# Removed to use the centralized utility instead
+
+
+def generate_kv_store_upsert(
+    table_name: str,
+    entity_key_field: str = "name",
+) -> str:
+    """
+    Generate trigger to maintain KV_STORE cache on entity changes.
+
+    Creates a trigger that:
+    1. Extracts entity_key from entity (e.g., name, key, label)
+    2. Updates KV_STORE on INSERT/UPDATE for O(1) lookups
+    3. Removes from KV_STORE on DELETE
+
+    Args:
+        table_name: Base table name
+        entity_key_field: Field to use as entity_key in KV_STORE
+
+    Returns:
+        SQL for trigger creation
+    """
+    trigger_name = f"trg_{table_name}_kv_store"
+    function_name = f"fn_{table_name}_kv_store_upsert"
+
+    return f"""
+-- Trigger function to maintain KV_STORE for {table_name}
+CREATE OR REPLACE FUNCTION {function_name}()
+RETURNS TRIGGER AS $$
+BEGIN
+    IF (TG_OP = 'DELETE') THEN
+        -- Remove from KV_STORE on delete
+        DELETE FROM kv_store
+        WHERE entity_id = OLD.id;
+        RETURN OLD;
+    ELSIF (TG_OP = 'INSERT' OR TG_OP = 'UPDATE') THEN
+        -- Upsert to KV_STORE (O(1) lookup by entity_key)
+        INSERT INTO kv_store (
+            entity_key,
+            entity_type,
+            entity_id,
+            tenant_id,
+            user_id,
+            metadata,
+            graph_edges,
+            updated_at
+        ) VALUES (
+            NEW.{entity_key_field}::VARCHAR,
+            '{table_name}',
+            NEW.id,
+            NEW.tenant_id,
+            NEW.user_id,
+            NEW.metadata,
+            COALESCE(NEW.graph_edges, '[]'::jsonb),
+            CURRENT_TIMESTAMP
+        )
+        ON CONFLICT (tenant_id, entity_key)
+        DO UPDATE SET
+            entity_id = EXCLUDED.entity_id,
+            user_id = EXCLUDED.user_id,
+            metadata = EXCLUDED.metadata,
+            graph_edges = EXCLUDED.graph_edges,
+            updated_at = CURRENT_TIMESTAMP;
+
+        RETURN NEW;
+    END IF;
+END;
+$$ LANGUAGE plpgsql;
+
+-- Create trigger
+DROP TRIGGER IF EXISTS {trigger_name} ON {table_name};
+CREATE TRIGGER {trigger_name}
+AFTER INSERT OR UPDATE OR DELETE ON {table_name}
+FOR EACH ROW EXECUTE FUNCTION {function_name}();
+""".strip()
+
+
+async def register_type(
+    model: Type[BaseModel],
+    table_name: str,
+    entity_key_field: str = "name",
+    tenant_scoped: bool = True,
+    create_embeddings: bool = True,
+    create_kv_trigger: bool = True,
+) -> dict[str, Any]:
+    """
+    Register a Pydantic model as a database schema.
+
+    Creates:
+    1. Primary table for entity storage
+    2. Embeddings table (if create_embeddings=True)
+    3. KV_STORE trigger (if create_kv_trigger=True)
+
+    Args:
+        model: Pydantic model class
+        table_name: Table name
+        entity_key_field: Field to use as natural key in KV_STORE
+        tenant_scoped: Add tenant_id column and indexes
+        create_embeddings: Create embeddings table
+        create_kv_trigger: Create KV_STORE trigger
+
+    Returns:
+        Dict with SQL statements and metadata
+    """
+    result = {
+        "table_name": table_name,
+        "model": model.__name__,
+        "sql": {},
+        "embeddable_fields": [],
+    }
+
+    # Generate primary table schema
+    table_sql = generate_table_schema(model, table_name, tenant_scoped)
+    result["sql"]["table"] = table_sql
+
+    # Generate embeddings schema
+    if create_embeddings:
+        embeddings_sql, embeddable_fields = generate_embeddings_schema(model, table_name)
+        result["sql"]["embeddings"] = embeddings_sql
+        result["embeddable_fields"] = embeddable_fields
+
+    # Generate KV_STORE trigger
+    if create_kv_trigger:
+        kv_trigger_sql = generate_kv_store_upsert(table_name, entity_key_field)
+        result["sql"]["kv_trigger"] = kv_trigger_sql
+
+    logger.info(f"Registered type {model.__name__} as table {table_name}")
+
+    return result

rem/services/postgres/repository.py

@@ -0,0 +1,337 @@
+"""Generic repository for entity persistence.
+
+Single repository class that works with any Pydantic model type.
+No need for model-specific repository classes.
+
+Usage:
+    from rem.models.entities import Message
+    from rem.services.repositories import Repository
+
+    repo = Repository(db, Message, table_name="messages")
+    message = await repo.upsert(message_instance)
+    messages = await repo.find({"session_id": "abc", "tenant_id": "xyz"})
+"""
+
+import json
+from typing import Any, Generic, Type, TypeVar, TYPE_CHECKING
+
+from loguru import logger
+from pydantic import BaseModel
+
+from .sql_builder import (
+    build_count,
+    build_delete,
+    build_insert,
+    build_select,
+    build_upsert,
+)
+from ...settings import settings
+
+if TYPE_CHECKING:
+    from .service import PostgresService
+
+
+def get_postgres_service() -> "PostgresService | None":
+    """
+    Get PostgresService instance with connection string from settings.
+
+    Returns None if Postgres is disabled.
+    """
+    if not settings.postgres.enabled:
+        return None
+    
+    from .service import PostgresService
+    return PostgresService()
+
+T = TypeVar("T", bound=BaseModel)
+
+# Known JSONB fields from CoreModel that need deserialization
+JSONB_FIELDS = {"graph_edges", "metadata"}
+
+
+class Repository(Generic[T]):
+    """Generic repository for any Pydantic model type."""
+
+    def __init__(
+        self,
+        model_class: Type[T],
+        table_name: str | None = None,
+        db: "PostgresService | None" = None,
+    ):
+        """
+        Initialize repository.
+
+        Args:
+            model_class: Pydantic model class (e.g., Message, Resource)
+            table_name: Optional table name (defaults to lowercase model name + 's')
+            db: Optional PostgresService instance (creates from settings if None)
+        """
+        self.db = db or get_postgres_service()
+        self.model_class = model_class
+        self.table_name = table_name or f"{model_class.__name__.lower()}s"
+
+    async def upsert(
+        self,
+        records: T | list[T],
+        embeddable_fields: list[str] | None = None,
+        generate_embeddings: bool = False,
+    ) -> T | list[T]:
+        """
+        Upsert single record or list of records (create or update on ID conflict).
+
+        Accepts both single items and lists - no need to distinguish batch vs non-batch.
+        Single items are coerced to lists internally for processing.
+
+        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
+
+        Returns:
+            Single record or list of records with generated IDs (matches input type)
+        """
+        # Coerce single item to list for uniform processing
+        is_single = not isinstance(records, list)
+        records_list: list[T]
+        if is_single:
+            records_list = [records]  # type: ignore[list-item]
+        else:
+            records_list = records  # Type narrowed by isinstance check
+
+        if not settings.postgres.enabled or not self.db:
+            logger.debug(f"Postgres disabled, skipping {self.model_class.__name__} upsert")
+            return records
+
+        # Ensure connection
+        if not self.db.pool:
+            await self.db.connect()
+
+        # Type guard: ensure pool is not None after connect
+        if not self.db.pool:
+            raise RuntimeError("Failed to establish database connection")
+
+        for record in records_list:
+            sql, params = build_upsert(record, self.table_name, conflict_field="id", return_id=True)
+            async with self.db.pool.acquire() as conn:
+                row = await conn.fetchrow(sql, *params)
+                if row and "id" in row:
+                    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:
+            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")
+
+        # Return single item or list to match input type
+        return records_list[0] if is_single else records_list
+
+    async def get_by_id(self, record_id: str, tenant_id: str) -> T | None:
+        """
+        Get a single record by ID.
+
+        Args:
+            record_id: Record identifier
+            tenant_id: Tenant identifier for multi-tenancy isolation
+
+        Returns:
+            Model instance or None if not found
+        """
+        if not settings.postgres.enabled or not self.db:
+            logger.debug(f"Postgres disabled, returning None for {self.model_class.__name__} get")
+            return None
+
+        # Ensure connection
+        if not self.db.pool:
+            await self.db.connect()
+
+        # Type guard: ensure pool is not None after connect
+        if not self.db.pool:
+            raise RuntimeError("Failed to establish database connection")
+
+        query = f"""
+            SELECT * FROM {self.table_name}
+            WHERE id = $1 AND tenant_id = $2 AND deleted_at IS NULL
+        """
+
+        async with self.db.pool.acquire() as conn:
+            row = await conn.fetchrow(query, record_id, tenant_id)
+
+        if not row:
+            return None
+
+        # PostgreSQL JSONB columns come back as strings, need to parse them
+        row_dict = dict(row)
+        return self.model_class.model_validate(row_dict)
+
+    async def find(
+        self,
+        filters: dict[str, Any],
+        order_by: str = "created_at ASC",
+        limit: int | None = None,
+        offset: int = 0,
+    ) -> list[T]:
+        """
+        Find records matching filters.
+
+        Args:
+            filters: Dict of field -> value filters (AND-ed together)
+            order_by: ORDER BY clause (default: "created_at ASC")
+            limit: Optional limit on number of records
+            offset: Offset for pagination
+
+        Returns:
+            List of model instances
+
+        Example:
+            messages = await repo.find({
+                "session_id": "abc-123",
+                "tenant_id": "acme-corp",
+                "user_id": "alice"
+            })
+        """
+        if not settings.postgres.enabled or not self.db:
+            logger.debug(f"Postgres disabled, returning empty {self.model_class.__name__} list")
+            return []
+
+        # Ensure connection
+        if not self.db.pool:
+            await self.db.connect()
+
+        # Type guard: ensure pool is not None after connect
+        if not self.db.pool:
+            raise RuntimeError("Failed to establish database connection")
+
+        sql, params = build_select(
+            self.model_class,
+            self.table_name,
+            filters,
+            order_by=order_by,
+            limit=limit,
+            offset=offset,
+        )
+
+        async with self.db.pool.acquire() as conn:
+            rows = await conn.fetch(sql, *params)
+
+        return [self.model_class.model_validate(dict(row)) for row in rows]
+
+    async def find_one(self, filters: dict[str, Any]) -> T | None:
+        """
+        Find single record matching filters.
+
+        Args:
+            filters: Dict of field -> value filters
+
+        Returns:
+            Model instance or None if not found
+        """
+        results = await self.find(filters, limit=1)
+        return results[0] if results else None
+
+    async def get_by_session(
+        self, session_id: str, tenant_id: str, user_id: str | None = None
+    ) -> list[T]:
+        """
+        Get all records for a session (convenience method for Message model).
+
+        Args:
+            session_id: Session identifier
+            tenant_id: Tenant identifier
+            user_id: Optional user identifier
+
+        Returns:
+            List of model instances ordered by created_at
+        """
+        filters = {"session_id": session_id, "tenant_id": tenant_id}
+        if user_id:
+            filters["user_id"] = user_id
+
+        return await self.find(filters, order_by="created_at ASC")
+
+    async def update(self, record: T) -> T:
+        """
+        Update a record (upsert).
+
+        Args:
+            record: Model instance to update
+
+        Returns:
+            Updated record
+        """
+        result = await self.upsert(record)
+        # upsert with single record returns single record
+        return result  # type: ignore[return-value]
+
+    async def delete(self, record_id: str, tenant_id: str) -> bool:
+        """
+        Soft delete a record (sets deleted_at).
+
+        Args:
+            record_id: Record identifier
+            tenant_id: Tenant identifier for multi-tenancy isolation
+
+        Returns:
+            True if deleted, False if not found
+        """
+        if not settings.postgres.enabled or not self.db:
+            logger.debug(f"Postgres disabled, skipping {self.model_class.__name__} deletion")
+            return False
+
+        # Ensure connection
+        if not self.db.pool:
+            await self.db.connect()
+
+        # Type guard: ensure pool is not None after connect
+        if not self.db.pool:
+            raise RuntimeError("Failed to establish database connection")
+
+        sql, params = build_delete(self.table_name, record_id, tenant_id)
+
+        async with self.db.pool.acquire() as conn:
+            row = await conn.fetchrow(sql, *params)
+
+        return row is not None
+
+    async def count(self, filters: dict[str, Any]) -> int:
+        """
+        Count records matching filters.
+
+        Args:
+            filters: Dict of field -> value filters
+
+        Returns:
+            Count of matching records
+        """
+        if not settings.postgres.enabled or not self.db:
+            return 0
+
+        # Ensure connection
+        if not self.db.pool:
+            await self.db.connect()
+
+        # Type guard: ensure pool is not None after connect
+        if not self.db.pool:
+            raise RuntimeError("Failed to establish database connection")
+
+        sql, params = build_count(self.table_name, filters)
+
+        async with self.db.pool.acquire() as conn:
+            row = await conn.fetchrow(sql, *params)
+
+        return row[0] if row else 0