WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

scripts/milvus_copy_acceptance.py

@@ -0,0 +1,96 @@
+#!/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())