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