remdb 0.3.0__py3-none-any.whl → 0.3.114__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

@@ -1,7 +1,12 @@
 """
 Schema generation utility from Pydantic models.
 
-Scans a directory of Pydantic models and generates complete database schemas including:
+Generates complete database schemas from:
+1. REM's core models (Resource, Moment, User, etc.)
+2. Models registered via rem.register_model() or rem.register_models()
+3. Models discovered from a directory scan
+
+Output includes:
 - Primary tables
 - Embeddings tables
 - KV_STORE triggers
@@ -11,8 +16,12 @@ Scans a directory of Pydantic models and generates complete database schemas inc
 Usage:
     from rem.services.postgres.schema_generator import SchemaGenerator
 
+    # Generate from registry (includes core + registered models)
     generator = SchemaGenerator()
-    schema = generator.generate_from_directory("src/rem/models/entities")
+    schema = await generator.generate_from_registry()
+
+    # Or generate from directory (legacy)
+    schema = await generator.generate_from_directory("src/rem/models/entities")
 
     # Write to file
     with open("src/rem/sql/schema.sql", "w") as f:
@@ -228,12 +237,65 @@ class SchemaGenerator:
         self.schemas[table_name] = schema
         return schema
 
+    async def generate_from_registry(
+        self, output_file: str | None = None, include_core: bool = True
+    ) -> str:
+        """
+        Generate complete schema from the model registry.
+
+        Includes:
+        1. REM's core models (if include_core=True)
+        2. Models registered via rem.register_model() or rem.register_models()
+
+        Args:
+            output_file: Optional output file path (relative to output_dir)
+            include_core: If True, include REM's core models (default: True)
+
+        Returns:
+            Complete SQL schema as string
+
+        Example:
+            import rem
+            from rem.models.core import CoreModel
+
+            # Register custom model
+            @rem.register_model
+            class CustomEntity(CoreModel):
+                name: str
+
+            # Generate schema (includes core + custom)
+            generator = SchemaGenerator()
+            schema = await generator.generate_from_registry()
+        """
+        from ...registry import get_model_registry
+
+        registry = get_model_registry()
+        models = registry.get_models(include_core=include_core)
+
+        logger.info(f"Generating schema from registry: {len(models)} models")
+
+        # Generate schemas for each model
+        for model_name, ext in models.items():
+            await self.generate_schema_for_model(
+                ext.model,
+                table_name=ext.table_name,
+                entity_key_field=ext.entity_key_field,
+            )
+
+        return self._generate_sql_output(
+            source="model registry",
+            output_file=output_file,
+        )
+
     async def generate_from_directory(
         self, directory: str | Path, output_file: str | None = None
     ) -> str:
         """
         Generate complete schema from all models in a directory.
 
+        Note: For most use cases, prefer generate_from_registry() which uses
+        the model registry pattern.
+
         Args:
             directory: Path to directory with Pydantic models
             output_file: Optional output file path (relative to output_dir)
@@ -248,12 +310,31 @@ class SchemaGenerator:
         for model_name, model in models.items():
             await self.generate_schema_for_model(model)
 
-        # Combine into single SQL file
+        return self._generate_sql_output(
+            source=f"directory: {directory}",
+            output_file=output_file,
+        )
+
+    def _generate_sql_output(
+        self, source: str, output_file: str | None = None
+    ) -> str:
+        """
+        Generate SQL output from accumulated schemas.
+
+        Args:
+            source: Description of schema source (for header comment)
+            output_file: Optional output file path (relative to output_dir)
+
+        Returns:
+            Complete SQL schema as string
+        """
+        import datetime
+
         sql_parts = [
             "-- REM Model Schema (install_models.sql)",
             "-- Generated from Pydantic models",
-            f"-- Source directory: {directory}",
-            "-- Generated at: " + __import__("datetime").datetime.now().isoformat(),
+            f"-- Source: {source}",
+            f"-- Generated at: {datetime.datetime.now().isoformat()}",
             "--",
             "-- DO NOT EDIT MANUALLY - Regenerate with: rem db schema generate",
             "--",

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/rate_limit.py

@@ -0,0 +1,113 @@
+"""
+Rate Limit Service - Postgres-backed rate limiting.
+
+Implements tenant-aware, tiered rate limiting using PostgreSQL UNLOGGED tables
+for high performance. Supports monthly quotas and short-term burst limits.
+"""
+
+import random
+from datetime import datetime, timezone
+from enum import Enum
+from typing import Optional
+
+from loguru import logger
+
+from ..models.entities.user import UserTier
+from .postgres.service import PostgresService
+
+
+class RateLimitService:
+    """
+    Service for tracking and enforcing API rate limits.
+    
+    Uses an UNLOGGED table `rate_limits` for performance.
+    Note: Counts in UNLOGGED tables may be lost on database crash/restart.
+    """
+
+    def __init__(self, db: PostgresService):
+        self.db = db
+        
+        # Rate limits configuration
+        # Format: (limit, period_seconds)
+        # This is a simple implementation. In production, move to settings.
+        self.TIER_CONFIG = {
+            UserTier.ANONYMOUS: {"limit": 1000, "period": 3600},  # 1000/hour (for testing)
+            UserTier.FREE: {"limit": 50, "period": 2592000},      # 50/month (~30 days)
+            UserTier.BASIC: {"limit": 10000, "period": 2592000},  # 10k/month
+            UserTier.PRO: {"limit": 100000, "period": 2592000},   # 100k/month
+        }
+
+    async def check_rate_limit(
+        self, 
+        tenant_id: str, 
+        identifier: str, 
+        tier: UserTier
+    ) -> tuple[bool, int, int]:
+        """
+        Check if request is allowed under the rate limit.
+
+        Args:
+            tenant_id: Tenant identifier
+            identifier: User ID or Anonymous ID
+            tier: User subscription tier
+
+        Returns:
+            Tuple (is_allowed, current_count, limit)
+        """
+        config = self.TIER_CONFIG.get(tier, self.TIER_CONFIG[UserTier.FREE])
+        limit = config["limit"]
+        period = config["period"]
+        
+        # Construct time-window key
+        now = datetime.now(timezone.utc)
+        
+        if period >= 2592000: # Monthly
+            time_key = now.strftime("%Y-%m") 
+        elif period >= 86400: # Daily
+            time_key = now.strftime("%Y-%m-%d")
+        elif period >= 3600: # Hourly
+            time_key = now.strftime("%Y-%m-%d-%H")
+        else: # Minute/Second (fallback)
+            time_key = int(now.timestamp() / period)
+
+        key = f"{tenant_id}:{identifier}:{tier.value}:{time_key}"
+        
+        # Calculate expiry (for cleanup)
+        expires_at = now.timestamp() + period
+
+        # Atomic UPSERT to increment counter
+        # Returns the new count
+        query = """
+            INSERT INTO rate_limits (key, count, expires_at)
+            VALUES ($1, 1, to_timestamp($2))
+            ON CONFLICT (key) DO UPDATE
+            SET count = rate_limits.count + 1
+            RETURNING count;
+        """
+        
+        try:
+            count = await self.db.fetchval(query, key, expires_at)
+        except Exception as e:
+            logger.error(f"Rate limit check failed: {e}")
+            # Fail open to avoid blocking users on DB error
+            return True, 0, limit
+
+        is_allowed = count <= limit
+        
+        # Probabilistic cleanup (1% chance)
+        if random.random() < 0.01:
+            await self.cleanup_expired()
+            
+        return is_allowed, count, limit
+
+    async def cleanup_expired(self):
+        """Remove expired rate limit keys."""
+        try:
+            # Use a small limit to avoid locking/long queries
+            query = """
+                DELETE FROM rate_limits 
+                WHERE expires_at < NOW()
+            """
+            await self.db.execute(query)
+        except Exception as e:
+            logger.warning(f"Rate limit cleanup failed: {e}")

rem/services/rem/README.md

@@ -302,3 +302,17 @@ See `tests/integration/test_rem_query_evolution.py` for stage-based validation a
 * **Unified View**: The underlying SQL function `rem_traverse` uses a view `all_graph_edges` that unions `graph_edges` from all entity tables (`resources`, `moments`, `users`, etc.). This enables polymorphic traversal without complex joins in the application layer.
 * **KV Store**: Edge destinations (`dst`) are resolved to entity IDs using the `kv_store`. This requires that all traversable entities have an entry in the `kv_store` (handled automatically by database triggers).
 * **Iterated Retrieval**: REM is architected for multi-turn retrieval where LLMs conduct conversational database exploration. Each query informs the next, enabling emergent information discovery without requiring upfront schema knowledge.
+
+## Scaling & Architectural Decisions
+
+### 1. Hybrid Adjacency List
+REM implements a **Hybrid Adjacency List** pattern to balance strict relational guarantees with graph flexibility:
+*   **Primary Storage (Source of Truth):** Standard PostgreSQL tables (`resources`, `moments`, etc.) enforce schema validation, constraints, and type safety.
+*   **Graph Overlay:** Relationships are stored as "inline edges" within a JSONB column (`graph_edges`) on each entity.
+*   **Performance Layer:** A denormalized `UNLOGGED` table (`kv_store`) acts as a high-speed cache, mapping human-readable keys to internal UUIDs and edges. This avoids the traditional "join bomb" of traversing normalized SQL tables while avoiding the operational complexity of a separate graph database (e.g., Neo4j).
+
+### 2. The Pareto Principle in Graph Algorithms
+We explicitly choose **Simplicity over Full-Scale Graph Analytics**.
+*   **Hypothesis:** For LLM Agent workloads, 80% of the value is derived from **local context retrieval** (1-3 hops via `LOOKUP` and `TRAVERSE`).
+*   **Diminishing Returns:** Global graph algorithms (PageRank, Community Detection) offer diminishing returns for real-time agentic retrieval tasks. Agents typically need to answer specific questions ("Who worked on file X?"), which is a local neighborhood problem, not a global cluster analysis problem.
+*   **Future Scaling:** If deeper analysis is needed, we prefer **Graph + Vector (RAG)** approaches (using semantic similarity to find implicit links) over complex explicit graph algorithms.

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,7 +166,8 @@ 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
         msg = Message(

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})"
         )