scripts/backfill_cve_completions.py

#!/usr/bin/env python3
"""scripts/backfill_cve_completions.py — generate teacher answers for prompted CVE rows.

Sister script to backfill_cve_prompts.py. That one generated the
`payload.prompt` (the question side of a training pair); this one
generates the `payload.completion` (the answer side).

Why a separate script
---------------------
Splitting prompt-generation from completion-generation lets each pass
run independently, restart cleanly on failure, and use a different
teacher per side if quality demands. For Stage-1 cybersec adapter
training the same Mistral mistral-medium-latest tier handles both — it
has the depth to write a senior-engineer answer and the Experiment
tier is free.

Selection criteria
------------------
  WHERE kind='cve' AND domain='cybersecurity'
    AND payload ? 'prompt'
    AND NOT (payload ? 'completion')
  ORDER BY KEV-flagged first, then CVSS severity, then recency.

KEV-first ordering matters because those are the CVEs adversaries are
ACTIVELY exploiting in the wild — the highest-signal training data we
have. If the run is interrupted partway, we still get the best rows.

Idempotent: re-run after Ctrl-C; the `NOT (payload ? 'completion')`
predicate is the resume cursor.

Throughput
----------
Mistral Experiment tier: 23 RPM account-wide. We pace at 20.
494 rows / 20 RPM = ~25 minutes for a full backfill.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path

try:
    from dotenv import load_dotenv

    load_dotenv(Path(__file__).resolve().parent.parent / ".env")
except ImportError:
    pass

import psycopg
from psycopg import rows as psycopg_rows

MISTRAL_ENDPOINT = "https://api.mistral.ai/v1/chat/completions"
DEFAULT_MODEL = "mistral-medium-latest"

RATE_LIMIT_RPM = 20
RATE_INTERVAL_S = 60.0 / RATE_LIMIT_RPM

# This system prompt mirrors the role we want the trained adapter to
# embody: a senior security engineer who explains root cause,
# exploitation, mitigation, and detection in concrete technical terms.
# Every completion becomes one (prompt, completion) training pair, so
# the answer style here directly shapes the adapter's voice.
SYSTEM_PROMPT = (
    "You are a senior offensive-and-defensive cybersecurity engineer "
    "answering for a peer-engineer audience. Given a vulnerability "
    "training prompt, write a concrete, technically rigorous answer "
    "that covers (1) root cause, (2) realistic exploitation pattern, "
    "(3) concrete mitigation/patch guidance, and (4) detection signals "
    "(log fields, EDR signatures, network markers). No marketing fluff, "
    "no generic 'apply security best practices' platitudes. Cite "
    "specific configuration keys, function names, or CWE IDs when "
    "relevant. 4-8 sentences total, dense and implementation-ready. "
    "Output the answer body only — no preface, no markdown headings, "
    "no JSON, no fences."
)


def strip_markdown_fences(s: str) -> str:
    """Some models wrap output in ```…``` even when asked not to."""
    s = s.strip()
    if not s.startswith("```"):
        return s
    parts = s.split("```")
    if len(parts) >= 3:
        inner = parts[1]
        if "\n" in inner:
            first, rest = inner.split("\n", 1)
            if not first.strip() or first.strip().isalpha():
                return rest.strip()
        return inner.strip()
    return s


def call_mistral(
    api_key: str,
    model: str,
    user_prompt: str,
    timeout_s: int = 90,
) -> tuple[str | None, str | None]:
    body = json.dumps(
        {
            "model": model,
            "messages": [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt},
            ],
            "max_tokens": 900,
            "temperature": 0.4,
        }
    ).encode("utf-8")
    req = urllib.request.Request(
        MISTRAL_ENDPOINT,
        data=body,
        method="POST",
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        },
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout_s) as resp:
            data = json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as e:
        if e.code == 429:
            return None, "429"
        msg = ""
        try:
            msg = e.read().decode("utf-8")[:200]
        except Exception:
            msg = ""
        print(f"  ! HTTP {e.code}: {msg}", file=sys.stderr)
        return None, "http_other"
    except Exception as e:
        print(f"  ! fetch error: {e}", file=sys.stderr)
        return None, "fetch"

    content = (data.get("choices") or [{}])[0].get("message", {}).get("content", "")
    content = strip_markdown_fences(content)
    if not content or len(content) < 80:
        return None, "empty"
    return content, None


def count_pending(conn) -> int:
    with conn.cursor() as cur:
        cur.execute(
            """
            SELECT count(*)
              FROM public.training_queue
             WHERE kind = 'cve'
               AND domain = 'cybersecurity'
               AND payload ? 'prompt'
               AND NOT (payload ? 'completion')
            """
        )
        return cur.fetchone()[0]


def fetch_rows(conn, limit: int) -> list[dict]:
    """KEV-flagged first (actively exploited), then CRITICAL/HIGH/MEDIUM,
    then most recent. Even a partial run captures the highest-signal data."""
    sql = """
        SELECT id, external_id, payload
          FROM public.training_queue
         WHERE kind = 'cve'
           AND domain = 'cybersecurity'
           AND payload ? 'prompt'
           AND NOT (payload ? 'completion')
         ORDER BY
            CASE WHEN (payload->>'kev')::boolean THEN 0 ELSE 1 END,
            CASE payload->>'cvss_severity'
                WHEN 'CRITICAL' THEN 1
                WHEN 'HIGH' THEN 2
                WHEN 'MEDIUM' THEN 3
                ELSE 9
            END,
            (payload->>'published') DESC NULLS LAST
         LIMIT %s
    """
    with conn.cursor(row_factory=psycopg_rows.dict_row) as cur:
        cur.execute(sql, (limit,))
        return list(cur.fetchall())


def update_row(conn, row_id: int, completion: str, model: str) -> None:
    sql = """
        UPDATE public.training_queue
           SET payload = payload
                         || jsonb_build_object('completion', %s::text)
                         || jsonb_build_object('completion_model', %s::text)
         WHERE id = %s
    """
    with conn.cursor() as cur:
        cur.execute(sql, (completion, model, row_id))
    conn.commit()


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    parser.add_argument("--limit", type=int, default=None, help="cap total rows enriched")
    parser.add_argument("--batch", type=int, default=50)
    parser.add_argument(
        "--model",
        default=os.environ.get("BEE_BACKFILL_MODEL", DEFAULT_MODEL),
    )
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    api_key = (os.environ.get("BEE_MISTRAL_API_KEY") or "").strip()
    if not api_key:
        print("ERROR: BEE_MISTRAL_API_KEY not set", file=sys.stderr)
        return 1
    pg_url = (os.environ.get("POSTGRES_URL_NON_POOLING") or "").strip()
    if not pg_url:
        print("ERROR: POSTGRES_URL_NON_POOLING not set", file=sys.stderr)
        return 1

    print(
        f"Completion backfill — model={args.model}  "
        f"batch={args.batch}  pace={RATE_LIMIT_RPM} req/min"
    )

    started = time.monotonic()
    enriched = 0
    skipped = 0
    rate_limited = 0
    last_call = 0.0

    with psycopg.connect(pg_url, autocommit=False) as conn:
        pending = count_pending(conn)
        print(f"  pending rows worth completing: {pending}")
        if args.dry_run:
            print("dry-run; exiting")
            return 0

        target = min(args.limit, pending) if args.limit else pending
        if target == 0:
            print("nothing to do")
            return 0
        print(f"  target this run: {target}")
        print()

        while enriched + skipped < target:
            remaining = target - enriched - skipped
            rows = fetch_rows(conn, min(args.batch, remaining))
            if not rows:
                break
            for row in rows:
                elapsed = time.monotonic() - last_call
                if elapsed < RATE_INTERVAL_S:
                    time.sleep(RATE_INTERVAL_S - elapsed)
                last_call = time.monotonic()

                content, err = call_mistral(api_key, args.model, row["payload"]["prompt"])
                if err == "429":
                    rate_limited += 1
                    print("  ! 429 — backing off 12s")
                    time.sleep(12.0)
                    continue
                if not content:
                    skipped += 1
                    continue

                update_row(conn, row["id"], content, args.model)
                enriched += 1

                if enriched % 10 == 0 or enriched == target:
                    elapsed_min = (time.monotonic() - started) / 60.0
                    rate = enriched / elapsed_min if elapsed_min > 0 else 0
                    eta_min = (target - enriched) / rate if rate > 0 else 0
                    print(
                        f"  enriched {enriched}/{target}  "
                        f"(skipped {skipped}, 429s {rate_limited}, "
                        f"~{rate:.1f}/min, ETA {eta_min:.1f}min)"
                    )

    elapsed_total = time.monotonic() - started
    print()
    print(
        f"Done. enriched={enriched}  skipped={skipped}  "
        f"rate_limited={rate_limited}  in {elapsed_total/60:.1f} min"
    )
    return 0


if __name__ == "__main__":
    sys.exit(main())