xinmi/scripts/milvus_copy_acceptance.py

#!/usr/bin/env python3
"""Copy existing vectors to an isolated BM25 collection, without re-embedding.

The manifest is an export of the acceptance DB's chunk IDs. Source collections
are read-only; acknowledgements and a manifest digest make retries resumable.
"""
import argparse
import hashlib
import json
import os
import time
import urllib.request
from pathlib import Path

FIELDS = ['id', 'embedding', 'content', 'source_id', 'source_type', 'chunk_id',
          'knowledge_id', 'knowledge_base_id', 'tag_id', 'is_enabled']


def post(base, path, payload):
    request = urllib.request.Request(base + '/v2/vectordb/' + path,
                                     data=json.dumps(payload).encode(),
                                     headers={'Content-Type': 'application/json'})
    with urllib.request.urlopen(request, timeout=120) as response:
        result = json.load(response)
    if result.get('code') != 0:
        raise RuntimeError(f"{path}: {result.get('message', result.get('code'))}")
    return result.get('data')


def save(path, data):
    temporary = path.with_suffix('.tmp')
    temporary.write_text(json.dumps(data, ensure_ascii=False, indent=2) + '\n')
    os.replace(temporary, path)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--manifest', type=Path, required=True)
    parser.add_argument('--state', type=Path, required=True)
    parser.add_argument('--target', default='yawan_acceptance_hybrid_512')
    parser.add_argument('--source', default='weknora_embeddings_512')
    args = parser.parse_args()
    if not args.target.startswith('yawan_acceptance_') or args.target == args.source:
        parser.error('target must be a distinct acceptance collection')
    base = 'http://127.0.0.1:19530'
    raw = args.manifest.read_bytes()
    manifest = json.loads(raw)
    digest = hashlib.sha256(raw).hexdigest()
    state = json.loads(args.state.read_text()) if args.state.exists() else {
        'manifest_sha256': digest, 'source': args.source, 'target': args.target,
        'offset': 0, 'copied': 0, 'missing': [], 'started_at': time.time(),
    }
    for key, value in [('manifest_sha256', digest), ('source', args.source), ('target', args.target)]:
        if state[key] != value:
            raise RuntimeError(f'{key} differs from checkpoint')
    description = post(base, 'collections/describe', {'collectionName': args.target})
    if 'content_sparse' not in {f['name'] for f in description['fields']}:
        raise RuntimeError('target must have BM25 schema before copying')
    for offset in range(state['offset'], len(manifest), 256):
        batch = manifest[offset:offset + 256]
        expected = {x['id']: x for x in batch}
        records = post(base, 'entities/query', {
            'collectionName': args.source,
            'filter': 'chunk_id in ' + json.dumps(list(expected)),
            'outputFields': FIELDS, 'limit': len(batch) * 2,
            'consistencyLevel': 'Strong',
        })
        found = set()
        for row in records:
            entry = expected.get(row['chunk_id'])
            if not entry or row['knowledge_base_id'] != entry['knowledge_base_id'] or row['knowledge_id'] != entry['knowledge_id']:
                raise RuntimeError('source scope differs from acceptance manifest')
            if row['chunk_id'] in found:
                raise RuntimeError('duplicate source vectors for chunk')
            found.add(row['chunk_id'])
        if records:
            result = post(base, 'entities/upsert', {'collectionName': args.target, 'data': records})
            if result.get('upsertCount') != len(records):
                raise RuntimeError('target did not acknowledge every vector')
        state['missing'].extend(sorted(expected.keys() - found))
        state['copied'] += len(records)
        state['offset'] = offset + len(batch)
        state['updated_at'] = time.time()
        save(args.state, state)
        if offset % (256 * 20) == 0 or state['offset'] == len(manifest):
            print(json.dumps({'processed': state['offset'], 'total': len(manifest),
                              'copied': state['copied'], 'missing': len(state['missing'])}), flush=True)
    state['complete'] = not state['missing']
    save(args.state, state)
    print(json.dumps({'complete': state['complete'], 'copied': state['copied'],
                      'missing': len(state['missing'])}), flush=True)
    return 0 if state['complete'] else 1


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