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

105 lines
6.9 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']]
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.')
from tools.graphrag.graph_search import requested_name_phrases
first_names={normalized(r['title']).split()[0] for r in catalog}
capitalized={normalized(match[0]) for match in re.finditer(
r'\b[A-ZÁÄČĎÉÍĹĽŇÓÔŔŠŤÚÝŽ][^\W\d_]+(?:\s+[A-ZÁÄČĎÉÍĹĽŇÓÔŔŠŤÚÝŽ][^\W\d_]+){1,3}\b',question)}
unknown=[]
for phrase in requested_name_phrases(question):
words=phrase.split()
named_companion=(named and phrase in capitalized and
re.search(r'\ba\s+'+re.escape(phrase)+r'\b',normalized(question)))
if (2<=len(words)<=4 and (words[0] in first_names or named_companion) and
not any(text_contains_value(phrase,r['title'],person=True) for r in catalog)):
unknown.append(phrase)
if unknown:
return clarify(payload,'unresolved_subject','Meno „'+unknown[0]+'“ sa nepodarilo spoľahlivo priradiť. '
'Uveď jeho presné celé meno; ostatné osoby nebudem používať ako náhradu.')
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}