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

rem/services/phoenix/client.py

@@ -53,7 +53,7 @@ from datetime import datetime
 from pathlib import Path
 from typing import Any, Callable, TYPE_CHECKING, cast
 
-import pandas as pd
+import polars as pl
 from loguru import logger
 
 from .config import PhoenixConfig
@@ -64,6 +64,95 @@ if TYPE_CHECKING:
     from phoenix.client.resources.experiments.types import RanExperiment
 
 
+def dataframe_to_phoenix_dataset(
+    client: "PhoenixClient",
+    df: pl.DataFrame,
+    dataset_name: str,
+    input_keys: list[str] | None = None,
+    output_keys: list[str] | None = None,
+    metadata_keys: list[str] | None = None,
+    description: str | None = None,
+) -> "Dataset":
+    """Convert a Polars DataFrame to a Phoenix Dataset.
+
+    This function transforms a Polars DataFrame into a Phoenix Dataset by:
+    1. Extracting input columns (what agents receive)
+    2. Extracting output columns (ground truth/expected output)
+    3. Extracting metadata columns (optional labels, difficulty, etc.)
+
+    If column keys are not specified, uses smart defaults:
+    - input_keys: columns containing 'input', 'query', 'question', or 'prompt'
+    - output_keys: columns containing 'output', 'expected', 'answer', or 'response'
+    - metadata_keys: remaining columns
+
+    Args:
+        client: PhoenixClient instance
+        df: Polars DataFrame with experiment data
+        dataset_name: Name for the created Phoenix dataset
+        input_keys: Optional list of column names for inputs
+        output_keys: Optional list of column names for outputs (ground truth)
+        metadata_keys: Optional list of column names for metadata
+        description: Optional dataset description
+
+    Returns:
+        Phoenix Dataset instance
+
+    Example:
+        >>> df = pl.read_csv("golden_set.csv")
+        >>> dataset = dataframe_to_phoenix_dataset(
+        ...     client=phoenix_client,
+        ...     df=df,
+        ...     dataset_name="my-golden-set",
+        ...     input_keys=["query"],
+        ...     output_keys=["expected_output"],
+        ...     metadata_keys=["difficulty"]
+        ... )
+    """
+    columns = df.columns
+
+    # Smart defaults for column detection
+    if input_keys is None:
+        input_keys = [c for c in columns if any(
+            k in c.lower() for k in ["input", "query", "question", "prompt"]
+        )]
+        if not input_keys:
+            # Fallback: first column
+            input_keys = [columns[0]] if columns else []
+
+    if output_keys is None:
+        output_keys = [c for c in columns if any(
+            k in c.lower() for k in ["output", "expected", "answer", "response", "reference"]
+        )]
+        if not output_keys:
+            # Fallback: second column
+            output_keys = [columns[1]] if len(columns) > 1 else []
+
+    if metadata_keys is None:
+        used_keys = set(input_keys) | set(output_keys)
+        metadata_keys = [c for c in columns if c not in used_keys]
+
+    logger.debug(
+        f"DataFrame to Phoenix Dataset: inputs={input_keys}, "
+        f"outputs={output_keys}, metadata={metadata_keys}"
+    )
+
+    # Convert to list of dicts
+    records = df.to_dicts()
+
+    inputs = [{k: row.get(k) for k in input_keys} for row in records]
+    outputs = [{k: row.get(k) for k in output_keys} for row in records]
+    metadata = [{k: row.get(k) for k in metadata_keys} for row in records] if metadata_keys else None
+
+    # Create Phoenix dataset
+    return client.create_dataset_from_data(
+        name=dataset_name,
+        inputs=inputs,
+        outputs=outputs,
+        metadata=metadata,
+        description=description,
+    )
+
+
 class PhoenixClient:
     """High-level Phoenix client for REM evaluation workflows.
 
@@ -260,19 +349,22 @@ class PhoenixClient:
             "SEARCH semantic AI engineer",sarah-chen,person,medium,SEARCH
         """
         try:
-            # Load CSV
-            df = pd.read_csv(csv_file_path)
+            # Load CSV with Polars
+            df = pl.read_csv(csv_file_path)
+
+            # Convert to list of dicts
+            records = df.to_dicts()
 
             # Extract inputs
-            inputs = cast(list[dict[str, Any]], df[input_keys].to_dict("records"))
+            inputs = [{k: row.get(k) for k in input_keys} for row in records]
 
             # Extract outputs
-            outputs = cast(list[dict[str, Any]], df[output_keys].to_dict("records"))
+            outputs = [{k: row.get(k) for k in output_keys} for row in records]
 
             # Extract metadata if specified
             metadata = None
             if metadata_keys:
-                metadata = cast(list[dict[str, Any]], df[metadata_keys].to_dict("records"))
+                metadata = [{k: row.get(k) for k in metadata_keys} for row in records]
 
             return self.create_dataset_from_data(
                 name=name,
@@ -331,13 +423,16 @@ class PhoenixClient:
 
     def run_experiment(
         self,
-        dataset: "Dataset" | str,
+        dataset: "Dataset" | str | pl.DataFrame,
         task: Callable[[Any], Any] | None = None,
         evaluators: list[Callable[[Any], Any]] | None = None,
         experiment_name: str | None = None,
         experiment_description: str | None = None,
         experiment_metadata: dict[str, Any] | None = None,
         experiment_config: Any | None = None,
+        input_keys: list[str] | None = None,
+        output_keys: list[str] | None = None,
+        metadata_keys: list[str] | None = None,
     ) -> "RanExperiment":
         """Run an evaluation experiment.
 
@@ -346,14 +441,22 @@ class PhoenixClient:
         2. Agent run: Provide task function to execute agents on dataset
         3. Evaluator run: Provide evaluators to score existing outputs
 
+        Dataset can be:
+        - Phoenix Dataset instance
+        - Dataset name (string) - will be loaded from Phoenix
+        - Polars DataFrame - will be converted to Phoenix Dataset
+
         Args:
-            dataset: Dataset instance or name (required unless experiment_config provided)
+            dataset: Dataset instance, name, or Polars DataFrame
             task: Optional task function to run on each example (agent execution)
             evaluators: Optional list of evaluator functions
             experiment_name: Optional experiment name
             experiment_description: Optional description
             experiment_metadata: Optional metadata dict
             experiment_config: Optional ExperimentConfig instance (overrides other params)
+            input_keys: Column names for inputs (required if dataset is DataFrame)
+            output_keys: Column names for outputs (required if dataset is DataFrame)
+            metadata_keys: Optional column names for metadata
 
         Returns:
             RanExperiment with results
@@ -369,6 +472,16 @@ class PhoenixClient:
             ...     experiment_name="rem-v1-baseline"
             ... )
 
+        Example - With Polars DataFrame:
+            >>> df = pl.read_csv("golden_set.csv")
+            >>> experiment = client.run_experiment(
+            ...     dataset=df,
+            ...     task=run_agent,
+            ...     experiment_name="rem-v1-baseline",
+            ...     input_keys=["query"],
+            ...     output_keys=["expected_output"]
+            ... )
+
         Example - Evaluator Run (Phase 2b):
             >>> experiment = client.run_experiment(
             ...     dataset=agent_results,
@@ -407,6 +520,21 @@ class PhoenixClient:
                     else:
                         dataset = dataset_ref.path
 
+            # Convert Polars DataFrame to Phoenix Dataset
+            if isinstance(dataset, pl.DataFrame):
+                dataset_name_for_phoenix = f"{experiment_name or 'experiment'}-dataset-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
+                logger.info(f"Converting Polars DataFrame to Phoenix Dataset: {dataset_name_for_phoenix}")
+                dataset = dataframe_to_phoenix_dataset(
+                    client=self,
+                    df=dataset,
+                    dataset_name=dataset_name_for_phoenix,
+                    input_keys=input_keys,
+                    output_keys=output_keys,
+                    metadata_keys=metadata_keys,
+                    description=f"Auto-created from DataFrame for experiment: {experiment_name}",
+                )
+                logger.info(f"✓ Created Phoenix Dataset: {dataset_name_for_phoenix}")
+
             # Load dataset if name provided
             if isinstance(dataset, str):
                 dataset = self.get_dataset(dataset)
@@ -454,7 +582,7 @@ class PhoenixClient:
         root_spans_only: bool = True,
         trace_id: str | None = None,
         span_id: str | None = None,
-    ) -> pd.DataFrame:
+    ) -> pl.DataFrame:
         """Query traces from Phoenix.
 
         Args:
@@ -467,7 +595,7 @@ class PhoenixClient:
             span_id: Filter by specific span ID
 
         Returns:
-            DataFrame with trace data
+            Polars DataFrame with trace data
 
         Example:
             >>> traces = client.get_traces(
@@ -492,8 +620,11 @@ class PhoenixClient:
             if span_id:
                 query_params["span_id"] = span_id
 
-            # Query traces
-            traces_df = self._client.query_spans(limit=limit, **query_params)  # type: ignore[attr-defined]
+            # Query traces (Phoenix returns pandas DataFrame)
+            pandas_df = self._client.query_spans(limit=limit, **query_params)  # type: ignore[attr-defined]
+
+            # Convert pandas to Polars
+            traces_df = pl.from_pandas(pandas_df)
 
             logger.debug(f"Retrieved {len(traces_df)} traces")
             return traces_df
@@ -535,7 +666,7 @@ class PhoenixClient:
             ... )
         """
         try:
-            # Query traces
+            # Query traces (returns Polars DataFrame)
             traces_df = self.get_traces(
                 project_name=project_name,
                 start_time=start_time,
@@ -547,12 +678,15 @@ class PhoenixClient:
             if len(traces_df) == 0:
                 raise ValueError("No traces found matching criteria")
 
+            # Convert to list of dicts for iteration
+            records = traces_df.to_dicts()
+
             # Extract inputs and outputs from traces
             inputs = []
             outputs = []
             metadata = []
 
-            for _, row in traces_df.iterrows():
+            for row in records:
                 # Extract input
                 span_input = row.get("attributes.input")
                 if span_input:

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