Subgraph Query

akms.graph.query_subgraph

query_subgraph.py — Subgraph Query Engine (§2.4 of system design).

Given a task description, seed tags, and agent role, extracts a ranked subgraph for loadout construction. Operates on the compiled graph.json — unaware of the global/local split.

Algorithm (12 steps): 1. Load query profile for agent_role from config 2. Find seed nodes matching any tag in domain_tags 3. Compute ego_graph of radius max_depth from each seed 4. Union all ego_graphs 5. Filter to LOADABLE_STATUSES (tentative, established) 6. Keep seeds + strict seed-anchored traversal via profile edge_types only 7. Apply prefer_domains boost (×1.5) 8. Apply exclude_domains filter 9. Exclude nodes below confidence threshold 10. Rank nodes by profile's rank_formula 11. Cap at MAX_NODES_PER_LOADOUT 12. Inject pitfall nodes (always included up to MAX_PITFALL_NODES)

query_subgraph

query_subgraph(
    G: DiGraph,
    domain_tags: list[str],
    agent_role: AgentRole | str,
    config: PropagationConfig | None = None,
    max_depth: int = 2,
) -> list[tuple[str, dict[str, Any]]]

Extract a ranked subgraph for loadout construction.

Parameters:

Name	Type	Description	Default
G	DiGraph	The compiled knowledge graph (from load_graph or build_graph).	required
domain_tags	list[str]	Tags to seed the search (e.g., ["taichi", "gpu"]).	required
agent_role	AgentRole | str	The agent role selecting a query profile.	required
config	PropagationConfig | None	PropagationConfig (uses defaults if None).	None
max_depth	int	Maximum graph traversal depth from seeds.	2

Returns:

Type	Description
list[tuple[str, dict[str, Any]]]	Ranked list of (node_id, node_data) tuples.
list[tuple[str, dict[str, Any]]]	Pitfall nodes are always included regardless of rank.

Source code in packages/akms/src/akms/graph/query_subgraph.py

@traced("akms.query_subgraph")
def query_subgraph(
    G: nx.DiGraph,
    domain_tags: list[str],
    agent_role: AgentRole | str,
    config: PropagationConfig | None = None,
    max_depth: int = 2,
) -> list[tuple[str, dict[str, Any]]]:
    """Extract a ranked subgraph for loadout construction.

    Args:
        G: The compiled knowledge graph (from load_graph or build_graph).
        domain_tags: Tags to seed the search (e.g., ["taichi", "gpu"]).
        agent_role: The agent role selecting a query profile.
        config: PropagationConfig (uses defaults if None).
        max_depth: Maximum graph traversal depth from seeds.

    Returns:
        Ranked list of (node_id, node_data) tuples.
        Pitfall nodes are always included regardless of rank.
    """
    if config is None:
        config = PropagationConfig()

    # Normalize agent_role to string
    role_key = (
        agent_role.value if isinstance(agent_role, AgentRole) else str(agent_role)
    )

    # ── Step 1: Load query profile ───────────────────────────────────
    profile = config.query_roles.get(role_key)
    if profile is None:
        logger.warning(
            "No query profile for role '%s', using implementer defaults",
            role_key,
        )
        profile = config.query_roles.get("implementer", QueryRoleProfile())

    loadout_config = config.loadout
    max_nodes = loadout_config.max_nodes_per_loadout
    max_pitfall = loadout_config.max_pitfall_nodes
    min_confidence = loadout_config.min_confidence_threshold

    logger.info(
        "query_subgraph: role=%s, tags=%s, depth=%d, max_nodes=%d",
        role_key,
        domain_tags,
        max_depth,
        max_nodes,
    )

    # ── Step 2: Find seed nodes ──────────────────────────────────────
    seeds = _find_seed_nodes(G, domain_tags)
    logger.info("Found %d seed nodes for tags %s", len(seeds), domain_tags)

    if not seeds:
        logger.warning("No seed nodes found for tags %s", domain_tags)
        return []

    # ── Steps 3-4: Ego graph union ───────────────────────────────────
    candidates = _extract_ego_union(G, seeds, max_depth)
    logger.info("Ego union: %d candidates", len(candidates))

    # ── Step 5: Filter by loadable status ────────────────────────────
    candidates = _filter_by_status(G, candidates)
    logger.info("After status filter: %d candidates", len(candidates))

    # Session nodes are non-loadable by contract (FR-G11).
    candidates = {n for n in candidates if _get_node_domain(G, n) != "session"}
    logger.info("After session-domain exclusion: %d candidates", len(candidates))

    # ── Step 6: Strict seed-anchored traversal via allowed edge types ─
    if profile.edge_types:
        candidates = _filter_edges_by_type(G, candidates, seeds, profile.edge_types)
        logger.info("After edge type filter: %d candidates", len(candidates))

    # ── Step 12 (early): Identify pitfall nodes ──────────────────────
    # We identify them early so they're preserved through subsequent filters
    pitfall_nodes = _find_pitfall_nodes(G, candidates)
    logger.info("Pitfall nodes: %d", len(pitfall_nodes))

    # ── Step 7: Apply prefer_domains boost ───────────────────────────
    # (Applied during ranking, not filtering — just track which nodes get boost)
    prefer_domains = set(profile.prefer_domains)

    # ── Step 8: Apply exclude_domains filter ─────────────────────────
    if profile.exclude_domains:
        exclude = set(profile.exclude_domains)
        # Never exclude pitfall nodes
        candidates = {
            n
            for n in candidates
            if _get_node_domain(G, n) not in exclude or n in pitfall_nodes
        }
        logger.info("After domain exclusion: %d candidates", len(candidates))

    # ── Step 9: Exclude nodes below confidence threshold ─────────────
    # Pitfall nodes are exempt from confidence threshold
    candidates = {
        n
        for n in candidates
        if _get_node_confidence(G, n) >= min_confidence or n in pitfall_nodes
    }
    logger.info(
        "After confidence threshold (%.2f): %d candidates",
        min_confidence,
        len(candidates),
    )

    # ── Step 10: Rank nodes ──────────────────────────────────────────
    ranked = []
    for node_id in candidates:
        rank = _compute_rank(G, node_id, profile.rank_formula)

        # Step 7: Prefer domains boost (×1.5)
        if prefer_domains and _get_node_domain(G, node_id) in prefer_domains:
            rank *= 1.5

        ranked.append((node_id, rank))

    # Sort by rank descending, then by id for determinism
    ranked.sort(key=lambda x: (-x[1], x[0]))

    # ── Step 11: Cap at max_nodes ────────────────────────────────────
    # Separate pitfall nodes from regular nodes
    regular_ranked = [(nid, r) for nid, r in ranked if nid not in pitfall_nodes]
    pitfall_ranked = [(nid, r) for nid, r in ranked if nid in pitfall_nodes]

    # Cap pitfall nodes
    pitfall_ranked = pitfall_ranked[:max_pitfall]

    # Cap regular nodes (leaving room for pitfalls)
    regular_slots = max(0, max_nodes - len(pitfall_ranked))
    regular_ranked = regular_ranked[:regular_slots]

    # ── Step 12: Merge and return ────────────────────────────────────
    # Pitfall nodes come first (always included), then ranked regular nodes
    result_ids = [nid for nid, _ in pitfall_ranked] + [nid for nid, _ in regular_ranked]

    # Deduplicate while preserving order
    seen = set()
    unique_ids = []
    for nid in result_ids:
        if nid not in seen:
            seen.add(nid)
            unique_ids.append(nid)

    result = [(nid, dict(G.nodes[nid])) for nid in unique_ids]
    result = [
        (nid, data) for nid, data in result if _get_node_domain(G, nid) != "session"
    ]

    # ── FR-G10: load_with co-activation hints ────────────────────────
    # Promote nodes that selected nodes are flagged to co-load with.
    # These hints are *pragmatic* (almost-always-co-loaded) and, by spec,
    # are distinct from semantic edges — so they bypass the edge-type
    # traversal (Step 6) and confidence threshold (Step 9), like pitfalls.
    # Non-transitive: only one hop out from the already-selected nodes.
    selected_ids = {nid for nid, _ in result}
    loadable_values = {s.value for s in LOADABLE_STATUSES}
    coactivated_ids: set[str] = set()
    for nid, _ in result:
        hints = G.nodes[nid].get("load_with", []) or []
        if isinstance(hints, str):
            hints = [hints]
        for target in hints:
            target = str(target)
            if (
                target in G
                and target not in selected_ids
                and target not in coactivated_ids
                and _get_node_status(G, target) in loadable_values
                and _get_node_domain(G, target) != "session"
            ):
                coactivated_ids.add(target)

    # Append deterministically; tag the copied node_data so downstream
    # rendering (generate_loadout) can distinguish co-activated nodes.
    for target in sorted(coactivated_ids):
        data = dict(G.nodes[target])
        data["_coactivated"] = True
        result.append((target, data))

    logger.info(
        "query_subgraph result: %d nodes (%d pitfall, %d regular, %d co-activated)",
        len(result),
        len(pitfall_ranked),
        len(regular_ranked),
        len(coactivated_ids),
    )

    return result

compute_query_hash

compute_query_hash(
    domain_tags: list[str], agent_role: str, max_depth: int
) -> str

Compute a deterministic hash for a query, used for caching.

Parameters:

Name	Type	Description	Default
domain_tags	list[str]	Sorted list of tags.	required
agent_role	str	Role string.	required
max_depth	int	Traversal depth.	required

Returns:

Type	Description
str	SHA256 hex digest of the query parameters.

Source code in packages/akms/src/akms/graph/query_subgraph.py

def compute_query_hash(
    domain_tags: list[str],
    agent_role: str,
    max_depth: int,
) -> str:
    """Compute a deterministic hash for a query, used for caching.

    Args:
        domain_tags: Sorted list of tags.
        agent_role: Role string.
        max_depth: Traversal depth.

    Returns:
        SHA256 hex digest of the query parameters.
    """
    # Normalize and sort tags for determinism
    normalized_tags = sorted(set(t.lower().strip() for t in domain_tags))
    key = f"{','.join(normalized_tags)}|{agent_role}|{max_depth}"
    return hashlib.sha256(key.encode()).hexdigest()