orchestrata_ultimate.knowledge_graph.context_pack.selector

Reference › orchestrata_ultimate › knowledge_graph › context_pack › selector

Deterministic selection and ranking, the second step of the context pack pipeline (SPEC-106 FR-003). For each TaskType there is one parameterised Cypher template (a frozen _Template with the relationships to expand and a hop limit, REQ-106-020), generalising the six older SPEC-009 inject_for_task handlers onto the canonical schema. Containment becomes canonical DEFINES expansion over Module, Class and Function; requirement traceability becomes TRACES_TO expansion from code to RequirementItem; and TESTED_BY never expands the subgraph but is a boolean ranking signal. Expansion is bounded to k hops from the frame's linked entities (default 2, REQ-106-021), with debug using an override of 1.

select takes a RequestFrame from frame and a synchronous repository, runs the template against the operational graph (orchkg on orchgraph) and returns a RankedSubgraph, best first. Each RankedNode records its template, hop distance, trace degree, tested_by, ambiguity flag and final rank_score. The result goes to packer and may be refined by relevance.

At a glance

Class diagram

Dependency diagram

API

Constants

Classes

_Template (dataclass)

Decorators: dataclass(frozen=True) · Line 68

One TaskType's selection template.

RankedNode (dataclass)

Decorators: dataclass(frozen=True) · Line 137

One selected node with its ranking inputs and final score.

RankedSubgraph (dataclass)

Decorators: dataclass(frozen=True) · Line 152

The result of :func:select — nodes already ranked (best first).

Functions

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/knowledge_graph/context_pack/selector.py

"""Deterministic selection — per-task-type templates + graph-signal ranking
(SPEC-106 FR-003).

One parameterized Cypher template per :class:`TaskType` (REQ-106-020),
generalizing the six SPEC-009 ``inject_for_task`` handlers
(``_inject_build_context`` .. ``_inject_review_context``) onto the FR-001
canon (REQ-106-025):

* containment — the legacy File/Class/Method AST-subtree walk becomes
  canonical ``DEFINES`` expansion (Module/Class/Function);
* requirement traceability — the legacy Spec/Requirement pull becomes
  canonical ``TRACES_TO`` expansion (code -> RequirementItem);
* test coverage — ``TESTED_BY`` never expands the subgraph (``Test`` is not
  a canonical FR-001 label — see ``canonical_schema``'s module docstring)
  but IS a ranking signal ("TESTED_BY presence", REQ-106-022), computed as a
  boolean per selected node.

Expansion is bounded to ``k`` hops from the frame's linked entities
(default 2, REQ-106-021); DEBUG uses a genuine per-template override of 1 —
debugging wants a tight, precise radius around the failing entity, mirroring
how narrowly-scoped ``_inject_debug_context`` already was.

Ranking (REQ-106-022) combines hop distance, trace-link degree, TESTED_BY
presence, recency (``updated_at`` when present — a graph-stored property,
never wall-clock, so this stays deterministic given fixed graph state), and
an ambiguity penalty for nodes reached through an ambiguous linked entity.
Ties break lexicographically on node id (REQ-106-023). No embedding/LLM
scoring participates anywhere (REQ-106-024).
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.knowledge_graph.canonical_schema import (
    CANONICAL_LABELS,
    REL_DEFINES,
    REL_TESTED_BY,
    REL_TRACES_TO,
)
from orchestrata_ultimate.knowledge_graph.context_injector import TaskType
from orchestrata_ultimate.knowledge_graph.context_pack.frame import (
    LINKABLE_LABELS,
    RequestFrame,
)

if TYPE_CHECKING:
    from orchestrata_ultimate.core.graph import SyncGraphRepository

DEFAULT_K_HOP = 2

# Ranking weights (REQ-106-022). Deliberately simple linear combination —
# hop distance dominates (it is the primary "how close to what was asked
# for" signal), then TESTED_BY presence, then trace-link degree, with
# recency as a light tiebreaker-ish nudge and the ambiguity penalty large
# enough to always sink an ambiguous-origin node below unambiguous ones at
# the same hop.
_HOP_WEIGHT = 10.0
_TRACE_DEGREE_WEIGHT = 1.0
_TESTED_BY_BONUS = 5.0
_RECENCY_WEIGHT = 0.001
_AMBIGUITY_PENALTY = 100.0


@dataclass(frozen=True)
class _Template:
    """One TaskType's selection template."""

    name: str
    expand_rels: tuple[str, ...]
    k_hop: int


# REQ-106-025: each template's comment names the SPEC-009 handler (and its
# traced REQs) it generalizes, preserving that coverage.
_TEMPLATES: dict[TaskType, _Template] = {
    # _inject_build_context (SPEC-009 REQ-009-US-001-003, REQ-017-005): pulled
    # Spec + Requirement context. Canonical analogue: requirement
    # traceability in both directions from the linked entities — PLUS
    # DEFINES containment (added 2026-07-21): with sparse TRACES_TO coverage
    # a Class root selected NOTHING but itself (measured: 148/4000 tokens
    # spent, zero structural context for a build brief). DEFINES gives the
    # entity's members/module — additive, loses no SPEC-009 coverage.
    TaskType.BUILD: _Template("build", (REL_TRACES_TO, REL_DEFINES), DEFAULT_K_HOP),
    # _inject_refactor_context (SPEC-009 REQ-009-US-002-001..007): walked one
    # file's AST subtree. Canonical analogue: DEFINES containment.
    TaskType.REFACTOR: _Template("refactor", (REL_DEFINES,), DEFAULT_K_HOP),
    # _inject_debug_context (SPEC-009 REQ-009-US-003-001..004): narrowly
    # test-focused — a real per-template k-hop override (REQ-106-021).
    TaskType.DEBUG: _Template("debug", (REL_DEFINES,), 1),
    # _inject_document_context (SPEC-009 REQ-009-US-004-001..004): containment
    # only (module/class/function tree for doc generation).
    TaskType.DOCUMENT: _Template("document", (REL_DEFINES,), DEFAULT_K_HOP),
    # _inject_test_context (SPEC-009 REQ-009-US-005-001..004): code paired
    # with the requirements it traces to, to test against.
    TaskType.TEST: _Template("test", (REL_DEFINES, REL_TRACES_TO), DEFAULT_K_HOP),
    # _inject_review_context (SPEC-009 REQ-009-US-006-001..005): the richest
    # handler — structure + traceability (+ TESTED_BY as a ranking signal,
    # universal below).
    TaskType.REVIEW: _Template("review", (REL_DEFINES, REL_TRACES_TO), DEFAULT_K_HOP),
}

# Root lookup is restricted to the linkable canonical labels so Neo4j can use
# each label's canonical-key index rather than a label-less property scan.
_ROOT_LABEL_FILTER = " OR ".join(f"root:`{label}`" for label in LINKABLE_LABELS)
# Degree via PATTERN COMPREHENSION — the only form BOTH engines accept.
# Neo4j 5.26 rejects a pattern expression inside size(); orchgraph/Memgraph
# rejects both COUNT{} and EXISTS{} ("Not yet implemented: Exists can only be
# used in ...", verified live 2026-07-26). size([...pattern... | 1]) works on
# both, so it is what the cutover left standing.
_TRACE_DEGREE_EXPR = f"size([(n)-[:`{REL_TRACES_TO}`]-() | 1])"
_TESTED_BY_EXPR = f"size([(n)-[:`{REL_TESTED_BY}`]->() | 1])"


def _build_cypher(template: _Template) -> str:
    rel_pattern = "|".join(f"`{rel}`" for rel in template.expand_rels)
    return (
        "UNWIND $root_ids AS rid "
        f"MATCH (root) WHERE ({_ROOT_LABEL_FILTER}) AND (root.id = rid OR root.alias = rid) "
        "WITH collect(DISTINCT root) AS roots "
        "UNWIND roots AS r "
        f"MATCH p = (r)-[:{rel_pattern}*0..{template.k_hop}]-(n) "
        "WITH n, min(length(p)) AS hop "
        f"RETURN n, hop, {_TRACE_DEGREE_EXPR} AS trace_degree, "
        f"{_TESTED_BY_EXPR} AS tested_by_count"
    )


_CYPHER_BY_TASK_TYPE: dict[TaskType, str] = {
    task_type: _build_cypher(template) for task_type, template in _TEMPLATES.items()
}


@dataclass(frozen=True)
class RankedNode:
    """One selected node with its ranking inputs and final score."""

    node_id: str
    label: str
    properties: dict[str, Any]
    template: str
    hop_distance: int
    trace_degree: int
    tested_by: bool
    ambiguous: bool
    rank_score: float


@dataclass(frozen=True)
class RankedSubgraph:
    """The result of :func:`select` — nodes already ranked (best first)."""

    task_type: TaskType
    nodes: tuple[RankedNode, ...] = ()


def _root_ids(frame: RequestFrame) -> list[str]:
    """Unambiguous root ids — the only roots the k-hop expansion grows from.

    Ambiguous entities' candidates are deliberately NOT expansion roots
    (2026-07-21): expanding every candidate of an ambiguous token filled
    packs with rank-penalized noise neighbors (measured: 3 of 4 sections in
    a build pack). Per REQ-106-013 the candidates are still RETAINED — they
    are fetched hop-0 (themselves only) by :func:`select`.
    """
    ids: list[str] = []
    for entity in frame.entities:
        if not entity.ambiguous and entity.node_id:
            ids.append(entity.node_id)
    return list(dict.fromkeys(ids))


def _ambiguous_ids(frame: RequestFrame) -> frozenset[str]:
    """Every node id that originated from an ambiguous linked entity."""
    return frozenset(
        node_id
        for entity in frame.entities
        if entity.ambiguous
        for _, node_id in entity.candidates
    )


def _pick_label(labels: list[str]) -> str:
    """Pick a deterministic, canonical label for a node with 1+ labels."""
    for label in sorted(labels):
        if label in CANONICAL_LABELS:
            return label
    return sorted(labels)[0] if labels else ""


def _score(
    hop: int,
    trace_degree: int,
    tested_by: bool,
    updated_at: Any,
    ambiguous: bool,
) -> float:
    """Combine the REQ-106-022 graph signals into one comparable score."""
    score = -_HOP_WEIGHT * hop + _TRACE_DEGREE_WEIGHT * trace_degree
    if tested_by:
        score += _TESTED_BY_BONUS
    if isinstance(updated_at, int | float):
        score += _RECENCY_WEIGHT * float(updated_at)
    if ambiguous:
        score -= _AMBIGUITY_PENALTY
    return score


@traces(
    spec="106",
    reqs=[
        "REQ-106-020",
        "REQ-106-021",
        "REQ-106-022",
        "REQ-106-023",
        "REQ-106-024",
        "REQ-106-025",
        "REQ-106-033",
        "REQ-106-034",
        "REQ-106-035",
    ],
)
def select(frame: RequestFrame, repo: SyncGraphRepository) -> RankedSubgraph:
    """Select and rank the subgraph rooted at ``frame``'s linked entities (FR-003).

    Args:
        frame: A :class:`RequestFrame` from :func:`analyze_request`.
        repo: A connected ``SyncGraphRepository``.

    Returns:
        A :class:`RankedSubgraph` — nodes sorted best-first (highest score,
        ties broken lexicographically on node id, REQ-106-023). Empty when
        the frame linked no entities at all (a legitimate, non-error result
        distinct from KG unavailability, which is the caller's concern).
    """
    root_ids = _root_ids(frame)
    ambiguous_ids = _ambiguous_ids(frame)
    if not root_ids and not ambiguous_ids:
        return RankedSubgraph(task_type=frame.task_type, nodes=())

    template = _TEMPLATES[frame.task_type]
    cypher = _CYPHER_BY_TASK_TYPE[frame.task_type]

    records: list[Any] = []
    if root_ids:
        records.extend(repo.execute_read(cypher, {"root_ids": root_ids}))
    if ambiguous_ids:
        # Ambiguous candidates: retained (REQ-106-013) but hop-0 ONLY —
        # they contribute themselves, never their neighborhoods. Reuse the
        # root-lookup shape so label indexes serve the match.
        hop0_cypher = (
            "UNWIND $root_ids AS rid "
            f"MATCH (n) WHERE ({_ROOT_LABEL_FILTER.replace('root:', 'n:')}) "
            "AND (n.id = rid OR n.alias = rid) "
            f"RETURN n, 0 AS hop, {_TRACE_DEGREE_EXPR} AS trace_degree, "
            f"{_TESTED_BY_EXPR} AS tested_by_count"
        )
        records.extend(
            repo.execute_read(hop0_cypher, {"root_ids": sorted(ambiguous_ids)})
        )

    nodes: list[RankedNode] = []
    seen_ids: set[str] = set()
    for record in records:
        n = record["n"]
        node_id = n.get("id") or n.get("alias")
        if not node_id:
            # Reached via a canonical rel but carries neither canonical
            # key — not a linkable/rankable entity (e.g. a non-canonical
            # label attached to a canonical node); skip it rather than
            # emit a section with no identity.
            continue
        if node_id in seen_ids:
            # Expansion records precede hop-0 ambiguous records, so a node
            # reached both ways keeps its real hop distance.
            continue
        seen_ids.add(node_id)
        hop = record["hop"]
        trace_degree = record["trace_degree"]
        tested_by = record["tested_by_count"] > 0
        properties = dict(n)
        label = _pick_label(list(n.labels))
        ambiguous = node_id in ambiguous_ids
        score = _score(
            hop, trace_degree, tested_by, properties.get("updated_at"), ambiguous
        )
        nodes.append(
            RankedNode(
                node_id=node_id,
                label=label,
                properties=properties,
                template=template.name,
                hop_distance=hop,
                trace_degree=trace_degree,
                tested_by=tested_by,
                ambiguous=ambiguous,
                rank_score=score,
            )
        )

    nodes.sort(key=lambda rn: (-rn.rank_score, rn.node_id))
    return RankedSubgraph(task_type=frame.task_type, nodes=tuple(nodes))