xinmi/scripts/apply_weknora_assets.py

#!/usr/bin/env python3
"""Checkpoint Agent/skill migration without writing private bodies to reports."""
from __future__ import annotations

import argparse
import datetime
import json
from pathlib import Path

from apply_weknora_identities import sql_run
from weknora_asset_etl import LEDGER, BINDING_LEDGER, asset_plan, asset_sql, safe_asset_summary
from weknora_category_etl import save_json
from weknora_etl import consistent_snapshot
from weknora_identity_etl import IDENTITY_LEDGER, TARGET_DATABASE, TARGET_TENANT
from weknora_model_etl import model_plan


def target_snapshot():
    return json.loads(sql_run("""SELECT json_build_object(
 'database',current_database(),
 'agents',(SELECT coalesce(json_agg(json_build_object('id',id,'tenant_id',tenant_id,
   'fingerprint',md5(to_jsonb(a)::text)) ORDER BY id),'[]') FROM custom_agents a),
 'skills',(SELECT coalesce(json_agg(json_build_object('id',id,'tenant_id',tenant_id,
   'fingerprint',md5(to_jsonb(s)::text)) ORDER BY id),'[]') FROM tenant_skill_catalog s),
 'models',(SELECT coalesce(json_agg(json_build_object('id',id,
   'fingerprint',md5(to_jsonb(m)::text)) ORDER BY id),'[]') FROM models m),
 'knowledge_count',(SELECT count(*) FROM knowledges),
 'chunk_count',(SELECT count(*) FROM chunks));"""))


def build_plan(db, mapping_file):
    identities = json.loads(sql_run(f"SELECT json_object_agg(source_user_id,target_user_id) FROM {IDENTITY_LEDGER} "
        f"WHERE source_system='yawan-workbench' AND source_organization_id=1 AND target_tenant_id={TARGET_TENANT};"))
    state = json.loads(mapping_file.read_text())
    if state.get('tenant_id') != TARGET_TENANT:
        raise ValueError('Knowledge mapping tenant mismatch')
    kbs = {key: item['id'] for key, item in state['knowledge_bases'].items()}
    return asset_plan(db, identities, kbs, model_plan(db, TARGET_TENANT))


def verification_sql():
    return f"""SELECT json_build_object(
 'agents',(SELECT count(*) FROM {LEDGER} WHERE target_tenant_id={TARGET_TENANT} AND source_kind='agent'),
 'skills',(SELECT count(*) FROM {LEDGER} WHERE target_tenant_id={TARGET_TENANT} AND source_kind='skill'),
 'bindings',(SELECT count(*) FROM {BINDING_LEDGER} WHERE target_tenant_id={TARGET_TENANT}),
 'effective_bindings',(SELECT coalesce(sum(jsonb_array_length(coalesce(a.config->'instruction_skill_ids','[]'))),0)
   FROM custom_agents a JOIN {LEDGER} l ON l.source_kind='agent' AND l.target_id=a.id WHERE l.target_tenant_id={TARGET_TENANT}),
 'disabled_agents',(SELECT count(*) FROM custom_agents a JOIN {LEDGER} l ON l.source_kind='agent' AND l.target_id=a.id
   WHERE l.target_tenant_id={TARGET_TENANT} AND NOT a.enabled),
 'private_agents',(SELECT count(*) FROM custom_agents a JOIN {LEDGER} l ON l.source_kind='agent' AND l.target_id=a.id
   WHERE l.target_tenant_id={TARGET_TENANT} AND a.visibility='private'),
 'private_skills',(SELECT count(*) FROM tenant_skill_catalog s JOIN {LEDGER} l ON l.source_kind='skill' AND l.target_id=s.id
   WHERE l.target_tenant_id={TARGET_TENANT} AND s.visibility='private'),
 'rules_match',(SELECT bool_and(a.config->>'system_prompt'=coalesce(l.source_snapshot->>'rules',''))
   FROM custom_agents a JOIN {LEDGER} l ON l.source_kind='agent' AND l.target_id=a.id WHERE l.target_tenant_id={TARGET_TENANT}),
 'skill_bodies_match',(SELECT bool_and(s.instructions=l.source_snapshot->>'instructions')
   FROM tenant_skill_catalog s JOIN {LEDGER} l ON l.source_kind='skill' AND l.target_id=s.id WHERE l.target_tenant_id={TARGET_TENANT}));"""


def transaction(query, commit=False):
    return f"""BEGIN ISOLATION LEVEL SERIALIZABLE;
SET LOCAL statement_timeout='30s';
SET LOCAL lock_timeout='5s';
CREATE TEMP TABLE _agents_before ON COMMIT DROP AS SELECT * FROM custom_agents;
CREATE TEMP TABLE _skills_before ON COMMIT DROP AS SELECT * FROM tenant_skill_catalog;
{query}
DO $preservation$
BEGIN
 IF EXISTS (SELECT 1 FROM _agents_before b LEFT JOIN custom_agents a ON a.id=b.id
            WHERE a.id IS NULL OR to_jsonb(a) IS DISTINCT FROM to_jsonb(b))
 OR EXISTS (SELECT 1 FROM _skills_before b LEFT JOIN tenant_skill_catalog s ON s.id=b.id
            WHERE s.id IS NULL OR to_jsonb(s) IS DISTINCT FROM to_jsonb(b))
 THEN RAISE EXCEPTION 'Asset import changed an existing row'; END IF;
END $preservation$;
SELECT json_build_object(
 'inserted_agents',(SELECT count(*) FROM custom_agents a WHERE NOT EXISTS (SELECT 1 FROM _agents_before b WHERE b.id=a.id)),
 'inserted_skills',(SELECT count(*) FROM tenant_skill_catalog s WHERE NOT EXISTS (SELECT 1 FROM _skills_before b WHERE b.id=s.id)),
 'existing_assets_unchanged',true);
{verification_sql()}
{'COMMIT' if commit else 'ROLLBACK'};
"""


def verify(plan, before, after, results):
    checks = {'acceptance_scope': before['database'] == after['database'] == TARGET_DATABASE}
    for kind in ('agents', 'skills'):
        old, new = ({v['id']: v for v in state[kind]} for state in (before, after))
        expected = {item['target']['id'] for item in plan[kind]}
        checks[kind + '_existing_unchanged'] = all(new.get(k) == v for k, v in old.items())
        checks[kind + '_planned_only'] = set(new) - set(old) == expected - set(old)
        checks[kind + '_all_present'] = expected <= set(new)
    checks['knowledge_and_models_unchanged'] = all(before[k] == after[k] for k in ('models', 'knowledge_count', 'chunk_count'))
    summary = next(item for item in results if 'effective_bindings' in item)
    checks['durable_mapping_counts'] = all(summary[k] == len(plan[k]) for k in ('agents', 'skills', 'bindings'))
    checks['effective_binding_count'] = summary['effective_bindings'] == sum(len(a['target']['config']['instruction_skill_ids']) for a in plan['agents'])
    checks['disabled_and_private_preserved'] = summary['disabled_agents'] == 1 and summary['private_agents'] == 2 and summary['private_skills'] == 2
    # Native edits may differ on future repeats; first import must preserve all bodies.
    if not any(x['target']['id'] in {a['id'] for a in before['agents']} for x in plan['agents']):
        checks['source_bodies_preserved'] = summary['rules_match'] and summary['skill_bodies_match']
    return {'passed': all(checks.values()), 'checks': checks}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('command', choices=('plan', 'apply'))
    parser.add_argument('--sqlite', type=Path, required=True)
    parser.add_argument('--kb-mapping', type=Path, required=True)
    parser.add_argument('--out', type=Path, required=True)
    args = parser.parse_args()
    suffixes = ('.before.json', '.manifest.json', '.attempt.json', '.committed.json', '.after.json')
    if any(p.exists() for p in (args.out, *(args.out.with_suffix(s) for s in suffixes))):
        parser.error('Checkpoint exists; use a new output name')
    with consistent_snapshot(args.sqlite) as (db, info):
        plan = build_plan(db, args.kb_mapping)
        if info['source_mutated']:
            raise RuntimeError('Source changed while taking the snapshot')
        plan['source_snapshot_sha256'] = info['sha256']
        query = asset_sql(plan)
    before = target_snapshot()
    save_json(args.out.with_suffix('.before.json'), before)
    save_json(args.out.with_suffix('.manifest.json'), {**safe_asset_summary(plan), 'source_snapshot_sha256': info['sha256']})
    rehearsal = [json.loads(line) for line in sql_run(transaction(query)).splitlines() if line.startswith('{')]
    result = {'checked_at': datetime.datetime.now(datetime.timezone.utc).isoformat(),
              'database': TARGET_DATABASE, 'tenant_id': TARGET_TENANT,
              'source_snapshot_sha256': info['sha256'], 'source_mutated': info['source_mutated'],
              'rehearsal': rehearsal, 'committed': False}
    if args.command == 'apply':
        save_json(args.out.with_suffix('.attempt.json'), {**result, 'commit_requested': True,
                  'recovery': 'Reconcile asset ledger and committed receipt; do not delete or recreate assets automatically'})
        results = [json.loads(line) for line in sql_run(transaction(query, True)).splitlines() if line.startswith('{')]
        result.update(committed=True, changes=results)
        save_json(args.out.with_suffix('.committed.json'), result)
        after = target_snapshot()
        save_json(args.out.with_suffix('.after.json'), after)
        result['verification'] = verify(plan, before, after, results)
    save_json(args.out, result)
    print(json.dumps(result))
    return 0 if result.get('verification', {}).get('passed', True) else 1


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