remdb 0.2.6__py3-none-any.whl → 0.3.103__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/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/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/user_service.py

@@ -0,0 +1,98 @@
+"""
+User Service - User account management.
+
+Handles user creation, profile updates, and session linking.
+"""
+
+from datetime import datetime
+from typing import Optional
+
+from loguru import logger
+
+from ..models.entities.user import User, UserTier
+from .postgres.repository import Repository
+from .postgres.service import PostgresService
+
+
+class UserService:
+    """
+    Service for managing user accounts and sessions.
+    """
+
+    def __init__(self, db: PostgresService):
+        self.db = db
+        self.repo = Repository(User, "users", db=db)
+
+    async def get_or_create_user(
+        self,
+        email: str,
+        tenant_id: str = "default",
+        name: str = "New User",
+        avatar_url: Optional[str] = None,
+    ) -> User:
+        """
+        Get existing user by email or create a new one.
+        """
+        users = await self.repo.find(filters={"email": email}, limit=1)
+        
+        if users:
+            user = users[0]
+            # Update profile if needed (e.g., name/avatar from OAuth)
+            updated = False
+            if name and user.name == "New User": # Only update if placeholder
+                user.name = name
+                updated = True
+            
+            # Store avatar in metadata if provided
+            if avatar_url:
+                user.metadata = user.metadata or {}
+                if user.metadata.get("avatar_url") != avatar_url:
+                    user.metadata["avatar_url"] = avatar_url
+                    updated = True
+            
+            if updated:
+                user.updated_at = datetime.utcnow()
+                await self.repo.upsert(user)
+            
+            return user
+        
+        # Create new user
+        user = User(
+            tenant_id=tenant_id,
+            user_id=email, # Use email as user_id for now? Or UUID?
+            # The User model has 'user_id' field but also 'id' UUID.
+            # Usually user_id is the external ID or email.
+            name=name,
+            email=email,
+            tier=UserTier.FREE,
+            created_at=datetime.utcnow(),
+            updated_at=datetime.utcnow(),
+            metadata={"avatar_url": avatar_url} if avatar_url else {},
+        )
+        await self.repo.upsert(user)
+        logger.info(f"Created new user: {email}")
+        return user
+
+    async def link_anonymous_session(self, user: User, anon_id: str) -> None:
+        """
+        Link an anonymous session ID to a user account.
+        
+        This allows merging history from the anonymous session into the user's profile.
+        """
+        if not anon_id:
+            return
+
+        # Check if already linked
+        if anon_id in user.anonymous_ids:
+            return
+
+        # Add to list
+        user.anonymous_ids.append(anon_id)
+        user.updated_at = datetime.utcnow()
+        
+        # Save
+        await self.repo.upsert(user)
+        logger.info(f"Linked anonymous session {anon_id} to user {user.email}")
+        
+        # TODO: Migrate/Merge actual data (rate limit counts, history) if needed.
+        # For now, we just link the IDs so future queries can include data from this anon_id.