dp-zp-agent/tools/answer_evaluation/rag_preparation.py

102 lines
6.6 KiB
Python

"""Resolve named source profiles and scope content before the shared token budget."""
from __future__ import annotations
import json
import re
import sqlite3
from contextlib import closing
from .context_policy import normalized, person_pattern, text_contains_value
from .question_intent import content_question
from .rag_evidence import PROFILE, academic_windows_metadata, requested_types, WORK_LABELS, project_number
from tools.retrieval.question_constraints import requested_years
def clarify(payload: dict, reason: str, message: str) -> dict:
return {**payload,'status':'ambiguous','reason':reason,'complete':False,
'clarification':message}
def prepare_named_content(db_file, question: str, payload: dict) -> dict:
from tools.retrieval.work_selector import parse_work_selector, label_key
selector=parse_work_selector(question)
if not content_question(question) and not selector.specific:
return payload
if re.search(r'\b(?:s tagom|s klucovym slovom|z kategorie|s autorom dokumentu)\b',normalized(question)):
return clarify(payload,'content_filter_requires_graph',
'Najprv over podmienku tagu, kategórie alebo autora cez /graphrag. '
'Potom môžem opísať obsah konkrétnej vyhovujúcej práce; túto podmienku nevynechám.')
from tools.rag.rag_utils import build_context_text
with closing(sqlite3.connect(db_file.resolve().as_uri()+'?mode=ro',uri=True)) as conn:
conn.row_factory=sqlite3.Row
if not {'documents','chunks'} <= {r[0] for r in conn.execute("SELECT name FROM sqlite_master WHERE type='table'")}:
return payload
catalog=[dict(r) for r in conn.execute('SELECT path,title,author,published FROM documents WHERE COALESCE(published,1)=1')
if PROFILE.fullmatch(r['path']) and r['title']]
from tools.retrieval.person_names import resolve_person_names
resolution=payload.get('name_resolution') or resolve_person_names(question,[
{'kind':'Person','name':r['title'],'source_document':r['path']} for r in catalog])
if resolution['status'] != 'ok':
return clarify(payload,'unresolved_subject',resolution['message'])
question=resolution['resolved_query']
payload={**payload,'name_resolution':resolution,'resolved_query':question}
identities = {normalized(m['name']): m['source_document']
for m in resolution.get('matches',[])
if m.get('source_document')}
if identities:
catalog = [r for r in catalog if identities.get(normalized(r['title'])) == r['path']]
named=[r for r in catalog if text_contains_value(question,r['title'],person=True)]
# Preserve question order rather than the original retrieval ranking.
named.sort(key=lambda r:re.search(person_pattern(r['title']),normalized(question)).start())
if len(named)>4:
return clarify(payload,'too_many_subjects','Pre podrobné obsahové porovnanie vyber najviac štyri osoby naraz.')
if not named:
return payload
types=requested_types(selector.positive_text,[r['title'] for r in named])
years=requested_years(question)
selected=[];scopes=[]
for person in named:
rows=list(conn.execute('SELECT text,heading_paths_json FROM chunks WHERE document_path=? '
'AND COALESCE(published,1)=1 ORDER BY chunk_index',(person['path'],)))
frames=[]; records={};seen=set()
for row in rows:
paths=json.loads(row['heading_paths_json'] or '[]')
headings=[' > '.join(p) if isinstance(p,list) else str(p) for p in paths]
first=headings[0] if headings else ''
for scope,label,text in academic_windows_metadata('Sekcia: '+first+'\n'+str(row['text'] or '')):
if not scope or (types.get(person['title']) and scope[0]!=types[person['title']]) or (years and scope[1] not in years):
continue
label=label or WORK_LABELS[scope[0]]+(f' {scope[1]}' if scope[1] is not None else '')
if not selector.matches(label):
continue
records[(scope,normalized(label))]=label
if (scope,normalized(label),text) not in seen:
seen.add((scope,normalized(label),text));frames.append('Sekcia: '+label+'\n'+text)
if selector.multiple and (any(label not in {label_key(v) for v in records.values()} for label in selector.labels) or
any(not any(project_number(v)==number for v in records.values()) for number in selector.numbers)):
return clarify(payload,'unresolved_work_selection',
'Niektoré z požadovaných označení prác osoby '+person['title']+' sa v indexe nenašli. '
'Uveď presné nadpisy zo zdroja; chýbajúcu prácu nevynechám z porovnania potichu.')
if len(records)>1 and not selector.multiple:
options='; '.join(dict.fromkeys(records.values()))
return clarify(payload,'ambiguous_work',f'Ktorú prácu osoby {person["title"]} máš na mysli? '
f'V indexe sú tieto možnosti: {options}. Uveď typ práce a rok alebo presné označenie projektu.')
scopes.extend({'person':person['title'],'work_type':scope[0],'year':scope[1],'work_label':label}
for (scope,_key),label in records.items())
from .content_selection import prioritize_frames
frames=prioritize_frames(frames,question)
source=next((s for s in payload.get('sources',[]) if s.get('document_path')==person['path']),None)
if source is None:
from tools.rag.search_core import make_source_url
source={'title':person['title'],'document_path':person['path'],'source_url':make_source_url(person['path']),
'author':person['author'],'published':person['published']}
selected.append({**source,'text':'\n\n'.join(frames)})
# Keep already-retrieved child documents of the requested profiles.
for source in payload.get('sources',[]):
if any(source.get('document_path','').startswith(p['path'].removesuffix('README.md'))
and source.get('document_path')!=p['path'] for p in named):
selected.append(source)
selected=[{**source,'source_id':f'S{i}'} for i,source in enumerate(selected,1)]
return {**payload,'sources':selected,'source_count':len(selected),'context':build_context_text(selected),
'named_subjects':[p['title'] for p in named],'work_scopes':scopes}