xinmi/scripts/weknora_model_etl.py

"""Plan legacy model configuration and emit guarded, in-memory SQL.

Callers provide a consistent SQLite snapshot and execute SQL in their own
acceptance transaction. This module never opens the source or executes SQL.
Only the safe plan may be persisted: credentials and SQL must stay in memory.
"""
from __future__ import annotations

import base64
import copy
import json
import re
import sqlite3
import uuid
from collections.abc import Callable, Mapping

SOURCE = "yawan-workbench"
TARGET_DATABASE = "yawan_native_acceptance_20260915"
TARGET_TENANT = 10059
SOURCE_ORGANIZATION = 1
LEDGER = "yawan_legacy_model_mappings"
NAMESPACE = uuid.UUID("d6f95f91-93f9-4afc-85a7-23d3fa0a7f6d")
REVIEWED_PROVIDER_ID = 3
REVIEWED_BASE_URL = "https://api.deepseek.com/v1"
REVIEWED_MODELS = {"deepseek-flash", "deepseek-v4-pro"}


def model_id(provider_id: int, model_name: str, tenant_id: int = TARGET_TENANT) -> str:
    """Stable tenant-scoped ID; use only IDs included in an approved plan."""
    return str(uuid.uuid5(NAMESPACE, json.dumps(
        [SOURCE, tenant_id, SOURCE_ORGANIZATION, provider_id, model_name, "KnowledgeQA"],
        separators=(",", ":"))))


def _list(value: str) -> list:
    parsed = json.loads(value)
    if not isinstance(parsed, list):
        raise ValueError("Source model list is invalid")
    return parsed


def model_plan(db: sqlite3.Connection, tenant_id: int = TARGET_TENANT) -> dict:
    """Produce safe metadata without reading any credential column value."""
    if type(tenant_id) is not int or tenant_id != TARGET_TENANT:
        raise ValueError("Model migration target tenant mismatch")
    providers = {}
    for row in db.execute("SELECT id,name,base_url,models,default_model,builtin,enabled,organization_id "
                          "FROM modelprovider ORDER BY id"):
        provider = dict(zip(("id", "name", "base_url", "models", "default_model", "builtin",
                             "enabled", "organization_id"), row))
        if provider["organization_id"] != SOURCE_ORGANIZATION:
            raise ValueError("Source model provider belongs to another organization")
        if provider["enabled"] not in (0, 1) or provider["builtin"] not in (0, 1):
            raise ValueError("Invalid source model provider flags")
        provider["models"] = _list(provider["models"])
        if not all(isinstance(name, str) and name.strip() == name and name for name in provider["models"]):
            raise ValueError("Invalid source model names")
        # Endpoint credentials must never enter a safe manifest.
        from urllib.parse import urlsplit
        url = urlsplit(provider["base_url"])
        if url.username or url.password or url.query or url.fragment:
            raise ValueError("Source provider URL contains unsupported private components")
        providers[provider["id"]] = provider
    catalogs = []
    for row in db.execute("SELECT id,provider_id,model_name,display_name,capabilities,capability_status,"
                          "enabled,builtin,rate_version,input_rate,output_rate FROM modelcatalog ORDER BY id"):
        catalog = dict(zip(("id", "provider_id", "model_name", "display_name", "capabilities",
                            "capability_status", "enabled", "builtin", "rate_version", "input_rate",
                            "output_rate"), row))
        if catalog["provider_id"] not in providers:
            raise ValueError("Model catalog has no source provider")
        catalog["capabilities"] = _list(catalog["capabilities"])
        if catalog["enabled"] not in (0, 1):
            raise ValueError("Invalid source catalog enabled flag")
        catalogs.append(catalog)

    imported, skipped, seen = [], [], set()
    for catalog in catalogs:
        provider = providers[catalog["provider_id"]]
        reason = None
        if not provider["enabled"] or not catalog["enabled"]:
            reason = "source_configuration_disabled"
        elif catalog["model_name"].startswith("gpt-image-"):
            reason = "image_generation_has_no_native_model_type"
        elif (provider["id"] != REVIEWED_PROVIDER_ID or provider["base_url"] != REVIEWED_BASE_URL
              or catalog["model_name"] not in REVIEWED_MODELS):
            reason = "model_protocol_not_reviewed"
        elif catalog["model_name"] not in provider["models"]:
            reason = "model_not_enabled_in_provider"
        if reason:
            skipped.append({"source_catalog": catalog, "reason": reason})
            continue
        identity = (provider["id"], catalog["model_name"])
        if identity in seen:
            raise ValueError("Duplicate enabled model catalog identity")
        seen.add(identity)
        imported.append({
            "source_catalog": catalog, "source_provider_id": provider["id"],
            "target_id": model_id(*identity), "name": catalog["model_name"],
            "display_name": catalog["display_name"] or catalog["model_name"],
            "type": "KnowledgeQA", "source": "remote", "status": "active",
            "is_default": False, "is_builtin": False, "managed_by": "",
            "parameters": {"base_url": provider["base_url"], "provider": "deepseek",
                           "interface_type": "", "supports_vision": False},
            "credential_state": "must_resolve_and_encrypt_before_insert",
            "compatibility": ["source_openai_chat_protocol_preserved", "no_paid_probe_performed",
                              "source_rates_retained_in_ledger_not_applied_to_native_billing",
                              "untested_source_capabilities_not_inferred"],
        })
    catalog_map = {str(item["source_catalog"]["id"]): item["target_id"] for item in imported}
    provider_model_map = {f"{item['source_provider_id']}:{item['name']}": item["target_id"] for item in imported}
    provider_defaults = {str(p["id"]): provider_model_map[f"{p['id']}:{p['default_model']}"]
                         for p in providers.values() if f"{p['id']}:{p['default_model']}" in provider_model_map}
    default_row = db.execute("SELECT value FROM appsetting WHERE key='default_model_provider_id'").fetchone()
    configured_default = int(default_row[0]) if default_row and str(default_row[0]).isdigit() else None
    enabled = sorted((p for p in providers.values() if p["enabled"]), key=lambda p: (-p["builtin"], p["id"]))
    chosen = providers.get(configured_default)
    if not chosen or not chosen["enabled"]:
        chosen = enabled[0] if enabled else None
    default_provider_id = chosen["id"] if chosen else None
    result = {
        "schema_version": 1, "source": SOURCE, "target_database": TARGET_DATABASE,
        "target_tenant_id": TARGET_TENANT, "source_organization_id": SOURCE_ORGANIZATION,
        "source_providers": list(providers.values()), "models": imported, "skipped": skipped,
        "catalog_map": catalog_map, "provider_model_map": provider_model_map,
        "provider_defaults": provider_defaults, "default_provider_id": default_provider_id,
        "default_gen_model_id": provider_defaults.get(str(default_provider_id)),
        "ledger": LEDGER, "repeat_policy": "preserve_existing_models_and_credentials",
        "credential_policy": "native_enc_v1_aes256gcm; memory_only; no_report_or_sql_file",
        "agent_resolution_policy": "legacy_agents.coordinator._provider_model",
    }
    _validate_plan(result)
    return result


def resolve_agent_model(plan: dict, agent: Mapping) -> dict:
    """Resolve the legacy coordinator's provider/default fallback explicitly."""
    _validate_plan(plan)
    providers = {p["id"]: p for p in plan["source_providers"]}
    provider = providers.get(agent.get("model_provider_id"))
    fallback = not provider or not provider["enabled"]
    if fallback:
        provider = providers.get(plan["default_provider_id"])
    if not provider:
        raise ValueError("Agent has no effective source model provider")
    requested = agent.get("model_name") or ""
    effective = requested if requested in provider["models"] else (
        provider["default_model"] or (provider["models"][0] if provider["models"] else ""))
    target_id = plan["provider_model_map"].get(f"{provider['id']}:{effective}")
    if not target_id:
        raise ValueError("Agent effective source model has no reviewed native mapping")
    return {"target_model_id": target_id, "effective_provider_id": provider["id"],
            "effective_model_name": effective, "requested_provider_id": agent.get("model_provider_id"),
            "requested_model_name": requested, "source_catalog_id": agent.get("model_catalog_id"),
            "provider_fallback": fallback, "model_fallback": requested != effective,
            "resolution_policy": plan["agent_resolution_policy"]}


def _validate_plan(plan: dict) -> None:
    if (plan.get("schema_version") != 1 or plan.get("source") != SOURCE
            or plan.get("target_database") != TARGET_DATABASE or plan.get("target_tenant_id") != TARGET_TENANT
            or plan.get("source_organization_id") != SOURCE_ORGANIZATION):
        raise ValueError("Model plan scope mismatch")
    seen = set()
    for item in plan.get("models", []):
        catalog = item["source_catalog"]
        if (item["source_provider_id"] != REVIEWED_PROVIDER_ID or item["name"] not in REVIEWED_MODELS
                or catalog["provider_id"] != item["source_provider_id"] or catalog["model_name"] != item["name"]
                or catalog["enabled"] != 1 or item["target_id"] != model_id(item["source_provider_id"], item["name"])
                or item["type"] != "KnowledgeQA" or item["source"] != "remote" or item["status"] != "active"
                or item["is_default"] is not False or item["is_builtin"] is not False or item["managed_by"] != ""
                or item["parameters"] != {"base_url": REVIEWED_BASE_URL, "provider": "deepseek",
                                          "interface_type": "", "supports_vision": False}):
            raise ValueError("Model plan changes reviewed model configuration")
        if not item["display_name"] or len(item["display_name"]) > 255 or item["target_id"] in seen:
            raise ValueError("Invalid or duplicate target model identity")
        seen.add(item["target_id"])
    expected = {f"{m['source_provider_id']}:{m['name']}": m["target_id"] for m in plan["models"]}
    if plan.get("provider_model_map") != expected:
        raise ValueError("Model mapping differs from approved models")
    if plan.get("catalog_map") != {str(m["source_catalog"]["id"]): m["target_id"] for m in plan["models"]}:
        raise ValueError("Catalog mapping differs from approved models")


def read_credentials(db: sqlite3.Connection, plan: dict,
                     resolver: Callable[[str, int], str] | None = None) -> dict[int, str]:
    """Read only imported providers' secrets; returned values must stay in memory."""
    _validate_plan(plan)
    result = {}
    for provider_id in sorted({m["source_provider_id"] for m in plan["models"]}):
        row = db.execute("SELECT api_key FROM modelprovider WHERE id=? AND organization_id=? AND enabled=1",
                         (provider_id, SOURCE_ORGANIZATION)).fetchone()
        if not row or not isinstance(row[0], str) or not row[0].strip():
            raise ValueError("Imported model provider credential is missing")
        value = resolver(row[0], provider_id) if resolver else row[0]
        if not isinstance(value, str) or not value.strip() or value.startswith(("keychain:", "enc:", "env:")):
            raise ValueError("Imported model provider credential requires a source resolver")
        result[provider_id] = value
    return result


def model_sql(plan: dict, credentials: Mapping[int, str], encrypt_secret: Callable[[str], str]) -> str:
    """Emit SQL with encrypted credentials, no transaction boundaries or execution.

    encrypt_secret must use the acceptance app's SYSTEM_AES_KEY and implement
    the native enc:v1: + base64url(nonce[12] + ciphertext + tag[16]) format.
    The caller must not log, print, or save the return value or SQL errors.
    """
    _validate_plan(plan)
    if not re.fullmatch(r"[a-f0-9]{64}", str(plan.get("source_snapshot_sha256", ""))):
        raise ValueError("Model SQL requires the caller's source snapshot fingerprint")
    encrypted = {}
    for provider_id in sorted({m["source_provider_id"] for m in plan["models"]}):
        raw = credentials.get(provider_id)
        if not isinstance(raw, str) or not raw.strip() or raw.startswith(("keychain:", "enc:", "env:")):
            raise ValueError("Model SQL requires resolved plaintext credentials in memory")
        try:
            value = encrypt_secret(raw)
            if not isinstance(value, str) or not value.startswith("enc:v1:") or value == raw:
                raise ValueError
            token = value[7:]
            if not re.fullmatch(r"[A-Za-z0-9_-]+", token):
                raise ValueError
            data = base64.urlsafe_b64decode(token + "=" * (-len(token) % 4))
            if len(data) != 12 + len(raw.encode()) + 16:
                raise ValueError
        except Exception:
            raise ValueError("Native model credential encryption failed") from None
        encrypted[provider_id] = value
    payload = copy.deepcopy(plan)
    for item in payload["models"]:
        item["stored_parameters"] = dict(item["parameters"], api_key=encrypted[item["source_provider_id"]])
    literal = "'" + json.dumps(payload, ensure_ascii=True, sort_keys=True, separators=(",", ":")).replace("'", "''") + "'::jsonb"
    # Do not derive identifiers or ledger metadata from credential material.
    return f"""SET LOCAL standard_conforming_strings = on;
SET LOCAL log_statement = 'none';
SET LOCAL log_min_error_statement = 'panic';
DO $yawan_models_v1$
DECLARE
    migration jsonb := {literal};
    item jsonb;
    mapped record;
    existing_model models%ROWTYPE;
BEGIN
    IF current_database() <> '{TARGET_DATABASE}' THEN
        RAISE EXCEPTION 'Model migration target database mismatch';
    END IF;
    PERFORM pg_advisory_xact_lock({TARGET_TENANT}, 11903);
    PERFORM 1 FROM tenants WHERE id={TARGET_TENANT} AND status='active' AND deleted_at IS NULL FOR SHARE;
    IF NOT FOUND THEN RAISE EXCEPTION 'Model migration requires the active acceptance tenant'; END IF;
    LOCK TABLE models IN SHARE ROW EXCLUSIVE MODE;
    CREATE TABLE IF NOT EXISTS {LEDGER} (
        source_system text NOT NULL,
        source_organization_id bigint NOT NULL,
        source_catalog_id bigint NOT NULL,
        source_provider_id bigint NOT NULL,
        target_tenant_id integer NOT NULL,
        target_model_id varchar(64) NOT NULL,
        source_snapshot jsonb NOT NULL,
        source_snapshot_sha256 varchar(64) NOT NULL CHECK (source_snapshot_sha256 ~ '^[a-f0-9]{{64}}$'),
        schema_version integer NOT NULL DEFAULT 1 CHECK (schema_version=1),
        created_at timestamptz NOT NULL DEFAULT CURRENT_TIMESTAMP,
        PRIMARY KEY (source_system,source_organization_id,source_catalog_id),
        UNIQUE (target_tenant_id,target_model_id)
    );
    LOCK TABLE {LEDGER} IN SHARE ROW EXCLUSIVE MODE;
    FOR item IN SELECT value FROM jsonb_array_elements(migration->'models') LOOP
        SELECT * INTO mapped FROM {LEDGER} WHERE source_system='{SOURCE}'
            AND source_organization_id={SOURCE_ORGANIZATION}
            AND source_catalog_id=(item->'source_catalog'->>'id')::bigint;
        IF FOUND THEN
            IF mapped.target_tenant_id <> {TARGET_TENANT} OR mapped.target_model_id <> item->>'target_id'
                OR mapped.source_provider_id <> (item->>'source_provider_id')::bigint
                OR mapped.source_snapshot IS DISTINCT FROM item - 'stored_parameters' THEN
                RAISE EXCEPTION 'Model migration ledger mapping differs from reviewed source';
            END IF;
            SELECT * INTO existing_model FROM models WHERE id=mapped.target_model_id;
            IF NOT FOUND OR existing_model.tenant_id <> {TARGET_TENANT} OR existing_model.deleted_at IS NOT NULL
                OR existing_model.name <> item->>'name' OR existing_model.type <> 'KnowledgeQA'
                OR existing_model.source <> 'remote' OR existing_model.status <> 'active' THEN
                RAISE EXCEPTION 'Previously migrated model is missing or incompatible; manual review required';
            END IF;
            -- Existing credentials and user edits are deliberately preserved.
            CONTINUE;
        END IF;
        PERFORM 1 FROM models WHERE id=item->>'target_id';
        IF FOUND THEN RAISE EXCEPTION 'Model ID collision without migration provenance'; END IF;
        INSERT INTO models (id,tenant_id,name,display_name,type,source,description,parameters,
                            is_default,is_builtin,managed_by,status,created_at,updated_at)
            VALUES (item->>'target_id',{TARGET_TENANT},item->>'name',item->>'display_name',
                    'KnowledgeQA','remote','Legacy model configuration; capabilities remain unverified.',
                    item->'stored_parameters',false,false,'','active',CURRENT_TIMESTAMP,CURRENT_TIMESTAMP);
        INSERT INTO {LEDGER} (source_system,source_organization_id,source_catalog_id,source_provider_id,
                             target_tenant_id,target_model_id,source_snapshot,source_snapshot_sha256)
            VALUES ('{SOURCE}',{SOURCE_ORGANIZATION},(item->'source_catalog'->>'id')::bigint,
                    (item->>'source_provider_id')::bigint,{TARGET_TENANT},item->>'target_id',
                    item - 'stored_parameters',migration->>'source_snapshot_sha256');
    END LOOP;
END;
$yawan_models_v1$;
"""


def model_verification_sql() -> str:
    """Read model identity/credential-presence metadata, never credential values."""
    return f"""SELECT json_build_object(
        'mapped', (SELECT count(*) FROM {LEDGER} WHERE target_tenant_id={TARGET_TENANT}),
        'models', (SELECT coalesce(json_agg(json_build_object('id',m.id,'name',m.name,
            'type',m.type,'status',m.status,'encrypted_credential',
            coalesce(m.parameters->>'api_key','') LIKE 'enc:v1:%') ORDER BY m.id),'[]')
            FROM {LEDGER} l JOIN models m ON m.id=l.target_model_id AND m.tenant_id=l.target_tenant_id
            WHERE l.target_tenant_id={TARGET_TENANT} AND m.deleted_at IS NULL));"""