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

215 lines
9.4 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 tools.answers.context_policy import normalized, person_pattern, text_contains_value
from tools.answers.question_intent import content_question
from tools.answers.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
from tools.answers.question_plan import QuestionPlan, build_question_plan
plan=QuestionPlan.model_validate(payload['question_plan']) if payload.get('question_plan') else None
selector = parse_work_selector(question)
metadata_scope=bool(plan and plan.intent=='metadata' and any(p.role=='student' for p in plan.people)
and (set(plan.fields)&{'title','year','work_type'} or plan.presentation=='count'))
if not content_question(question) and not selector.specific and not metadata_scope:
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
if plan is None:
bound=build_question_plan(question,[{'name':r['title'],'kind':'Person','source_document':r['path']} for r in catalog])
if bound.clarification:return clarify(payload,bound.reason,bound.clarification)
plan=bound.spec
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 not plan.matches_work(person["title"],label,scope[0],scope[1])
):
continue
label = label or WORK_LABELS[scope[0]] + (
f" {scope[1]}" if scope[1] is not None else ""
)
if not selector.matches(label):
continue
if metadata_scope and (scope,normalized(label)) in records:
continue
records[(scope, normalized(label))] = label
if (scope, normalized(label), text) not in seen:
seen.add((scope, normalized(label), text))
# Exact metadata needs every record's source heading
# and introduction, not all subsequent meeting notes.
body='\n'.join(line for line in text.splitlines() if line.strip()) if metadata_scope else text
if metadata_scope:body='\n'.join(body.splitlines()[:6])
frames.append("Sekcia: " + label + "\n" + body)
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 plan.scope == "single_work":
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 tools.answers.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,
"question_plan": plan.model_dump(),
"named_scope_complete": True,
}