remdb 0.3.14__py3-none-any.whl → 0.3.133__py3-none-any.whl

rem/services/postgres/repository.py

@@ -335,3 +335,135 @@ class Repository(Generic[T]):
             row = await conn.fetchrow(sql, *params)
 
         return row[0] if row else 0
+
+    async def find_paginated(
+        self,
+        filters: dict[str, Any],
+        page: int = 1,
+        page_size: int = 50,
+        order_by: str = "created_at DESC",
+        partition_by: str | None = None,
+    ) -> dict[str, Any]:
+        """
+        Find records with page-based pagination using CTE with ROW_NUMBER().
+
+        Uses a CTE with ROW_NUMBER() OVER (PARTITION BY ... ORDER BY ...) for
+        efficient pagination with total count in a single query.
+
+        Args:
+            filters: Dict of field -> value filters (AND-ed together)
+            page: Page number (1-indexed)
+            page_size: Number of records per page
+            order_by: ORDER BY clause for row numbering (default: "created_at DESC")
+            partition_by: Optional field to partition by (e.g., "user_id").
+                         If None, uses global row numbering.
+
+        Returns:
+            Dict containing:
+            - data: List of model instances for the page
+            - total: Total count of records matching filters
+            - page: Current page number
+            - page_size: Records per page
+            - total_pages: Total number of pages
+            - has_next: Whether there are more pages
+            - has_previous: Whether there are previous pages
+
+        Example:
+            result = await repo.find_paginated(
+                {"tenant_id": "acme", "user_id": "alice"},
+                page=2,
+                page_size=20,
+                order_by="created_at DESC",
+                partition_by="user_id"
+            )
+            # result = {
+            #     "data": [...],
+            #     "total": 150,
+            #     "page": 2,
+            #     "page_size": 20,
+            #     "total_pages": 8,
+            #     "has_next": True,
+            #     "has_previous": True
+            # }
+        """
+        if not settings.postgres.enabled or not self.db:
+            logger.debug(f"Postgres disabled, returning empty {self.model_class.__name__} pagination")
+            return {
+                "data": [],
+                "total": 0,
+                "page": page,
+                "page_size": page_size,
+                "total_pages": 0,
+                "has_next": False,
+                "has_previous": 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")
+
+        # Build WHERE clause from filters
+        where_conditions = ["deleted_at IS NULL"]
+        params: list[Any] = []
+        param_idx = 1
+
+        for field, value in filters.items():
+            where_conditions.append(f"{field} = ${param_idx}")
+            params.append(value)
+            param_idx += 1
+
+        where_clause = " AND ".join(where_conditions)
+
+        # Build PARTITION BY clause
+        partition_clause = f"PARTITION BY {partition_by}" if partition_by else ""
+
+        # Build the CTE query with ROW_NUMBER() and COUNT() window functions
+        # This gives us pagination + total count in a single query
+        sql = f"""
+        WITH numbered AS (
+            SELECT *,
+                   ROW_NUMBER() OVER ({partition_clause} ORDER BY {order_by}) as _row_num,
+                   COUNT(*) OVER ({partition_clause}) as _total_count
+            FROM {self.table_name}
+            WHERE {where_clause}
+        )
+        SELECT * FROM numbered
+        WHERE _row_num > ${param_idx} AND _row_num <= ${param_idx + 1}
+        ORDER BY _row_num
+        """
+
+        # Calculate row range for the page
+        start_row = (page - 1) * page_size
+        end_row = page * page_size
+        params.extend([start_row, end_row])
+
+        async with self.db.pool.acquire() as conn:
+            rows = await conn.fetch(sql, *params)
+
+        # Extract total from first row (all rows have the same _total_count)
+        total = rows[0]["_total_count"] if rows else 0
+
+        # Remove internal columns and convert to models
+        data = []
+        for row in rows:
+            row_dict = dict(row)
+            row_dict.pop("_row_num", None)
+            row_dict.pop("_total_count", None)
+            data.append(self.model_class.model_validate(row_dict))
+
+        # Calculate pagination metadata
+        total_pages = (total + page_size - 1) // page_size if total > 0 else 0
+
+        return {
+            "data": data,
+            "total": total,
+            "page": page,
+            "page_size": page_size,
+            "total_pages": total_pages,
+            "has_next": page < total_pages,
+            "has_previous": page > 1,
+        }

rem/services/postgres/schema_generator.py

@@ -12,6 +12,7 @@ Output includes:
 - KV_STORE triggers
 - Indexes (foreground and background)
 - Migrations
+- Schema table entries (for agent-like table access)
 
 Usage:
     from rem.services.postgres.schema_generator import SchemaGenerator
@@ -30,14 +31,192 @@ Usage:
 
 import importlib.util
 import inspect
+import json
+import uuid
 from pathlib import Path
-from typing import Type
+from typing import Any, Type
 
 from loguru import logger
 from pydantic import BaseModel
 
 from ...settings import settings
-from .register_type import register_type
+from ...utils.sql_paths import get_package_sql_dir
+from .register_type import register_type, should_embed_field
+
+# Namespace UUID for generating deterministic UUIDs from model names
+# Using UUID5 with this namespace ensures same model always gets same UUID
+REM_SCHEMA_NAMESPACE = uuid.UUID("6ba7b810-9dad-11d1-80b4-00c04fd430c8")  # DNS namespace
+
+
+def generate_model_uuid(fully_qualified_name: str) -> uuid.UUID:
+    """
+    Generate deterministic UUID from fully qualified model name.
+
+    Uses UUID5 (SHA-1 hash) with REM namespace for reproducibility.
+    Same fully qualified name always produces same UUID.
+
+    Args:
+        fully_qualified_name: Full module path, e.g., "rem.models.entities.Resource"
+
+    Returns:
+        Deterministic UUID for this model
+    """
+    return uuid.uuid5(REM_SCHEMA_NAMESPACE, fully_qualified_name)
+
+
+def extract_model_schema_metadata(
+    model: Type[BaseModel],
+    table_name: str,
+    entity_key_field: str,
+    include_search_tool: bool = True,
+) -> dict[str, Any]:
+    """
+    Extract schema metadata from a Pydantic model for schemas table.
+
+    Args:
+        model: Pydantic model class
+        table_name: Database table name
+        entity_key_field: Field used as entity key in kv_store
+        include_search_tool: If True, add search_rem tool for querying this table
+
+    Returns:
+        Dict with schema metadata ready for schemas table insert
+    """
+    # Get fully qualified name
+    fqn = f"{model.__module__}.{model.__name__}"
+
+    # Generate deterministic UUID
+    schema_id = generate_model_uuid(fqn)
+
+    # Get JSON schema from Pydantic
+    json_schema = model.model_json_schema()
+
+    # Find embedding fields
+    embedding_fields = []
+    for field_name, field_info in model.model_fields.items():
+        if should_embed_field(field_name, field_info):
+            embedding_fields.append(field_name)
+
+    # Build description with search capability note
+    base_description = model.__doc__ or f"Schema for {model.__name__}"
+    search_note = (
+        f"\n\nThis agent can search the `{table_name}` table using the `search_rem` tool. "
+        f"Use REM query syntax: LOOKUP for exact match, FUZZY for typo-tolerant search, "
+        f"SEARCH for semantic similarity, or SQL for complex queries."
+    ) if include_search_tool else ""
+
+    # Build spec with table metadata and tools
+    # Note: default_search_table is used by create_agent to append a description
+    # suffix to the search_rem tool when loading it dynamically
+    has_embeddings = bool(embedding_fields)
+
+    spec = {
+        "type": "object",
+        "description": base_description + search_note,
+        "properties": json_schema.get("properties", {}),
+        "required": json_schema.get("required", []),
+        "json_schema_extra": {
+            "table_name": table_name,
+            "entity_key_field": entity_key_field,
+            "embedding_fields": embedding_fields,
+            "fully_qualified_name": fqn,
+            "tools": ["search_rem"] if include_search_tool else [],
+            "default_search_table": table_name,
+            "has_embeddings": has_embeddings,
+        },
+    }
+
+    # Build content (documentation)
+    content = f"""# {model.__name__}
+
+{base_description}
+
+## Overview
+
+The `{model.__name__}` entity is stored in the `{table_name}` table. Each record is uniquely
+identified by its `{entity_key_field}` field for lookups and graph traversal.
+
+## Search Capabilities
+
+This schema includes the `search_rem` tool which supports:
+- **LOOKUP**: O(1) exact match by {entity_key_field} (e.g., `LOOKUP "entity-name"`)
+- **FUZZY**: Typo-tolerant search (e.g., `FUZZY "partial" THRESHOLD 0.3`)
+- **SEARCH**: Semantic vector search on {', '.join(embedding_fields) if embedding_fields else 'content'} (e.g., `SEARCH "concept" FROM {table_name} LIMIT 10`)
+- **SQL**: Complex queries (e.g., `SELECT * FROM {table_name} WHERE ...`)
+
+## Table Info
+
+| Property | Value |
+|----------|-------|
+| Table | `{table_name}` |
+| Entity Key | `{entity_key_field}` |
+| Embedding Fields | {', '.join(f'`{f}`' for f in embedding_fields) if embedding_fields else 'None'} |
+| Tools | {', '.join(['`search_rem`'] if include_search_tool else ['None'])} |
+
+## Fields
+
+"""
+    for field_name, field_info in model.model_fields.items():
+        field_type = str(field_info.annotation) if field_info.annotation else "Any"
+        field_desc = field_info.description or ""
+        required = "Required" if field_info.is_required() else "Optional"
+        content += f"### `{field_name}`\n"
+        content += f"- **Type**: `{field_type}`\n"
+        content += f"- **{required}**\n"
+        if field_desc:
+            content += f"- {field_desc}\n"
+        content += "\n"
+
+    return {
+        "id": str(schema_id),
+        "name": model.__name__,
+        "table_name": table_name,
+        "entity_key_field": entity_key_field,
+        "embedding_fields": embedding_fields,
+        "fqn": fqn,
+        "spec": spec,
+        "content": content,
+        "category": "entity",
+    }
+
+
+def generate_schema_upsert_sql(schema_metadata: dict[str, Any]) -> str:
+    """
+    Generate SQL UPSERT statement for schemas table.
+
+    Uses ON CONFLICT DO UPDATE for idempotency.
+
+    Args:
+        schema_metadata: Dict from extract_model_schema_metadata()
+
+    Returns:
+        SQL INSERT ... ON CONFLICT statement
+    """
+    # Escape single quotes in content and spec
+    content_escaped = schema_metadata["content"].replace("'", "''")
+    spec_json = json.dumps(schema_metadata["spec"]).replace("'", "''")
+
+    sql = f"""
+-- Schema entry for {schema_metadata['name']} ({schema_metadata['table_name']})
+INSERT INTO schemas (id, tenant_id, name, content, spec, category, metadata)
+VALUES (
+    '{schema_metadata['id']}'::uuid,
+    'system',
+    '{schema_metadata['name']}',
+    '{content_escaped}',
+    '{spec_json}'::jsonb,
+    'entity',
+    '{{"table_name": "{schema_metadata['table_name']}", "entity_key_field": "{schema_metadata['entity_key_field']}", "embedding_fields": {json.dumps(schema_metadata['embedding_fields'])}, "fqn": "{schema_metadata['fqn']}"}}'::jsonb
+)
+ON CONFLICT (id) DO UPDATE SET
+    name = EXCLUDED.name,
+    content = EXCLUDED.content,
+    spec = EXCLUDED.spec,
+    category = EXCLUDED.category,
+    metadata = EXCLUDED.metadata,
+    updated_at = CURRENT_TIMESTAMP;
+"""
+    return sql.strip()
 
 
 class SchemaGenerator:
@@ -56,9 +235,9 @@ class SchemaGenerator:
         Initialize schema generator.
 
         Args:
-            output_dir: Optional directory for output files (defaults to settings.sql_dir)
+            output_dir: Optional directory for output files (defaults to package sql dir)
         """
-        self.output_dir = output_dir or Path(settings.sql_dir)
+        self.output_dir = output_dir or get_package_sql_dir()
         self.schemas: dict[str, dict] = {}
 
     def discover_models(self, directory: str | Path) -> dict[str, Type[BaseModel]]:
@@ -234,6 +413,14 @@ class SchemaGenerator:
             create_kv_trigger=True,
         )
 
+        # Extract schema metadata for schemas table entry
+        schema_metadata = extract_model_schema_metadata(
+            model=model,
+            table_name=table_name,
+            entity_key_field=entity_key_field,
+        )
+        schema["schema_metadata"] = schema_metadata
+
         self.schemas[table_name] = schema
         return schema
 
@@ -343,6 +530,7 @@ class SchemaGenerator:
             "-- 2. Embeddings tables (embeddings_<table>)",
             "-- 3. KV_STORE triggers for cache maintenance",
             "-- 4. Indexes (foreground only, background indexes separate)",
+            "-- 5. Schema table entries (for agent-like table access)",
             "",
             "-- ============================================================================",
             "-- PREREQUISITES CHECK",
@@ -388,6 +576,19 @@ class SchemaGenerator:
                 sql_parts.append(schema["sql"]["kv_trigger"])
                 sql_parts.append("")
 
+        # Add schema table entries (every entity table is also an "agent")
+        sql_parts.append("-- ============================================================================")
+        sql_parts.append("-- SCHEMA TABLE ENTRIES")
+        sql_parts.append("-- Every entity table gets a schemas entry for agent-like access")
+        sql_parts.append("-- ============================================================================")
+        sql_parts.append("")
+
+        for table_name, schema in self.schemas.items():
+            if "schema_metadata" in schema:
+                schema_upsert = generate_schema_upsert_sql(schema["schema_metadata"])
+                sql_parts.append(schema_upsert)
+                sql_parts.append("")
+
         # Add migration record
         sql_parts.append("-- ============================================================================")
         sql_parts.append("-- RECORD MIGRATION")

rem/services/postgres/service.py

@@ -190,19 +190,19 @@ class PostgresService:
 
     async def connect(self) -> None:
         """Establish database connection pool."""
-        logger.info(f"Connecting to PostgreSQL with pool size {self.pool_size}")
+        logger.debug(f"Connecting to PostgreSQL with pool size {self.pool_size}")
         self.pool = await asyncpg.create_pool(
             self.connection_string,
             min_size=1,
             max_size=self.pool_size,
             init=self._init_connection,  # Configure JSONB codec on each connection
         )
-        logger.info("PostgreSQL connection pool established")
+        logger.debug("PostgreSQL connection pool established")
 
         # Start embedding worker if available
         if self.embedding_worker and hasattr(self.embedding_worker, "start"):
             await self.embedding_worker.start()
-            logger.info("Embedding worker started")
+            logger.debug("Embedding worker started")
 
     async def disconnect(self) -> None:
         """Close database connection pool."""
@@ -211,10 +211,10 @@ class PostgresService:
         # The worker will be stopped explicitly when the application shuts down
 
         if self.pool:
-            logger.info("Closing PostgreSQL connection pool")
+            logger.debug("Closing PostgreSQL connection pool")
             await self.pool.close()
             self.pool = None
-            logger.info("PostgreSQL connection pool closed")
+            logger.debug("PostgreSQL connection pool closed")
 
     async def execute(
         self,
@@ -631,7 +631,7 @@ class PostgresService:
         table_name: str,
         embedding: list[float],
         limit: int = 10,
-        min_similarity: float = 0.7,
+        min_similarity: float = 0.3,
         tenant_id: Optional[str] = None,
     ) -> list[dict[str, Any]]:
         """

rem/services/rem/parser.py

@@ -50,9 +50,36 @@ class RemQueryParser:
         params: Dict[str, Any] = {}
         positional_args: List[str] = []
 
-        # Process remaining tokens
-        for token in tokens[1:]:
-            if "=" in token:
+        # For SQL queries, preserve the raw query (keywords like LIMIT are SQL keywords)
+        if query_type == QueryType.SQL:
+            # Everything after "SQL" is the raw SQL query
+            raw_sql = query_string[3:].strip()  # Skip "SQL" prefix
+            params["raw_query"] = raw_sql
+            return query_type, params
+
+        # Process remaining tokens, handling REM keywords
+        i = 1
+        while i < len(tokens):
+            token = tokens[i]
+            token_upper = token.upper()
+
+            # Handle REM keywords that take a value
+            if token_upper in ("LIMIT", "DEPTH", "THRESHOLD", "TYPE", "FROM", "WITH"):
+                if i + 1 < len(tokens):
+                    keyword_map = {
+                        "LIMIT": "limit",
+                        "DEPTH": "max_depth",
+                        "THRESHOLD": "threshold",
+                        "TYPE": "edge_types",
+                        "FROM": "initial_query",
+                        "WITH": "initial_query",
+                    }
+                    key = keyword_map[token_upper]
+                    value = tokens[i + 1]
+                    params[key] = self._convert_value(key, value)
+                    i += 2
+                    continue
+            elif "=" in token:
                 # It's a keyword argument
                 key, value = token.split("=", 1)
                 # Handle parameter aliases
@@ -61,6 +88,7 @@ class RemQueryParser:
             else:
                 # It's a positional argument part
                 positional_args.append(token)
+            i += 1
 
         # Map positional arguments to specific fields based on QueryType
         self._map_positional_args(query_type, positional_args, params)
@@ -133,13 +161,20 @@ class RemQueryParser:
             params["query_text"] = combined_value
 
         elif query_type == QueryType.SEARCH:
-            params["query_text"] = combined_value
+            # SEARCH expects: SEARCH <table> <query_text> [LIMIT n]
+            # First positional arg is table name, rest is query text
+            if len(positional_args) >= 2:
+                params["table_name"] = positional_args[0]
+                params["query_text"] = " ".join(positional_args[1:])
+            elif len(positional_args) == 1:
+                # Could be table name or query text - assume query text if no table
+                params["query_text"] = positional_args[0]
+            # If no positional args, params stays empty
 
         elif query_type == QueryType.TRAVERSE:
             params["initial_query"] = combined_value
         
-        # SQL typically requires named arguments (table=...), but if we supported 
-        # SQL SELECT * FROM ..., we might handle it differently. 
-        # For now, RemService expects table=...
-        # If there are positional args for SQL, we might ignore or raise, 
-        # but current service doesn't use them.
+        elif query_type == QueryType.SQL:
+            # SQL with positional args means "SQL SELECT * FROM ..." form
+            # Treat the combined positional args as the raw SQL query
+            params["raw_query"] = combined_value

rem/services/rem/service.py

@@ -13,6 +13,31 @@ Design:
 - All queries pushed down to Postgres for performance
 - Model schema inspection for validation only
 - Exceptions for missing fields/embeddings
+
+TODO: Staged Plan Execution
+- Implement execute_staged_plan() method for multi-stage query execution
+- Each stage can be:
+  1. Static query (query field): Execute REM dialect directly
+  2. Dynamic query (intent field): LLM interprets intent + previous results to build query
+- Flow for dynamic stages:
+  1. Gather results from depends_on stages (from previous_results or current execution)
+  2. Pass intent + previous results to LLM (like ask_rem but with context)
+  3. LLM generates REM query based on what it learned from previous stages
+  4. Execute generated query
+  5. Store results in stage_results for client to use in continuation
+- Multi-turn continuation:
+  - Client passes previous_results back from response's stage_results
+  - Client sets resume_from_stage to skip already-executed stages
+  - Server uses previous_results as context for depends_on lookups
+- Use cases:
+  - LOOKUP "Sarah" → intent: "find her team members" (LLM sees Sarah's graph_edges, builds TRAVERSE)
+  - SEARCH "API docs" → intent: "get authors" (LLM extracts author refs, builds LOOKUP)
+  - Complex graph exploration with LLM-driven navigation
+- API: POST /api/v1/query with:
+  - mode="staged-plan"
+  - plan=[{stage, query|intent, name, depends_on}]
+  - previous_results=[{stage, name, query_executed, results, count}] (for continuation)
+  - resume_from_stage=N (to skip completed stages)
 """
 
 from typing import Any
@@ -309,17 +334,26 @@ class RemService:
         )
 
         # Execute vector search via rem_search() PostgreSQL function
+        min_sim = params.min_similarity if params.min_similarity is not None else 0.3
+        limit = params.limit or 10
         query_params = get_search_params(
             query_embedding,
             table_name,
             field_name,
             tenant_id,
             provider,
-            params.min_similarity or 0.7,
-            params.limit or 10,
+            min_sim,
+            limit,
             tenant_id, # Use tenant_id (query.user_id) as user_id
         )
+        logger.debug(
+            f"SEARCH params: table={table_name}, field={field_name}, "
+            f"tenant_id={tenant_id}, provider={provider}, "
+            f"min_similarity={min_sim}, limit={limit}, "
+            f"embedding_dims={len(query_embedding)}"
+        )
         results = await self.db.execute(SEARCH_QUERY, query_params)
+        logger.debug(f"SEARCH results: {len(results)} rows")
 
         return {
             "query_type": "SEARCH",

rem/services/session/compression.py

@@ -14,6 +14,21 @@ from typing import Any
 
 from loguru import logger
 
+# Max length for entity keys (kv_store.entity_key is varchar(255))
+MAX_ENTITY_KEY_LENGTH = 255
+
+
+def truncate_key(key: str, max_length: int = MAX_ENTITY_KEY_LENGTH) -> str:
+    """Truncate a key to max length, preserving useful suffix if possible."""
+    if len(key) <= max_length:
+        return key
+    # Keep first part and add hash suffix for uniqueness
+    import hashlib
+    hash_suffix = hashlib.md5(key.encode()).hexdigest()[:8]
+    truncated = key[:max_length - 9] + "-" + hash_suffix
+    logger.warning(f"Truncated key from {len(key)} to {len(truncated)} chars: {key[:50]}...")
+    return truncated
+
 from rem.models.entities import Message
 from rem.services.postgres import PostgresService, Repository
 from rem.settings import settings
@@ -151,15 +166,20 @@ class SessionMessageStore:
             return f"msg-{message_index}"
 
         # Create entity key for REM LOOKUP: session-{session_id}-msg-{index}
-        entity_key = f"session-{session_id}-msg-{message_index}"
+        # Truncate to avoid exceeding kv_store.entity_key varchar(255) limit
+        entity_key = truncate_key(f"session-{session_id}-msg-{message_index}")
 
         # Create Message entity for assistant response
+        # Use pre-generated id from message dict if available (for frontend feedback)
         msg = Message(
+            id=message.get("id"),  # Use pre-generated ID if provided
             content=message.get("content", ""),
             message_type=message.get("role", "assistant"),
             session_id=session_id,
             tenant_id=self.user_id,  # Set tenant_id to user_id (application scoped to user)
             user_id=user_id or self.user_id,
+            trace_id=message.get("trace_id"),
+            span_id=message.get("span_id"),
             metadata={
                 "message_index": message_index,
                 "entity_key": entity_key,  # Store entity key for LOOKUP
@@ -268,11 +288,14 @@ class SessionMessageStore:
                 # Short assistant messages, user messages, and system messages stored as-is
                 # Store ALL messages in database for full audit trail
                 msg = Message(
+                    id=message.get("id"),  # Use pre-generated ID if provided
                     content=content,
                     message_type=message.get("role", "user"),
                     session_id=session_id,
                     tenant_id=self.user_id,  # Set tenant_id to user_id (application scoped to user)
                     user_id=user_id or self.user_id,
+                    trace_id=message.get("trace_id"),
+                    span_id=message.get("span_id"),
                     metadata={
                         "message_index": idx,
                         "timestamp": message.get("timestamp"),

rem/services/session/reload.py

@@ -65,7 +65,7 @@ async def reload_session(
             session_id=session_id, user_id=user_id, decompress=decompress_messages
         )
 
-        logger.info(
+        logger.debug(
             f"Reloaded {len(messages)} messages for session {session_id} "
             f"(decompressed={decompress_messages})"
         )