pgbee 0.1.0__py3-none-any.whl

pgbee/__init__.py

@@ -0,0 +1 @@
+"""pgbee: installer and reference worker for pgbee derived columns."""

pgbee/cli.py

@@ -0,0 +1,231 @@
+"""`pgbee` command line: install the extension, run the worker, inspect derived columns."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import signal
+import sys
+import time
+from pathlib import Path
+from typing import Annotated, Any
+
+import psycopg
+import structlog
+import typer
+from psycopg.rows import dict_row
+
+from pgbee import extension
+from pgbee.db import WORKER_BACKENDS, Contract
+from pgbee.installer import InstalledAsExtension, install
+from pgbee.providers import OpenAICompatibleProvider, OpenRouterProvider, Provider
+from pgbee.settings import Settings, load_settings
+from pgbee.worker import Worker, WorkerPool
+
+app = typer.Typer(no_args_is_help=True, add_completion=False)
+
+
+def _configure_logging(level: str) -> None:
+    numeric = logging.getLevelName(level.upper())
+    if not isinstance(numeric, int):
+        numeric = logging.INFO
+    structlog.configure(
+        wrapper_class=structlog.make_filtering_bound_logger(numeric),
+        processors=[
+            structlog.processors.TimeStamper(fmt="%H:%M:%S"),
+            structlog.processors.add_log_level,
+            structlog.dev.ConsoleRenderer(colors=sys.stderr.isatty()),
+        ],
+    )
+
+
+@app.command("install")
+def install_cmd() -> None:
+    """Apply the SQL files in sql/ that the database has not seen yet."""
+    settings = load_settings()
+    with psycopg.connect(settings.database_url) as conn:
+        try:
+            applied = install(conn)
+        except InstalledAsExtension as exc:
+            raise typer.BadParameter(str(exc)) from exc
+    if applied:
+        typer.echo(f"applied: {', '.join(str(v) for v in applied)}")
+    else:
+        typer.echo("up to date")
+
+
+@app.command("extension-files")
+def extension_files_cmd(
+    out_dir: Annotated[Path, typer.Argument(help="Where to write pgbee.control and the scripts.")],
+) -> None:
+    """Write the files for CREATE EXTENSION pgbee, to copy into `pg_config --sharedir`/extension."""
+    for path in extension.build(out_dir):
+        typer.echo(path)
+
+
+@app.command("status")
+def status_cmd() -> None:
+    """List derived columns with their version and queue counters."""
+    settings = load_settings()
+    with psycopg.connect(settings.database_url, row_factory=dict_row) as conn:
+        rows = conn.execute(
+            "SELECT c.table_schema, c.table_name, c.column_name, c.version, c.backend, c.model,"
+            " c.enabled, c.pending, c.claimed, c.done, c.dead, c.stale, c.human_overrides,"
+            " c.backfill_pending, c.backfill_scanned,"
+            " b.budget_usd, b.budget_period, b.spent_usd, b.exhausted"
+            " FROM bee.columns c JOIN bee.budgets b ON b.column_def_id = c.id"
+            " ORDER BY c.table_schema, c.table_name, c.column_name"
+        ).fetchall()
+    if not rows:
+        typer.echo("no derived columns")
+        return
+    for r in rows:
+        state = "on" if r["enabled"] else "off"
+        typer.echo(
+            f"{r['table_schema']}.{r['table_name']}.{r['column_name']}"
+            f"  v{r['version']} {r['backend']} {r['model']} [{state}]"
+            f"  pending={r['pending']} claimed={r['claimed']} done={r['done']} dead={r['dead']}"
+            f" stale={r['stale']} human={r['human_overrides']}"
+            f"{_backfill(r)}  {_spend(r)}"
+        )
+
+
+def _backfill(r: dict[str, Any]) -> str:
+    if not r["backfill_pending"]:
+        return ""
+    return f" backfilling (scanned {r['backfill_scanned']} rows)"
+
+
+def _spend(r: dict[str, Any]) -> str:
+    spent = f"${r['spent_usd']:.4f}"
+    if r["budget_usd"] is None:
+        return f"spent {spent} this {r['budget_period']}, no cap"
+    flag = " EXHAUSTED" if r["exhausted"] else ""
+    return f"spent {spent} of ${r['budget_usd']} per {r['budget_period']}{flag}"
+
+
+OPENAI_DEFAULT_URL = "https://api.openai.com/v1"
+
+
+def build_provider(settings: Settings) -> Provider:
+    try:
+        kind = settings.provider_kind()
+    except ValueError as exc:
+        raise typer.BadParameter(str(exc)) from exc
+    if kind == "openrouter":
+        assert settings.openrouter_api_key is not None
+        return OpenRouterProvider(
+            settings.openrouter_api_key,
+            settings.openrouter_base_url,
+            timeout=settings.request_timeout_seconds,
+        )
+    # Local servers (Ollama, vLLM) take no key, but the SDK wants a non empty one.
+    return OpenAICompatibleProvider(
+        settings.openai_api_key or "not-needed",
+        settings.openai_base_url or OPENAI_DEFAULT_URL,
+        timeout=settings.request_timeout_seconds,
+    )
+
+
+@app.command("run")
+def run_cmd(
+    once: bool = typer.Option(False, "--once", help="Process one batch and exit."),
+    drain: bool = typer.Option(
+        False,
+        "--drain",
+        help="Process until nothing is ready, then exit: for cron and serverless schedulers.",
+    ),
+    max_seconds: float | None = typer.Option(
+        None,
+        "--max-seconds",
+        min=1,
+        help="With --drain, stop claiming after this many seconds (the batch in flight finishes).",
+    ),
+    batch_size: int = typer.Option(20, "--batch-size", min=1),
+    backends: str = typer.Option(
+        "",
+        "--backends",
+        help="Comma separated backends this worker serves, one lane each. Default: all the"
+        f" provider supports ({','.join(WORKER_BACKENDS)} on OpenRouter).",
+    ),
+) -> None:
+    """Consume the queue: call the models and write the results back."""
+    if once and drain:
+        raise typer.BadParameter("--once and --drain are alternatives")
+    if max_seconds is not None and not drain:
+        raise typer.BadParameter("--max-seconds works with --drain")
+    settings = load_settings()
+    provider = build_provider(settings)
+    _configure_logging(settings.log_level)
+    chosen = [b.strip() for b in backends.split(",") if b.strip()] or list(provider.backends)
+    if set(chosen) - set(WORKER_BACKENDS):
+        raise typer.BadParameter(f"--backends takes a subset of {','.join(WORKER_BACKENDS)}")
+    unsupported = set(chosen) - set(provider.backends)
+    if unsupported:
+        raise typer.BadParameter(
+            f"{','.join(sorted(unsupported))} not available with {settings.provider_kind()}:"
+            " the decision backend needs OpenRouter"
+        )
+    asyncio.run(
+        _run(
+            settings,
+            provider,
+            once=once,
+            drain=drain,
+            max_seconds=max_seconds,
+            batch_size=batch_size,
+            backends=chosen,
+        )
+    )
+
+
+async def _run(
+    settings: Settings,
+    provider: Provider,
+    *,
+    once: bool,
+    drain: bool,
+    max_seconds: float | None,
+    batch_size: int,
+    backends: list[str],
+) -> None:
+    structlog.get_logger("pgbee").info("provider", kind=settings.provider_kind(), backends=backends)
+    options: dict[str, Any] = {
+        "batch_size": batch_size,
+        "poll_interval": settings.poll_interval_seconds,
+        "claim_timeout_seconds": settings.claim_timeout_seconds,
+        "maintenance_interval": settings.maintenance_interval_seconds,
+    }
+    if once:
+        contract = await Contract.connect(settings.database_url)
+        try:
+            worker = Worker(
+                contract, provider, worker_id=settings.worker_id, backends=backends, **options
+            )
+            stats = await worker.run_once()
+        finally:
+            await contract.close()
+        typer.echo(
+            f"claimed={stats.claimed} failed={stats.failed} "
+            + " ".join(f"{k}={v}" for k, v in stats.outcomes.items())
+        )
+        return
+    pool = WorkerPool(
+        lambda: Contract.connect(settings.database_url),
+        provider,
+        worker_id=settings.worker_id,
+        backends=backends,
+        **options,
+    )
+    loop = asyncio.get_running_loop()
+    for sig in (signal.SIGINT, signal.SIGTERM):
+        loop.add_signal_handler(sig, pool.stop)
+    if not drain:
+        await pool.run_forever()
+        return
+    deadline = None if max_seconds is None else time.monotonic() + max_seconds
+    for backend, stats in (await pool.drain(deadline)).items():
+        typer.echo(
+            f"{backend}: claimed={stats.claimed} failed={stats.failed} "
+            + " ".join(f"{k}={v}" for k, v in stats.outcomes.items())
+        )

pgbee/db.py

@@ -0,0 +1,114 @@
+"""The worker's side of the SQL contract: claim, complete, fail, reclaim, listen."""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Sequence
+from typing import Any
+
+import psycopg
+from psycopg.rows import DictRow, dict_row
+
+from pgbee.jobs import Job
+
+WORKER_BACKENDS = ["llm", "decision", "embedding"]
+
+
+class Contract:
+    def __init__(self, conn: psycopg.AsyncConnection[DictRow]) -> None:
+        self._conn = conn
+
+    @classmethod
+    async def connect(cls, database_url: str) -> Contract:
+        conn = await psycopg.AsyncConnection.connect(
+            database_url, autocommit=True, row_factory=dict_row
+        )
+        return cls(conn)
+
+    async def close(self) -> None:
+        await self._conn.close()
+
+    @property
+    def broken(self) -> bool:
+        """True when the connection is lost: the worker must stop, not fail jobs."""
+        return self._conn.broken or self._conn.closed
+
+    async def claim(
+        self, worker_id: str, batch_size: int, backends: Sequence[str] = WORKER_BACKENDS
+    ) -> list[Job]:
+        cur = await self._conn.execute(
+            "SELECT * FROM bee.claim_jobs(%s, %s, %s::bee.backend[])"
+            " ORDER BY column_def_id, job_id",
+            (worker_id, batch_size, list(backends)),
+        )
+        return [Job.from_row(row) for row in await cur.fetchall()]
+
+    async def complete(
+        self,
+        job: Job,
+        value: Any,
+        *,
+        confidence: float | None,
+        model: str,
+        usage: dict[str, Any],
+        latency_ms: int,
+        details: dict[str, Any] | None = None,
+    ) -> str:
+        cur = await self._conn.execute(
+            "SELECT bee.complete_job(%s::bigint, %s, %s::jsonb, %s::real, %s, %s::jsonb, %s,"
+            " %s::jsonb) AS outcome",
+            (
+                job.job_id,
+                job.source_hash,
+                json.dumps(value),
+                confidence,
+                model,
+                json.dumps(usage),
+                latency_ms,
+                None if details is None else json.dumps(details),
+            ),
+        )
+        row = await cur.fetchone()
+        assert row is not None
+        return str(row["outcome"])
+
+    async def fail(self, job: Job, error: str, *, retryable: bool) -> str | None:
+        cur = await self._conn.execute(
+            "SELECT bee.fail_job(%s::bigint, %s, %s) AS status",
+            (job.job_id, error[:2000], retryable),
+        )
+        row = await cur.fetchone()
+        assert row is not None
+        return None if row["status"] is None else str(row["status"])
+
+    async def reclaim_stale(self, timeout_seconds: int) -> int:
+        cur = await self._conn.execute(
+            "SELECT bee.reclaim_stale(make_interval(secs => %s)) AS n", (timeout_seconds,)
+        )
+        row = await cur.fetchone()
+        assert row is not None
+        return int(row["n"])
+
+    async def prune(self, limit: int) -> tuple[int, int]:
+        """One maintenance batch: (results, jobs) deleted, each at most `limit`."""
+        cur = await self._conn.execute("SELECT results, jobs FROM bee.prune(%s)", (limit,))
+        row = await cur.fetchone()
+        assert row is not None
+        return int(row["results"]), int(row["jobs"])
+
+    async def listen(self) -> None:
+        await self._conn.execute("LISTEN bee_jobs")
+
+    async def wait_for_notify(self, max_wait: float) -> bool:
+        """Block until a NOTIFY on bee_jobs arrives or max_wait seconds pass. True on notify.
+
+        Notifications received while the worker was busy are queued by psycopg; they are all
+        drained here, since one claim serves them all.
+        """
+        received = False
+        async for _ in self._conn.notifies(timeout=max_wait, stop_after=1):
+            received = True
+        if received:
+            async for _ in self._conn.notifies(timeout=0):
+                pass
+        return received

pgbee/extension.py

@@ -0,0 +1,87 @@
+"""Build the files for `CREATE EXTENSION pgbee` from the same versioned SQL files the installer
+applies, so the two ways of installing cannot drift apart.
+
+Version 0.N is file NNNN: `pgbee--0.1.sql` is file 0001 and every later file becomes an update
+script `pgbee--0.(N-1)--0.N.sql`. Postgres chains them, both on a fresh CREATE EXTENSION and on
+ALTER EXTENSION pgbee UPDATE.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from pgbee.installer import SqlFile, sql_files
+
+NAME = "pgbee"
+COMMENT = "Derived columns computed by models: queue, versions, lineage, human overrides, budgets"
+
+GUARD = (
+    '\\echo Use "CREATE EXTENSION pgbee" or "ALTER EXTENSION pgbee UPDATE" to load this file.'
+    " \\quit\n"
+)
+
+# Tables and sequences of the extension hold user data (definitions, jobs, lineage, spend):
+# pg_dump must dump their rows. schema_version is left out because the scripts fill it.
+CONFIG_DUMP = """
+DO $pgbee$
+DECLARE
+  r record;
+BEGIN
+  FOR r IN
+    SELECT c.oid::regclass AS rel
+    FROM pg_catalog.pg_class c
+    JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
+    JOIN pg_catalog.pg_depend d ON d.classid = 'pg_catalog.pg_class'::regclass AND d.objid = c.oid
+                               AND d.deptype = 'e'
+    JOIN pg_catalog.pg_extension e ON e.oid = d.refobjid AND e.extname = 'pgbee'
+    WHERE n.nspname = 'bee' AND c.relkind IN ('r', 'S')
+      AND c.relname <> 'schema_version'
+      AND NOT (c.oid = ANY (coalesce(e.extconfig, '{}')))
+  LOOP
+    PERFORM pg_catalog.pg_extension_config_dump(r.rel, '');
+  END LOOP;
+END $pgbee$;
+"""
+
+
+def version_of(sql_file: SqlFile) -> str:
+    return f"0.{sql_file.version}"
+
+
+def script(sql_file: SqlFile) -> str:
+    body = sql_file.path.read_text(encoding="utf-8")
+    return (
+        f"{GUARD}\n{body.rstrip()}\n{CONFIG_DUMP}\n"
+        f"INSERT INTO bee.schema_version (version) VALUES ({sql_file.version});\n"
+    )
+
+
+def control(default_version: str) -> str:
+    return (
+        f"comment = '{COMMENT}'\n"
+        f"default_version = '{default_version}'\n"
+        "relocatable = false\n"
+        "superuser = true\n"
+    )
+
+
+def build(out_dir: Path, directory: Path | None = None) -> list[Path]:
+    """Write the control file and one script per SQL file into out_dir. Returns the paths."""
+    files = sql_files(directory)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    written: list[Path] = []
+    previous: SqlFile | None = None
+    for sql_file in files:
+        if previous is None:
+            name = f"{NAME}--{version_of(sql_file)}.sql"
+        else:
+            name = f"{NAME}--{version_of(previous)}--{version_of(sql_file)}.sql"
+        path = out_dir / name
+        path.write_text(script(sql_file), encoding="utf-8")
+        written.append(path)
+        previous = sql_file
+    assert previous is not None
+    control_path = out_dir / f"{NAME}.control"
+    control_path.write_text(control(version_of(previous)), encoding="utf-8")
+    written.append(control_path)
+    return written

pgbee/installer.py

@@ -0,0 +1,84 @@
+"""Apply the versioned SQL files in `sql/` to a database and track them in bee.schema_version."""
+
+from __future__ import annotations
+
+import os
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import psycopg
+from psycopg.rows import tuple_row
+
+SQL_FILE_RE = re.compile(r"^(\d{4})_[a-z0-9_]+\.sql$")
+
+
+class InstalledAsExtension(RuntimeError):
+    """The database has pgbee as a Postgres extension; the file installer must not touch it."""
+
+
+@dataclass(frozen=True)
+class SqlFile:
+    version: int
+    path: Path
+
+
+def sql_dir() -> Path:
+    """Directory holding the extension SQL files, shipped inside the package.
+
+    Overridable with PGBEE_SQL_DIR. The repository's top level `sql` links here.
+    """
+    env = os.environ.get("PGBEE_SQL_DIR")
+    if env:
+        return Path(env)
+    return Path(__file__).resolve().parent / "sql"
+
+
+def sql_files(directory: Path | None = None) -> list[SqlFile]:
+    directory = directory or sql_dir()
+    files: list[SqlFile] = []
+    for path in sorted(directory.glob("*.sql")):
+        match = SQL_FILE_RE.match(path.name)
+        if not match:
+            raise ValueError(f"unexpected SQL file name {path.name}: want NNNN_name.sql")
+        files.append(SqlFile(version=int(match.group(1)), path=path))
+    if not files:
+        raise FileNotFoundError(f"no SQL files in {directory}")
+    return files
+
+
+def applied_versions(conn: psycopg.Connection[Any]) -> set[int]:
+    exists = conn.execute(
+        "SELECT 1 FROM pg_tables WHERE schemaname = 'bee' AND tablename = 'schema_version'"
+    ).fetchone()
+    if not exists:
+        return set()
+    cur = conn.cursor(row_factory=tuple_row)
+    rows = cur.execute("SELECT version FROM bee.schema_version").fetchall()
+    return {int(row[0]) for row in rows}
+
+
+def install(conn: psycopg.Connection[Any], directory: Path | None = None) -> list[int]:
+    """Apply every SQL file not yet recorded, one transaction each. Returns applied versions.
+
+    Refuses a database where pgbee was installed as an extension: there the files arrive through
+    ALTER EXTENSION pgbee UPDATE, and applying them here would detach objects from it.
+    """
+    if conn.execute("SELECT 1 FROM pg_extension WHERE extname = 'pgbee'").fetchone():
+        raise InstalledAsExtension(
+            "pgbee is installed as the extension pgbee:"
+            " upgrade it with ALTER EXTENSION pgbee UPDATE"
+        )
+    applied = applied_versions(conn)
+    done: list[int] = []
+    for sql_file in sql_files(directory):
+        if sql_file.version in applied:
+            continue
+        with conn.transaction():
+            conn.execute(sql_file.path.read_text(encoding="utf-8"))
+            conn.execute(
+                "INSERT INTO bee.schema_version (version) VALUES (%s)", (sql_file.version,)
+            )
+        done.append(sql_file.version)
+    return done

pgbee/jobs.py

@@ -0,0 +1,58 @@
+"""Job payloads as returned by bee.claim_jobs, and the text the models see."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass(frozen=True)
+class Job:
+    job_id: int
+    column_def_id: int
+    column_version_id: int
+    backend: str
+    model: str
+    prompt: str | None
+    output_type: str
+    output_schema: Any
+    backend_config: dict[str, Any]
+    config: dict[str, Any]
+    row_pk: dict[str, Any]
+    source_hash: bytes
+    source: dict[str, Any]
+    attempts: int
+
+    @classmethod
+    def from_row(cls, row: dict[str, Any]) -> Job:
+        return cls(
+            job_id=int(row["job_id"]),
+            column_def_id=int(row["column_def_id"]),
+            column_version_id=int(row["column_version_id"]),
+            backend=str(row["backend"]),
+            model=str(row["model"]),
+            prompt=row["prompt"],
+            output_type=str(row["output_type"]),
+            output_schema=row["output_schema"],
+            backend_config=dict(row["backend_config"] or {}),
+            config=dict(row["config"] or {}),
+            row_pk=dict(row["row_pk"]),
+            source_hash=bytes(row["source_hash"]),
+            source=dict(row["source"]),
+            attempts=int(row["attempts"]),
+        )
+
+    def source_text(self) -> str:
+        """Render the sources for a model: raw value for one column, labeled block otherwise."""
+        if len(self.source) == 1:
+            return _as_text(next(iter(self.source.values())))
+        return "\n".join(f"{name}: {_as_text(value)}" for name, value in self.source.items())
+
+
+def _as_text(value: Any) -> str:
+    if value is None:
+        return ""
+    if isinstance(value, str):
+        return value
+    return json.dumps(value, ensure_ascii=False)