erdos-renyi 0.0.0__tar.gz

erdos_renyi-0.0.0/PKG-INFO

@@ -0,0 +1,70 @@
+Metadata-Version: 2.4
+Name: erdos-renyi
+Version: 0.0.0
+Summary: MCP connector: run Erdos queries on Reyni-registered EHR exports.
+Author-email: Krv Labs <team@krv.ai>
+License: Apache-2.0
+Requires-Python: <3.15,>=3.11
+Description-Content-Type: text/markdown
+Requires-Dist: fastmcp>=2.0
+Requires-Dist: polars>=1.0.0
+Requires-Dist: sqlglot>=23.0.0
+Requires-Dist: marimo>=0.8.0
+Requires-Dist: pyarrow>=15.0.0
+Requires-Dist: pydantic>=2.0.0
+Provides-Extra: dev
+Requires-Dist: pytest>=8.3; extra == "dev"
+Requires-Dist: ruff>=0.9; extra == "dev"
+
+# erdos-renyi
+
+`erdos-renyi` is an MCP server providing the privacy-preserving handshake between Erdos and Reyni. Clinicians register fine-grained EHR cohort exports with Reyni, and agents rewrite Erdos predictive SQL queries against tokenized schema names (`tok_…`) without ever accessing raw cell values or PHI. Results are staged locally in Arrow IPC format, gated by k-anonymity for agent-visible counts, and rendered in a private localhost Marimo dashboard for clinician review.
+
+## Quickstart
+
+```bash
+# Install dependencies
+uv sync --extra dev
+
+# Run test suite
+uv run pytest
+
+# Run the MCP server over stdio
+uv run er-mcp
+```
+
+## Reyni Gateway Setup
+
+Add `erdos-renyi` to `~/Library/Application Support/reyni/servers.toml` (Windows: `%APPDATA%\reyni\servers.toml`):
+
+```toml
+[[server]]
+name       = "erdos"
+launch     = ["uv", "tool", "run", "--python", "3.12", "--from", "erdos-renyi", "er-mcp"]
+phi_access = true
+```
+
+During local development, point `--from` to the local repository path:
+
+```toml
+[[server]]
+name       = "erdos"
+launch     = ["uv", "tool", "run", "--python", "3.12", "--from", "/path/to/erdos-renyi", "er-mcp"]
+phi_access = true
+```
+
+Agent-visible Erdos tools: `erdos__get_schema`, `erdos__stage_query`. After staging, call `reyni__register_dashboard` (Reyni gateway). `launch_dashboard` is gateway-internal only.
+
+## Agent Workflow
+
+1. **`reyni__list_datasets`**: Obtain the `safe_path` for the tokenized twin export.
+2. **`erdos__get_schema(safe_path)`**: Inspect column names and data types (no row values or samples).
+3. **Rewrite SQL**: Adapt the Erdos predictive query to target the local schema. The table name in Polars SQL is always **`cohort`**.
+4. **`erdos__stage_query(safe_path, sql)`**: Sanitize and execute SQL. If the match count is below the k-anonymity threshold ($k < 5$), the exact count is withheld (`below_threshold`), but the session is staged for the clinician.
+5. **`reyni__register_dashboard(session_id)`**: Reyni launches the Marimo dashboard internally, reads a clinician-only sidecar (`session_{id}.launch.json`), and stores the URL encrypted — the agent receives only a `dashboard_ref`.
+
+## Security Guarantees
+
+- **Zero cell egress**: MCP tools only return column metadata and aggregate match counts.
+- **SQL Sanitization**: Enforces read-only `SELECT` queries, rejects external filesystem functions (`read_csv`, `eval`), and clamps `LIMIT` to 1000 rows.
+- **Localhost Enclave**: Staging files use Arrow IPC; dashboards bind to `127.0.0.1`. Launch URLs are written to mode-`0600` sidecars, not returned over MCP.

erdos_renyi-0.0.0/README.md

@@ -0,0 +1,52 @@
+# erdos-renyi
+
+`erdos-renyi` is an MCP server providing the privacy-preserving handshake between Erdos and Reyni. Clinicians register fine-grained EHR cohort exports with Reyni, and agents rewrite Erdos predictive SQL queries against tokenized schema names (`tok_…`) without ever accessing raw cell values or PHI. Results are staged locally in Arrow IPC format, gated by k-anonymity for agent-visible counts, and rendered in a private localhost Marimo dashboard for clinician review.
+
+## Quickstart
+
+```bash
+# Install dependencies
+uv sync --extra dev
+
+# Run test suite
+uv run pytest
+
+# Run the MCP server over stdio
+uv run er-mcp
+```
+
+## Reyni Gateway Setup
+
+Add `erdos-renyi` to `~/Library/Application Support/reyni/servers.toml` (Windows: `%APPDATA%\reyni\servers.toml`):
+
+```toml
+[[server]]
+name       = "erdos"
+launch     = ["uv", "tool", "run", "--python", "3.12", "--from", "erdos-renyi", "er-mcp"]
+phi_access = true
+```
+
+During local development, point `--from` to the local repository path:
+
+```toml
+[[server]]
+name       = "erdos"
+launch     = ["uv", "tool", "run", "--python", "3.12", "--from", "/path/to/erdos-renyi", "er-mcp"]
+phi_access = true
+```
+
+Agent-visible Erdos tools: `erdos__get_schema`, `erdos__stage_query`. After staging, call `reyni__register_dashboard` (Reyni gateway). `launch_dashboard` is gateway-internal only.
+
+## Agent Workflow
+
+1. **`reyni__list_datasets`**: Obtain the `safe_path` for the tokenized twin export.
+2. **`erdos__get_schema(safe_path)`**: Inspect column names and data types (no row values or samples).
+3. **Rewrite SQL**: Adapt the Erdos predictive query to target the local schema. The table name in Polars SQL is always **`cohort`**.
+4. **`erdos__stage_query(safe_path, sql)`**: Sanitize and execute SQL. If the match count is below the k-anonymity threshold ($k < 5$), the exact count is withheld (`below_threshold`), but the session is staged for the clinician.
+5. **`reyni__register_dashboard(session_id)`**: Reyni launches the Marimo dashboard internally, reads a clinician-only sidecar (`session_{id}.launch.json`), and stores the URL encrypted — the agent receives only a `dashboard_ref`.
+
+## Security Guarantees
+
+- **Zero cell egress**: MCP tools only return column metadata and aggregate match counts.
+- **SQL Sanitization**: Enforces read-only `SELECT` queries, rejects external filesystem functions (`read_csv`, `eval`), and clamps `LIMIT` to 1000 rows.
+- **Localhost Enclave**: Staging files use Arrow IPC; dashboards bind to `127.0.0.1`. Launch URLs are written to mode-`0600` sidecars, not returned over MCP.

erdos_renyi-0.0.0/dashboard/app.py

@@ -0,0 +1,76 @@
+import os
+import sys
+import tempfile
+from pathlib import Path
+from typing import Any
+
+import marimo
+import polars as pl
+
+__generated_with = "0.24.0"
+app = marimo.App(width="full", app_title="Erdos-Renyi Cohort Review")
+
+SCORE_CANDIDATES = ("score", "rank", "fit", "probability", "prob")
+
+
+def _extract_session_id() -> str | None:
+    for arg in sys.argv:
+        if arg.startswith("--session="):
+            return arg.split("=", 1)[1]
+    return None
+
+
+def _find_sort_col(columns: list[str]) -> str | None:
+    for col in columns:
+        if any(c in col.lower() for c in SCORE_CANDIDATES):
+            return col
+    return None
+
+
+def _load_and_render(session_id: str, staging_dir: Path, mo: Any) -> Any:
+    ipc_path = staging_dir / f"session_{session_id}.arrow"
+    if not ipc_path.is_file():
+        return mo.md(f"### ⚠️ Session `{session_id}` data file not found or expired.")
+    try:
+        df = pl.read_ipc(ipc_path)
+        matched = _find_sort_col(df.columns)
+        if matched:
+            df = df.sort(matched, descending=True)
+        sort_info = (
+            f"- **Sorted by:** `{matched}` (descending)"
+            if matched
+            else "- **Sorted by:** (original order)"
+        )
+        header = mo.md(
+            f"""
+# 🔬 Erdos-Renyi Cohort Review
+- **Session ID:** `{session_id}`
+- **Matched Records:** `{len(df)}`
+{sort_info}
+"""
+        )
+        return mo.vstack([header, mo.ui.table(df, selection=None)])
+    except Exception as e:
+        return mo.md(f"### ⚠️ Failed to read session dataset: {e}")
+
+
+@app.cell
+def _():
+    import marimo as mo
+
+    session_id = _extract_session_id()
+    if not session_id:
+        output = mo.md("### ⚠️ No session ID provided to dashboard.")
+    else:
+        staging_dir = Path(
+            os.getenv("ERDOS_RENYI_STAGING_DIR")
+            or Path(tempfile.gettempdir()) / "erdos-renyi-sessions"
+        )
+        output = _load_and_render(session_id, staging_dir, mo)
+
+    output
+    return
+
+
+if __name__ == "__main__":
+    app.run()

erdos_renyi-0.0.0/er/__init__.py

@@ -0,0 +1,3 @@
+"""erdos-renyi: MCP connector to run Erdos queries on Reyni-registered EHR exports."""
+
+__version__ = "0.1.0"

erdos_renyi-0.0.0/er/config.py

@@ -0,0 +1,40 @@
+"""Configuration and settings for erdos-renyi."""
+
+from __future__ import annotations
+
+import os
+import tempfile
+from pathlib import Path
+
+from pydantic import BaseModel, Field
+
+
+class Settings(BaseModel):
+    """Runtime configuration settings."""
+
+    max_row_limit: int = Field(
+        default_factory=lambda: int(os.getenv("ERDOS_RENYI_MAX_ROWS", "1000"))
+    )
+    k_anonymity: int = Field(default_factory=lambda: int(os.getenv("ERDOS_RENYI_K", "5")))
+    session_ttl_minutes: int = Field(
+        default_factory=lambda: int(os.getenv("ERDOS_RENYI_TTL_MIN", "15"))
+    )
+    dashboard_host: str = Field(
+        default_factory=lambda: os.getenv("ERDOS_RENYI_DASHBOARD_HOST", "127.0.0.1")
+    )
+    secret: str = Field(default_factory=lambda: os.getenv("ERDOS_RENYI_SECRET", ""))
+    staging_dir: Path = Field(
+        default_factory=lambda: Path(
+            os.getenv("ERDOS_RENYI_STAGING_DIR")
+            or Path(tempfile.gettempdir()) / "erdos-renyi-sessions"
+        )
+    )
+
+    model_config = {"frozen": True}
+
+
+def get_settings() -> Settings:
+    """Return application settings, ensuring staging directory exists."""
+    settings = Settings()
+    settings.staging_dir.mkdir(parents=True, exist_ok=True)
+    return settings

erdos_renyi-0.0.0/er/dashboard.py

@@ -0,0 +1,250 @@
+"""Localhost Marimo dashboard launcher with HMAC authentication and TTL."""
+
+from __future__ import annotations
+
+import atexit
+import hashlib
+import json
+import hmac
+import secrets
+import socket
+from subprocess import DEVNULL, Popen, SubprocessError
+import sys
+import time
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+from er.config import Settings, get_settings
+from er.data import ArrowSessionStore, drop_session
+
+_LOCAL_SECRET: str = secrets.token_hex(32)
+
+
+def resolve_dashboard_app() -> Path:
+    """Resolve path to the Marimo app.py file."""
+    candidates = [
+        Path(__file__).resolve().parent.parent / "dashboard" / "app.py",
+        Path(__file__).resolve().parent / "dashboard" / "app.py",
+        Path(__file__).resolve().parent / "app.py",
+    ]
+    for c in candidates:
+        if c.is_file():
+            return c
+    raise FileNotFoundError(
+        f"Marimo dashboard script app.py not found in {[str(c) for c in candidates]}"
+    )
+
+
+def find_free_port() -> int:
+    """Find an available port on 127.0.0.1."""
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("127.0.0.1", 0))
+        return s.getsockname()[1]
+
+
+class ActiveSession(BaseModel):
+    """Metadata and process reference for an active dashboard session."""
+
+    session_id: str
+    port: int
+    token: str
+    process: Popen[Any] = Field(exclude=True)
+    expires_at: float
+
+    model_config = {"arbitrary_types_allowed": True}
+
+
+class DashboardRunner(ABC):
+    """Abstract base class for running localhost dashboards."""
+
+    @abstractmethod
+    def launch(self, session_id: str) -> dict[str, Any]:
+        """Launch a dashboard instance for the given session ID."""
+
+    @abstractmethod
+    def cleanup_expired(self) -> None:
+        """Clean up expired dashboard processes and staged data."""
+
+    @abstractmethod
+    def stop_all(self) -> None:
+        """Stop all running dashboard processes."""
+
+
+
+def write_launch_sidecar(staging_dir: Path, session_id: str, url: str, token: str, expires_at: float) -> None:
+    """Write clinician-only launch metadata; never returned over MCP."""
+    staging_dir.mkdir(parents=True, exist_ok=True)
+    sidecar = staging_dir / f"session_{session_id}.launch.json"
+    payload = {
+        "dashboard_url": url,
+        "token": token,
+        "expires_at": expires_at,
+    }
+    sidecar.write_text(json.dumps(payload), encoding="utf-8")
+    try:
+        sidecar.chmod(0o600)
+    except OSError:
+        pass
+
+
+def agent_launch_response(
+    session_id: str, expires_in_minutes: int
+) -> dict[str, Any]:
+    """MCP-safe launch payload — no localhost URL."""
+    return {
+        "status": "launched",
+        "session_id": session_id,
+        "expires_in_minutes": expires_in_minutes,
+        "instructions": "Call reyni__register_dashboard so the clinician can open this review in the Reyni app.",
+    }
+
+
+class MarimoDashboardRunner(DashboardRunner):
+    """Manages lifecycle of ephemeral Marimo dashboard subprocesses."""
+
+    def __init__(self, settings: Settings | None = None) -> None:
+        self.settings = settings or get_settings()
+        self.sessions: dict[str, ActiveSession] = {}
+        atexit.register(self.stop_all)
+
+    def _get_secret(self) -> str:
+        return self.settings.secret if self.settings.secret else _LOCAL_SECRET
+
+    def _generate_token(self, session_id: str, expires_at: float) -> str:
+        secret = self._get_secret()
+        msg = f"{session_id}:{int(expires_at)}".encode()
+        return hmac.new(secret.encode(), msg, hashlib.sha256).hexdigest()
+
+    def _terminate_and_drop(self, session: ActiveSession) -> None:
+        try:
+            session.process.terminate()
+            session.process.wait(timeout=2)
+        except (OSError, SubprocessError):
+            try:
+                session.process.kill()
+            except (OSError, SubprocessError):
+                pass
+        drop_session(self.settings.staging_dir, session.session_id)
+
+    def cleanup_expired(self) -> None:
+        now = time.time()
+        expired_ids = [
+            sid
+            for sid, s in self.sessions.items()
+            if now >= s.expires_at or s.process.poll() is not None
+        ]
+        for sid in expired_ids:
+            session = self.sessions.pop(sid, None)
+            if session:
+                self._terminate_and_drop(session)
+
+    def stop_all(self) -> None:
+        for sid in list(self.sessions.keys()):
+            session = self.sessions.pop(sid, None)
+            if session:
+                self._terminate_and_drop(session)
+
+    def launch(self, session_id: str) -> dict[str, Any]:
+        self.cleanup_expired()
+
+        store = ArrowSessionStore(self.settings.staging_dir)
+        if not store.exists(session_id):
+            return {
+                "status": "error",
+                "message": f"Session data not found or expired: {session_id}",
+            }
+
+        # If already running, return existing URL if not expired
+        if session_id in self.sessions:
+            existing = self.sessions[session_id]
+            if existing.process.poll() is None and time.time() < existing.expires_at:
+                url = (
+                    f"http://{self.settings.dashboard_host}:{existing.port}/?token={existing.token}"
+                )
+                write_launch_sidecar(
+                    self.settings.staging_dir,
+                    session_id,
+                    url,
+                    existing.token,
+                    existing.expires_at,
+                )
+                return agent_launch_response(session_id, self.settings.session_ttl_minutes)
+            else:
+                self.sessions.pop(session_id, None)
+
+        try:
+            app_py = resolve_dashboard_app()
+        except FileNotFoundError as exc:
+            return {"status": "error", "message": str(exc)}
+
+        port = find_free_port()
+        expires_at = time.time() + (self.settings.session_ttl_minutes * 60)
+        token = self._generate_token(session_id, expires_at)
+
+        cmd = [
+            sys.executable,
+            "-m",
+            "marimo",
+            "run",
+            str(app_py),
+            "--host",
+            "127.0.0.1",
+            "--port",
+            str(port),
+            "--headless",
+            "--",
+            f"--session={session_id}",
+            f"--token={token}",
+        ]
+
+        try:
+            proc = Popen(
+                cmd,
+                stdout=DEVNULL,
+                stderr=DEVNULL,
+            )
+        except (OSError, SubprocessError) as exc:
+            return {
+                "status": "error",
+                "message": f"Failed to spawn dashboard process: {exc}",
+            }
+
+        active = ActiveSession(
+            session_id=session_id,
+            port=port,
+            token=token,
+            process=proc,
+            expires_at=expires_at,
+        )
+        self.sessions[session_id] = active
+
+        url = f"http://{self.settings.dashboard_host}:{port}/?token={token}"
+        write_launch_sidecar(
+            self.settings.staging_dir,
+            session_id,
+            url,
+            token,
+            expires_at,
+        )
+        return agent_launch_response(session_id, self.settings.session_ttl_minutes)
+
+
+# Default global dashboard runner instance
+_runner: MarimoDashboardRunner | None = None
+
+
+def get_dashboard_runner() -> MarimoDashboardRunner:
+    """Return singleton Marimo dashboard runner."""
+    global _runner
+    if _runner is None:
+        _runner = MarimoDashboardRunner()
+    return _runner
+
+
+def launch_dashboard(session_id: str, settings: Settings | None = None) -> dict[str, Any]:
+    """Launch dashboard for session_id."""
+    runner = MarimoDashboardRunner(settings) if settings else get_dashboard_runner()
+    return runner.launch(session_id)

erdos_renyi-0.0.0/er/data.py

@@ -0,0 +1,126 @@
+"""Dataset scanning, SQL execution on Polars, and Arrow session staging."""
+
+from __future__ import annotations
+
+import secrets
+from abc import ABC, abstractmethod
+from pathlib import Path
+
+import polars as pl
+
+
+class DatasetScanner(ABC):
+    """Abstract base class for scanning local datasets."""
+
+    @abstractmethod
+    def scan(self, path: Path) -> pl.LazyFrame:
+        """Scan a dataset file into a LazyFrame without loading all rows into memory."""
+
+
+class CSVDatasetScanner(DatasetScanner):
+    """Scanner for CSV datasets (Reyni tokenized twins)."""
+
+    def scan(self, path: Path) -> pl.LazyFrame:
+        return pl.scan_csv(path, infer_schema_length=1000)
+
+
+class ParquetDatasetScanner(DatasetScanner):
+    """Scanner for Parquet datasets."""
+
+    def scan(self, path: Path) -> pl.LazyFrame:
+        return pl.scan_parquet(path)
+
+
+def get_scanner(path: str | Path) -> DatasetScanner:
+    """Return the appropriate dataset scanner for a given path."""
+    p = Path(path)
+    if p.suffix.lower() == ".parquet":
+        return ParquetDatasetScanner()
+    return CSVDatasetScanner()
+
+
+def scan_path(path: str | Path) -> pl.LazyFrame:
+    """Scan a dataset path into a Polars LazyFrame.
+
+    Raises:
+        FileNotFoundError: If the file does not exist or is not a regular file.
+    """
+    p = Path(path)
+    if not p.is_file():
+        raise FileNotFoundError(f"Dataset file not found: {path}")
+    scanner = get_scanner(p)
+    return scanner.scan(p)
+
+
+def schema_of(path: str | Path) -> list[dict[str, str]]:
+    """Return column names and dtypes of the dataset without reading cell values."""
+    sch = scan_path(path).collect_schema()
+    return [{"name": name, "dtype": str(dtype)} for name, dtype in sch.items()]
+
+
+class SessionStore(ABC):
+    """Abstract base class for staging cohort session results."""
+
+    @abstractmethod
+    def stage(self, session_id: str, df: pl.DataFrame) -> Path:
+        """Stage query results under the session ID."""
+
+    @abstractmethod
+    def get_path(self, session_id: str) -> Path | None:
+        """Get path to staged session data if it exists."""
+
+    @abstractmethod
+    def exists(self, session_id: str) -> bool:
+        """Check if staged session data exists."""
+
+    @abstractmethod
+    def drop(self, session_id: str) -> None:
+        """Drop staged session data."""
+
+
+class ArrowSessionStore(SessionStore):
+    """Stores query results as Arrow IPC files."""
+
+    def __init__(self, staging_dir: Path) -> None:
+        self.staging_dir = Path(staging_dir)
+
+    def _file_path(self, session_id: str) -> Path:
+        return self.staging_dir / f"session_{session_id}.arrow"
+
+    def stage(self, session_id: str, df: pl.DataFrame) -> Path:
+        self.staging_dir.mkdir(parents=True, exist_ok=True)
+        dest = self._file_path(session_id)
+        df.write_ipc(dest)
+        return dest
+
+    def get_path(self, session_id: str) -> Path | None:
+        p = self._file_path(session_id)
+        return p if p.is_file() else None
+
+    def exists(self, session_id: str) -> bool:
+        return self._file_path(session_id).is_file()
+
+    def drop(self, session_id: str) -> None:
+        p = self._file_path(session_id)
+        if p.exists():
+            p.unlink()
+
+
+def execute_and_stage(path: str | Path, sql: str, staging_dir: Path) -> tuple[str, int]:
+    """Execute SQL query on the dataset and stage results into an Arrow IPC file.
+
+    Returns:
+        tuple[str, int]: (session_id, row_count)
+    """
+    session_id = secrets.token_hex(8)
+    ctx = pl.SQLContext(cohort=scan_path(path), eager=False)
+    df = ctx.execute(sql).collect()
+    store = ArrowSessionStore(staging_dir)
+    store.stage(session_id, df)
+    return session_id, len(df)
+
+
+def drop_session(staging_dir: Path, session_id: str) -> None:
+    """Drop staged Arrow IPC session file."""
+    store = ArrowSessionStore(staging_dir)
+    store.drop(session_id)

erdos_renyi-0.0.0/er/privacy.py

@@ -0,0 +1,15 @@
+"""Privacy rules and k-anonymity gates."""
+
+from __future__ import annotations
+
+
+def agent_count(row_count: int, k: int = 5) -> int | str:
+    """Gate the agent-visible count with k-anonymity.
+
+    If 0 < row_count < k, returns 'below_threshold' to prevent
+    identifying small cohort sizes (e.g. rare disease queries).
+    If count is 0 or >= k, returns the exact integer count.
+    """
+    if 0 < row_count < k:
+        return "below_threshold"
+    return row_count

erdos_renyi-0.0.0/er/sanitizer.py

@@ -0,0 +1,102 @@
+"""SQL sanitization and preparation for Polars execution."""
+
+from __future__ import annotations
+
+import sqlglot
+from sqlglot import exp
+
+
+class SQLSanitizerError(Exception):
+    """Raised when an SQL query fails sanitization or validation."""
+
+
+FORBIDDEN_EXPRESSIONS = (
+    exp.Insert,
+    exp.Update,
+    exp.Delete,
+    exp.Drop,
+    exp.Alter,
+    exp.Create,
+    exp.Command,
+    exp.Pragma,
+    exp.Transaction,
+    exp.ReadCSV,
+)
+
+FORBIDDEN_FUNCTIONS = {
+    "read_csv",
+    "read_parquet",
+    "scan_parquet",
+    "scan_csv",
+    "scan_ndjson",
+    "read_json",
+    "read_ndjson",
+    "write_file",
+    "eval",
+    "readcsv",
+    "readparquet",
+    "scanparquet",
+    "scancsv",
+    "writefile",
+    "readjson",
+}
+
+
+def _check_forbidden_nodes(parsed: exp.Expression) -> None:
+    for node in parsed.walk():
+        if isinstance(node, FORBIDDEN_EXPRESSIONS):
+            raise SQLSanitizerError(f"Forbidden SQL statement type: {type(node).__name__}")
+        if isinstance(node, (exp.Anonymous, exp.Func)):
+            name = (getattr(node, "name", "") or "").lower()
+            cls = type(node).__name__.lower()
+            if name in FORBIDDEN_FUNCTIONS or cls in FORBIDDEN_FUNCTIONS or cls.replace("_", "") in FORBIDDEN_FUNCTIONS:
+                raise SQLSanitizerError(f"Forbidden SQL function: {name or cls}")
+
+
+def _clamp_limit(parsed: exp.Select, max_row_limit: int) -> exp.Select:
+    limit_node = parsed.args.get("limit")
+    if limit_node is None or limit_node.expression is None:
+        return parsed.limit(max_row_limit)
+
+    expr = limit_node.expression
+    if isinstance(expr, exp.Literal) and expr.is_int:
+        val = int(expr.this)
+        if 0 < val <= max_row_limit:
+            return parsed
+
+    limit_node.set("expression", exp.Literal.number(max_row_limit))
+    return parsed
+
+
+def _parse_select(sql: str) -> exp.Select:
+    if not sql or not sql.strip():
+        raise SQLSanitizerError("Empty SQL query")
+    try:
+        parsed = sqlglot.parse_one(sql, read="postgres")
+    except Exception as exc:
+        raise SQLSanitizerError(f"SQL parsing error: {exc}") from exc
+
+    if not isinstance(parsed, exp.Select):
+        raise SQLSanitizerError(
+            f"Query root must be a SELECT statement, got {type(parsed).__name__}"
+        )
+    return parsed
+
+
+def sanitize_and_prepare_sql(sql: str, max_row_limit: int = 1000) -> str:
+    """Sanitize and clamp an SQL query for safe local execution.
+
+    Args:
+        sql: Raw SQL text provided by the agent.
+        max_row_limit: Maximum allowed rows in the LIMIT clause.
+
+    Returns:
+        Cleaned SQL string in postgres dialect with clamped LIMIT.
+
+    Raises:
+        SQLSanitizerError: If query violates safety constraints.
+    """
+    parsed = _parse_select(sql)
+    _check_forbidden_nodes(parsed)
+    clamped = _clamp_limit(parsed, max_row_limit)
+    return clamped.sql(dialect="duckdb")

erdos_renyi-0.0.0/er/server.py

@@ -0,0 +1,167 @@
+"""FastMCP server entrypoint for erdos-renyi."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import polars as pl
+from fastmcp import FastMCP
+from pydantic import BaseModel
+
+from er.config import get_settings
+from er.dashboard import launch_dashboard as do_launch_dashboard
+from er.data import execute_and_stage, schema_of
+from er.privacy import agent_count
+from er.sanitizer import SQLSanitizerError, sanitize_and_prepare_sql
+
+
+class ColumnInfo(BaseModel):
+    """Metadata for a single column."""
+
+    name: str
+    dtype: str
+
+
+class SchemaResponse(BaseModel):
+    """Schema inspection return model."""
+
+    table: str = "cohort"
+    path: str
+    n_columns: int
+    columns: list[ColumnInfo]
+    sql_notes: str = (
+        "FROM cohort. Identifier columns are Reyni tokens (tok_…), not raw IDs. "
+        "Rewrite the Erdos query against these names. Do not SELECT rows back through MCP — call stage_query."
+    )
+
+
+class StagedResponse(BaseModel):
+    """Successful query staging return model."""
+
+    status: str = "staged"
+    session_id: str
+    matched_records: int
+    message: str = "Staged. Call reyni__register_dashboard with this session_id. Do not request rows."
+
+
+class StagedSuppressedResponse(BaseModel):
+    """Suppressed count staging return model (k-anonymity)."""
+
+    status: str = "staged_suppressed"
+    session_id: str
+    matched_records: str = "below_threshold"
+    message: str = (
+        "Cohort is below the k-anonymity threshold. Count withheld. "
+        "Call reyni__register_dashboard so the clinician can review in the Reyni app."
+    )
+
+
+class ErrorResponse(BaseModel):
+    """Error return model."""
+
+    status: str = "error"
+    error_type: str | None = None
+    message: str
+
+
+mcp = FastMCP(
+    "er-mcp",
+    instructions=(
+        "Run Erdos trial-matching queries on a Reyni-registered export. "
+        "1) Reyni list_datasets → safe_path. "
+        "2) get_schema(safe_path). "
+        "3) Rewrite the Erdos SQL against table cohort using those column names. "
+        "4) stage_query(safe_path, sql). "
+        "5) reyni__register_dashboard(session_id) so the clinician opens the review in the Reyni app. "
+        "Never pass the original file path. Never return or request patient rows. "
+        "Never fetch dashboard URLs."
+    ),
+)
+
+
+@mcp.tool()
+def get_schema(safe_path: str) -> dict[str, Any]:
+    """Lazy-scan the dataset at safe_path and return column names and dtypes.
+
+    Never exposes raw rows or cell values.
+
+    Args:
+        safe_path: Path to the registered dataset (from Reyni list_datasets).
+
+    Returns:
+        Dictionary with table name, column count, and column name/dtype pairs.
+    """
+    try:
+        cols_raw = schema_of(safe_path)
+        cols = [ColumnInfo(**c) for c in cols_raw]
+        resp = SchemaResponse(
+            path=str(safe_path),
+            n_columns=len(cols),
+            columns=cols,
+        )
+        return resp.model_dump()
+    except (FileNotFoundError, pl.exceptions.PolarsError, OSError, ValueError) as exc:
+        return ErrorResponse(message=str(exc)).model_dump(exclude_none=True)
+
+
+@mcp.tool()
+def stage_query(safe_path: str, sql: str) -> dict[str, Any]:
+    """Sanitize, execute SQL against safe_path dataset as table 'cohort', and stage results.
+
+    Stages the resulting cohort into a private Arrow IPC session for the localhost dashboard.
+    Applies k-anonymity gating to the matched record count reported to the agent.
+
+    Args:
+        safe_path: Path to the registered dataset (from Reyni list_datasets).
+        sql: SQL query to run against the 'cohort' table.
+
+    Returns:
+        Dictionary containing staging status, session_id, and matched records count.
+    """
+    settings = get_settings()
+
+    try:
+        sanitized_sql = sanitize_and_prepare_sql(sql, max_row_limit=settings.max_row_limit)
+    except SQLSanitizerError as exc:
+        return ErrorResponse(error_type="SQLSanitizerError", message=str(exc)).model_dump()
+
+    try:
+        session_id, row_count = execute_and_stage(
+            safe_path, sanitized_sql, staging_dir=settings.staging_dir
+        )
+    except (pl.exceptions.PolarsError, OSError, ValueError, RuntimeError) as exc:
+        return ErrorResponse(
+            error_type="ExecutionError",
+            message=f"Query execution failed: {exc}",
+        ).model_dump()
+
+    gated_count = agent_count(row_count, k=settings.k_anonymity)
+    if gated_count == "below_threshold":
+        return StagedSuppressedResponse(session_id=session_id).model_dump()
+
+    return StagedResponse(session_id=session_id, matched_records=row_count).model_dump()
+
+
+@mcp.tool()
+def launch_dashboard(session_id: str) -> dict[str, Any]:
+    """Spawn Marimo dashboard bound to 127.0.0.1 (Reyni gateway-internal).
+
+    Writes launch metadata to a clinician-only sidecar; MCP response has no URL.
+
+    Args:
+        session_id: Session ID returned by stage_query.
+
+    Returns:
+        Dictionary with status, session_id, TTL, and instructions (no dashboard_url).
+    """
+    settings = get_settings()
+    return do_launch_dashboard(session_id, settings=settings)
+
+
+def main() -> None:
+    """Run FastMCP stdio server."""
+    mcp.run()
+
+
+if __name__ == "__main__":
+    main()

erdos_renyi-0.0.0/erdos_renyi.egg-info/PKG-INFO

@@ -0,0 +1,70 @@
+Metadata-Version: 2.4
+Name: erdos-renyi
+Version: 0.0.0
+Summary: MCP connector: run Erdos queries on Reyni-registered EHR exports.
+Author-email: Krv Labs <team@krv.ai>
+License: Apache-2.0
+Requires-Python: <3.15,>=3.11
+Description-Content-Type: text/markdown
+Requires-Dist: fastmcp>=2.0
+Requires-Dist: polars>=1.0.0
+Requires-Dist: sqlglot>=23.0.0
+Requires-Dist: marimo>=0.8.0
+Requires-Dist: pyarrow>=15.0.0
+Requires-Dist: pydantic>=2.0.0
+Provides-Extra: dev
+Requires-Dist: pytest>=8.3; extra == "dev"
+Requires-Dist: ruff>=0.9; extra == "dev"
+
+# erdos-renyi
+
+`erdos-renyi` is an MCP server providing the privacy-preserving handshake between Erdos and Reyni. Clinicians register fine-grained EHR cohort exports with Reyni, and agents rewrite Erdos predictive SQL queries against tokenized schema names (`tok_…`) without ever accessing raw cell values or PHI. Results are staged locally in Arrow IPC format, gated by k-anonymity for agent-visible counts, and rendered in a private localhost Marimo dashboard for clinician review.
+
+## Quickstart
+
+```bash
+# Install dependencies
+uv sync --extra dev
+
+# Run test suite
+uv run pytest
+
+# Run the MCP server over stdio
+uv run er-mcp
+```
+
+## Reyni Gateway Setup
+
+Add `erdos-renyi` to `~/Library/Application Support/reyni/servers.toml` (Windows: `%APPDATA%\reyni\servers.toml`):
+
+```toml
+[[server]]
+name       = "erdos"
+launch     = ["uv", "tool", "run", "--python", "3.12", "--from", "erdos-renyi", "er-mcp"]
+phi_access = true
+```
+
+During local development, point `--from` to the local repository path:
+
+```toml
+[[server]]
+name       = "erdos"
+launch     = ["uv", "tool", "run", "--python", "3.12", "--from", "/path/to/erdos-renyi", "er-mcp"]
+phi_access = true
+```
+
+Agent-visible Erdos tools: `erdos__get_schema`, `erdos__stage_query`. After staging, call `reyni__register_dashboard` (Reyni gateway). `launch_dashboard` is gateway-internal only.
+
+## Agent Workflow
+
+1. **`reyni__list_datasets`**: Obtain the `safe_path` for the tokenized twin export.
+2. **`erdos__get_schema(safe_path)`**: Inspect column names and data types (no row values or samples).
+3. **Rewrite SQL**: Adapt the Erdos predictive query to target the local schema. The table name in Polars SQL is always **`cohort`**.
+4. **`erdos__stage_query(safe_path, sql)`**: Sanitize and execute SQL. If the match count is below the k-anonymity threshold ($k < 5$), the exact count is withheld (`below_threshold`), but the session is staged for the clinician.
+5. **`reyni__register_dashboard(session_id)`**: Reyni launches the Marimo dashboard internally, reads a clinician-only sidecar (`session_{id}.launch.json`), and stores the URL encrypted — the agent receives only a `dashboard_ref`.
+
+## Security Guarantees
+
+- **Zero cell egress**: MCP tools only return column metadata and aggregate match counts.
+- **SQL Sanitization**: Enforces read-only `SELECT` queries, rejects external filesystem functions (`read_csv`, `eval`), and clamps `LIMIT` to 1000 rows.
+- **Localhost Enclave**: Staging files use Arrow IPC; dashboards bind to `127.0.0.1`. Launch URLs are written to mode-`0600` sidecars, not returned over MCP.

erdos_renyi-0.0.0/erdos_renyi.egg-info/SOURCES.txt

@@ -0,0 +1,21 @@
+README.md
+pyproject.toml
+dashboard/app.py
+er/__init__.py
+er/config.py
+er/dashboard.py
+er/data.py
+er/privacy.py
+er/sanitizer.py
+er/server.py
+erdos_renyi.egg-info/PKG-INFO
+erdos_renyi.egg-info/SOURCES.txt
+erdos_renyi.egg-info/dependency_links.txt
+erdos_renyi.egg-info/entry_points.txt
+erdos_renyi.egg-info/requires.txt
+erdos_renyi.egg-info/top_level.txt
+tests/test_dashboard.py
+tests/test_privacy.py
+tests/test_sanitizer.py
+tests/test_schema.py
+tests/test_stage.py

erdos_renyi-0.0.0/erdos_renyi.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

erdos_renyi-0.0.0/erdos_renyi.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+er-mcp = er.server:main

erdos_renyi-0.0.0/erdos_renyi.egg-info/requires.txt

@@ -0,0 +1,10 @@
+fastmcp>=2.0
+polars>=1.0.0
+sqlglot>=23.0.0
+marimo>=0.8.0
+pyarrow>=15.0.0
+pydantic>=2.0.0
+
+[dev]
+pytest>=8.3
+ruff>=0.9

erdos_renyi-0.0.0/erdos_renyi.egg-info/top_level.txt

@@ -0,0 +1,2 @@
+dashboard
+er

erdos_renyi-0.0.0/pyproject.toml

@@ -0,0 +1,43 @@
+[project]
+name = "erdos-renyi"
+version = "0.0.0"
+description = "MCP connector: run Erdos queries on Reyni-registered EHR exports."
+readme = "README.md"
+requires-python = ">=3.11,<3.15"
+authors = [{ name = "Krv Labs", email = "team@krv.ai" }]
+license = { text = "Apache-2.0" }
+dependencies = [
+    "fastmcp>=2.0",
+    "polars>=1.0.0",
+    "sqlglot>=23.0.0",
+    "marimo>=0.8.0",
+    "pyarrow>=15.0.0",
+    "pydantic>=2.0.0",
+]
+
+[project.optional-dependencies]
+dev = ["pytest>=8.3", "ruff>=0.9"]
+
+[project.scripts]
+er-mcp = "er.server:main"
+
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools.packages.find]
+include = ["er", "er.*", "dashboard"]
+
+[tool.setuptools]
+include-package-data = true
+
+[tool.setuptools.package-data]
+er = ["*.py", "dashboard/*.py"]
+dashboard = ["*.py"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
+[tool.ruff]
+line-length = 100
+target-version = "py312"

erdos_renyi-0.0.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

erdos_renyi-0.0.0/tests/test_dashboard.py

@@ -0,0 +1,59 @@
+"""Unit tests for localhost dashboard lifecycle."""
+
+from pathlib import Path
+
+from er.config import Settings
+from er.dashboard import MarimoDashboardRunner
+from er.server import launch_dashboard
+
+
+def test_launch_dashboard_missing_session():
+    result = launch_dashboard("nonexistent_session_id")
+    assert result["status"] == "error"
+    assert "not found" in result["message"].lower()
+
+
+def test_launch_dashboard_success(sample_csv_path: Path, tmp_path: Path):
+    settings = Settings(staging_dir=tmp_path / "staging")
+    sql = "SELECT patient_id, score FROM cohort WHERE bmi > 30"
+    runner = MarimoDashboardRunner(settings=settings)
+    # Stage session into runner's staging dir for test
+    from er.data import execute_and_stage
+    from er.sanitizer import sanitize_and_prepare_sql
+
+    staged_sid, _ = execute_and_stage(
+        sample_csv_path, sanitize_and_prepare_sql(sql), settings.staging_dir
+    )
+
+    res = runner.launch(staged_sid)
+    try:
+        assert res["status"] == "launched"
+        assert res["session_id"] == staged_sid
+        assert "dashboard_url" not in res
+        assert res["expires_in_minutes"] == 15
+        assert "register_dashboard" in res["instructions"].lower()
+        sidecar = settings.staging_dir / f"session_{staged_sid}.launch.json"
+        assert sidecar.is_file()
+        assert "127.0.0.1" in sidecar.read_text()
+    finally:
+        runner.stop_all()
+
+
+def test_dashboard_table_with_polars(sample_csv_path: Path, tmp_path: Path):
+    import marimo as mo
+    import polars as pl
+
+    from er.data import execute_and_stage
+    from er.sanitizer import sanitize_and_prepare_sql
+
+    staging_dir = tmp_path / "staging"
+    sql = "SELECT patient_id, score FROM cohort WHERE bmi > 30"
+    staged_sid, _ = execute_and_stage(
+        sample_csv_path, sanitize_and_prepare_sql(sql), staging_dir
+    )
+
+    ipc_path = staging_dir / f"session_{staged_sid}.arrow"
+    df = pl.read_ipc(ipc_path)
+    table = mo.ui.table(df, selection=None)
+    assert table is not None
+

erdos_renyi-0.0.0/tests/test_privacy.py

@@ -0,0 +1,19 @@
+"""Unit tests for k-anonymity privacy gate."""
+
+from er.privacy import agent_count
+
+
+def test_privacy_below_threshold():
+    assert agent_count(1, 5) == "below_threshold"
+    assert agent_count(3, 5) == "below_threshold"
+    assert agent_count(4, 5) == "below_threshold"
+
+
+def test_privacy_zero_allowed():
+    assert agent_count(0, 5) == 0
+
+
+def test_privacy_at_and_above_threshold():
+    assert agent_count(5, 5) == 5
+    assert agent_count(10, 5) == 10
+    assert agent_count(100, 5) == 100

erdos_renyi-0.0.0/tests/test_sanitizer.py

@@ -0,0 +1,63 @@
+"""Unit tests for SQL sanitizer."""
+
+import pytest
+
+from er.sanitizer import SQLSanitizerError, sanitize_and_prepare_sql
+
+
+def test_select_passes():
+    sql = "SELECT patient_id, bmi FROM cohort WHERE bmi > 30"
+    sanitized = sanitize_and_prepare_sql(sql)
+    assert "SELECT" in sanitized
+    assert "LIMIT 1000" in sanitized
+
+
+def test_insert_rejected():
+    sql = "INSERT INTO cohort (patient_id, bmi) VALUES ('x', 25.0)"
+    with pytest.raises(SQLSanitizerError):
+        sanitize_and_prepare_sql(sql)
+
+
+def test_drop_rejected():
+    sql = "DROP TABLE cohort"
+    with pytest.raises(SQLSanitizerError):
+        sanitize_and_prepare_sql(sql)
+
+
+def test_update_rejected():
+    sql = "UPDATE cohort SET bmi = 20.0 WHERE patient_id = 'x'"
+    with pytest.raises(SQLSanitizerError):
+        sanitize_and_prepare_sql(sql)
+
+
+def test_read_csv_rejected():
+    sql = "SELECT * FROM read_csv('secret.csv')"
+    with pytest.raises(SQLSanitizerError):
+        sanitize_and_prepare_sql(sql)
+
+
+def test_scan_parquet_rejected():
+    sql = "SELECT * FROM scan_parquet('secret.parquet')"
+    with pytest.raises(SQLSanitizerError):
+        sanitize_and_prepare_sql(sql)
+
+
+def test_limit_clamping():
+    sql = "SELECT * FROM cohort LIMIT 10000"
+    sanitized = sanitize_and_prepare_sql(sql, max_row_limit=1000)
+    assert "LIMIT 1000" in sanitized
+    assert "10000" not in sanitized
+
+
+def test_missing_limit_added():
+    sql = "SELECT * FROM cohort"
+    sanitized = sanitize_and_prepare_sql(sql, max_row_limit=500)
+    assert "LIMIT 500" in sanitized
+
+
+def test_negative_or_zero_limit_clamped():
+    sql_zero = "SELECT * FROM cohort LIMIT 0"
+    assert "LIMIT 1000" in sanitize_and_prepare_sql(sql_zero, max_row_limit=1000)
+
+    sql_neg = "SELECT * FROM cohort LIMIT -1"
+    assert "LIMIT 1000" in sanitize_and_prepare_sql(sql_neg, max_row_limit=1000)

erdos_renyi-0.0.0/tests/test_schema.py

@@ -0,0 +1,32 @@
+"""Unit tests for schema inspection tool."""
+
+import json
+from pathlib import Path
+
+from er.server import get_schema
+
+
+def test_get_schema_structure(sample_csv_path: Path):
+    result = get_schema(str(sample_csv_path))
+
+    assert result["table"] == "cohort"
+    assert result["path"] == str(sample_csv_path)
+    assert result["n_columns"] == 4
+
+    col_names = [c["name"] for c in result["columns"]]
+    assert col_names == ["patient_id", "icd10", "bmi", "score"]
+    assert all(set(c.keys()) == {"name", "dtype"} for c in result["columns"])
+
+
+def test_get_schema_no_cell_leak(sample_csv_path: Path):
+    result = get_schema(str(sample_csv_path))
+    dumped = json.dumps(result)
+    assert "secret-mrn-1" not in dumped
+    assert "E11.9" not in dumped
+    assert "32.5" not in dumped
+
+
+def test_get_schema_nonexistent_file():
+    result = get_schema("/nonexistent/path/cohort.csv")
+    assert result["status"] == "error"
+    assert "not found" in result["message"].lower()

erdos_renyi-0.0.0/tests/test_stage.py

@@ -0,0 +1,69 @@
+"""Unit tests for query staging and privacy gating."""
+
+import json
+from pathlib import Path
+
+import polars as pl
+
+from er.config import get_settings
+from er.server import stage_query
+
+
+def test_stage_query_happy_path(sample_csv_path: Path):
+    sql = "SELECT patient_id, score FROM cohort WHERE bmi > 30 ORDER BY score DESC"
+    result = stage_query(str(sample_csv_path), sql)
+
+    assert result["status"] == "staged"
+    assert "session_id" in result
+    assert isinstance(result["matched_records"], int)
+    assert result["matched_records"] == 11
+
+    # Ensure no patient data/cells leaked in MCP return
+    dumped = json.dumps(result)
+    assert "secret-mrn-1" not in dumped
+    assert "32.5" not in dumped
+
+    # Verify staged Arrow IPC exists and holds the rows
+    staging_dir = get_settings().staging_dir
+    session_id = result["session_id"]
+    ipc_file = staging_dir / f"session_{session_id}.arrow"
+    assert ipc_file.exists()
+
+    df = pl.read_ipc(ipc_file)
+    assert len(df) == 11
+    assert set(df.columns) == {"patient_id", "score"}
+
+
+def test_stage_query_k_anonymity_suppression(sample_csv_path: Path):
+    # Matches 2 rows in fixture (row 7: 26.4, row 13: 27.1)
+    sql = "SELECT patient_id, score FROM cohort WHERE icd10 = 'C34.9' AND bmi < 28"
+    result = stage_query(str(sample_csv_path), sql)
+
+    assert result["status"] == "staged_suppressed"
+    assert "session_id" in result
+    assert result["matched_records"] == "below_threshold"
+
+    # Arrow IPC must still be written for the clinician
+    staging_dir = get_settings().staging_dir
+    session_id = result["session_id"]
+    ipc_file = staging_dir / f"session_{session_id}.arrow"
+    assert ipc_file.exists()
+
+    df = pl.read_ipc(ipc_file)
+    assert len(df) == 2
+
+
+def test_stage_query_sanitizer_error(sample_csv_path: Path):
+    sql = "DROP TABLE cohort"
+    result = stage_query(str(sample_csv_path), sql)
+
+    assert result["status"] == "error"
+    assert result["error_type"] == "SQLSanitizerError"
+
+
+def test_stage_query_execution_error(sample_csv_path: Path):
+    sql = "SELECT nonexistent_column FROM cohort"
+    result = stage_query(str(sample_csv_path), sql)
+
+    assert result["status"] == "error"
+    assert result["error_type"] == "ExecutionError"