671 lines
20 KiB
Python
671 lines
20 KiB
Python
"""Versioned task-shaped retrieval plans and immutable context capsules."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, replace
|
|
from typing import Literal, cast
|
|
|
|
from .errors import DocForgeError
|
|
from .models import Edge, Node, ProjectDescriptor
|
|
from .pagination import canonical_hash
|
|
|
|
TaskKind = Literal[
|
|
"change",
|
|
"implementation",
|
|
"failure",
|
|
"ownership",
|
|
"test",
|
|
"operation",
|
|
"release",
|
|
]
|
|
TASK_KINDS: tuple[TaskKind, ...] = (
|
|
"change",
|
|
"implementation",
|
|
"failure",
|
|
"ownership",
|
|
"test",
|
|
"operation",
|
|
"release",
|
|
)
|
|
|
|
RelationCategory = Literal[
|
|
"structure",
|
|
"implementation",
|
|
"dependency",
|
|
"execution",
|
|
"data",
|
|
"evidence",
|
|
"context",
|
|
"unclassified",
|
|
]
|
|
BASE_RELATION_CATEGORIES: tuple[RelationCategory, ...] = (
|
|
"structure",
|
|
"implementation",
|
|
"dependency",
|
|
"execution",
|
|
"data",
|
|
"evidence",
|
|
"context",
|
|
)
|
|
|
|
RELATION_CATEGORIES: dict[RelationCategory, tuple[str, ...]] = {
|
|
"structure": ("contains", "defined_in", "defines", "owns"),
|
|
"implementation": (
|
|
"implemented_by",
|
|
"implements",
|
|
"inherits",
|
|
"inherits_from",
|
|
),
|
|
"dependency": ("depends_on", "imports"),
|
|
"execution": ("activates", "calls", "dispatches_to", "launches"),
|
|
"data": ("reads", "writes"),
|
|
"evidence": ("documents", "governs", "proves", "tested_by", "verifies"),
|
|
"context": ("relates_to",),
|
|
"unclassified": (),
|
|
}
|
|
|
|
TASK_REQUIREMENTS: dict[TaskKind, tuple[RelationCategory, ...]] = {
|
|
"change": ("dependency",),
|
|
"implementation": ("implementation",),
|
|
"failure": ("execution",),
|
|
"ownership": ("structure",),
|
|
"test": ("evidence",),
|
|
"operation": ("execution",),
|
|
"release": ("evidence",),
|
|
}
|
|
|
|
PLANNER_ID = "docforge.core.task-context"
|
|
PLANNER_VERSION = 1
|
|
MAX_TASK_EVIDENCE = 1_000
|
|
MAX_TASK_CANDIDATE_EDGES = 100_000
|
|
MAX_TASK_QUERY_CHARS = 10_000
|
|
PROVENANCE_LIMITATIONS = (
|
|
"evidence_type_unavailable",
|
|
"extractor_identity_unavailable",
|
|
"relationship_provenance_unavailable",
|
|
"observation_time_unavailable",
|
|
)
|
|
|
|
_RELATION_TO_CATEGORY = {
|
|
relation: category
|
|
for category, relations in RELATION_CATEGORIES.items()
|
|
for relation in relations
|
|
}
|
|
|
|
|
|
def relation_category(relation: str) -> RelationCategory:
|
|
"""Classify only versioned known aliases; preserve every other relation."""
|
|
|
|
return cast(RelationCategory, _RELATION_TO_CATEGORY.get(relation, "unclassified"))
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RetrievalStepV1:
|
|
step_id: str
|
|
operation: Literal["exact", "search", "outgoing", "incoming", "metadata"]
|
|
relation_scope: Literal["none", "project_allowed"]
|
|
relation_set_hash: str | None
|
|
direction: Literal["none", "outgoing", "incoming"]
|
|
depth: int
|
|
limit: int
|
|
required: bool
|
|
evidence_role: str
|
|
|
|
def as_dict(self) -> dict[str, object]:
|
|
return {
|
|
"step_id": self.step_id,
|
|
"operation": self.operation,
|
|
"relation_scope": self.relation_scope,
|
|
"relation_set_hash": self.relation_set_hash,
|
|
"direction": self.direction,
|
|
"depth": self.depth,
|
|
"limit": self.limit,
|
|
"required": self.required,
|
|
"evidence_role": self.evidence_role,
|
|
}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RetrievalRequirementV1:
|
|
requirement_id: str
|
|
category: RelationCategory
|
|
|
|
def as_dict(self) -> dict[str, object]:
|
|
return {
|
|
"requirement_id": self.requirement_id,
|
|
"check": "selected_relation_category",
|
|
"category": self.category,
|
|
"required": True,
|
|
}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RetrievalPlanV1:
|
|
schema_version: Literal[1]
|
|
planner_id: str
|
|
planner_version: int
|
|
task_kind: TaskKind
|
|
task_query: str
|
|
focus_node_id: str | None
|
|
request_hash: str
|
|
effective_policy_hash: str
|
|
max_evidence: int
|
|
max_tokens: int
|
|
max_depth: int
|
|
max_candidate_edges: int
|
|
category_order: tuple[RelationCategory, ...]
|
|
steps: tuple[RetrievalStepV1, ...]
|
|
requirements: tuple[RetrievalRequirementV1, ...]
|
|
plan_hash: str
|
|
|
|
def as_dict(self) -> dict[str, object]:
|
|
return self.payload(include_hash=True)
|
|
|
|
def payload(self, *, include_hash: bool) -> dict[str, object]:
|
|
result: dict[str, object] = {
|
|
"schema_version": self.schema_version,
|
|
"planner": {
|
|
"id": self.planner_id,
|
|
"version": self.planner_version,
|
|
},
|
|
"task_kind": self.task_kind,
|
|
"request_hash": self.request_hash,
|
|
"effective_policy_hash": self.effective_policy_hash,
|
|
"focus_node_id": self.focus_node_id,
|
|
"limits": {
|
|
"max_evidence": self.max_evidence,
|
|
"max_tokens": self.max_tokens,
|
|
"max_depth": self.max_depth,
|
|
"max_candidate_edges": self.max_candidate_edges,
|
|
},
|
|
"category_order": list(self.category_order),
|
|
"steps": [step.as_dict() for step in self.steps],
|
|
"requirements": [requirement.as_dict() for requirement in self.requirements],
|
|
}
|
|
if include_hash:
|
|
result["plan_hash"] = self.plan_hash
|
|
return result
|
|
|
|
|
|
def build_retrieval_plan(
|
|
descriptor: ProjectDescriptor,
|
|
*,
|
|
task_kind: str,
|
|
task: str,
|
|
focus_node_id: str | None,
|
|
budget: int | None,
|
|
limit: int | None,
|
|
effective_policy: dict[str, object],
|
|
) -> RetrievalPlanV1:
|
|
"""Derive one fixed plan from bounded inputs rather than accepting caller operations."""
|
|
|
|
return _build_retrieval_plan(
|
|
descriptor,
|
|
task_kind=task_kind,
|
|
task=task,
|
|
focus_node_id=focus_node_id,
|
|
budget=budget,
|
|
limit=limit,
|
|
effective_policy_hash=canonical_hash(effective_policy),
|
|
)
|
|
|
|
|
|
def validate_retrieval_plan(
|
|
plan: RetrievalPlanV1,
|
|
descriptor: ProjectDescriptor,
|
|
) -> RetrievalPlanV1:
|
|
"""Reject forged, stale-shape, or internally inconsistent public plan objects."""
|
|
|
|
try:
|
|
expected = _build_retrieval_plan(
|
|
descriptor,
|
|
task_kind=plan.task_kind,
|
|
task=plan.task_query,
|
|
focus_node_id=plan.focus_node_id,
|
|
budget=plan.max_tokens,
|
|
limit=plan.max_evidence,
|
|
effective_policy_hash=plan.effective_policy_hash,
|
|
)
|
|
except (AttributeError, TypeError, DocForgeError) as error:
|
|
raise DocForgeError(
|
|
"invalid_retrieval_plan",
|
|
"Task retrieval plan is malformed or outside the fixed version-1 contract",
|
|
) from error
|
|
if plan != expected:
|
|
raise DocForgeError(
|
|
"invalid_retrieval_plan",
|
|
"Task retrieval plan does not match its fixed version-1 derivation",
|
|
)
|
|
return plan
|
|
|
|
|
|
def _build_retrieval_plan(
|
|
descriptor: ProjectDescriptor,
|
|
*,
|
|
task_kind: str,
|
|
task: str,
|
|
focus_node_id: str | None,
|
|
budget: int | None,
|
|
limit: int | None,
|
|
effective_policy_hash: str,
|
|
) -> RetrievalPlanV1:
|
|
if task_kind not in TASK_KINDS:
|
|
raise DocForgeError(
|
|
"invalid_task_kind",
|
|
"Task context kind is unsupported",
|
|
task_kind=task_kind,
|
|
allowed=list(TASK_KINDS),
|
|
)
|
|
selected_kind: TaskKind = task_kind # type: ignore[assignment]
|
|
normalized_task = task.strip()
|
|
if not normalized_task or len(normalized_task) > min(
|
|
descriptor.limits.max_query_chars,
|
|
MAX_TASK_QUERY_CHARS,
|
|
):
|
|
raise DocForgeError(
|
|
"invalid_task_focus",
|
|
"Task description is empty or exceeds the configured query limit",
|
|
)
|
|
if focus_node_id is not None and (not focus_node_id or len(focus_node_id) > 256):
|
|
raise DocForgeError("invalid_task_focus", "Task focus node ID is invalid")
|
|
selected_budget = _bounded_value(
|
|
budget,
|
|
default=min(8_000, descriptor.limits.max_context_tokens),
|
|
maximum=descriptor.limits.max_context_tokens,
|
|
code="invalid_budget",
|
|
)
|
|
selected_limit = _bounded_value(
|
|
limit,
|
|
default=min(20, descriptor.limits.max_results),
|
|
maximum=min(descriptor.limits.max_results, MAX_TASK_EVIDENCE),
|
|
code="invalid_limit",
|
|
)
|
|
if not _is_sha256(effective_policy_hash):
|
|
raise DocForgeError(
|
|
"invalid_retrieval_plan",
|
|
"Effective policy identity is not a SHA-256 value",
|
|
)
|
|
selected_depth = min(2, descriptor.limits.max_traversal_depth)
|
|
requirements = tuple(
|
|
RetrievalRequirementV1(
|
|
requirement_id=f"{selected_kind}.{category}",
|
|
category=category,
|
|
)
|
|
for category in TASK_REQUIREMENTS[selected_kind]
|
|
)
|
|
category_order: tuple[RelationCategory, ...] = (
|
|
*TASK_REQUIREMENTS[selected_kind],
|
|
*(
|
|
category
|
|
for category in BASE_RELATION_CATEGORIES
|
|
if category not in TASK_REQUIREMENTS[selected_kind]
|
|
),
|
|
"unclassified",
|
|
)
|
|
relation_set_hash = canonical_hash(sorted(descriptor.allowed_relations))
|
|
focus_operation: Literal["exact", "search"] = "exact" if focus_node_id else "search"
|
|
steps = (
|
|
RetrievalStepV1(
|
|
step_id="focus",
|
|
operation=focus_operation,
|
|
relation_scope="none",
|
|
relation_set_hash=None,
|
|
direction="none",
|
|
depth=0,
|
|
limit=1,
|
|
required=True,
|
|
evidence_role="focus",
|
|
),
|
|
RetrievalStepV1(
|
|
step_id="outgoing",
|
|
operation="outgoing",
|
|
relation_scope="project_allowed",
|
|
relation_set_hash=relation_set_hash,
|
|
direction="outgoing",
|
|
depth=selected_depth,
|
|
limit=selected_limit,
|
|
required=False,
|
|
evidence_role="related",
|
|
),
|
|
RetrievalStepV1(
|
|
step_id="incoming",
|
|
operation="incoming",
|
|
relation_scope="project_allowed",
|
|
relation_set_hash=relation_set_hash,
|
|
direction="incoming",
|
|
depth=selected_depth,
|
|
limit=selected_limit,
|
|
required=False,
|
|
evidence_role="related",
|
|
),
|
|
RetrievalStepV1(
|
|
step_id="metadata",
|
|
operation="metadata",
|
|
relation_scope="none",
|
|
relation_set_hash=None,
|
|
direction="none",
|
|
depth=0,
|
|
limit=selected_limit,
|
|
required=True,
|
|
evidence_role="provenance",
|
|
),
|
|
)
|
|
request_hash = canonical_hash(
|
|
{
|
|
"task_kind": selected_kind,
|
|
"task": normalized_task,
|
|
"focus_node_id": focus_node_id,
|
|
"budget": selected_budget,
|
|
"limit": selected_limit,
|
|
}
|
|
)
|
|
placeholder = RetrievalPlanV1(
|
|
schema_version=1,
|
|
planner_id=PLANNER_ID,
|
|
planner_version=PLANNER_VERSION,
|
|
task_kind=selected_kind,
|
|
task_query=normalized_task,
|
|
focus_node_id=focus_node_id,
|
|
request_hash=request_hash,
|
|
effective_policy_hash=effective_policy_hash,
|
|
max_evidence=selected_limit,
|
|
max_tokens=selected_budget,
|
|
max_depth=selected_depth,
|
|
max_candidate_edges=min(
|
|
(selected_limit + 1) ** 2,
|
|
MAX_TASK_CANDIDATE_EDGES,
|
|
),
|
|
category_order=category_order,
|
|
steps=steps,
|
|
requirements=requirements,
|
|
plan_hash="",
|
|
)
|
|
return replace(
|
|
placeholder,
|
|
plan_hash=canonical_hash(placeholder.payload(include_hash=False)),
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class CapsuleRelationshipV1:
|
|
source_id: str
|
|
relation: str
|
|
target_id: str
|
|
direction: Literal["outgoing", "incoming"]
|
|
category: RelationCategory
|
|
|
|
def as_dict(self) -> dict[str, object]:
|
|
return {
|
|
"source_id": self.source_id,
|
|
"relation": self.relation,
|
|
"target_id": self.target_id,
|
|
"direction": self.direction,
|
|
"category": self.category,
|
|
"provenance": "validated_graph_edge_without_source_provenance",
|
|
}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class CapsuleEvidenceV1:
|
|
evidence_hash: str
|
|
role: Literal["focus", "related"]
|
|
reason_code: Literal["exact_focus", "lexical_focus", "relationship_path"]
|
|
node: Node
|
|
depth: int
|
|
relationship_path: tuple[CapsuleRelationshipV1, ...]
|
|
relationship_reasons: tuple[CapsuleRelationshipV1, ...]
|
|
estimated_tokens: int
|
|
|
|
def as_dict(self) -> dict[str, object]:
|
|
result = self.payload()
|
|
return {"evidence_hash": self.evidence_hash, **result}
|
|
|
|
def payload(self) -> dict[str, object]:
|
|
result: dict[str, object] = {
|
|
"role": self.role,
|
|
"reason_code": self.reason_code,
|
|
"node_id": self.node.node_id,
|
|
"title": self.node.title,
|
|
"family": self.node.family,
|
|
"authority": self.node.authority,
|
|
"status": self.node.status,
|
|
"tags": list(self.node.tags),
|
|
"summary": self.node.summary,
|
|
"text": _node_text(self.node),
|
|
"estimated_tokens": self.estimated_tokens,
|
|
"source": {
|
|
"path": self.node.source_path,
|
|
"anchor": self.node.source_anchor,
|
|
"content_hash": self.node.content_hash,
|
|
},
|
|
"depth": self.depth,
|
|
"relationship_path": [
|
|
relationship.as_dict() for relationship in self.relationship_path
|
|
],
|
|
"relationship_reasons": [
|
|
relationship.as_dict() for relationship in self.relationship_reasons
|
|
],
|
|
"provenance_limitations": list(PROVENANCE_LIMITATIONS),
|
|
}
|
|
return result
|
|
|
|
|
|
def capsule_evidence(
|
|
*,
|
|
role: Literal["focus", "related"],
|
|
reason_code: Literal["exact_focus", "lexical_focus", "relationship_path"],
|
|
node: Node,
|
|
depth: int,
|
|
relationship_path: tuple[CapsuleRelationshipV1, ...],
|
|
relationship_reasons: tuple[CapsuleRelationshipV1, ...],
|
|
) -> CapsuleEvidenceV1:
|
|
tokens = estimate_tokens(_node_text(node))
|
|
placeholder = CapsuleEvidenceV1(
|
|
evidence_hash="",
|
|
role=role,
|
|
reason_code=reason_code,
|
|
node=node,
|
|
depth=depth,
|
|
relationship_path=relationship_path,
|
|
relationship_reasons=relationship_reasons,
|
|
estimated_tokens=tokens,
|
|
)
|
|
return replace(
|
|
placeholder,
|
|
evidence_hash=canonical_hash(placeholder.payload()),
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class EvidenceGapV1:
|
|
code: Literal[
|
|
"focus_not_found",
|
|
"focus_ambiguous",
|
|
"category_not_declared",
|
|
"no_selected_evidence",
|
|
"evidence_incomplete",
|
|
"unclassified_relation",
|
|
]
|
|
requirement_id: str
|
|
category: RelationCategory | None
|
|
state: Literal["missing", "incomplete", "blocked", "limitation"]
|
|
check_complete: bool
|
|
detail: str
|
|
|
|
def as_dict(self) -> dict[str, object]:
|
|
return {
|
|
"code": self.code,
|
|
"requirement_id": self.requirement_id,
|
|
"category": self.category,
|
|
"state": self.state,
|
|
"check_complete": self.check_complete,
|
|
"detail": self.detail,
|
|
}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class CapsuleOmissionV1:
|
|
code: Literal[
|
|
"result_limit",
|
|
"token_budget",
|
|
"response_limit",
|
|
"edge_examination_limit",
|
|
"unclassified_relation_limit",
|
|
]
|
|
subject: str
|
|
detail_hash: str
|
|
|
|
def as_dict(self) -> dict[str, str]:
|
|
return {
|
|
"code": self.code,
|
|
"subject": self.subject,
|
|
"detail_hash": self.detail_hash,
|
|
}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ContextCapsuleV1:
|
|
schema_version: Literal[1]
|
|
state: Literal["complete", "incomplete", "blocked"]
|
|
task_kind: TaskKind
|
|
generation: tuple[tuple[str, object], ...]
|
|
plan: RetrievalPlanV1
|
|
focus_state: Literal["resolved", "not_found", "ambiguous"]
|
|
focus_node_id: str | None
|
|
focus_candidate_count: int
|
|
evidence: tuple[CapsuleEvidenceV1, ...]
|
|
gaps: tuple[EvidenceGapV1, ...]
|
|
omissions: tuple[CapsuleOmissionV1, ...]
|
|
selected_count: int
|
|
examined_edge_count: int
|
|
estimated_tokens: int
|
|
unclassified_relations: tuple[str, ...]
|
|
collection_hash: str
|
|
capsule_hash: str
|
|
|
|
def as_dict(self) -> dict[str, object]:
|
|
return self.payload(include_hashes=True)
|
|
|
|
def payload(self, *, include_hashes: bool) -> dict[str, object]:
|
|
evidence = [item.as_dict() for item in self.evidence]
|
|
gaps = [gap.as_dict() for gap in self.gaps]
|
|
omissions = [omission.as_dict() for omission in self.omissions]
|
|
result: dict[str, object] = {
|
|
"schema_version": self.schema_version,
|
|
"state": self.state,
|
|
"task_kind": self.task_kind,
|
|
"generation": dict(self.generation),
|
|
"plan": self.plan.as_dict(),
|
|
"focus": {
|
|
"state": self.focus_state,
|
|
"node_id": self.focus_node_id,
|
|
"candidate_count": self.focus_candidate_count,
|
|
},
|
|
"evidence": evidence,
|
|
"gaps": gaps,
|
|
"omissions": omissions,
|
|
"summary": {
|
|
"evidence_count": len(evidence),
|
|
"gap_count": len(gaps),
|
|
"omission_count": len(omissions),
|
|
"selected_count": self.selected_count,
|
|
"examined_edge_count": self.examined_edge_count,
|
|
"estimated_tokens": self.estimated_tokens,
|
|
"unclassified_relations": list(self.unclassified_relations),
|
|
},
|
|
}
|
|
if include_hashes:
|
|
result["collection_hash"] = self.collection_hash
|
|
result["capsule_hash"] = self.capsule_hash
|
|
return result
|
|
|
|
|
|
def finalize_capsule(
|
|
*,
|
|
state: Literal["complete", "incomplete", "blocked"],
|
|
task_kind: TaskKind,
|
|
generation: dict[str, object],
|
|
plan: RetrievalPlanV1,
|
|
focus_state: Literal["resolved", "not_found", "ambiguous"],
|
|
focus_node_id: str | None,
|
|
focus_candidate_count: int,
|
|
evidence: tuple[CapsuleEvidenceV1, ...],
|
|
gaps: tuple[EvidenceGapV1, ...],
|
|
omissions: tuple[CapsuleOmissionV1, ...],
|
|
selected_count: int,
|
|
examined_edge_count: int,
|
|
estimated_tokens: int,
|
|
unclassified_relations: tuple[str, ...],
|
|
) -> ContextCapsuleV1:
|
|
collection_hash = canonical_hash(
|
|
{
|
|
"generation": generation,
|
|
"plan_hash": plan.plan_hash,
|
|
"evidence": [item.evidence_hash for item in evidence],
|
|
"gaps": [gap.as_dict() for gap in gaps],
|
|
"omissions": [omission.as_dict() for omission in omissions],
|
|
}
|
|
)
|
|
placeholder = ContextCapsuleV1(
|
|
schema_version=1,
|
|
state=state,
|
|
task_kind=task_kind,
|
|
generation=tuple(generation.items()),
|
|
plan=plan,
|
|
focus_state=focus_state,
|
|
focus_node_id=focus_node_id,
|
|
focus_candidate_count=focus_candidate_count,
|
|
evidence=evidence,
|
|
gaps=gaps,
|
|
omissions=omissions,
|
|
selected_count=selected_count,
|
|
examined_edge_count=examined_edge_count,
|
|
estimated_tokens=estimated_tokens,
|
|
unclassified_relations=unclassified_relations,
|
|
collection_hash=collection_hash,
|
|
capsule_hash="",
|
|
)
|
|
capsule_hash = canonical_hash(
|
|
{
|
|
**placeholder.payload(include_hashes=False),
|
|
"collection_hash": collection_hash,
|
|
}
|
|
)
|
|
return replace(placeholder, capsule_hash=capsule_hash)
|
|
|
|
|
|
def estimate_tokens(text: str) -> int:
|
|
return max(1, (len(text) + 3) // 4)
|
|
|
|
|
|
def edge_tuple(edge: Edge) -> tuple[str, str, str]:
|
|
return edge.source_id, edge.relation, edge.target_id
|
|
|
|
|
|
def _node_text(node: Node) -> str:
|
|
return (
|
|
f"ID: {node.node_id}\nTitle: {node.title}\nFamily: {node.family}\n"
|
|
f"Authority: {node.authority}\nStatus: {node.status}\nSource: {node.source_path}\n"
|
|
f"Summary: {node.summary}\n\n{node.content}"
|
|
)
|
|
|
|
|
|
def _bounded_value(
|
|
value: int | None,
|
|
*,
|
|
default: int,
|
|
maximum: int,
|
|
code: str,
|
|
) -> int:
|
|
selected = default if value is None else value
|
|
if type(selected) is not int or selected < 1 or selected > maximum:
|
|
raise DocForgeError(code, "Task context limit is outside the configured range")
|
|
return selected
|
|
|
|
|
|
def _is_sha256(value: object) -> bool:
|
|
return (
|
|
isinstance(value, str)
|
|
and len(value) == 64
|
|
and all(character in "0123456789abcdef" for character in value)
|
|
)
|