dp-zp-agent/tools/rag/search_scope.py

92 lines
5.2 KiB
Python

"""Hard question filters, independent of lexical/semantic relevance scores."""
from __future__ import annotations
from dataclasses import dataclass
import json
from tools.answers.question_plan import QuestionPlan, build_question_plan
from tools.answers.rag_evidence import academic_windows_metadata
from tools.answers.context_policy import normalized
from tools.retrieval.person_names import sqlite_people
from tools.retrieval.retrieval_contract import result_fields
@dataclass
class SearchScope:
plan: QuestionPlan
failure: dict | None = None
eligible_ids: set[str] | None = None
def project(self, row: dict) -> dict | None:
"""Clip mixed-work chunks; a date/name elsewhere is not a work binding."""
students = [p for p in self.plan.people if p.role == 'student']
path = row.get('document_path', '')
owners = [p for p in students if p.source_document and (
path == p.source_document or path.startswith(p.source_document.removesuffix('README.md')))]
if students and not owners:
return None
authors = [p.name for p in self.plan.people if p.role == 'document_author']
if authors and normalized(row.get('author') or '') not in {normalized(p) for p in authors}:
return None
constrained = any(w.work_type or w.years or w.labels or w.excluded_labels or
w.numbers or w.excluded_numbers or w.unnumbered for w in self.plan.works)
constrained |= bool(self.plan.period and self.plan.period.field == 'work_year')
if not constrained:
return dict(row)
text = row.get('text') or ''
windows = academic_windows_metadata(text)
if not any(scope for scope, _, _ in windows):
# Chunk ancestry can be the only explicit work heading. Never use
# a profile directory or taxonomy year as an academic work year.
paths = row.get('heading_paths')
if paths is None:
paths = json.loads(row.get('heading_paths_json') or '[]')
first = paths[0] if paths else ''
heading = ' > '.join(first) if isinstance(first, list) else str(first)
body = '\n'.join(line for line in text.splitlines() if not line.startswith('Dokument:'))
windows = academic_windows_metadata('Sekcia: ' + heading + '\n' + body)
names = [p.name for p in owners] or [None]
selected = [(scope, label, body) for scope, label, body in windows if scope and any(
self.plan.matches_work(name, label, scope[0], scope[1]) for name in names)]
if not selected:
return None
# Preserve scoped source headings for the evidence verifier, but omit
# other works even if one chunk or an expansion spans several works.
clipped = '\n\n'.join('Sekcia: ' + label + '\n' + body for _, label, body in selected)
if 'start_year' in self.plan.fields:
# Study metadata belongs to the profile, independently of work
# sections. Keep the explicit line when both values were requested.
study_lines = [line for scope, _, body in windows if scope is None
for line in body.splitlines() if normalized(line).startswith('rok zaciatku studia ')]
clipped = '\n'.join(study_lines) + '\n\n' + clipped if study_lines else clipped
headings = [[label] for _, label, _ in selected]
return {**row, 'text': clipped, 'text_length': len(clipped),
'snippet': clipped[:320], 'heading_paths': headings,
'heading_paths_json': json.dumps(headings, ensure_ascii=False)}
def project_rows(self, rows):
return [value for row in rows if (value := self.project(row)) is not None]
def prepare_search_scope(conn, db_file, query, *, published_only, question_plan=None):
if question_plan is None:
bound = build_question_plan(query, sqlite_people(db_file, published_only=published_only))
plan = bound.spec
if bound.clarification:
status = 'no_answer' if bound.resolution.get('status') == 'no_answer' and bound.reason == bound.resolution.get('reason') else 'ambiguous'
return SearchScope(plan, failure={
**result_fields(status=status, sources=[], reason=bound.reason, complete=False),
'engine': 'hybrid_fts5_embeddings', 'results': [], 'strategies': [],
'clarification': bound.clarification, 'name_resolution': bound.resolution})
else:
plan = question_plan if isinstance(question_plan, QuestionPlan) else QuestionPlan.model_validate(question_plan)
if normalized(plan.question) != normalized(query) or plan.unsupported:
raise ValueError('Search plan does not match the supported question')
scope = SearchScope(plan)
if plan.people or any(w.work_type or w.years or w.labels or w.excluded_labels or
w.numbers or w.excluded_numbers or w.unnumbered for w in plan.works) or plan.period:
rows = conn.execute('SELECT * FROM chunks WHERE (?=0 OR COALESCE(published,1)=1)',
(int(published_only),)).fetchall()
scope.eligible_ids = {row['chunk_id'] for row in rows if scope.project(dict(row)) is not None}
return scope