388 lines
12 KiB
Python
388 lines
12 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import re
|
||
|
|
from dataclasses import dataclass
|
||
|
|
|
||
|
|
from django.db.models import Q
|
||
|
|
|
||
|
|
from control_plane.authoring.models import (
|
||
|
|
DocumentAuthority,
|
||
|
|
SourceDocumentVersion,
|
||
|
|
SourcePassage,
|
||
|
|
Work,
|
||
|
|
WorkType,
|
||
|
|
)
|
||
|
|
from control_plane.authoring.state_management import json_sha256
|
||
|
|
|
||
|
|
STOP_WORDS = {
|
||
|
|
"and",
|
||
|
|
"are",
|
||
|
|
"about",
|
||
|
|
"after",
|
||
|
|
"again",
|
||
|
|
"also",
|
||
|
|
"before",
|
||
|
|
"being",
|
||
|
|
"between",
|
||
|
|
"but",
|
||
|
|
"could",
|
||
|
|
"for",
|
||
|
|
"from",
|
||
|
|
"has",
|
||
|
|
"her",
|
||
|
|
"him",
|
||
|
|
"his",
|
||
|
|
"have",
|
||
|
|
"into",
|
||
|
|
"its",
|
||
|
|
"must",
|
||
|
|
"not",
|
||
|
|
"scene",
|
||
|
|
"she",
|
||
|
|
"should",
|
||
|
|
"that",
|
||
|
|
"the",
|
||
|
|
"their",
|
||
|
|
"them",
|
||
|
|
"then",
|
||
|
|
"there",
|
||
|
|
"they",
|
||
|
|
"this",
|
||
|
|
"through",
|
||
|
|
"what",
|
||
|
|
"when",
|
||
|
|
"where",
|
||
|
|
"which",
|
||
|
|
"while",
|
||
|
|
"with",
|
||
|
|
"would",
|
||
|
|
"was",
|
||
|
|
"were",
|
||
|
|
"write",
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass(frozen=True)
|
||
|
|
class RankedPassage:
|
||
|
|
passage: SourcePassage
|
||
|
|
score: float
|
||
|
|
reason: str
|
||
|
|
|
||
|
|
|
||
|
|
def query_terms(query: str, limit: int = 16) -> list[str]:
|
||
|
|
counts: dict[str, int] = {}
|
||
|
|
for token in re.findall(r"[a-zA-Z][a-zA-Z0-9']{2,}", query.lower()):
|
||
|
|
if token in STOP_WORDS:
|
||
|
|
continue
|
||
|
|
counts[token] = counts.get(token, 0) + 1
|
||
|
|
ordered = sorted(
|
||
|
|
counts.items(), key=lambda item: (-item[1], -len(item[0]), item[0])
|
||
|
|
)
|
||
|
|
return [token for token, _count in ordered[:limit]]
|
||
|
|
|
||
|
|
|
||
|
|
def retrieve_scene_passages(
|
||
|
|
*,
|
||
|
|
work: Work,
|
||
|
|
query: str,
|
||
|
|
authorities: list[str] | None = None,
|
||
|
|
pinned_document_keys: list[str] | None = None,
|
||
|
|
limit: int = 24,
|
||
|
|
) -> list[RankedPassage]:
|
||
|
|
authorities = authorities or [DocumentAuthority.CANON]
|
||
|
|
invalid = sorted(set(authorities) - set(DocumentAuthority.values))
|
||
|
|
if invalid:
|
||
|
|
raise ValueError(f"unsupported document authorities: {', '.join(invalid)}")
|
||
|
|
pinned = {value.strip() for value in (pinned_document_keys or []) if value.strip()}
|
||
|
|
terms = query_terms(query)
|
||
|
|
visible_work_ids = list(
|
||
|
|
Work.objects.filter(series=work.series, work_type=WorkType.SERIES_REFERENCE).values_list(
|
||
|
|
"id", flat=True
|
||
|
|
)
|
||
|
|
)
|
||
|
|
visible_work_ids.append(work.id)
|
||
|
|
base = SourcePassage.objects.select_related(
|
||
|
|
"document_version__document"
|
||
|
|
).filter(
|
||
|
|
document_version__document__work_id__in=visible_work_ids,
|
||
|
|
document_version__authority__in=authorities,
|
||
|
|
document_version__superseded_by__isnull=True,
|
||
|
|
)
|
||
|
|
|
||
|
|
candidates: dict[object, SourcePassage] = {}
|
||
|
|
if terms:
|
||
|
|
term_filter = Q()
|
||
|
|
for term in terms:
|
||
|
|
term_filter |= Q(content__icontains=term)
|
||
|
|
term_filter |= Q(document_version__document__title__icontains=term)
|
||
|
|
for passage in base.filter(term_filter)[:4000]:
|
||
|
|
candidates[passage.id] = passage
|
||
|
|
if pinned:
|
||
|
|
for passage in base.filter(document_version__document__logical_key__in=pinned)[:2000]:
|
||
|
|
candidates[passage.id] = passage
|
||
|
|
if not candidates:
|
||
|
|
for passage in base[:300]:
|
||
|
|
candidates[passage.id] = passage
|
||
|
|
|
||
|
|
ranked: list[RankedPassage] = []
|
||
|
|
lowered_query = query.lower()
|
||
|
|
for passage in candidates.values():
|
||
|
|
document = passage.document_version.document
|
||
|
|
haystack = passage.content.lower()
|
||
|
|
identity = f"{document.logical_key} {document.title}".lower()
|
||
|
|
score = 0.0
|
||
|
|
matched = []
|
||
|
|
for term in terms:
|
||
|
|
occurrences = haystack.count(term)
|
||
|
|
if occurrences:
|
||
|
|
score += 1.0 + min(occurrences, 4) * 0.5
|
||
|
|
matched.append(term)
|
||
|
|
if term in identity:
|
||
|
|
score += 4.0
|
||
|
|
if document.logical_key in pinned:
|
||
|
|
score += 100.0
|
||
|
|
if passage.document_version.authority == DocumentAuthority.CANON:
|
||
|
|
score += 2.0
|
||
|
|
if document.title.lower() in lowered_query:
|
||
|
|
score += 5.0
|
||
|
|
reason = "pinned" if document.logical_key in pinned else "terms: " + ", ".join(matched[:6])
|
||
|
|
ranked.append(RankedPassage(passage=passage, score=score, reason=reason.strip()))
|
||
|
|
ranked.sort(
|
||
|
|
key=lambda item: (
|
||
|
|
-item.score,
|
||
|
|
item.passage.document_version.document.logical_key,
|
||
|
|
item.passage.ordinal,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
requested_limit = max(1, limit)
|
||
|
|
if not pinned:
|
||
|
|
if len(authorities) == 1:
|
||
|
|
return ranked[:requested_limit]
|
||
|
|
return _select_across_authorities(ranked, authorities, requested_limit)
|
||
|
|
|
||
|
|
pinned_ranked = [
|
||
|
|
item
|
||
|
|
for item in ranked
|
||
|
|
if item.passage.document_version.document.logical_key in pinned
|
||
|
|
]
|
||
|
|
other_ranked = [
|
||
|
|
item
|
||
|
|
for item in ranked
|
||
|
|
if item.passage.document_version.document.logical_key not in pinned
|
||
|
|
]
|
||
|
|
if not other_ranked:
|
||
|
|
return ranked[:requested_limit]
|
||
|
|
|
||
|
|
pinned_budget = min(len(pinned_ranked), max(1, requested_limit * 2 // 3))
|
||
|
|
found_pinned_keys = {
|
||
|
|
item.passage.document_version.document.logical_key for item in pinned_ranked
|
||
|
|
}
|
||
|
|
per_document_limit = max(
|
||
|
|
1,
|
||
|
|
(pinned_budget + len(found_pinned_keys) - 1) // max(1, len(found_pinned_keys)),
|
||
|
|
)
|
||
|
|
selected: list[RankedPassage] = []
|
||
|
|
pinned_counts: dict[str, int] = {}
|
||
|
|
for item in pinned_ranked:
|
||
|
|
document_key = item.passage.document_version.document.logical_key
|
||
|
|
if len(selected) >= pinned_budget:
|
||
|
|
break
|
||
|
|
if pinned_counts.get(document_key, 0) >= per_document_limit:
|
||
|
|
continue
|
||
|
|
selected.append(item)
|
||
|
|
pinned_counts[document_key] = pinned_counts.get(document_key, 0) + 1
|
||
|
|
|
||
|
|
selected.extend(
|
||
|
|
_select_across_authorities(
|
||
|
|
other_ranked,
|
||
|
|
authorities,
|
||
|
|
requested_limit - len(selected),
|
||
|
|
)
|
||
|
|
)
|
||
|
|
selected_ids = {item.passage.id for item in selected}
|
||
|
|
for item in ranked:
|
||
|
|
if len(selected) >= requested_limit:
|
||
|
|
break
|
||
|
|
if item.passage.id not in selected_ids:
|
||
|
|
selected.append(item)
|
||
|
|
selected_ids.add(item.passage.id)
|
||
|
|
return selected
|
||
|
|
|
||
|
|
|
||
|
|
def _select_across_authorities(
|
||
|
|
ranked: list[RankedPassage], authorities: list[str], limit: int
|
||
|
|
) -> list[RankedPassage]:
|
||
|
|
authority_groups = {
|
||
|
|
authority: [
|
||
|
|
item
|
||
|
|
for item in ranked
|
||
|
|
if item.passage.document_version.authority == authority
|
||
|
|
]
|
||
|
|
for authority in authorities
|
||
|
|
}
|
||
|
|
offsets = {authority: 0 for authority in authorities}
|
||
|
|
selected: list[RankedPassage] = []
|
||
|
|
while len(selected) < limit:
|
||
|
|
added = False
|
||
|
|
for authority in authorities:
|
||
|
|
offset = offsets[authority]
|
||
|
|
group = authority_groups[authority]
|
||
|
|
if offset >= len(group):
|
||
|
|
continue
|
||
|
|
selected.append(group[offset])
|
||
|
|
offsets[authority] += 1
|
||
|
|
added = True
|
||
|
|
if len(selected) >= limit:
|
||
|
|
break
|
||
|
|
if not added:
|
||
|
|
break
|
||
|
|
selected_ids = {item.passage.id for item in selected}
|
||
|
|
for item in ranked:
|
||
|
|
if len(selected) >= limit:
|
||
|
|
break
|
||
|
|
if item.passage.id not in selected_ids:
|
||
|
|
selected.append(item)
|
||
|
|
selected_ids.add(item.passage.id)
|
||
|
|
return selected
|
||
|
|
|
||
|
|
|
||
|
|
def build_scene_context_pack(
|
||
|
|
*,
|
||
|
|
work: Work,
|
||
|
|
query: str,
|
||
|
|
authorities: list[str] | None = None,
|
||
|
|
pinned_document_keys: list[str] | None = None,
|
||
|
|
governing_document_keys: list[str] | None = None,
|
||
|
|
limit: int = 24,
|
||
|
|
max_chars: int = 50000,
|
||
|
|
) -> tuple[dict, list[RankedPassage]]:
|
||
|
|
authorities = authorities or [DocumentAuthority.CANON]
|
||
|
|
governing_keys = list(
|
||
|
|
dict.fromkeys(
|
||
|
|
value.strip() for value in (governing_document_keys or []) if value.strip()
|
||
|
|
)
|
||
|
|
)
|
||
|
|
ranked = retrieve_scene_passages(
|
||
|
|
work=work,
|
||
|
|
query=query,
|
||
|
|
authorities=authorities,
|
||
|
|
pinned_document_keys=pinned_document_keys,
|
||
|
|
limit=limit,
|
||
|
|
)
|
||
|
|
citations = []
|
||
|
|
rendered = []
|
||
|
|
used_chars = 0
|
||
|
|
kept: list[RankedPassage] = []
|
||
|
|
if governing_keys:
|
||
|
|
visible_work_ids = list(
|
||
|
|
Work.objects.filter(
|
||
|
|
series=work.series,
|
||
|
|
work_type=WorkType.SERIES_REFERENCE,
|
||
|
|
).values_list("id", flat=True)
|
||
|
|
)
|
||
|
|
visible_work_ids.append(work.id)
|
||
|
|
versions = list(
|
||
|
|
SourceDocumentVersion.objects.select_related("document")
|
||
|
|
.filter(
|
||
|
|
document__work_id__in=visible_work_ids,
|
||
|
|
document__logical_key__in=governing_keys,
|
||
|
|
authority__in=authorities,
|
||
|
|
superseded_by__isnull=True,
|
||
|
|
)
|
||
|
|
.order_by("document__logical_key")
|
||
|
|
)
|
||
|
|
versions_by_key: dict[str, list[SourceDocumentVersion]] = {}
|
||
|
|
for version in versions:
|
||
|
|
versions_by_key.setdefault(version.document.logical_key, []).append(version)
|
||
|
|
missing = [key for key in governing_keys if key not in versions_by_key]
|
||
|
|
ambiguous = [key for key, values in versions_by_key.items() if len(values) > 1]
|
||
|
|
if missing:
|
||
|
|
raise ValueError("governing documents not found: " + ", ".join(missing))
|
||
|
|
if ambiguous:
|
||
|
|
raise ValueError("governing document keys are ambiguous: " + ", ".join(ambiguous))
|
||
|
|
for key in governing_keys:
|
||
|
|
version = versions_by_key[key][0]
|
||
|
|
document = version.document
|
||
|
|
label = f"SRC-{len(citations) + 1:02d}"
|
||
|
|
end_line = version.content.count("\n") + 1
|
||
|
|
block = (
|
||
|
|
f"[{label}] authority={version.authority} source={document.logical_key} "
|
||
|
|
f"version={version.version} lines=1-{end_line} scope=governing-document\n"
|
||
|
|
f"{version.content}"
|
||
|
|
)
|
||
|
|
if used_chars + len(block) > max_chars:
|
||
|
|
raise ValueError("governing documents exceed the context character budget")
|
||
|
|
used_chars += len(block)
|
||
|
|
rendered.append(block)
|
||
|
|
citations.append(
|
||
|
|
{
|
||
|
|
"id": label,
|
||
|
|
"kind": "governing_document",
|
||
|
|
"passage_id": None,
|
||
|
|
"document_version_id": str(version.id),
|
||
|
|
"document_key": document.logical_key,
|
||
|
|
"document_title": document.title,
|
||
|
|
"document_version": version.version,
|
||
|
|
"authority": version.authority,
|
||
|
|
"source_path": version.source_path,
|
||
|
|
"start_line": 1,
|
||
|
|
"end_line": end_line,
|
||
|
|
"start_char": 0,
|
||
|
|
"end_char": len(version.content),
|
||
|
|
"sha256": version.source_sha256,
|
||
|
|
"score": None,
|
||
|
|
"reason": "governing document supplied in full",
|
||
|
|
}
|
||
|
|
)
|
||
|
|
governing_set = set(governing_keys)
|
||
|
|
for item in ranked:
|
||
|
|
passage = item.passage
|
||
|
|
version = passage.document_version
|
||
|
|
document = version.document
|
||
|
|
if document.logical_key in governing_set:
|
||
|
|
continue
|
||
|
|
excerpt = passage.content[:2500]
|
||
|
|
label = f"SRC-{len(citations) + 1:02d}"
|
||
|
|
block = (
|
||
|
|
f"[{label}] authority={version.authority} source={document.logical_key} "
|
||
|
|
f"version={version.version} lines={passage.start_line}-{passage.end_line}\n{excerpt}"
|
||
|
|
)
|
||
|
|
if rendered and used_chars + len(block) > max_chars:
|
||
|
|
continue
|
||
|
|
used_chars += len(block)
|
||
|
|
kept.append(item)
|
||
|
|
citations.append(
|
||
|
|
{
|
||
|
|
"id": label,
|
||
|
|
"passage_id": str(passage.id),
|
||
|
|
"document_key": document.logical_key,
|
||
|
|
"document_title": document.title,
|
||
|
|
"document_version": version.version,
|
||
|
|
"authority": version.authority,
|
||
|
|
"source_path": version.source_path,
|
||
|
|
"start_line": passage.start_line,
|
||
|
|
"end_line": passage.end_line,
|
||
|
|
"start_char": passage.start_char,
|
||
|
|
"end_char": passage.end_char,
|
||
|
|
"sha256": passage.sha256,
|
||
|
|
"score": item.score,
|
||
|
|
"reason": item.reason,
|
||
|
|
}
|
||
|
|
)
|
||
|
|
rendered.append(block)
|
||
|
|
pack = {
|
||
|
|
"schema_version": 1,
|
||
|
|
"work_id": str(work.id),
|
||
|
|
"query": query,
|
||
|
|
"authorities": authorities,
|
||
|
|
"governing_document_keys": governing_keys,
|
||
|
|
"citations": citations,
|
||
|
|
"rendered_context": (
|
||
|
|
"\n\n".join(rendered)
|
||
|
|
if rendered
|
||
|
|
else "(No matching approved source passages.)"
|
||
|
|
),
|
||
|
|
}
|
||
|
|
pack["sha256"] = json_sha256(pack)
|
||
|
|
return pack, kept
|