business-assistant/src/step2_knowledge_ingest/ingest.py

"""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