xinmi/scripts/verify_model_etl.py

#!/usr/bin/env python3
"""Exercise model migration conflicts in PostgreSQL, always rolling back.

An optional single Pro model debug request records only a status, finish reason,
token counts, and elapsed time. SQL, credentials, and response text never leave
memory or appear in the report.
"""
from __future__ import annotations

import argparse
import datetime
import json
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path

from apply_weknora_identities import sql_run
from apply_weknora_models import native_encryptor, resolve_source_credential, target_snapshot
from verify_identity_etl import rejected_case
from weknora_category_etl import literal, save_json
from weknora_etl import consistent_snapshot
from weknora_model_etl import (LEDGER, TARGET_DATABASE, TARGET_TENANT, model_id,
                               model_plan, model_sql, read_credentials)


def ledger_snapshot():
    return json.loads(sql_run(f"SELECT coalesce(json_agg(to_jsonb(l) ORDER BY source_catalog_id),'[]') "
                             f"FROM {LEDGER} l;"))


def regression_sql(plan, query):
    uid = literal(plan["models"][0]["target_id"])
    cases = [
        ("mapped_model_wrong_tenant_rejected", f"UPDATE models SET tenant_id={TARGET_TENANT + 1} WHERE id={uid};",
         "Previously migrated model is missing or incompatible"),
        ("ledger_wrong_tenant_rejected", f"UPDATE {LEDGER} SET target_tenant_id={TARGET_TENANT + 1} WHERE target_model_id={uid};",
         "Model migration ledger mapping differs"),
        ("ledger_wrong_model_id_rejected", f"UPDATE {LEDGER} SET target_model_id='rollback-only-conflicting-id' WHERE target_model_id={uid};",
         "Model migration ledger mapping differs"),
        ("ledger_wrong_provider_rejected", f"UPDATE {LEDGER} SET source_provider_id=999 WHERE target_model_id={uid};",
         "Model migration ledger mapping differs"),
        ("ledger_changed_source_rejected", f"UPDATE {LEDGER} SET source_snapshot='{{}}'::jsonb WHERE target_model_id={uid};",
         "Model migration ledger mapping differs"),
        ("existing_model_id_without_ledger_rejected", f"DELETE FROM {LEDGER} WHERE target_model_id={uid};",
         "Model ID collision without migration provenance"),
        ("deleted_model_not_revived", f"UPDATE models SET deleted_at=now() WHERE id={uid};",
         "Previously migrated model is missing or incompatible"),
        ("wrong_model_type_rejected", f"UPDATE models SET type='Embedding' WHERE id={uid};",
         "Previously migrated model is missing or incompatible"),
        ("inactive_target_tenant_rejected", f"UPDATE tenants SET status='suspended' WHERE id={TARGET_TENANT};",
         "Model migration requires the active acceptance tenant"),
    ]
    sql = f"""BEGIN ISOLATION LEVEL SERIALIZABLE;
SET LOCAL statement_timeout='30s';
SET LOCAL lock_timeout='5s';
SET LOCAL log_statement='none';
SET LOCAL log_min_error_statement='panic';
CREATE TEMP TABLE _models_before ON COMMIT DROP AS SELECT * FROM models;
CREATE TEMP TABLE _model_ledger_before ON COMMIT DROP AS SELECT * FROM {LEDGER};
{query}
DO $model_repeat$
BEGIN
 IF EXISTS (SELECT * FROM models EXCEPT SELECT * FROM _models_before)
 OR EXISTS (SELECT * FROM _models_before EXCEPT SELECT * FROM models)
 OR EXISTS (SELECT * FROM {LEDGER} EXCEPT SELECT * FROM _model_ledger_before)
 OR EXISTS (SELECT * FROM _model_ledger_before EXCEPT SELECT * FROM {LEDGER})
 THEN RAISE EXCEPTION 'Exact model migration repeat changed stored rows'; END IF;
END $model_repeat$;
SELECT json_build_object('check','exact_repeat_is_noop','passed',true);
SELECT json_build_object('check','existing_embedding_and_credentials_unchanged','passed',
 EXISTS (SELECT 1 FROM _models_before WHERE type='Embedding') AND NOT EXISTS (
  SELECT 1 FROM _models_before b LEFT JOIN models m ON m.id=b.id
  WHERE to_jsonb(m) IS DISTINCT FROM to_jsonb(b)));
"""
    for name, mutation, expected in cases:
        sql += rejected_case(name, mutation, expected, query)
    sql += f"""
UPDATE models SET display_name='Rollback-only native display edit',
 parameters=jsonb_set(parameters,'{{api_key}}','"enc:v1:rollback-only-rotation"'::jsonb)
 WHERE id={uid};
CREATE TEMP TABLE _model_rotated_before ON COMMIT DROP AS SELECT * FROM models;
{query}
DO $model_rotation_preservation$
BEGIN
 IF EXISTS (SELECT * FROM models EXCEPT SELECT * FROM _model_rotated_before)
 OR EXISTS (SELECT * FROM _model_rotated_before EXCEPT SELECT * FROM models)
 THEN RAISE EXCEPTION 'Model migration overwrote native edits or rotated credentials'; END IF;
END $model_rotation_preservation$;
SELECT json_build_object('check','native_display_and_rotated_credential_preserved','passed',true);
SELECT json_build_object('check','all_model_ledger_rows_secret_free','passed',
 NOT EXISTS (SELECT 1 FROM {LEDGER} WHERE source_snapshot::text LIKE '%"api_key"%'
 OR source_snapshot::text LIKE '%"stored_parameters"%'
 OR source_snapshot::text LIKE '%enc:v1:%'));
ROLLBACK;
"""
    return sql


def pro_api_check(credentials_path):
    credentials = json.loads(credentials_path.read_text())
    opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))
    login = urllib.request.Request("http://127.0.0.1:18080/api/v1/auth/login",
        data=json.dumps({key: credentials[key] for key in ("email", "password")}).encode(),
        headers={"Content-Type": "application/json"})
    try:
        with opener.open(login, timeout=20) as response:
            auth = json.load(response)
    except (urllib.error.URLError, ValueError):
        raise RuntimeError("Acceptance verification login failed") from None
    if auth.get("active_tenant", {}).get("id") != TARGET_TENANT or not auth.get("token"):
        raise RuntimeError("Verification login is outside the acceptance tenant")
    body = urllib.parse.urlencode({"input": "\u4ec5\u56de\u590dOK", "options": json.dumps({"max_tokens": 256})}).encode()
    request = urllib.request.Request(
        f"http://127.0.0.1:18080/api/v1/models/{model_id(3, 'deepseek-v4-pro')}/debug",
        data=body, headers={"Content-Type": "application/x-www-form-urlencoded",
                            "Authorization": "Bearer " + auth["token"]})
    try:
        with opener.open(request, timeout=55) as response:
            payload = json.load(response)
    except (urllib.error.URLError, ValueError):
        return {"ok": False, "finish_reason": "request_failed", "prompt_tokens": None,
                "completion_tokens": None, "total_tokens": None, "elapsed_ms": None}
    data = payload.get("data") or {}
    raw = data.get("raw_response") or {}
    usage = raw.get("usage") or {}
    finish = raw.get("finish_reason")
    # Whitelist the only model text field permitted in the saved result.
    allowed_finish = {"stop", "length", "tool_calls", "content_filter", "function_call", None, ""}
    return {"ok": bool(payload.get("success") and data.get("ok") and str(raw.get("content", "")).strip() == "OK"),
            "finish_reason": finish if finish in allowed_finish else "unrecognized",
            "prompt_tokens": usage.get("prompt_tokens"), "completion_tokens": usage.get("completion_tokens"),
            "total_tokens": usage.get("total_tokens"), "elapsed_ms": data.get("elapsed_ms")}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--sqlite", type=Path, required=True)
    parser.add_argument("--environment", type=Path, required=True)
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--credentials", type=Path)
    args = parser.parse_args()
    if args.out.exists():
        parser.error("Output already exists")
    manifest = json.loads(args.manifest.read_text())
    with consistent_snapshot(args.sqlite) as (db, info):
        fresh = model_plan(db)
        expected = {k: v for k, v in manifest.items() if k != "source_snapshot_sha256"}
        if fresh != expected or info["source_mutated"]:
            raise RuntimeError("Model source differs from reviewed manifest")
        query = model_sql(manifest, read_credentials(db, manifest, resolve_source_credential),
                          native_encryptor(args.environment))
    before, ledger_before = target_snapshot(), ledger_snapshot()
    output = sql_run(regression_sql(manifest, query))
    checks = [json.loads(line) for line in output.splitlines() if line.startswith("{")]
    after, ledger_after = target_snapshot(), ledger_snapshot()
    checks.append({"check": "all_test_changes_rolled_back", "passed": before == after and ledger_before == ledger_after})
    checks.append({"check": "acceptance_scope", "passed": before["database"] == after["database"] == TARGET_DATABASE})
    report = {"checked_at": datetime.datetime.now(datetime.timezone.utc).isoformat(),
              "database": TARGET_DATABASE, "tenant_id": TARGET_TENANT,
              "committed": False, "checks": checks, "passed": all(c["passed"] for c in checks)}
    if args.credentials:
        report["pro_debug_max_tokens_256"] = pro_api_check(args.credentials)
    save_json(args.out, report)
    print(json.dumps(report))
    return 0 if report["passed"] else 1


if __name__ == "__main__":
    raise SystemExit(main())