lemma-db 0.1.0__py3-none-any.whl

db_extension/AGENT.md

@@ -0,0 +1,125 @@
+# Lemma OpenRouter agent
+
+Schema-driven **Verus** optimizer (not SSB-only). The host transpiles SQL with a caller-supplied
+or catalog-resolved schema via `verus_transpiler`; agents edit `runquery_agent.rs` against
+`/context/ro/spec.rs`. Legacy Dafny paths live under `research_loop/dafny_legacy/` and are not
+used by the default optimizer.
+
+## Schema
+
+Pass an explicit schema dict to `run_optimization_loop(..., schema={...})`, set
+`LEMMA_SCHEMA_JSON=/path/to/schema.json`, or use `DatabaseCatalog` against the DuckDB
+file that holds the workload tables (`LEMMA_DUCKDB_PATH`). Unknown tables fail loudly —
+there is no hardcoded SSB/TPCH schema invent path.
+
+## Architecture
+
+```
+Host
+  ├─ OpenRouter ReAct (LLM on host)     [default]
+  ├─ EgressBridge allowlist UDS         [CLI-in-Docker: Cursor/Anthropic/…]
+  ├─ McpSocketServer UDS → measure_core
+  └─ tool/CLI container (--network none)
+        ├─ tools_worker / agent CLI
+        ├─ mcp_proxy (stdio MCP → MCP sock)
+        └─ egress sidecar HTTPS_PROXY → only vendor API hosts
+```
+
+**CLI:** mount `~/.cursor` (or `AGENT_CREDENTIALS_DIR`), pass `AGENT_ENV` keys, egress
+allowlist by profile (`cursor`, `anthropic`, …). Non-API web is blocked by the proxy.
+With egress on, Cursor also gets `CURSOR_FORCED_SHELL_EGRESS=1` so **WebSearch/WebFetch**
+are disabled automatically (they run in Cursor cloud via `api2`, not as local curls).
+See `research_loop/AGENT_SANDBOX.md`.
+
+**Real MCP:** `mcp_tool_specs.py` → host FastMCP, sandbox `mcp_proxy`, OpenRouter defs.
+
+### Tools
+
+| Tool | Action |
+|------|--------|
+| `validate_runquery` | Body lint |
+| `run_runquery` | Host harness; `run_id` + metrics |
+| `submit_runquery` / `submit` | Mark only |
+| `get_submit_result` / `list_runs` / `mcp_health` | Read back |
+
+### Register MCP (CLI on host)
+
+```json
+{
+  "mcpServers": {
+    "lemma-host": {
+      "command": "uv",
+      "args": ["run", "python", "-m", "db_extension.agent.mcp_host", "--transport", "stdio"],
+      "env": {
+        "LEMMA_AGENT_WORKSPACE": "/ABS/PATH/research_loop/agent_workspace",
+        "LEMMA_QUERY_ID": "1"
+      }
+    }
+  }
+}
+```
+
+Inside Docker CLI mode, entrypoint writes `.cursor/mcp.json` → `python -m lemma_agent.mcp_proxy`.
+
+## CLI-in-Docker (Cursor / other CLIs)
+
+`--network none` + **allowlisted egress** (not open internet):
+
+- Mount `AGENT_CREDENTIALS_DIR` (default `~/.cursor`) → `/root/.cursor`
+- Pass `AGENT_ENV` keys (e.g. `CURSOR_API_KEY`)
+- Host `EgressBridge` allowlist by `AGENT_EGRESS_PROFILE` (`cursor`, `anthropic`, …)
+- Sandbox `HTTPS_PROXY=http://127.0.0.1:8118` → only those API hosts
+- MCP measure via Unix socket / `mcp_proxy`
+
+See `research_loop/AGENT_SANDBOX.md`. Rebuild: `docker build -t lemma-agent:latest -f docker/agent/Dockerfile .`
+
+## Run
+
+```bash
+docker build -t lemma-agent:latest -f docker/agent/Dockerfile .
+MOCK_AGENT=0 OPENROUTER_API_KEY=sk-or-... uv run python -m db_extension.run_optimizer "SELECT ..."
+```
+
+Custom schema example:
+
+```bash
+export LEMMA_SCHEMA_JSON=/path/to/schema.json   # {"V":"int"} or multi-table JSON
+uv run python -c "
+from db_extension.optimizer import run_optimization_loop
+print(run_optimization_loop('SELECT SUM(v) FROM t', schema={'V':'int'}, use_mock=True))
+"
+```
+
+## Tests
+
+```bash
+uv run python -m pytest db_extension/tests/test_measure_core.py db_extension/tests/test_mcp_socket.py db_extension/tests/test_mcp_host.py db_extension/tests/test_mcp_registry_sync.py db_extension/tests/test_model_bridge.py db_extension/tests/test_tools_worker_paths.py -q
+```
+
+## Run artifacts (GCP harvest)
+
+Each optimizer job **with ``LEMMA_RESEARCH_LOG=1``** writes one harvest directory under
+`research_loop/runs/<id>/`:
+
+```
+manifest.json
+result.json
+history.json
+workspace/          # agent files, mcp_results/, egress_*.jsonl, trace.jsonl
+logs/pipeline.log
+logs/pipeline.jsonl
+logs/docker_agent.*
+logs/harness_iter*.json
+logs/openrouter_meta_iter*.json
+```
+
+`research_loop/runs/LATEST` points at the most recent run (absolute path). Set automatically via
+`LEMMA_RUN_DIR` and `LEMMA_AGENT_WORKSPACE` when research logging is on.
+
+```bash
+LEMMA_RESEARCH_LOG=1  # required for harvest dirs (default 0)
+```
+
+Archive the whole directory for GCP: `tar czf run.tgz -C research_loop/runs <id>/` or `gsutil cp -r research_loop/runs/<id>/ gs://bucket/path/`.
+
+For Spot/preemptible VMs, also set `LEMMA_EXPERIMENT_EVENT_URL` to a durable receiver (`research_loop/scripts/experiment_event_receiver.py`; see `holdout/EXPERIMENTS_GENDB.md`) so `query_end` summaries stream off-box as each query finishes.

db_extension/DEMO.md

@@ -0,0 +1,59 @@
+# Lemma interactive demo
+
+## Live demo (real agent)
+
+```bash
+chmod +x scripts/demo.sh
+./scripts/demo.sh
+```
+
+Right pane (split terminal):
+
+```bash
+./scripts/demo_view/follow-agent-log.sh
+```
+
+This script:
+
+1. Clears **all** cached optimized queries
+2. Sets `LEMMA_DEMO=1`, `MOCK_AGENT=0`, `LEMMA_DATASET_SIZE=100000`
+3. Drops you into the **official DuckDB CLI** (downloaded to `build/duckdb` on first run)
+
+Requires the Cursor \`agent\` CLI on PATH (same auth as running \`agent\` in your shell).
+
+## Mock demo (offline / no API key)
+
+```bash
+./scripts/mockdemo.sh
+```
+
+Same UX, but seeds a hardcoded RunQuery body (`MOCK_AGENT=1`, `LEMMA_DATASET_SIZE=2000000`, no LLM).
+
+Override query: `DEMO_QUERY_ID=5 ./scripts/demo.sh`
+
+## In the DuckDB shell
+
+```sql
+-- (1) Vanilla (DuckDB prints its own timing after the result)
+SELECT SUM(...) FROM lineorder_flat WHERE ...;
+
+-- (2) Lemma (demo steps on stdout, then result)
+SELECT lemma('SELECT SUM(...) FROM lineorder_flat WHERE ...');
+
+-- (3) Run (2) again → cached 💾 path
+```
+
+## Env toggles
+
+| Variable | `demo.sh` | `mockdemo.sh` |
+|----------|-----------|---------------|
+| `LEMMA_DEMO` | on | on |
+| `MOCK_AGENT` | **0** (real agent) | **1** (fixture) |
+| `LEMMA_DATASET_SIZE` | 100,000 | 2,000,000 |
+| `LEMMA_LOG_LEVEL` | WARN | WARN |
+
+Production `config.env` keeps demo **off**.
+
+## Re-run fresh pipeline
+
+Run the script again — it clears cache every time.

db_extension/__init__.py

@@ -0,0 +1,3 @@
+from .catalog import DatabaseCatalog
+
+__all__ = ["DatabaseCatalog"]

db_extension/agent/__init__.py

@@ -0,0 +1,41 @@
+"""Sandboxed OpenRouter agent for Lemma RunQuery optimization."""
+from __future__ import annotations
+
+# ``config`` imports the host's ``lemma_config``, which the container image does not have (this package is mounted into it): the three
+# config names load on first use, so ``import lemma_agent.mcp_proxy`` works in the container.
+_CONFIG_NAMES = ("AgentFlags", "load_agent_flags", "truthy")
+
+__all__ = [
+    "AgentFlags",
+    "extract_marked_body",
+    "load_agent_flags",
+    "run_openrouter_agent_iteration",
+    "truthy",
+    "wrap_body_with_markers",
+]
+
+
+def __getattr__(name: str):
+    if name in _CONFIG_NAMES:
+        from . import config
+
+        return getattr(config, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def extract_marked_body(*args, **kwargs):
+    from .extract import extract_marked_body as _fn
+
+    return _fn(*args, **kwargs)
+
+
+def wrap_body_with_markers(*args, **kwargs):
+    from .extract import wrap_body_with_markers as _fn
+
+    return _fn(*args, **kwargs)
+
+
+def run_openrouter_agent_iteration(*args, **kwargs):
+    from .harness import run_openrouter_agent_iteration as _run
+
+    return _run(*args, **kwargs)

db_extension/agent/claude_stream.py

@@ -0,0 +1,157 @@
+"""Convert Claude Code ``--output-format stream-json`` into the Cursor-style transcript.
+
+Runs inside the agent container as the tail of the agent command:
+``claude -p ... | python -m lemma_agent.claude_stream --raw <file>``.
+stdout becomes ``agent_stream.jsonl`` (the entrypoint tees it), so the trace-reading
+rules (thinking deltas, tool_call started/completed) work for both agents.
+The untouched Claude lines go to ``--raw``.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+import time
+from typing import IO, Iterable, Iterator
+
+_TOOL_KEYS = {
+    "Read": "readToolCall",
+    "Edit": "editToolCall",
+    "MultiEdit": "editToolCall",
+    "Write": "writeToolCall",
+    "Glob": "globToolCall",
+    "Grep": "grepToolCall",
+    "Bash": "shellToolCall",
+}
+
+
+def _tool_call(name: str, tool_input: dict) -> tuple[str, dict]:
+    """Return (cursor tool_call key, args) for one Claude ``tool_use`` block."""
+    if name.startswith("mcp__"):
+        _, server, tool = name.split("__", 2)
+        return "mcpToolCall", {"server": server, "name": tool, "args": tool_input}
+    key = _TOOL_KEYS.get(name, f"{name[0].lower()}{name[1:]}ToolCall")
+    args = dict(tool_input)
+    if "file_path" in args:
+        args["path"] = args.pop("file_path")
+    if key == "globToolCall" and "pattern" in args:
+        args["globPattern"] = args.pop("pattern")
+    return key, args
+
+
+def _result_text(content) -> str:
+    if isinstance(content, str):
+        return content
+    return "".join(b.get("text", "") for b in content if b.get("type") == "text")
+
+
+def convert_events(events: Iterable[dict]) -> Iterator[dict]:
+    """Yield Cursor-style events for Claude Code stream-json events."""
+    session_id = ""
+    pending: dict[str, tuple[str, dict]] = {}
+    for ev in events:
+        now_ms = int(time.time() * 1000)
+        kind = ev.get("type")
+        session_id = ev.get("session_id", session_id)
+        if kind == "system":
+            if ev.get("subtype") == "init":
+                yield {
+                    "type": "system",
+                    "subtype": "init",
+                    "apiKeySource": ev.get("apiKeySource"),
+                    "cwd": ev.get("cwd"),
+                    "session_id": session_id,
+                    "model": ev.get("model"),
+                    "permissionMode": ev.get("permissionMode"),
+                    "mcp_servers": ev.get("mcp_servers"),
+                }
+        elif kind == "assistant":
+            for block in ev["message"]["content"]:
+                btype = block["type"]
+                if btype == "thinking":
+                    yield {
+                        "type": "thinking",
+                        "subtype": "delta",
+                        "text": block["thinking"],
+                        "session_id": session_id,
+                        "timestamp_ms": now_ms,
+                    }
+                    yield {
+                        "type": "thinking",
+                        "subtype": "completed",
+                        "session_id": session_id,
+                        "timestamp_ms": now_ms,
+                    }
+                elif btype == "text":
+                    yield {
+                        "type": "assistant",
+                        "message": {"role": "assistant", "content": [{"type": "text", "text": block["text"]}]},
+                        "session_id": session_id,
+                        "timestamp_ms": now_ms,
+                    }
+                elif btype == "tool_use":
+                    key, args = _tool_call(block["name"], block["input"])
+                    pending[block["id"]] = (key, args)
+                    yield {
+                        "type": "tool_call",
+                        "subtype": "started",
+                        "call_id": block["id"],
+                        "tool_call": {key: {"args": args}},
+                        "session_id": session_id,
+                        "timestamp_ms": now_ms,
+                    }
+        elif kind == "user":
+            content = ev["message"]["content"]
+            if isinstance(content, str):
+                continue
+            for block in content:
+                if block.get("type") != "tool_result":
+                    continue
+                key, args = pending.pop(block["tool_use_id"])
+                text = _result_text(block["content"])
+                outcome = {"error": {"message": text}} if block.get("is_error") else {"success": {"content": text}}
+                yield {
+                    "type": "tool_call",
+                    "subtype": "completed",
+                    "call_id": block["tool_use_id"],
+                    "tool_call": {key: {"args": args, "result": outcome}},
+                    "session_id": session_id,
+                    "timestamp_ms": now_ms,
+                }
+        elif kind == "result":
+            yield {
+                "type": "result",
+                "subtype": ev["subtype"],
+                "is_error": ev.get("is_error"),
+                "result": ev.get("result", ""),
+                "duration_ms": ev.get("duration_ms"),
+                "num_turns": ev.get("num_turns"),
+                "total_cost_usd": ev.get("total_cost_usd"),
+                "session_id": session_id,
+            }
+
+
+def convert_stream(src: IO[str], dst: IO[str], raw: IO[str]) -> None:
+    def parsed() -> Iterator[dict]:
+        for line in src:
+            raw.write(line)
+            raw.flush()
+            if line.strip():
+                yield json.loads(line)
+
+    for out in convert_events(parsed()):
+        dst.write(json.dumps(out) + "\n")
+        dst.flush()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--raw", required=True)
+    args = parser.parse_args()
+    with open(args.raw, "a", encoding="utf-8") as raw:
+        convert_stream(sys.stdin, sys.stdout, raw)
+
+
+if __name__ == "__main__":
+    main()

db_extension/agent/config.py

@@ -0,0 +1,84 @@
+"""Agent configuration: run settings from lemma.toml (lemma_config), the imperative/OpenRouter ones from research_loop/config.env and process env."""
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from pathlib import Path
+
+import lemma_config
+
+ROOT = Path(__file__).resolve().parents[2]
+DEFAULT_CONFIG_ENV = ROOT / "research_loop" / "config.env"
+
+
+def truthy(s: str | None) -> bool:
+    if s is None:
+        return False
+    return s.strip().lower() not in ("", "0", "false", "no", "off")
+
+
+def _load_config_file(path: Path) -> dict[str, str]:
+    out: dict[str, str] = {}
+    if not path.is_file():
+        return out
+    with open(path) as f:
+        for line in f:
+            line = line.strip()
+            if not line or line.startswith("#") or "=" not in line:
+                continue
+            k, v = line.split("=", 1)
+            out[k.strip()] = v.strip()
+    return out
+
+
+@dataclass(frozen=True)
+class AgentFlags:
+    backend: str = "openrouter"
+    openrouter_api_key: str = ""
+    openrouter_model: str = "anthropic/claude-sonnet-4"
+    openrouter_base_url: str = "https://openrouter.ai/api/v1"
+    agent_timeout_sec: int = 600
+    agent_network: bool = False
+    agent_data_mode: str = "stats"
+    agent_workload_hint: bool = True
+    agent_submit_ends_session: bool = False
+    agent_image: str = "lemma-agent:latest"
+    agent_max_turns: int = 40
+
+    @classmethod
+    def from_mapping(cls, cfg: dict[str, str]) -> AgentFlags:
+        return cls(
+            backend=cfg.get("LEMMA_AGENT_BACKEND", "openrouter"),
+            openrouter_api_key=cfg.get("OPENROUTER_API_KEY", ""),
+            openrouter_model=cfg.get("OPENROUTER_MODEL", "anthropic/claude-sonnet-4"),
+            openrouter_base_url=cfg.get("OPENROUTER_BASE_URL", "https://openrouter.ai/api/v1"),
+            agent_timeout_sec=int(cfg.get("AGENT_TIMEOUT_SEC", "600")),
+            agent_network=truthy(cfg.get("AGENT_NETWORK", "0")),
+            agent_data_mode=cfg.get("AGENT_DATA_MODE", "stats"),
+            agent_workload_hint=truthy(cfg.get("AGENT_WORKLOAD_HINT", "1")),
+            agent_submit_ends_session=truthy(cfg.get("AGENT_SUBMIT_ENDS_SESSION", "0")),
+            agent_image=cfg.get("AGENT_IMAGE", "lemma-agent:latest"),
+            agent_max_turns=int(cfg.get("AGENT_MAX_TURNS", "40")),
+        )
+
+
+def load_agent_flags(config_env: Path | None = None) -> AgentFlags:
+    """Load flags from config.env; process environment overrides file values."""
+    cfg = _load_config_file(config_env or DEFAULT_CONFIG_ENV)
+    keys = (
+        "LEMMA_AGENT_BACKEND",
+        "OPENROUTER_API_KEY",
+        "OPENROUTER_MODEL",
+        "OPENROUTER_BASE_URL",
+        "AGENT_NETWORK",
+        "AGENT_WORKLOAD_HINT",
+        "AGENT_IMAGE",
+        "AGENT_MAX_TURNS",
+    )
+    for key in keys:
+        if key in os.environ:
+            cfg[key] = os.environ[key]
+    cfg["AGENT_TIMEOUT_SEC"] = str(lemma_config.get("agent_timeout_sec"))
+    cfg["AGENT_DATA_MODE"] = lemma_config.get("agent_data_mode")
+    cfg["AGENT_SUBMIT_ENDS_SESSION"] = "1" if lemma_config.get("agent_submit_ends_session") else "0"
+    return AgentFlags.from_mapping(cfg)

db_extension/agent/docker_runner.py

@@ -0,0 +1,130 @@
+"""Docker tool-container runner for the OpenRouter agent."""
+from __future__ import annotations
+
+import json
+import subprocess
+import threading
+import uuid
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from db_extension.agent.config import AgentFlags
+from research_loop.agent_image_check import agent_code_mount_args
+from research_loop.job_label import job_label_args
+from research_loop.sandbox_hide import shadow_mount_args
+
+ROOT = Path(__file__).resolve().parents[2]
+DOCKERFILE = ROOT / "docker" / "agent" / "Dockerfile"
+
+
+@dataclass
+class ContainerSession:
+    proc: subprocess.Popen[str]
+    _lock: threading.Lock = field(default_factory=threading.Lock)
+    _counter: int = 0
+
+    def call_tool(self, name: str, args: dict | None = None) -> dict:
+        with self._lock:
+            self._counter += 1
+            req_id = self._counter
+        req = {"id": req_id, "tool": name, "args": args or {}}
+        assert self.proc.stdin is not None
+        assert self.proc.stdout is not None
+        self.proc.stdin.write(json.dumps(req) + "\n")
+        self.proc.stdin.flush()
+        line = self.proc.stdout.readline()
+        if not line:
+            raise RuntimeError("tool container closed stdout")
+        return json.loads(line)
+
+    def close(self) -> None:
+        if self.proc.poll() is None:
+            self.proc.terminate()
+            try:
+                self.proc.wait(timeout=5)
+            except subprocess.TimeoutExpired:
+                self.proc.kill()
+                self.proc.wait()
+
+
+def docker_image_built(image: str) -> bool:
+    return subprocess.run(
+        ["docker", "image", "inspect", image],
+        capture_output=True,
+    ).returncode == 0
+
+
+def ensure_image(image: str) -> None:
+    if docker_image_built(image):
+        return
+    if not DOCKERFILE.is_file():
+        raise FileNotFoundError(f"Dockerfile not found: {DOCKERFILE}")
+    subprocess.run(
+        ["docker", "build", "-t", image, "-f", str(DOCKERFILE), str(ROOT)],
+        check=True,
+    )
+
+
+def start_tool_container(
+    workspace: Path,
+    context_ro: Path,
+    data_dir: Path | None,
+    flags: AgentFlags,
+    *,
+    data_file_name: str | None = None,
+    mcp_sock: Path | None = None,
+    query_id: int | None = None,
+) -> ContainerSession:
+    if mcp_sock is None:
+        raise RuntimeError(
+            "Host MCP Unix socket path is required for tool container measure/submit tools. "
+            "Start McpSocketServer and pass mcp_sock= to start_tool_container."
+        )
+    ensure_image(flags.agent_image)
+    workspace = workspace.resolve()
+    context_ro = context_ro.resolve()
+    name = f"lemma-agent-{uuid.uuid4().hex[:12]}"
+    cmd = [
+        "docker",
+        "run",
+        "--rm",
+        "-i",
+        "--name",
+        name,
+        *job_label_args(),
+        "--cap-drop",
+        "ALL",
+        "-v",
+        f"{workspace}:/workspace:rw",
+        *shadow_mount_args(workspace),
+        "-v",
+        f"{context_ro}:/context/ro:ro",
+        *agent_code_mount_args(),  # the commit's agent code over the image's baked copy
+        "-e",
+        f"AGENT_DATA_MODE={flags.agent_data_mode}",
+        "-e",
+        "PYTHONPATH=/app",
+        "-w",
+        "/workspace",
+    ]
+    if not flags.agent_network:
+        cmd.extend(["--network", "none"])
+    if data_dir is not None and data_dir.is_dir():
+        cmd.extend(["-v", f"{data_dir.resolve()}:/data:ro"])
+    if data_file_name:
+        cmd.extend(["-e", f"AGENT_DATA_FILE={data_file_name}"])
+    sock = mcp_sock.resolve()
+    cmd.extend(["-v", f"{sock}:/lemma-mcp.sock"])
+    cmd.extend(["-e", "LEMMA_MCP_SOCK=/lemma-mcp.sock"])
+    if query_id is not None:
+        cmd.extend(["-e", f"LEMMA_QUERY_ID={query_id}"])
+    cmd.append(flags.agent_image)
+    proc = subprocess.Popen(
+        cmd,
+        stdin=subprocess.PIPE,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+        bufsize=1,
+    )
+    return ContainerSession(proc=proc)