|
|
@@ -8,6 +8,7 @@ import os
|
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
|
from threading import RLock
|
|
|
from .timing import slot, measured, mapped
|
|
|
+from .title_recall import TITLE_LIMIT, title_queries, title_match, select_titles
|
|
|
|
|
|
from src.step2_knowledge_ingest.ingest import (
|
|
|
CONTRACT, DIM, TYPES, Embedder, canonical, config, model_matches_manifest, restore, sha,
|
|
|
@@ -46,7 +47,7 @@ def related_guide_queries(question, records):
|
|
|
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):
|
|
|
+def choose_candidates(hits_by_type, query_hits, limit=MAX_CANDIDATES, title_hits=()):
|
|
|
"""Keep bank coverage, then each query's best hit before score-only filling."""
|
|
|
chosen, seen = [], set()
|
|
|
def add(kind, hit):
|
|
|
@@ -54,6 +55,8 @@ def choose_candidates(hits_by_type, query_hits, limit=MAX_CANDIDATES):
|
|
|
if len(chosen) < limit and key not in seen:
|
|
|
chosen.append((kind, hit))
|
|
|
seen.add(key)
|
|
|
+ for kind, hit in title_hits:
|
|
|
+ add(kind, hit)
|
|
|
for kind, hits in hits_by_type.items():
|
|
|
if hits:
|
|
|
add(kind, hits[0])
|
|
|
@@ -101,6 +104,32 @@ class ReadStore:
|
|
|
def get(self, name, ids):
|
|
|
return self.client.get(name, ids=ids, output_fields=['*'], consistency_level='Strong', timeout=self.timeout)
|
|
|
|
|
|
+ def titles(self, name):
|
|
|
+ # Existing snapshots index title + newline + summary (QA indexes its question).
|
|
|
+ # Cache only title/hash/identity, never all summaries or document bodies.
|
|
|
+ iterator = self.client.query_iterator(collection_name=name, batch_size=256,
|
|
|
+ output_fields=['id', 'index_text', 'record_hash'],
|
|
|
+ consistency_level='Strong', timeout=self.timeout)
|
|
|
+ rows, seen = [], set()
|
|
|
+ try:
|
|
|
+ while batch := iterator.next():
|
|
|
+ for row in batch:
|
|
|
+ identity = row['id']
|
|
|
+ if (not isinstance(identity, str) or not re.fullmatch('k_[a-f0-9]{24}', identity)
|
|
|
+ or identity in seen or not isinstance(row['index_text'], str)
|
|
|
+ or not re.fullmatch('[a-f0-9]{64}', row['record_hash'])):
|
|
|
+ raise RetrievalError('invalid_title_index')
|
|
|
+ seen.add(identity)
|
|
|
+ rows.append({'id': identity, 'title': row['index_text'].split('\n', 1)[0],
|
|
|
+ 'record_hash': row['record_hash']})
|
|
|
+ finally:
|
|
|
+ iterator.close()
|
|
|
+ return rows
|
|
|
+
|
|
|
+ def index_vectors(self, name, ids):
|
|
|
+ return self.client.get(name, ids=ids, output_fields=['vector', 'record_hash'],
|
|
|
+ consistency_level='Strong', timeout=self.timeout)
|
|
|
+
|
|
|
def close(self):
|
|
|
self.client.close()
|
|
|
|
|
|
@@ -132,6 +161,8 @@ class Retriever:
|
|
|
self._stores = {}
|
|
|
self._warmed = False
|
|
|
self.lock = RLock()
|
|
|
+ self._title_locks = {kind: RLock() for kind in SEARCH_TYPES}
|
|
|
+ self._title_cache = {}
|
|
|
|
|
|
def _initialize_model(self):
|
|
|
if self.embedder is None:
|
|
|
@@ -169,6 +200,62 @@ class Retriever:
|
|
|
for store in self._stores.values():
|
|
|
store.close()
|
|
|
self._stores.clear()
|
|
|
+ self._title_cache.clear()
|
|
|
+
|
|
|
+ def _titles_for(self, kind):
|
|
|
+ with self._title_locks[kind]:
|
|
|
+ if kind not in self._title_cache:
|
|
|
+ rows = self._store_for(kind).titles(self.manifest['databases'][kind]['index_collection'])
|
|
|
+ self._title_cache[kind] = tuple({**row, 'type': kind} for row in rows)
|
|
|
+ return self._title_cache[kind]
|
|
|
+
|
|
|
+ def _recall_titles(self, queries, vectors):
|
|
|
+ if not queries:
|
|
|
+ return [], {'status': 'skipped', 'limit': TITLE_LIMIT, 'matched_count': 0, 'failed_databases': []}
|
|
|
+ failures, entries = [], []
|
|
|
+ def read(kind):
|
|
|
+ try:
|
|
|
+ return kind, self._titles_for(kind), False
|
|
|
+ except Exception:
|
|
|
+ return kind, (), True
|
|
|
+ with ThreadPoolExecutor(max_workers=self.bank_workers) as pool:
|
|
|
+ for kind, rows, failed in mapped(pool, read, SEARCH_TYPES):
|
|
|
+ entries.extend(rows)
|
|
|
+ if failed:
|
|
|
+ failures.append(kind)
|
|
|
+ selected, count = select_titles(entries, queries)
|
|
|
+ hits = []
|
|
|
+ for kind in SEARCH_TYPES:
|
|
|
+ group = [row for row in selected if row['type'] == kind]
|
|
|
+ if not group:
|
|
|
+ continue
|
|
|
+ try:
|
|
|
+ values = self._store_for(kind).index_vectors(
|
|
|
+ self.manifest['databases'][kind]['index_collection'], [row['id'] for row in group])
|
|
|
+ by_id = {row['id']: row for row in values}
|
|
|
+ if len(by_id) != len(group):
|
|
|
+ raise RetrievalError('missing_title_vector')
|
|
|
+ bank_hits = []
|
|
|
+ for row in group:
|
|
|
+ value = by_id[row['id']]
|
|
|
+ vector = value['vector']
|
|
|
+ if (value['record_hash'] != row['record_hash'] or len(vector) != len(vectors[0])
|
|
|
+ or not all(math.isfinite(v) for v in vector) or not any(vector)):
|
|
|
+ raise RetrievalError('invalid_title_vector')
|
|
|
+ scores = [sum(a * b for a, b in zip(vector, q)) /
|
|
|
+ (math.sqrt(sum(a*a for a in vector)) * math.sqrt(sum(b*b for b in q)))
|
|
|
+ for q in vectors]
|
|
|
+ bank_hits.append((kind, {'id': row['id'], 'distance': max(scores),
|
|
|
+ 'entity': {'record_hash': row['record_hash']}, 'title_match': row['title_match']}))
|
|
|
+ hits.extend(bank_hits)
|
|
|
+ except Exception:
|
|
|
+ if kind not in failures:
|
|
|
+ failures.append(kind)
|
|
|
+ order = {(row['type'], row['id']): index for index, row in enumerate(selected)}
|
|
|
+ hits.sort(key=lambda pair: order[(pair[0], pair[1]['id'])])
|
|
|
+ return hits, {'status': 'partial' if failures else 'complete', 'limit': TITLE_LIMIT,
|
|
|
+ 'matched_count': count, 'selected_count': len(hits),
|
|
|
+ 'overflow': count > TITLE_LIMIT, 'failed_databases': failures}
|
|
|
|
|
|
def search(self, question, expansions=()):
|
|
|
# GPU inference is serialized, but network retrieval can overlap across users.
|
|
|
@@ -184,6 +271,7 @@ class Retriever:
|
|
|
if not isinstance(question, str) or not question.strip() or len(question) > 8000:
|
|
|
raise RetrievalError('invalid_question')
|
|
|
queries = [question]
|
|
|
+ named_queries = title_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)
|
|
|
@@ -229,7 +317,9 @@ class Retriever:
|
|
|
failed_queries = []
|
|
|
query_hits = [[] for _ in queries]
|
|
|
successful_searches = 0
|
|
|
- with ThreadPoolExecutor(max_workers=min(self.query_workers, len(queries))) as pool:
|
|
|
+ with ThreadPoolExecutor(max_workers=1) as title_pool, ThreadPoolExecutor(
|
|
|
+ max_workers=min(self.query_workers, len(queries))) as pool:
|
|
|
+ title_future = title_pool.submit(self._recall_titles, named_queries, vectors)
|
|
|
# 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:
|
|
|
@@ -242,13 +332,14 @@ class Retriever:
|
|
|
previous = merged[kind].get(hit['id'])
|
|
|
if previous is None or hit['distance'] > previous['distance']:
|
|
|
merged[kind][hit['id']] = hit
|
|
|
+ title_hits, title_diagnostics = title_future.result()
|
|
|
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)
|
|
|
+ chosen = choose_candidates(hits_by_type, query_hits, title_hits=title_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}
|
|
|
@@ -293,10 +384,14 @@ class Retriever:
|
|
|
record = restore(parts)
|
|
|
if record['id'] != hit['id'] or sha(canonical(record)) != hit['entity']['record_hash']:
|
|
|
raise RetrievalError('index_content_mismatch')
|
|
|
+ match = title_match(record['title'], named_queries) if hit.get('title_match') else None
|
|
|
+ if hit.get('title_match') and not match:
|
|
|
+ raise RetrievalError('title_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']})
|
|
|
+ 'snapshot': self.manifest['snapshot'],
|
|
|
+ **({'title_match': {**match, 'verified': True}} if match else {})})
|
|
|
except Exception:
|
|
|
content_failures.append(hit['id'])
|
|
|
related = related_guide_queries(question, records)
|
|
|
@@ -348,7 +443,9 @@ class Retriever:
|
|
|
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,
|
|
|
+ 'complete': not failures and not content_failures and not related_guide_failures
|
|
|
+ and not title_diagnostics['failed_databases'],
|
|
|
+ 'title_recall': title_diagnostics,
|
|
|
'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,
|