business-assistant/src/step3_question_answer/retrieval.py

"""Read-only, manifest-pinned, cross-database retrieval with verified full records."""

import json
import math
import re
import logging
import os
from concurrent.futures import ThreadPoolExecutor
from threading import RLock
from .timing import slot, measured, mapped

from src.step2_knowledge_ingest.ingest import (
    CONTRACT, DIM, TYPES, Embedder, canonical, config, model_matches_manifest, restore, sha,
)

SEARCH_TYPES = tuple(kind for kind in TYPES if kind != 'catalog')
PER_BANK_CANDIDATES = 3
MAX_CANDIDATES = 12
MAX_RELATED_GUIDES = 8
RELATED_GUIDE_SEARCH_LIMIT = 8


def related_guide_queries(question, records):
    """Expand an explicit list request from named sibling scenarios, never from prose."""
    if not re.search(r'列举|几个|哪些|所有|全部|一览|清单', question) or not re.search(
            r'办理事项|办事事项|办事指南|办理情形', question):
        return []
    groups, present = {}, set()
    for row in records:
        if row['type'] != 'service_guide':
            continue
        fields = (row['record'].get('data') or {}).get('structured_content') or {}
        if not isinstance(fields, dict):
            continue
        matter, district, scenario = (fields.get(key) for key in ('事项名称', '区划名称', '情形名称'))
        if not all(isinstance(value, str) and value.strip() for value in (matter, district, scenario)):
            continue
        group = (matter.strip(), district.strip())
        present.add((*group, scenario.strip()))
        basic = fields.get('基本信息')
        related = basic.get('其他业务办理项') if isinstance(basic, dict) else None
        names = related if isinstance(related, list) else related.split() if isinstance(related, str) else []
        for name in names:
            if isinstance(name, str) and 4 <= len(name.strip()) <= 120:
                groups.setdefault((*group, name.strip()), None)
    return [key for key in groups if key not in present][:MAX_RELATED_GUIDES]


def choose_candidates(hits_by_type, query_hits, limit=MAX_CANDIDATES):
    """Keep bank coverage, then each query's best hit before score-only filling."""
    chosen, seen = [], set()
    def add(kind, hit):
        key = (kind, hit['id'])
        if len(chosen) < limit and key not in seen:
            chosen.append((kind, hit))
            seen.add(key)
    for kind, hits in hits_by_type.items():
        if hits:
            add(kind, hits[0])
    for hits in query_hits:
        if hits:
            kind, hit = max(hits, key=lambda pair: pair[1]['distance'])
            add(kind, hit)
    for kind, hit in sorted(((kind, hit) for kind, hits in hits_by_type.items() for hit in hits),
                            key=lambda pair: pair[1]['distance'], reverse=True):
        add(kind, hit)
    return chosen


def query_concurrency(path, environ=None, key='RETRIEVAL_QUERY_CONCURRENCY', default=4, maximum=64):
    value = str(default)
    if path.exists():
        for line in path.read_text(encoding='utf-8-sig').splitlines():
            line = line.strip().removeprefix('export ')
            if '=' in line and not line.startswith('#'):
                name, raw = line.split('=', 1)
                if name.strip() == key:
                    value = raw.split(' #', 1)[0].strip().strip('\"\'')
    value = (os.environ if environ is None else environ).get(key, value)
    if not re.fullmatch(r'[0-9]+', value) or not 1 <= int(value) <= maximum:
        raise ValueError('invalid_retrieval_query_concurrency')
    return int(value)


class RetrievalError(Exception):
    pass


def query_text(text):
    return '为这个句子生成表示以用于检索相关文章：' + text


class ReadStore:
    """One client pinned to one database; never switches the active database."""
    def __init__(self, settings, database):
        from pymilvus import MilvusClient
        logging.getLogger('pymilvus').setLevel(logging.CRITICAL)
        self.timeout = settings['MILVUS_TIMEOUT_SECONDS']
        self.client = MilvusClient(uri=settings['MILVUS_URI'], db_name=database, timeout=self.timeout)

    def get(self, name, ids):
        return self.client.get(name, ids=ids, output_fields=['*'], consistency_level='Strong', timeout=self.timeout)

    def close(self):
        self.client.close()


class Retriever:
    def __init__(self, env_path, manifest_path):
        self.query_workers = query_concurrency(env_path)
        self.bank_workers = query_concurrency(env_path, key='RETRIEVAL_BANK_CONCURRENCY', default=6, maximum=6)
        self.content_workers = query_concurrency(env_path, key='RETRIEVAL_CONTENT_CONCURRENCY', default=8, maximum=32)
        self.settings = config(env_path)
        self.manifest = json.loads(manifest_path.read_text(encoding='utf-8'))
        m = self.manifest
        if m.get('complete') is not True or m.get('contract') != CONTRACT or m.get('dimension') != DIM:
            raise RetrievalError('invalid_manifest')
        if not re.fullmatch('[a-f0-9]{64}', m.get('snapshot', '')):
            raise RetrievalError('invalid_snapshot')
        for kind in SEARCH_TYPES:
            spec = m['databases'][kind]
            if spec['database'] != self.settings['MILVUS_DB_PREFIX'] + '_' + kind:
                raise RetrievalError('database_mismatch')
            version = spec.get('snapshot', m['snapshot'])
            if not isinstance(version, str) or not re.fullmatch('[a-f0-9]{64}', version):
                raise RetrievalError('invalid_library_snapshot')
            for key in ('content_collection', 'index_collection'):
                expected = ('content_' if key == 'content_collection' else 'index_') + version[:20]
                if spec[key] != expected:
                    raise RetrievalError('collection_mismatch')
        self.embedder = None
        self._stores = {}
        self._warmed = False
        self.lock = RLock()

    def _initialize_model(self):
        if self.embedder is None:
            model_path = self.settings['EMBEDDING_MODEL_PATH']
            if not model_matches_manifest(model_path, self.manifest):
                raise RetrievalError('embedding_model_changed')
            self.embedder = Embedder(self.settings)

    def warmup(self):
        """Load and run one local inference before accepting requests; safe to repeat."""
        with self.lock:
            if self._warmed:
                return
            try:
                self._initialize_model()
                self.embedder.encode([query_text('政策与公共服务查询')])
            except RetrievalError:
                raise
            except Exception:
                raise RetrievalError('embedding_warmup_failed') from None
            self._warmed = True

    def _initialize(self):
        self._initialize_model()

    def _store_for(self, kind):
        # One pinned client per bank, shared only for read-only operations.
        with self.lock:
            if kind not in self._stores:
                self._stores[kind] = ReadStore(self.settings, self.manifest['databases'][kind]['database'])
            return self._stores[kind]

    def close(self):
        with self.lock:
            for store in self._stores.values():
                store.close()
            self._stores.clear()

    def search(self, question, expansions=()):
        # GPU inference is serialized, but network retrieval can overlap across users.
        # Hosts must drain searches before close(), and bound concurrent requests.
        try:
            return self._search(question, expansions)
        except RetrievalError:
            raise
        except Exception:
            raise RetrievalError('retrieval_service_failure') from None

    def _search(self, question, expansions):
        if not isinstance(question, str) or not question.strip() or len(question) > 8000:
            raise RetrievalError('invalid_question')
        queries = [question]
        for text in expansions:
            if isinstance(text, str) and text.strip() and text not in queries and len(text) <= 500:
                queries.append(text)
        with slot(self.lock, 'embedding'):
            self._initialize()
            vectors = self.embedder.encode([query_text(text) for text in queries])
        # Initialize each pinned client before query tasks share its read-only searches.
        unavailable = set()
        for kind in SEARCH_TYPES:
            try:
                self._store_for(kind)
            except Exception:
                unavailable.add(kind)

        def recall(kind, vector):
            spec = self.manifest['databases'][kind]
            try:
                if kind in unavailable:
                    raise RetrievalError('database_unavailable')
                store = self._stores[kind]
                groups = store.client.search(spec['index_collection'], data=[vector],
                    limit=PER_BANK_CANDIDATES, output_fields=['record_hash'],
                    search_params={'metric_type': 'COSINE', 'params': {'ef': 64}},
                    consistency_level='Strong', timeout=store.timeout)
                unique = {}
                for group in groups:
                    for hit in group:
                        record_id = hit['id']
                        if not re.fullmatch('k_[a-f0-9]{24}', record_id) or not math.isfinite(hit['distance']):
                            raise RetrievalError('invalid_search_hit')
                        previous = unique.get(record_id)
                        if previous is None or hit['distance'] > previous['distance']:
                            unique[record_id] = hit
                return kind, sorted(unique.values(), key=lambda hit: hit['distance'], reverse=True), False
            except Exception:
                return kind, [], True
        @measured('vector_search')
        def recall_query(index):
            with ThreadPoolExecutor(max_workers=self.bank_workers) as banks:
                return index, list(banks.map(lambda kind: recall(kind, vectors[index]), SEARCH_TYPES))

        merged = {kind: {} for kind in SEARCH_TYPES}
        failed_queries = []
        query_hits = [[] for _ in queries]
        successful_searches = 0
        with ThreadPoolExecutor(max_workers=min(self.query_workers, len(queries))) as pool:
            # Every query is scheduled; map preserves query order for deterministic merging.
            for index, rows in mapped(pool, recall_query, range(len(queries))):
                for kind, hits, failed in rows:
                    if failed:
                        failed_queries.append({'query_index': index, 'type': kind})
                        continue
                    successful_searches += 1
                    query_hits[index].extend((kind, hit) for hit in hits)
                    for hit in hits:
                        previous = merged[kind].get(hit['id'])
                        if previous is None or hit['distance'] > previous['distance']:
                            merged[kind][hit['id']] = hit
        failures = [kind for kind in SEARCH_TYPES if any(row['type'] == kind for row in failed_queries)]
        hits_by_type = {kind: sorted(hits.values(), key=lambda hit: hit['distance'], reverse=True)
                        for kind, hits in merged.items()}
        if not successful_searches:
            raise RetrievalError('all_databases_failed')
        # Preserve cross-bank coverage before filling the original 12-candidate budget.
        chosen = choose_candidates(hits_by_type, query_hits)
        query_best = [max(hits, key=lambda pair: pair[1]['distance']) if hits else None for hits in query_hits]
        selected_by_type = {kind: [hit for selected_kind, hit in chosen if selected_kind == kind]
                            for kind in SEARCH_TYPES}
        # Batch all segments, including long records, then verify each full record.
        segments = {}
        @measured('content_read')
        def fetch_segments(kind, ids, bank_segments):
            for start in range(0, len(ids), 128):
                batch_ids = ids[start:start + 128]
                try:
                    rows = self._store_for(kind).get(self.manifest['databases'][kind]['content_collection'], batch_ids)
                    for row in rows:
                        if row.get('id') in batch_ids:
                            bank_segments.setdefault((kind, row['id']), []).append(row)
                except Exception:
                    # Missing segments become explicit per-record failures below.
                    continue
        def fetch_bank(kind):
            # Each worker owns its segments; first segments must arrive before
            # discovering the rest of that bank's records.
            bank_segments = {}
            fetch_segments(kind, [hit['id'] + ':0' for hit in selected_by_type[kind]], bank_segments)
            remaining = []
            for hit in selected_by_type[kind]:
                first = bank_segments.get((kind, hit['id'] + ':0'), [])
                if len(first) == 1 and type(first[0].get('total')) is int and 1 <= first[0]['total'] <= 10000:
                    remaining.extend(f"{hit['id']}:{i}" for i in range(1, first[0]['total']))
            fetch_segments(kind, remaining, bank_segments)
            return bank_segments
        with ThreadPoolExecutor(max_workers=self.content_workers) as pool:
            for bank_segments in mapped(pool, fetch_bank, (kind for kind in SEARCH_TYPES if selected_by_type[kind])):
                segments.update(bank_segments)
        records, content_failures = [], []
        for kind, hit in chosen:
            spec = self.manifest['databases'][kind]
            try:
                first = segments.get((kind, hit['id'] + ':0'), [])
                if len(first) != 1 or type(first[0]['total']) is not int or not 1 <= first[0]['total'] <= 10000:
                    raise RetrievalError('missing_content')
                parts = [part for i in range(first[0]['total'])
                         for part in segments.get((kind, f"{hit['id']}:{i}"), [])]
                record = restore(parts)
                if record['id'] != hit['id'] or sha(canonical(record)) != hit['entity']['record_hash']:
                    raise RetrievalError('index_content_mismatch')
                records.append({'record': record, 'type': kind, 'score': hit['distance'],
                                'best_query_indices': [i for i, best in enumerate(query_best)
                                                       if best and best[0] == kind and best[1]['id'] == hit['id']],
                                'snapshot': self.manifest['snapshot']})
            except Exception:
                content_failures.append(hit['id'])
        related = related_guide_queries(question, records)
        related_guide_failures = 0
        if related and 'service_guide' not in unavailable:
            spec = self.manifest['databases']['service_guide']
            store = self._stores['service_guide']
            with slot(self.lock, 'embedding'):
                vectors = self.embedder.encode([query_text(name) for _, _, name in related])
            known = {row['record']['id'] for row in records}
            for (matter, district, scenario), vector in zip(related, vectors):
                try:
                    found = False
                    groups = store.client.search(spec['index_collection'], data=[vector],
                        limit=RELATED_GUIDE_SEARCH_LIMIT, output_fields=['record_hash'],
                        search_params={'metric_type': 'COSINE', 'params': {'ef': 64}},
                        consistency_level='Strong', timeout=store.timeout)
                    for group in groups:
                        for hit in group:
                            record_id = hit['id']
                            if record_id in known or not re.fullmatch('k_[a-f0-9]{24}', record_id):
                                continue
                            first = store.get(spec['content_collection'], [record_id + ':0'])
                            if (len(first) != 1 or first[0].get('id') != record_id + ':0'
                                    or type(first[0].get('total')) is not int or not 1 <= first[0]['total'] <= 10000):
                                continue
                            ids = [f'{record_id}:{i}' for i in range(1, first[0]['total'])]
                            later = {part.get('id'): part for start in range(0, len(ids), 128)
                                     for part in store.get(spec['content_collection'], ids[start:start + 128])}
                            if any(identity not in later for identity in ids):
                                continue
                            parts = first + [later[identity] for identity in ids]
                            record = restore(parts)
                            fields = (record.get('data') or {}).get('structured_content') or {}
                            if (record.get('id') == record_id and sha(canonical(record)) == hit['entity']['record_hash']
                                    and isinstance(fields, dict) and fields.get('事项名称') == matter
                                    and fields.get('区划名称') == district and fields.get('情形名称') == scenario):
                                records.append({'record': record, 'type': 'service_guide', 'score': hit['distance'],
                                                'best_query_indices': [], 'snapshot': self.manifest['snapshot']})
                                known.add(record_id)
                                found = True
                                break
                        if found:
                            break
                except Exception:
                    # Supplementary lookup must not hide the original verified results.
                    related_guide_failures += 1
                    continue
        return {'candidates': records, 'queries': queries, 'failed_databases': failures,
                'failed_queries': failed_queries,
                'content_failures': content_failures, 'related_guide_failures': related_guide_failures,
                'complete': not failures and not content_failures and not related_guide_failures,
                'snapshot': self.manifest['snapshot'], 'searched_types': list(SEARCH_TYPES),
                'candidate_limit': MAX_CANDIDATES + (MAX_RELATED_GUIDES if related else 0),
                'base_candidate_limit': MAX_CANDIDATES,
                'per_type_limit': PER_BANK_CANDIDATES,
                'related_guide_limit': MAX_RELATED_GUIDES}