|
|
@@ -2,20 +2,26 @@
|
|
|
|
|
|
import argparse
|
|
|
from collections import Counter
|
|
|
+import copy
|
|
|
import hashlib
|
|
|
import json
|
|
|
from pathlib import Path
|
|
|
import re
|
|
|
+import shutil
|
|
|
import sys
|
|
|
+import tempfile
|
|
|
+import uuid
|
|
|
from zipfile import ZIP_DEFLATED, ZipFile
|
|
|
|
|
|
from src.step1_knowledge_prepare.summarize import atomic_write, utc_now
|
|
|
from src.step2_knowledge_ingest.ingest import (
|
|
|
CONTRACT, DIM, KINDS, TYPES, Store, canonical, config, content_parts,
|
|
|
- index_text, load, model_matches_manifest, restore, route, sha, valid_vector,
|
|
|
+ index_text, load, model_matches_manifest, publication_lock, restore, route, sha, valid_vector,
|
|
|
)
|
|
|
|
|
|
FORMAT = 'business-assistant-milvus-snapshot-v1'
|
|
|
+LIBRARY_FORMAT = 'business-assistant-milvus-library-v1'
|
|
|
+LIBRARIES = tuple(kind for kind in TYPES if kind != 'catalog')
|
|
|
SOURCE_FILES = tuple(f'{kind}{suffix}.jsonl' for kind in KINDS
|
|
|
for suffix in ('', '.with_summary'))
|
|
|
ARCHIVE_META = 'archive.json'
|
|
|
@@ -50,9 +56,12 @@ def _check_manifest(manifest):
|
|
|
for kind, spec in databases.items():
|
|
|
if not isinstance(spec, dict) or not re.fullmatch(r'[A-Za-z][A-Za-z0-9_]{0,99}', spec.get('database', '')):
|
|
|
raise ValueError('invalid_snapshot_manifest')
|
|
|
- if spec.get('content_collection') != 'content_' + manifest['snapshot'][:20]:
|
|
|
+ version = spec.get('snapshot', manifest['snapshot'])
|
|
|
+ if not isinstance(version, str) or not re.fullmatch(r'[a-f0-9]{64}', version):
|
|
|
+ raise ValueError('invalid_library_snapshot')
|
|
|
+ if spec.get('content_collection') != 'content_' + version[:20]:
|
|
|
raise ValueError('invalid_snapshot_manifest')
|
|
|
- if spec.get('index_collection') != (None if kind == 'catalog' else 'index_' + manifest['snapshot'][:20]):
|
|
|
+ if spec.get('index_collection') != (None if kind == 'catalog' else 'index_' + version[:20]):
|
|
|
raise ValueError('invalid_snapshot_manifest')
|
|
|
|
|
|
|
|
|
@@ -74,11 +83,29 @@ def _snapshot(records, manifest):
|
|
|
raise ValueError('snapshot_hash_mismatch')
|
|
|
|
|
|
|
|
|
-def export(prepared, output, settings, store_factory=Store):
|
|
|
- manifest = json.loads((prepared / 'ingestion/manifest.json').read_text(encoding='utf-8'))
|
|
|
- _check_manifest(manifest)
|
|
|
+def export(prepared, output, settings, store_factory=Store, database=None):
|
|
|
+ with publication_lock(prepared):
|
|
|
+ return _export(prepared, output, settings, store_factory, database)
|
|
|
+
|
|
|
+
|
|
|
+def _export(prepared, output, settings, store_factory, database):
|
|
|
+ source_manifest = json.loads((prepared / 'ingestion/manifest.json').read_text(encoding='utf-8'))
|
|
|
+ _check_manifest(source_manifest)
|
|
|
records = load(prepared)
|
|
|
- _snapshot(records, manifest)
|
|
|
+ _snapshot(records, source_manifest)
|
|
|
+ if database is not None:
|
|
|
+ if database not in LIBRARIES or database not in source_manifest['counts']:
|
|
|
+ raise ValueError('unknown_or_empty_library')
|
|
|
+ records = [record for record in records if route(record) == database]
|
|
|
+ manifest = copy.deepcopy(source_manifest)
|
|
|
+ manifest['counts'] = dict(Counter(route(record) for record in records))
|
|
|
+ manifest['snapshot'] = sha(canonical({'contract': CONTRACT, 'model': manifest['model_hash'], 'records': records}))
|
|
|
+ # Transfer packages always use one version, even when exporting a live
|
|
|
+ # manifest assembled from several independently updated libraries.
|
|
|
+ manifest['databases'] = {kind: {'database': source_manifest['databases'][kind]['database'],
|
|
|
+ 'content_collection': 'content_' + manifest['snapshot'][:20],
|
|
|
+ 'index_collection': None if kind == 'catalog' else 'index_' + manifest['snapshot'][:20]}
|
|
|
+ for kind in manifest['counts']}
|
|
|
if output.exists():
|
|
|
raise ValueError('archive_already_exists')
|
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
|
@@ -92,13 +119,17 @@ def export(prepared, output, settings, store_factory=Store):
|
|
|
archive.writestr(MANIFEST_NAME, manifest_bytes)
|
|
|
digests = {MANIFEST_NAME: hashlib.sha256(manifest_bytes).hexdigest()}
|
|
|
for name in SOURCE_FILES:
|
|
|
- archive.write(prepared / name, 'prepared/' + name)
|
|
|
- with (prepared / name).open('rb') as source:
|
|
|
- digests['prepared/' + name] = hashlib.file_digest(source, 'sha256').hexdigest()
|
|
|
+ content = (prepared / name).read_bytes()
|
|
|
+ if database is not None:
|
|
|
+ selected_ids = {record['id'] for record in records}
|
|
|
+ content = b''.join(line for line in content.splitlines(keepends=True)
|
|
|
+ if json.loads(line)['id'] in selected_ids)
|
|
|
+ archive.writestr('prepared/' + name, content)
|
|
|
+ digests['prepared/' + name] = hashlib.sha256(content).hexdigest()
|
|
|
for kind in TYPES:
|
|
|
if kind not in manifest['databases']:
|
|
|
continue
|
|
|
- spec = manifest['databases'][kind]
|
|
|
+ spec = source_manifest['databases'][kind]
|
|
|
if store.use(kind, create=False) != spec['database']:
|
|
|
raise ValueError('source_database_mismatch')
|
|
|
subset = [record for record in records if route(record) == kind]
|
|
|
@@ -138,8 +169,11 @@ def export(prepared, output, settings, store_factory=Store):
|
|
|
if actual != expected:
|
|
|
raise ValueError('source_count_mismatch')
|
|
|
print(json.dumps({'type': kind, 'exported': len(subset)}), flush=True)
|
|
|
- archive.writestr(ARCHIVE_META, canonical({'format': FORMAT, 'created_at': utc_now(),
|
|
|
- 'snapshot': manifest['snapshot'], 'counts': manifest['counts'], 'sha256': digests}) + '\n')
|
|
|
+ metadata = {'format': LIBRARY_FORMAT if database else FORMAT, 'created_at': utc_now(),
|
|
|
+ 'snapshot': manifest['snapshot'], 'counts': manifest['counts'], 'sha256': digests}
|
|
|
+ if database:
|
|
|
+ metadata['database'] = database
|
|
|
+ archive.writestr(ARCHIVE_META, canonical(metadata) + '\n')
|
|
|
verify_archive(partial)
|
|
|
partial.replace(output)
|
|
|
except BaseException:
|
|
|
@@ -160,8 +194,11 @@ def verify_archive(path):
|
|
|
_check_manifest(manifest)
|
|
|
expected = {MANIFEST_NAME, *(f'prepared/{name}' for name in SOURCE_FILES),
|
|
|
*(f'records/{kind}.jsonl' for kind in manifest['counts'])}
|
|
|
- if meta.get('format') != FORMAT or set(names) != expected | {ARCHIVE_META} or set(meta.get('sha256', {})) != expected:
|
|
|
+ if meta.get('format') not in (FORMAT, LIBRARY_FORMAT) or set(names) != expected | {ARCHIVE_META} or set(meta.get('sha256', {})) != expected:
|
|
|
raise ValueError('invalid_archive_members')
|
|
|
+ if meta['format'] == LIBRARY_FORMAT:
|
|
|
+ if meta.get('database') not in LIBRARIES or set(manifest['counts']) != {meta['database']}:
|
|
|
+ raise ValueError('invalid_library_archive')
|
|
|
if meta.get('snapshot') != manifest['snapshot'] or meta.get('counts') != manifest['counts']:
|
|
|
raise ValueError('archive_manifest_mismatch')
|
|
|
for name in expected:
|
|
|
@@ -188,7 +225,16 @@ def verify_archive(path):
|
|
|
source_records = []
|
|
|
for source_kind in KINDS:
|
|
|
with archive.open(f'prepared/{source_kind}.with_summary.jsonl') as source:
|
|
|
- source_records.extend(json.loads(line) for line in source)
|
|
|
+ enhanced = [json.loads(line) for line in source]
|
|
|
+ with archive.open(f'prepared/{source_kind}.jsonl') as source:
|
|
|
+ original = [json.loads(line) for line in source]
|
|
|
+ if len(original) != len(enhanced) or any(
|
|
|
+ r['kind'] != source_kind or any(k not in r or r[k] != v for k, v in raw.items())
|
|
|
+ for raw, r in zip(original, enhanced, strict=True)):
|
|
|
+ raise ValueError('archive_original_mismatch')
|
|
|
+ source_records.extend(enhanced)
|
|
|
+ if len(source_records) != len(records) or len({r['id'] for r in source_records}) != len(source_records):
|
|
|
+ raise ValueError('archive_prepared_mismatch')
|
|
|
if {r['id']: r for r in source_records} != {r['id']: r for r in records}:
|
|
|
raise ValueError('archive_prepared_mismatch')
|
|
|
_snapshot(source_records, manifest)
|
|
|
@@ -197,23 +243,109 @@ def verify_archive(path):
|
|
|
return manifest, len(records)
|
|
|
|
|
|
|
|
|
-def import_archive(path, prepared, settings, store_factory=Store):
|
|
|
- manifest, total = verify_archive(path)
|
|
|
- if not model_matches_manifest(settings['EMBEDDING_MODEL_PATH'], manifest):
|
|
|
- raise ValueError('target_model_mismatch')
|
|
|
- target_manifest = prepared / 'ingestion/manifest.json'
|
|
|
- if target_manifest.exists():
|
|
|
- existing = json.loads(target_manifest.read_text(encoding='utf-8'))
|
|
|
- if existing.get('complete') and existing.get('snapshot') != manifest['snapshot']:
|
|
|
- raise ValueError('target_has_different_snapshot')
|
|
|
- with ZipFile(path) as archive:
|
|
|
- for name in SOURCE_FILES:
|
|
|
- dest = prepared / name
|
|
|
- if dest.exists():
|
|
|
- with dest.open('rb') as source:
|
|
|
- existing_hash = hashlib.file_digest(source, 'sha256').hexdigest()
|
|
|
- if existing_hash != _digest(archive, 'prepared/' + name):
|
|
|
- raise ValueError('target_prepared_conflict')
|
|
|
+def _stage_sources(archive, stage, prepared, manifest, database, settings):
|
|
|
+ """Assemble and validate all local records before any remote writes."""
|
|
|
+ existing = None
|
|
|
+ if database:
|
|
|
+ path = prepared / 'ingestion/manifest.json'
|
|
|
+ if not path.is_file():
|
|
|
+ raise ValueError('single_library_requires_existing_snapshot')
|
|
|
+ existing = json.loads(path.read_text(encoding='utf-8'))
|
|
|
+ _check_manifest(existing)
|
|
|
+ _snapshot(load(prepared), existing)
|
|
|
+ if not model_matches_manifest(settings['EMBEDDING_MODEL_PATH'], existing):
|
|
|
+ raise ValueError('existing_model_mismatch')
|
|
|
+ # Both manifests must describe identical vector spaces; OS path spelling
|
|
|
+ # in a model fingerprint is not a different model.
|
|
|
+ normalize = lambda m: {k.replace('\\', '/'): v for k, v in m.get('model_files', {}).items()}
|
|
|
+ if existing['model_hash'] != manifest['model_hash'] and normalize(existing) != normalize(manifest):
|
|
|
+ raise ValueError('library_model_mismatch')
|
|
|
+ for kind, spec in existing['databases'].items():
|
|
|
+ if spec['database'] != settings['MILVUS_DB_PREFIX'] + '_' + kind:
|
|
|
+ raise ValueError('existing_database_mismatch')
|
|
|
+ for name in SOURCE_FILES:
|
|
|
+ content = archive.read('prepared/' + name)
|
|
|
+ if database:
|
|
|
+ # Pair IDs with the validated enhanced rows, rather than trusting
|
|
|
+ # raw input fields to implement the routing rules.
|
|
|
+ enhanced_name = name.split('.')[0] + '.with_summary.jsonl'
|
|
|
+ kept = {r['id'] for r in map(json.loads, (prepared / enhanced_name).read_text(encoding='utf-8').splitlines())
|
|
|
+ if route(r) != database}
|
|
|
+ old_lines = (prepared / name).read_bytes().splitlines(keepends=True)
|
|
|
+ prefix = b''.join(line.rstrip(b'\r\n') + b'\n' for line in old_lines if json.loads(line)['id'] in kept)
|
|
|
+ content = prefix + content
|
|
|
+ (stage / name).write_bytes(content)
|
|
|
+ combined = load(stage) # Includes global duplicate-ID and raw/enhanced checks.
|
|
|
+ result = copy.deepcopy(existing or manifest)
|
|
|
+ result['snapshot'] = sha(canonical({'contract': CONTRACT, 'model': result['model_hash'], 'records': combined}))
|
|
|
+ result['counts'] = dict(Counter(route(r) for r in combined))
|
|
|
+ if existing:
|
|
|
+ for spec in result['databases'].values():
|
|
|
+ spec.setdefault('snapshot', existing['snapshot'])
|
|
|
+ for kind, spec in manifest['databases'].items():
|
|
|
+ result['databases'][kind] = {**spec, 'database': settings['MILVUS_DB_PREFIX'] + '_' + kind,
|
|
|
+ 'snapshot': spec.get('snapshot', manifest['snapshot'])}
|
|
|
+ _check_manifest(result)
|
|
|
+ _snapshot(combined, result)
|
|
|
+ return result
|
|
|
+
|
|
|
+
|
|
|
+def _publish(stage, prepared):
|
|
|
+ """Same-filesystem directory switch, with the prior snapshot kept for recovery."""
|
|
|
+ parent = prepared.parent.resolve()
|
|
|
+ backup = parent / (prepared.name + '.previous-' + uuid.uuid4().hex)
|
|
|
+ # Never move a symlink or a computed path outside the selected workspace.
|
|
|
+ if prepared.is_symlink() or stage.is_symlink() or not stage.resolve().is_relative_to(parent):
|
|
|
+ raise ValueError('unsafe_publication_path')
|
|
|
+ had_previous = prepared.exists()
|
|
|
+ if had_previous:
|
|
|
+ prepared.rename(backup)
|
|
|
+ try:
|
|
|
+ stage.rename(prepared)
|
|
|
+ except BaseException:
|
|
|
+ if had_previous:
|
|
|
+ backup.rename(prepared)
|
|
|
+ raise
|
|
|
+ return backup if had_previous else None
|
|
|
+
|
|
|
+
|
|
|
+def import_archive(path, prepared, settings, store_factory=Store, database=None):
|
|
|
+ if database is not None and database not in LIBRARIES:
|
|
|
+ raise ValueError('unknown_library')
|
|
|
+ prepared = Path(prepared).absolute()
|
|
|
+ if prepared.is_symlink():
|
|
|
+ raise ValueError('unsafe_publication_path')
|
|
|
+ prepared.parent.mkdir(parents=True, exist_ok=True)
|
|
|
+ with publication_lock(prepared), tempfile.TemporaryDirectory(prefix='.' + prepared.name + '-import-', dir=prepared.parent) as work:
|
|
|
+ work = Path(work)
|
|
|
+ # Pin the exact bytes being verified and imported, even if the caller
|
|
|
+ # replaces the original ZIP while this operation is running.
|
|
|
+ pinned = work / 'source.zip'
|
|
|
+ shutil.copyfile(path, pinned)
|
|
|
+ manifest, total = verify_archive(pinned)
|
|
|
+ if not model_matches_manifest(settings['EMBEDDING_MODEL_PATH'], manifest):
|
|
|
+ raise ValueError('target_model_mismatch')
|
|
|
+ stage = work / 'prepared'
|
|
|
+ stage.mkdir()
|
|
|
+ with ZipFile(pinned) as archive:
|
|
|
+ meta = _json_entry(archive, ARCHIVE_META)
|
|
|
+ if database is None and meta['format'] != FORMAT:
|
|
|
+ raise ValueError('single_library_requires_database')
|
|
|
+ if database and (meta['format'] != LIBRARY_FORMAT or meta.get('database') != database):
|
|
|
+ raise ValueError('library_archive_mismatch')
|
|
|
+ result = _stage_sources(archive, stage, prepared, manifest, database, settings)
|
|
|
+ _restore(pinned, manifest, settings, store_factory)
|
|
|
+ result['verified_at'] = utc_now()
|
|
|
+ (stage / 'ingestion').mkdir()
|
|
|
+ atomic_write(stage / 'ingestion/manifest.json', json.dumps(result, ensure_ascii=False, indent=2) + '\n')
|
|
|
+ backup = _publish(stage, prepared)
|
|
|
+ print(json.dumps({'published': result['snapshot'], 'imported_library': database,
|
|
|
+ 'total_records': sum(result['counts'].values()),
|
|
|
+ 'backup': str(backup) if backup else None}), flush=True)
|
|
|
+ return result['snapshot'], total
|
|
|
+
|
|
|
+
|
|
|
+def _restore(path, manifest, settings, store_factory):
|
|
|
store = store_factory(settings)
|
|
|
try:
|
|
|
with ZipFile(path) as archive:
|
|
|
@@ -265,19 +397,8 @@ def import_archive(path, prepared, settings, store_factory=Store):
|
|
|
raise ValueError('target_vector_mismatch')
|
|
|
manifest['databases'][kind]['database'] = database
|
|
|
print(json.dumps({'type': kind, 'restored': len(records)}), flush=True)
|
|
|
- for name in SOURCE_FILES:
|
|
|
- dest = prepared / name
|
|
|
- dest.parent.mkdir(parents=True, exist_ok=True)
|
|
|
- if not dest.exists():
|
|
|
- temp = dest.with_name(dest.name + '.partial')
|
|
|
- temp.write_bytes(archive.read('prepared/' + name))
|
|
|
- temp.replace(dest)
|
|
|
- manifest['verified_at'] = utc_now()
|
|
|
- target_manifest.parent.mkdir(parents=True, exist_ok=True)
|
|
|
- atomic_write(target_manifest, json.dumps(manifest, ensure_ascii=False, indent=2) + '\n')
|
|
|
finally:
|
|
|
store.client.close()
|
|
|
- return manifest['snapshot'], total
|
|
|
|
|
|
|
|
|
def main():
|
|
|
@@ -286,7 +407,10 @@ def main():
|
|
|
parser.add_argument('archive', type=Path)
|
|
|
parser.add_argument('--prepared', type=Path, default=Path(__file__).resolve().parents[2] / 'knowledge/prepared')
|
|
|
parser.add_argument('--env-file', type=Path, default=Path(__file__).resolve().parents[2] / '.env')
|
|
|
+ parser.add_argument('--database', choices=LIBRARIES, help='Export/import one complete library; omitted means a full snapshot')
|
|
|
args = parser.parse_args()
|
|
|
+ if args.mode == 'verify' and args.database:
|
|
|
+ parser.error('verify detects the package type automatically; do not pass --database')
|
|
|
try:
|
|
|
if args.mode == 'verify':
|
|
|
snapshot, total = verify_archive(args.archive)
|
|
|
@@ -294,7 +418,7 @@ def main():
|
|
|
else:
|
|
|
settings = config(args.env_file)
|
|
|
operation = export if args.mode == 'export' else import_archive
|
|
|
- snapshot, total = operation(args.prepared, args.archive, settings) if args.mode == 'export' else operation(args.archive, args.prepared, settings)
|
|
|
+ snapshot, total = operation(args.prepared, args.archive, settings, database=args.database) if args.mode == 'export' else operation(args.archive, args.prepared, settings, database=args.database)
|
|
|
print(json.dumps({'mode': args.mode, 'snapshot': snapshot, 'records': total, 'complete': True}), flush=True)
|
|
|
return 0
|
|
|
except Exception as exc:
|
