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

rem/__init__.py

@@ -1,2 +1,129 @@
-def hello() -> str:
-    return "Hello from rem!"
+"""
+REM - Resources, Entities, Moments.
+
+A bio-inspired memory system for agentic AI, built on FastAPI.
+
+Usage (API mode):
+    from rem import create_app
+
+    # Create REM app (FastAPI with MCP server pre-configured)
+    app = create_app()
+
+    # Extend like any FastAPI app
+    @app.get("/my-endpoint")
+    async def my_endpoint():
+        return {"custom": True}
+
+    # Add routers
+    app.include_router(my_router)
+
+    # Access MCP server directly (FastMCP instance)
+    @app.mcp_server.tool()
+    async def my_custom_tool(query: str) -> dict:
+        '''Custom MCP tool for my application.'''
+        return {"result": "..."}
+
+    @app.mcp_server.resource("custom://config")
+    async def get_config() -> str:
+        '''Custom resource.'''
+        return '{"setting": "value"}'
+
+Usage (model registration - works with or without API):
+    import rem
+    from rem.models.core import CoreModel
+
+    @rem.register_model
+    class CustomEntity(CoreModel):
+        name: str
+        custom_field: str
+
+    # Or register multiple:
+    rem.register_models(ModelA, ModelB)
+
+    # Then schema generation includes your models:
+    # rem db schema generate
+"""
+
+from .registry import (
+    # Model registration
+    register_model,
+    register_models,
+    get_model_registry,
+    clear_model_registry,
+    # Schema path registration
+    register_schema_path,
+    register_schema_paths,
+    get_schema_paths,
+    get_schema_path_registry,
+    clear_schema_path_registry,
+)
+
+
+def create_app():
+    """
+    Create and return a FastAPI application with REM features pre-configured.
+
+    The returned app has:
+    - MCP server mounted at /api/v1/mcp
+    - Chat completions endpoint at /api/v1/chat/completions
+    - Health check at /health
+    - OpenAPI docs at /docs
+
+    The app exposes `app.mcp_server` (FastMCP instance) for adding custom
+    tools, resources, and prompts.
+
+    Returns:
+        FastAPI application with .mcp_server attribute
+
+    Example:
+        from rem import create_app
+
+        app = create_app()
+
+        # Add custom endpoint
+        @app.get("/custom")
+        async def custom():
+            return {"custom": True}
+
+        # Add custom MCP tool
+        @app.mcp_server.tool()
+        async def my_tool(query: str) -> dict:
+            return {"result": query}
+    """
+    from .api.main import create_app as _create_app
+    return _create_app()
+
+
+# Lazy app instance - created on first access
+_app = None
+
+
+def get_app():
+    """
+    Get or create the default REM app instance.
+
+    For most cases, use create_app() to get a fresh instance.
+    This is provided for convenience in simple scripts.
+    """
+    global _app
+    if _app is None:
+        _app = create_app()
+    return _app
+
+
+__all__ = [
+    # App creation
+    "create_app",
+    "get_app",
+    # Model registration
+    "register_model",
+    "register_models",
+    "get_model_registry",
+    "clear_model_registry",
+    # Schema path registration
+    "register_schema_path",
+    "register_schema_paths",
+    "get_schema_paths",
+    "get_schema_path_registry",
+    "clear_schema_path_registry",
+]

rem/agentic/context.py

@@ -72,7 +72,7 @@ class AgentContext(BaseModel):
     def get_user_id_or_default(
         user_id: str | None,
         source: str = "context",
-        default: str = "default",
+        default: str | None = None,
     ) -> str:
         """
         Get user_id or fallback to default with logging.
@@ -83,10 +83,10 @@ class AgentContext(BaseModel):
         Args:
             user_id: User identifier (may be None)
             source: Source of the call (for logging clarity)
-            default: Default value to use (default: "default")
+            default: Default value to use (default: settings.test.effective_user_id)
 
         Returns:
-            user_id if provided, otherwise default
+            user_id if provided, otherwise default from settings
 
         Example:
             # In MCP tool
@@ -105,8 +105,10 @@ class AgentContext(BaseModel):
             )
         """
         if user_id is None:
-            logger.debug(f"No user_id provided from {source}, using '{default}'")
-            return default
+            from rem.settings import settings
+            effective_default = default or settings.test.effective_user_id
+            logger.debug(f"No user_id provided from {source}, using '{effective_default}'")
+            return effective_default
         return user_id
 
     @classmethod

rem/agentic/providers/phoenix.py

@@ -128,15 +128,16 @@ def sanitize_tool_name(tool_name: str) -> str:
 
 
 def load_evaluator_schema(evaluator_name: str) -> dict[str, Any]:
-    """Load evaluator schema from schemas/evaluators/ directory.
+    """Load evaluator schema using centralized schema loader.
 
-    Searches for evaluator schema in rem/schemas/evaluators/
-    Supports .json, .yaml, and .yml files.
+    Uses the same unified search logic as agent schemas:
+    - "hello-world/default" → schemas/evaluators/hello-world/default.yaml
+    - "lookup-correctness" → schemas/evaluators/rem/lookup-correctness.yaml
+    - "rem-lookup-correctness" → schemas/evaluators/rem/lookup-correctness.yaml
 
     Args:
-        evaluator_name: Evaluator name (with or without extension)
-                       e.g., "rem-lookup-correctness" or
-                             "rem-lookup-correctness.yaml"
+        evaluator_name: Evaluator name or path
+                       e.g., "hello-world/default", "lookup-correctness"
 
     Returns:
         Evaluator schema dictionary with keys:
@@ -150,43 +151,13 @@ def load_evaluator_schema(evaluator_name: str) -> dict[str, Any]:
         FileNotFoundError: If evaluator schema not found
 
     Example:
-        >>> schema = load_evaluator_schema("rem-lookup-correctness")
+        >>> schema = load_evaluator_schema("hello-world/default")
         >>> print(schema["description"])
     """
-    # Get schemas directory (rem/schemas/evaluators/)
-    # rem.__file__ = rem/src/rem/__init__.py
-    # We need rem/schemas/evaluators/
-    import rem
-    rem_module_dir = Path(rem.__file__).parent  # rem/src/rem
-    rem_package_root = rem_module_dir.parent.parent  # rem/src/rem -> rem/src -> rem
-    schema_dir = rem_package_root / "schemas" / "evaluators"
-
-    # Try .yaml first (preferred format)
-    yaml_path = schema_dir / f"{evaluator_name}.yaml"
-    if yaml_path.exists():
-        logger.debug(f"Loading evaluator schema from {yaml_path}")
-        with open(yaml_path) as f:
-            return yaml.safe_load(f)
-
-    # Try .yml
-    yml_path = schema_dir / f"{evaluator_name}.yml"
-    if yml_path.exists():
-        logger.debug(f"Loading evaluator schema from {yml_path}")
-        with open(yml_path) as f:
-            return yaml.safe_load(f)
-
-    # Try .json
-    json_path = schema_dir / f"{evaluator_name}.json"
-    if json_path.exists():
-        logger.debug(f"Loading evaluator schema from {json_path}")
-        with open(json_path) as f:
-            return json.load(f)
-
-    raise FileNotFoundError(
-        f"Evaluator schema not found: {evaluator_name}\n"
-        f"Searched in: {schema_dir}\n"
-        f"Supported formats: .yaml, .yml, .json"
-    )
+    from ...utils.schema_loader import load_agent_schema
+
+    # Use centralized schema loader (searches evaluator paths too)
+    return load_agent_schema(evaluator_name)
 
 
 # =============================================================================
@@ -338,6 +309,22 @@ def create_evaluator_from_schema(
         # Already a dict
         schema = evaluator_schema_path
 
+    # Extract model from schema's provider_configs if not explicitly provided
+    if model_name is None:
+        json_schema_extra = schema.get("json_schema_extra", {})
+        provider_configs = json_schema_extra.get("provider_configs", [])
+        if provider_configs:
+            # Use first provider config
+            first_provider = provider_configs[0]
+            provider_name = first_provider.get("provider_name", "openai")
+            schema_model_name = first_provider.get("model_name", "gpt-4o-mini")
+            # Format as "provider:model" if not OpenAI (OpenAI is default)
+            if provider_name == "openai":
+                model_name = schema_model_name
+            else:
+                model_name = f"{provider_name}:{schema_model_name}"
+            logger.debug(f"Using model from schema provider_configs: {model_name}")
+
     # Create evaluator config
     evaluator_config = create_phoenix_evaluator(
         evaluator_schema=schema,
@@ -361,7 +348,8 @@ def create_evaluator_from_schema(
         Returns:
             Evaluation result with score, label, explanation
         """
-        logger.debug(f"Evaluating example: {example.get('input', '')[:100]}...")
+        input_preview = str(example.get('input', ''))[:100]
+        logger.debug(f"Evaluating example: {input_preview}...")
 
         # Phoenix llm_classify() expects a flat dict with string values
         # Build evaluation input by flattening nested dicts
@@ -393,6 +381,7 @@ def create_evaluator_from_schema(
 
         try:
             # Create single-row DataFrame for llm_classify
+            # Note: Phoenix's llm_classify requires pandas DataFrame (imported above)
             df = pd.DataFrame([eval_input])
 
             # Call Phoenix llm_classify
@@ -404,7 +393,7 @@ def create_evaluator_from_schema(
                 provide_explanation=True,
             )
 
-            # Extract result
+            # Extract result (results_df is pandas DataFrame from Phoenix)
             if not results_df.empty:
                 row = results_df.iloc[0]
                 label = row.get("label", "error")

rem/api/README.md

@@ -392,6 +392,29 @@ Middleware runs in reverse order of addition:
 
 ## Error Responses
 
+### 429 - Rate Limit Exceeded
+
+When a user exceeds their rate limit (based on their tier), the API returns a 429 status code with a structured error body. The frontend should intercept this error to prompt the user to sign in or upgrade.
+
+```json
+{
+  "error": {
+    "code": "rate_limit_exceeded",
+    "message": "You have exceeded your rate limit. Please sign in or upgrade to continue.",
+    "details": {
+      "limit": 50,
+      "tier": "anonymous",
+      "retry_after": 60
+    }
+  }
+}
+```
+
+**Handling Strategy:**
+1.  **Intercept 429s:** API client should listen for `status === 429`.
+2.  **Check Code:** If `error.code === 'rate_limit_exceeded'` AND `error.details.tier === 'anonymous'`, trigger "Login / Sign Up" flow.
+3.  **Authenticated Users:** If `tier !== 'anonymous'`, prompt to upgrade plan.
+
 ### 500 - Agent Schema Not Found
 
 ```json

rem/api/main.py

@@ -163,7 +163,22 @@ async def lifespan(app: FastAPI):
 
 def create_app() -> FastAPI:
     """
-    Create and configure the FastAPI application.
+    Create and configure the FastAPI application with MCP server.
+
+    The returned app exposes `app.mcp_server` (FastMCP instance) for adding
+    custom tools, resources, and prompts:
+
+        app = create_app()
+
+        @app.mcp_server.tool()
+        async def my_tool(query: str) -> dict:
+            '''Custom MCP tool.'''
+            return {"result": query}
+
+        @app.mcp_server.resource("custom://data")
+        async def my_resource() -> str:
+            '''Custom resource.'''
+            return '{"data": "value"}'
 
     Design Pattern:
     1. Create MCP server
@@ -174,9 +189,10 @@ def create_app() -> FastAPI:
     6. Define health endpoints
     7. Register API routers
     8. Mount MCP app
+    9. Expose mcp_server on app for extension
 
     Returns:
-        Configured FastAPI application
+        Configured FastAPI application with .mcp_server attribute
     """
     # Create MCP server and get HTTP app
     # path="/" creates routes at root, then mount at /api/v1/mcp
@@ -228,6 +244,11 @@ def create_app() -> FastAPI:
 
     # Add SSE buffering middleware (for MCP SSE transport)
     app.add_middleware(SSEBufferingMiddleware)
+    
+    # Add Anonymous Tracking & Rate Limiting (Runs AFTER Auth if Auth is enabled)
+    # Must be added BEFORE AuthMiddleware in code to be INNER in the stack
+    from .middleware.tracking import AnonymousTrackingMiddleware
+    app.add_middleware(AnonymousTrackingMiddleware)
 
     # Add authentication middleware (if enabled)
     if settings.auth.enabled:
@@ -305,6 +326,10 @@ def create_app() -> FastAPI:
     # Mount MCP app at /api/v1/mcp
     app.mount("/api/v1/mcp", mcp_app)
 
+    # Expose MCP server on app for extension
+    # Users can add tools/resources/prompts via app.mcp_server
+    app.mcp_server = mcp_server  # type: ignore[attr-defined]
+
     return app
 
 

rem/api/middleware/tracking.py

@@ -0,0 +1,172 @@
+"""
+Anonymous User Tracking & Rate Limiting Middleware.
+
+Handles:
+1. Anonymous Identity: Generates/Validates 'rem_anon_id' cookie.
+2. Context Injection: Sets request.state.anon_id.
+3. Rate Limiting: Enforces tenant-aware tiered limits via RateLimitService.
+"""
+
+import hmac
+import hashlib
+import uuid
+import secrets
+from typing import Optional
+
+from fastapi import Request, Response
+from fastapi.responses import JSONResponse
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.types import ASGIApp
+
+from ...services.postgres.service import PostgresService
+from ...services.rate_limit import RateLimitService
+from ...models.entities.user import UserTier
+from ...settings import settings
+
+
+class AnonymousTrackingMiddleware(BaseHTTPMiddleware):
+    """
+    Middleware for anonymous user tracking and rate limiting.
+    
+    Design Pattern:
+    - Uses a secure, signed cookie for anonymous ID.
+    - Enforces rate limits before request processing.
+    - Injects anon_id into request state.
+    """
+    
+    def __init__(self, app: ASGIApp):
+        super().__init__(app)
+        # Secret for signing cookies (should be in settings, fallback for safety)
+        self.secret_key = settings.auth.session_secret or "fallback-secret-change-me"
+        self.cookie_name = "rem_anon_id"
+        
+        # Dedicated DB service for this middleware (one pool per app instance)
+        self.db = PostgresService()
+        self.rate_limiter = RateLimitService(self.db)
+        
+        # Excluded paths (health checks, static assets, auth callbacks)
+        self.excluded_paths = {
+            "/health", 
+            "/docs", 
+            "/openapi.json", 
+            "/favicon.ico",
+            "/api/auth", # Don't rate limit auth flow heavily
+        }
+
+    async def dispatch(self, request: Request, call_next):
+        # 0. Skip excluded paths
+        if any(request.url.path.startswith(p) for p in self.excluded_paths):
+            return await call_next(request)
+
+        # 1. Lazy DB Connection
+        if not self.db.pool:
+            # Note: simple lazy init. In high concurrency startup, might trigger multiple connects
+            # followed by disconnects, but asyncpg pool handles this gracefully usually.
+            # Ideally hook into lifespan, but middleware is separate.
+            if settings.postgres.enabled:
+                await self.db.connect()
+
+        # 2. Identification (Cookie Strategy)
+        anon_id = request.cookies.get(self.cookie_name)
+        is_new_anon = False
+        
+        if not anon_id or not self._validate_signature(anon_id):
+            anon_id = self._generate_signed_id()
+            is_new_anon = True
+            
+        # Strip signature for internal use
+        raw_anon_id = anon_id.split(".")[0]
+        request.state.anon_id = raw_anon_id
+        
+        # 3. Determine User Tier & ID for Rate Limiting
+        # Check if user is authenticated (set by AuthMiddleware usually, but that runs AFTER?)
+        # Actually middleware runs in reverse order of addition. 
+        # If AuthMiddleware adds user to request.session, we might need to access session directly.
+        # request.user is standard.
+        
+        user = getattr(request.state, "user", None)
+        if user:
+            # Authenticated User
+            identifier = user.get("id") # Assuming user dict or object
+            # Determine tier from user object
+            tier_str = user.get("tier", UserTier.FREE.value)
+            try:
+                tier = UserTier(tier_str)
+            except ValueError:
+                tier = UserTier.FREE
+            tenant_id = user.get("tenant_id", "default")
+        else:
+            # Anonymous User
+            identifier = raw_anon_id
+            tier = UserTier.ANONYMOUS
+            # Tenant ID from header or default
+            tenant_id = request.headers.get("X-Tenant-Id", "default")
+
+        # 4. Rate Limiting
+        if settings.postgres.enabled:
+            is_allowed, current, limit = await self.rate_limiter.check_rate_limit(
+                tenant_id=tenant_id,
+                identifier=identifier,
+                tier=tier
+            )
+            
+            if not is_allowed:
+                return JSONResponse(
+                    status_code=429,
+                    content={
+                        "error": {
+                            "code": "rate_limit_exceeded",
+                            "message": "You have exceeded your rate limit. Please sign in or upgrade to continue.",
+                            "details": {
+                                "limit": limit,
+                                "tier": tier.value,
+                                "retry_after": 60
+                            }
+                        }
+                    },
+                    headers={"Retry-After": "60"}
+                )
+        
+        # 5. Process Request
+        response = await call_next(request)
+        
+        # 6. Set Cookie if new
+        if is_new_anon:
+            response.set_cookie(
+                key=self.cookie_name,
+                value=anon_id,
+                max_age=31536000, # 1 year
+                httponly=True,
+                samesite="lax",
+                secure=settings.environment == "production"
+            )
+            
+        # Add Rate Limit headers
+        if settings.postgres.enabled and 'limit' in locals():
+            response.headers["X-RateLimit-Limit"] = str(limit)
+            response.headers["X-RateLimit-Remaining"] = str(max(0, limit - current))
+            
+        return response
+
+    def _generate_signed_id(self) -> str:
+        """Generate a UUID4 signed with HMAC."""
+        val = str(uuid.uuid4())
+        sig = hmac.new(
+            self.secret_key.encode(), 
+            val.encode(), 
+            hashlib.sha256
+        ).hexdigest()[:12] # Short signature
+        return f"{val}.{sig}"
+
+    def _validate_signature(self, signed_val: str) -> bool:
+        """Validate the HMAC signature."""
+        try:
+            val, sig = signed_val.split(".")
+            expected_sig = hmac.new(
+                self.secret_key.encode(), 
+                val.encode(), 
+                hashlib.sha256
+            ).hexdigest()[:12]
+            return secrets.compare_digest(sig, expected_sig)
+        except ValueError:
+            return False

rem/api/routers/auth.py

@@ -49,6 +49,8 @@ from authlib.integrations.starlette_client import OAuth
 from loguru import logger
 
 from ...settings import settings
+from ...services.postgres.service import PostgresService
+from ...services.user_service import UserService
 
 router = APIRouter(prefix="/api/auth", tags=["auth"])
 
@@ -168,6 +170,53 @@ async def callback(provider: str, request: Request):
         if not user_info:
             # Fetch from userinfo endpoint if not in ID token
             user_info = await client.userinfo(token=token)
+            
+        # --- REM Integration Start ---
+        if settings.postgres.enabled:
+            # Connect to DB
+            db = PostgresService()
+            try:
+                await db.connect()
+                user_service = UserService(db)
+                
+                # Get/Create User
+                user_entity = await user_service.get_or_create_user(
+                    email=user_info.get("email"),
+                    name=user_info.get("name", "New User"),
+                    avatar_url=user_info.get("picture"),
+                    tenant_id="default", # Single tenant for now
+                )
+                
+                # Link Anonymous Session
+                # TrackingMiddleware sets request.state.anon_id
+                anon_id = getattr(request.state, "anon_id", None)
+                # Fallback to cookie if middleware didn't run or state missing
+                if not anon_id:
+                    # Attempt to parse cookie manually if needed, but middleware 
+                    # usually handles the signature logic.
+                    # Just check raw cookie for simple case (not recommended if signed)
+                    pass 
+                
+                if anon_id:
+                    await user_service.link_anonymous_session(user_entity, anon_id)
+                    
+                # Enrich session user with DB info
+                db_info = {
+                    "id": str(user_entity.id),
+                    "tenant_id": user_entity.tenant_id,
+                    "tier": user_entity.tier.value if user_entity.tier else "free",
+                    "roles": [user_entity.role] if user_entity.role else [],
+                }
+                
+            except Exception as db_e:
+                logger.error(f"Database error during auth callback: {db_e}")
+                # Continue login even if DB fails, but warn
+                db_info = {"id": "db_error", "tier": "free"}
+            finally:
+                await db.disconnect()
+        else:
+            db_info = {"id": "no_db", "tier": "free"}
+        # --- REM Integration End ---
 
         # Store user info in session
         request.session["user"] = {
@@ -176,6 +225,11 @@ async def callback(provider: str, request: Request):
             "email": user_info.get("email"),
             "name": user_info.get("name"),
             "picture": user_info.get("picture"),
+            # Add DB info
+            "id": db_info.get("id"),
+            "tenant_id": db_info.get("tenant_id", "default"),
+            "tier": db_info.get("tier"),
+            "roles": db_info.get("roles", []),
         }
 
         # Store tokens in session for API access

rem/api/routers/chat/completions.py

@@ -251,7 +251,7 @@ async def chat_completions(body: ChatCompletionRequest, request: Request):
         }
 
         # Store messages with compression
-        store = SessionMessageStore(user_id=context.user_id or "default")
+        store = SessionMessageStore(user_id=context.user_id or settings.test.effective_user_id)
 
         await store.store_session_messages(
             session_id=context.session_id,