"""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) )