"""Typed, versioned ingestion. Full records and index text have separate storage.""" import argparse from collections import Counter from contextlib import contextmanager import hashlib import json import logging import math import os from pathlib import Path import re import sqlite3 import sys from src.step1_knowledge_prepare.prepare import ROOT, strict_json from src.step1_knowledge_prepare.summary_input import scenario from src.step1_knowledge_prepare.summarize import atomic_write, run_lock, utc_now from src.knowledge_profiles import AGENTS, check_records, for_agent, prepared_path KINDS = ('policy', 'finance', 'general', 'qa') TYPES = ('policy', 'finance', 'qa', 'article', 'service_guide', 'service_catalog', 'catalog') DEFAULTS = { 'MILVUS_URI': 'http://192.168.2.69:19530', 'MILVUS_DB_PREFIX': 'business_knowledge', 'TALENT_MILVUS_DB_PREFIX': 'talent_knowledge', 'MILVUS_TIMEOUT_SECONDS': '60', 'EMBEDDING_MODEL_PATH': 'models/bge-large-zh-v1.5', 'EMBEDDING_DEVICE': 'cpu', 'EMBEDDING_BATCH_SIZE': '8', 'EMBEDDING_CPU_THREADS': '16', } DIM = 1024 CONTRACT = 'title-summary-bge-zh-cls-l2-fp32-token-weighted-windows-v1' def canonical(value): return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(',', ':'), allow_nan=False) def sha(value): return hashlib.sha256(value.encode('utf-8')).hexdigest() def file_sha(path): with path.open('rb') as stream: return hashlib.file_digest(stream, 'sha256').hexdigest() def model_fingerprint(model_dir): files = {str(p.relative_to(model_dir)): file_sha(p) for p in sorted(model_dir.rglob('*')) if p.is_file() and p.suffix in ('.json', '.safetensors', '.bin', '.txt')} if not ({'model.safetensors', 'pytorch_model.bin'} & files.keys()) or 'tokenizer.json' not in files: raise ValueError('missing_local_model') return files def model_matches_manifest(model_dir, manifest): """Accept path separator differences while requiring identical model files.""" actual = model_fingerprint(model_dir) if sha(canonical(actual)) == manifest['model_hash']: return True expected = manifest.get('model_files') if not isinstance(expected, dict) or sha(canonical(expected)) != manifest['model_hash']: return False normalize = lambda files: {name.replace('\\', '/'): digest for name, digest in files.items()} return len(normalize(actual)) == len(actual) and len(normalize(expected)) == len(expected) and normalize(actual) == normalize(expected) def config(path, environ=None): values = dict(DEFAULTS) if path.exists(): for line in path.read_text(encoding='utf-8-sig').splitlines(): line = line.strip().removeprefix('export ') if '=' not in line or line.startswith('#'): continue key, value = (item.strip() for item in line.split('=', 1)) if key not in DEFAULTS: continue if value.startswith(('"', "'")): end = value.find(value[0], 1) if end < 0 or value[end + 1:].strip() and not value[end + 1:].strip().startswith('#'): raise ValueError('invalid_ingest_config_quote') value = value[1:end] else: value = value.split(' #', 1)[0].strip() values[key] = value env = os.environ if environ is None else environ values.update({key: env[key] for key in DEFAULTS if key in env}) if not re.fullmatch(r'[A-Za-z][A-Za-z0-9_]{0,99}', values['MILVUS_DB_PREFIX']): raise ValueError('invalid_database_prefix') for key in ('MILVUS_TIMEOUT_SECONDS', 'EMBEDDING_BATCH_SIZE', 'EMBEDDING_CPU_THREADS'): values[key] = int(values[key]) if values[key] <= 0: raise ValueError('nonpositive_ingest_config') model = Path(values['EMBEDDING_MODEL_PATH']) values['EMBEDDING_MODEL_PATH'] = model if model.is_absolute() else ROOT / model return values def route(record): if record['kind'] == 'qa': return 'qa' return scenario(record) or 'catalog' def index_text(record): kind = route(record) if kind == 'catalog': return None if kind == 'qa': value = record['data']['question'] if record.get('summary_status') != 'not_required_qa': raise ValueError('qa_status_mismatch') else: if record.get('summary_status') != 'generated' or not isinstance(record.get('summary'), str) or not record['summary'].strip(): raise ValueError('missing_complete_summary') value = record['title'] + '\n' + record['summary'] if not isinstance(value, str) or not value.strip(): raise ValueError('empty_index_text') if record.get('retrieval_text') != value: raise ValueError('retrieval_text_mismatch') if len(value.encode('utf-8')) > 65535: raise ValueError('index_text_exceeds_milvus_capacity_no_truncation') return value def load(prepared): records, seen = [], set() for kind in KINDS: original = [strict_json(line) for line in (prepared / f'{kind}.jsonl').read_text(encoding='utf-8').splitlines()] enhanced = [strict_json(line) for line in (prepared / f'{kind}.with_summary.jsonl').read_text(encoding='utf-8').splitlines()] if len(original) != len(enhanced): raise ValueError('source_count_mismatch') for source, record in zip(original, enhanced, strict=True): if any(key not in record or record[key] != value for key, value in source.items()): raise ValueError('original_record_changed') if record['kind'] != kind or record['id'] in seen or not re.fullmatch(r'k_[a-f0-9]{24}', record['id']): raise ValueError('invalid_id_or_kind') seen.add(record['id']) index_text(record) records.append(record) if not records: raise ValueError('empty_knowledge') return records def content_parts(record): raw = canonical(record) # 12,000 Unicode characters occupy at most 48,000 UTF-8 bytes. parts = [raw[i:i + 12000] for i in range(0, len(raw), 12000)] return [{'id': f"{record['id']}:{i}", 'record_id': record['id'], 'part': i, 'total': len(parts), 'payload': part, 'record_hash': sha(raw), 'vector': [1.0, 0.0]} for i, part in enumerate(parts)] def restore(parts): if not parts: raise ValueError('missing_content') parts = sorted(parts, key=lambda part: part['part']) if any(p['part'] != i or p['total'] != len(parts) or p['record_id'] != parts[0]['record_id'] or p['record_hash'] != parts[0]['record_hash'] for i, p in enumerate(parts)): raise ValueError('incomplete_content_parts') raw = ''.join(part['payload'] for part in parts) if sha(raw) != parts[0]['record_hash']: raise ValueError('content_hash_mismatch') return strict_json(raw) def valid_vector(vector): return len(vector) == DIM and all(math.isfinite(v) for v in vector) and abs(sum(v * v for v in vector) - 1) < 0.001 class Embedder: def __init__(self, settings): import torch from transformers import AutoModel, AutoTokenizer torch.set_num_threads(settings['EMBEDDING_CPU_THREADS']) self.torch = torch path = str(settings['EMBEDDING_MODEL_PATH']) self.tokenizer = AutoTokenizer.from_pretrained(path, local_files_only=True, padding_side='right') self.model = AutoModel.from_pretrained(path, local_files_only=True, trust_remote_code=False, dtype=torch.float32).to(settings['EMBEDDING_DEVICE']).eval() self.limit = min(512, self.model.config.max_position_embeddings) self.batch_size = settings['EMBEDDING_BATCH_SIZE'] if self.model.config.hidden_size != DIM or self.model.config.model_type != 'bert': raise ValueError('model_dimension_mismatch') def preflight(self, texts): lengths = [len(self.tokenizer(text, truncation=False)['input_ids']) for text in texts] return max(lengths, default=0) def encode(self, texts): # Preserve every token. Long index/query texts use contiguous 510-token # windows; aggregate normalized CLS vectors by payload-token count. windows, owners, weights = [], [], [] for owner, text in enumerate(texts): ids = self.tokenizer(text, add_special_tokens=False, truncation=False)['input_ids'] for start in range(0, max(1, len(ids)), self.limit - 2): part = ids[start:start + self.limit - 2] windows.append({'input_ids': [self.tokenizer.cls_token_id] + part + [self.tokenizer.sep_token_id]}) owners.append(owner) weights.append(max(1, len(part))) totals = self.torch.zeros((len(texts), DIM), dtype=self.torch.float32) for start in range(0, len(windows), self.batch_size): inputs = self.tokenizer.pad(windows[start:start + self.batch_size], padding=True, return_tensors='pt').to(self.model.device) with self.torch.inference_mode(): hidden = self.model(**inputs).last_hidden_state[:, 0] chunks = self.torch.nn.functional.normalize(hidden.float(), p=2, dim=1).cpu() for offset, vector in enumerate(chunks): i = start + offset totals[owners[i]] += vector * weights[i] vectors = self.torch.nn.functional.normalize(totals, p=2, dim=1).tolist() if len(vectors) != len(texts) or not all(valid_vector(v) for v in vectors): raise ValueError('invalid_embedding') return vectors class Store: def __init__(self, settings): from pymilvus import MilvusClient logging.getLogger('pymilvus').setLevel(logging.CRITICAL) self.settings = settings self.timeout = settings['MILVUS_TIMEOUT_SECONDS'] self.client = MilvusClient(uri=settings['MILVUS_URI'], timeout=settings['MILVUS_TIMEOUT_SECONDS']) def use(self, kind, create): database = self.settings['MILVUS_DB_PREFIX'] + '_' + kind if database not in self.client.list_databases(timeout=self.timeout): if not create: raise ValueError('missing_database') self.client.create_database(database, timeout=self.timeout) self.client.use_database(database, timeout=self.timeout) return database def ensure(self, name, content=False): from pymilvus import DataType if self.client.has_collection(name, timeout=self.timeout): info = self.client.describe_collection(name, timeout=self.timeout) if info.get('description') != CONTRACT: raise ValueError('existing_collection_contract_mismatch') return schema = self.client.create_schema(auto_id=False, enable_dynamic_field=False, description=CONTRACT) schema.add_field('id', DataType.VARCHAR, is_primary=True, max_length=128) schema.add_field('vector', DataType.FLOAT_VECTOR, dim=2 if content else DIM) schema.add_field('record_hash', DataType.VARCHAR, max_length=64) if content: schema.add_field('record_id', DataType.VARCHAR, max_length=128) schema.add_field('part', DataType.INT64) schema.add_field('total', DataType.INT64) schema.add_field('payload', DataType.VARCHAR, max_length=65535) else: schema.add_field('index_text', DataType.VARCHAR, max_length=65535) indexes = self.client.prepare_index_params() indexes.add_index(field_name='vector', index_type='FLAT' if content else 'HNSW', metric_type='COSINE', params={} if content else {'M': 16, 'efConstruction': 200}) self.client.create_collection(name, schema=schema, index_params=indexes, consistency_level='Strong', timeout=self.timeout) def get(self, name, ids): return self.client.get(name, ids=ids, output_fields=['*'], consistency_level='Strong', timeout=self.timeout) def verify(self, records, content_name, index_name): for start in range(0, len(records), 32): batch = records[start:start + 32] ids = [part['id'] for record in batch for part in content_parts(record)] grouped = {} for offset in range(0, len(ids), 128): for part in self.get(content_name, ids[offset:offset + 128]): grouped.setdefault(part['record_id'], []).append(part) indexed_ids = [r['id'] for r in batch if index_text(r) is not None] rows = {row['id']: row for row in self.get(index_name, indexed_ids)} if indexed_ids else {} for record in batch: restored = restore(grouped.get(record['id'], [])) if restored != record: raise ValueError('remote_content_mismatch') text = index_text(record) if text is not None: row = rows.get(record['id']) if row is None or row['record_hash'] != sha(canonical(record)) or row['index_text'] != text or not valid_vector(row['vector']): raise ValueError('remote_index_mismatch') for name, expected in ((content_name, sum(len(content_parts(r)) for r in records)), (index_name, sum(index_text(r) is not None for r in records))): if expected: count = self.client.query(name, filter='', output_fields=['count(*)'], consistency_level='Strong', timeout=self.timeout)[0]['count(*)'] if count != expected: raise ValueError('remote_count_mismatch') indexed = next((r for r in records if index_text(r) is not None), None) if indexed: vector = self.get(index_name, [indexed['id']])[0]['vector'] hits = self.client.search(index_name, data=[vector], limit=1, search_params={'metric_type': 'COSINE', 'params': {'ef': 64}}, timeout=self.timeout)[0] if not hits or hits[0]['distance'] < 0.999: raise ValueError('vector_search_smoke_failed') @contextmanager def publication_lock(prepared): """Stable lock outside the directory that an import may replace.""" try: with run_lock(prepared.parent / ('.' + prepared.name + '.publish.lock')): yield except ValueError as exc: if str(exc) == 'another_summary_run_is_active': raise ValueError('another_knowledge_operation_is_active') from None raise def execute(args): with publication_lock(args.prepared): return _execute(args) def _execute(args): settings = for_agent(config(args.env_file), getattr(args, 'agent', 'business'), getattr(args, 'db_prefix', None)) records = load(args.prepared) check_records(records, settings) model_dir = settings['EMBEDDING_MODEL_PATH'] model_files = model_fingerprint(model_dir) model_hash = sha(canonical(model_files)) snapshot = sha(canonical({'contract': CONTRACT, 'model': model_hash, 'records': records})) content_name, index_name = 'content_' + snapshot[:20], 'index_' + snapshot[:20] counts = Counter(route(r) for r in records) published = None if args.verify_only: published = json.loads((args.prepared / 'ingestion/manifest.json').read_text(encoding='utf-8')) if (published.get('complete') is not True or published.get('contract') != CONTRACT or published.get('dimension') != DIM or published.get('counts') != dict(counts) or set(published.get('databases', {})) != set(counts) or not model_matches_manifest(model_dir, published)): raise ValueError('invalid_published_manifest') model_hash, model_files = published['model_hash'], published['model_files'] snapshot = sha(canonical({'contract': CONTRACT, 'model': model_hash, 'records': records})) if published.get('snapshot') != snapshot: raise ValueError('published_snapshot_mismatch') print(json.dumps({'counts': counts, 'snapshot': snapshot, 'mode': 'verify' if args.verify_only else 'dry-run' if args.dry_run else 'ingest'}), flush=True) if args.dry_run: return output = args.prepared / 'ingestion' output.mkdir(exist_ok=True) with run_lock(output / '.ingest.lock'): store = Store(settings) cache = sqlite3.connect(output / 'embedding_cache.sqlite3') try: cache.execute('CREATE TABLE IF NOT EXISTS vectors (key TEXT PRIMARY KEY, vector TEXT NOT NULL)') embedder = None if not args.verify_only: embedder = Embedder(settings) print(json.dumps({'max_index_tokens': embedder.preflight([index_text(r) for r in records if index_text(r) is not None])}), flush=True) manifest = {'complete': False, 'snapshot': snapshot, 'contract': CONTRACT, 'model_hash': model_hash, 'model_files': model_files, 'dimension': DIM, 'counts': dict(counts), 'databases': {}, 'verified_at': None} generated, reused = 0, 0 for kind in TYPES: subset = [r for r in records if route(r) == kind] if not subset: continue database = store.use(kind, create=not args.verify_only) if published: spec = published['databases'][kind] version = spec.get('snapshot', snapshot) if (not isinstance(version, str) or not re.fullmatch(r'[a-f0-9]{64}', version) or spec.get('database') != database or spec.get('content_collection') != 'content_' + version[:20] or spec.get('index_collection') != (None if kind == 'catalog' else 'index_' + version[:20])): raise ValueError('invalid_published_collection') content_name, index_name = spec['content_collection'], spec['index_collection'] manifest['databases'][kind] = {'database': database, 'content_collection': content_name, 'index_collection': index_name if kind != 'catalog' else None} if published and 'snapshot' in spec: manifest['databases'][kind]['snapshot'] = spec['snapshot'] if not args.verify_only: store.ensure(content_name, content=True) if kind != 'catalog': store.ensure(index_name) batch_size = settings['EMBEDDING_BATCH_SIZE'] for start in range(0, len(subset), batch_size): batch = subset[start:start + batch_size] texts = [index_text(r) for r in batch] keys = [sha(model_hash + CONTRACT + text) if text is not None else None for text in texts] vectors, missing = {}, [] for i, key in enumerate(keys): if key is None: continue saved = cache.execute('SELECT vector FROM vectors WHERE key=?', (key,)).fetchone() if saved and valid_vector(json.loads(saved[0])): vectors[i] = json.loads(saved[0]) reused += 1 else: missing.append(i) if missing: generated += len(missing) for i, vector in zip(missing, embedder.encode([texts[i] for i in missing]), strict=True): vectors[i] = vector cache.execute('INSERT OR REPLACE INTO vectors VALUES (?, ?)', (keys[i], canonical(vector))) cache.commit() parts = [part for r in batch for part in content_parts(r)] for offset in range(0, len(parts), 32): store.client.upsert(content_name, data=parts[offset:offset + 32], timeout=store.timeout) rows = [{'id': r['id'], 'vector': vectors[i], 'index_text': texts[i], 'record_hash': sha(canonical(r))} for i, r in enumerate(batch) if texts[i] is not None] if rows: store.client.upsert(index_name, data=rows, timeout=store.timeout) if start % (batch_size * 10) == 0 or start + batch_size >= len(subset): print(json.dumps({'type': kind, 'written': min(start + batch_size, len(subset)), 'total': len(subset)}), flush=True) store.client.flush(content_name, timeout=store.timeout) if kind != 'catalog': store.client.flush(index_name, timeout=store.timeout) store.verify(subset, content_name, index_name) print(json.dumps({'type': kind, 'verified': len(subset)}), flush=True) manifest['complete'] = True manifest['verified_at'] = utc_now() atomic_write(output / 'manifest.json', json.dumps(manifest, ensure_ascii=False, indent=2) + '\n') print(json.dumps({'complete': True, 'records_verified': len(records), 'vectors_generated': generated, 'vectors_reused': reused}), flush=True) finally: cache.close() store.client.close() def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument('--agent', choices=AGENTS, default='business') parser.add_argument('--db-prefix', help='Explicit target prefix; validated against the other agent') parser.add_argument('--prepared', type=Path) parser.add_argument('--env-file', type=Path, default=ROOT / '.env') mode = parser.add_mutually_exclusive_group() mode.add_argument('--dry-run', action='store_true') mode.add_argument('--verify-only', action='store_true') args = parser.parse_args() args.prepared = args.prepared or prepared_path(ROOT, args.agent) try: execute(args) return 0 except Exception as exc: # SDK exceptions may contain request payloads/addresses; do not echo them. code = str(exc) if isinstance(exc, ValueError) and re.fullmatch(r'[a-z_]+', str(exc)) else type(exc).__name__ print('knowledge_ingest_failed:' + code, file=sys.stderr) return 1