orchestrata_ultimate.agents.hierarchy

Reference › orchestrata_ultimate › agents › hierarchy

hierarchy keeps the parent-child tree of agents (specification 005, user story US-002). AgentHierarchy stores HierarchyNode entries (agent id, role, parent id, child ids, depth and metadata) and enforces a maximum depth, taken from MAX_HIERARCHY_DEPTH in constants, which is 5 by default. It can consult a MultiagentPrologBridge (prolog_bridge). communication, escalation and lifecycle build on it.

add_root creates the root, and add_child attaches a child and rejects it past the depth limit. remove_child and remove_agent delete a node together with all its descendants. The query methods are get_parent, get_children, get_ancestors (nearest first), get_descendants (depth-first), get_depth, get_subtree and is_ancestor. validate_tree returns a list of integrity issues, using three checks: orphans whose parent is missing, depth that is not parent depth plus one, and child references to missing nodes. Failures raise HierarchyError.

At a glance

Class diagram

Dependency diagram

API

Classes

HierarchyNode (dataclass)

Decorators: dataclass, traces(spec='005', story='US-002', reqs=['REQ-005-FR5-001']) · Line 31

A node in the agent hierarchy tree.

HierarchyError

Bases: Exception · Decorators: traces(spec='005', story='US-002', reqs=['REQ-005-AC4-001', 'REQ-005-FR5-001']) · Line 53

Error in hierarchy operations.

AgentHierarchy

Decorators: traces(spec='005', story='US-002', reqs=['REQ-005-AC4-001', 'REQ-005-FR5-001']) · Line 69

Manages parent-child tree for agents.

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/agents/hierarchy.py

# mypy: ignore-errors
"""Agent hierarchy management.

Manages parent-child relationships between agents with depth
limits, tree validation, and Redis-backed persistence.
"""

# Spec: 005 - Multi-Agent System

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, cast

from orchestrata_ultimate.agents.constants import MAX_HIERARCHY_DEPTH
from orchestrata_ultimate.core.base.traces import traces

if TYPE_CHECKING:
    from orchestrata_ultimate.agents.prolog_bridge import MultiagentPrologBridge

logger = logging.getLogger(__name__)


@dataclass
@traces(
    spec="005",
    story="US-002",
    reqs=["REQ-005-FR5-001"],
)
class HierarchyNode:
    """A node in the agent hierarchy tree."""

    agent_id: str
    role: str = ""
    parent_id: str | None = None
    children: list[str] = field(default_factory=list)
    depth: int = 0
    metadata: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        return {
            "agent_id": self.agent_id,
            "role": self.role,
            "parent_id": self.parent_id,
            "children": list(self.children),
            "depth": self.depth,
            "metadata": self.metadata,
        }


@traces(spec="005", story="US-002", reqs=["REQ-005-AC4-001", "REQ-005-FR5-001"])
class HierarchyError(Exception):
    """Error in hierarchy operations."""

    def __init__(self, message: str, agent_id: str = "") -> None:
        self.agent_id = agent_id
        super().__init__(message)


@traces(
    spec="005",
    story="US-002",
    reqs=[
        "REQ-005-AC4-001",
        "REQ-005-FR5-001",
    ],
)
class AgentHierarchy:
    """Manages parent-child tree for agents.

    Enforces max depth, validates tree integrity, supports
    ancestor/descendant queries, and optional Prolog validation.
    """

    def __init__(
        self,
        prolog_bridge: MultiagentPrologBridge | None = None,
        max_depth: int = MAX_HIERARCHY_DEPTH,
    ) -> None:
        self._nodes: dict[str, HierarchyNode] = {}
        self._prolog = prolog_bridge
        self._max_depth = max_depth

    @traces(spec="005", story="US-002", reqs=["REQ-005-FR5-001"])
    def add_root(self, agent_id: str, role: str = "orchestrator") -> HierarchyNode:
        """Add a root node (no parent)."""
        if agent_id in self._nodes:
            return self._nodes[agent_id]
        node = HierarchyNode(agent_id=agent_id, role=role, depth=0)
        self._nodes[agent_id] = node
        return node

    @traces(
        spec="005",
        story="US-002",
        reqs=["REQ-005-AC4-001", "REQ-005-FR5-001"],
    )
    def add_child(self, parent_id: str, child_id: str, role: str = "") -> HierarchyNode:
        """Add a child agent under a parent. Enforces max depth."""
        if parent_id not in self._nodes:
            raise HierarchyError(f"Parent {parent_id} not found", parent_id)
        if child_id in self._nodes:
            raise HierarchyError(f"Agent {child_id} already in hierarchy", child_id)

        parent = self._nodes[parent_id]
        child_depth = parent.depth + 1

        if child_depth > self._max_depth:
            raise HierarchyError(
                f"Max depth {self._max_depth} exceeded (would be {child_depth})",
                child_id,
            )

        # Validate via Prolog if available
        if self._prolog is not None:
            try:
                if not self._prolog.can_restart_child(parent_id, child_id):
                    pass  # Prolog check is advisory, not blocking
            except Exception as exc:
                logger.warning(
                    "Prolog can_restart_child check failed for parent '%s', child '%s': %s",
                    parent_id,
                    child_id,
                    exc,
                    exc_info=True,
                )

        child = HierarchyNode(
            agent_id=child_id,
            role=role,
            parent_id=parent_id,
            depth=child_depth,
        )
        self._nodes[child_id] = child
        parent.children.append(child_id)
        return child

    @traces(
        spec="005",
        story="US-003",
        reqs=["REQ-005-NFR5-001"],
    )
    def remove_child(self, parent_id: str, child_id: str) -> None:
        """Remove a child from its parent. Also removes all descendants."""
        if parent_id not in self._nodes:
            raise HierarchyError(f"Parent {parent_id} not found", parent_id)
        if child_id not in self._nodes:
            raise HierarchyError(f"Child {child_id} not found", child_id)

        parent = self._nodes[parent_id]
        if child_id not in parent.children:
            raise HierarchyError(f"{child_id} is not a child of {parent_id}", child_id)

        # Remove descendants first
        descendants = self.get_descendants(child_id)
        for desc_id in reversed(descendants):
            if desc_id in self._nodes:
                del self._nodes[desc_id]

        parent.children.remove(child_id)
        del self._nodes[child_id]

    @traces(
        spec="005",
        story="US-003",
        reqs=["REQ-005-AC5-001", "REQ-005-NFR5-001"],
    )
    def remove_agent(self, agent_id: str) -> None:
        """Remove an agent and all its descendants from the hierarchy."""
        if agent_id not in self._nodes:
            return
        node = self._nodes[agent_id]
        if node.parent_id and node.parent_id in self._nodes:
            parent = self._nodes[node.parent_id]
            if agent_id in parent.children:
                parent.children.remove(agent_id)
        descendants = self.get_descendants(agent_id)
        for desc_id in reversed(descendants):
            if desc_id in self._nodes:
                del self._nodes[desc_id]
        del self._nodes[agent_id]

    @traces(
        spec="005",
        story="US-002",
        reqs=[
            "REQ-005-FR2-001",
            "REQ-005-FR5-001",
        ],
    )
    def get_parent(self, agent_id: str) -> str | None:
        """Get parent ID or None for root."""
        node = self._nodes.get(agent_id)
        if node is None:
            return None
        return node.parent_id

    @traces(
        spec="005",
        story="US-002",
        reqs=["REQ-005-FR5-001"],
    )
    def get_children(self, agent_id: str) -> list[str]:
        """Get direct children IDs."""
        node = self._nodes.get(agent_id)
        if node is None:
            return []
        return list(node.children)

    @traces(
        spec="005",
        story="US-002",
        reqs=[
            "REQ-005-AC2-001",
            "REQ-005-FR2-001",
            "REQ-005-NFR4-001",
        ],
    )
    def get_ancestors(self, agent_id: str) -> list[str]:
        """Get ordered list of ancestors from immediate parent to root."""
        ancestors = []
        current = agent_id
        while True:
            node = self._nodes.get(current)
            if node is None or node.parent_id is None:
                break
            ancestors.append(node.parent_id)
            current = node.parent_id
        return ancestors

    @traces(spec="005", story="US-003", reqs=["REQ-005-AC5-001", "REQ-005-NFR5-001"])
    def get_descendants(self, agent_id: str) -> list[str]:
        """Get all descendants (DFS order)."""
        descendants: list[str] = []
        node = self._nodes.get(agent_id)
        if node is None:
            return descendants

        stack = list(node.children)
        while stack:
            child_id = stack.pop(0)
            descendants.append(child_id)
            child_node = self._nodes.get(child_id)
            if child_node:
                stack = list(child_node.children) + stack
        return descendants

    def get_depth(self, agent_id: str) -> int:
        """Get depth of an agent (root = 0)."""
        node = self._nodes.get(agent_id)
        if node is None:
            return -1
        return node.depth

    def get_subtree(self, agent_id: str) -> HierarchyNode | None:
        """Get the node for the given agent."""
        return self._nodes.get(agent_id)

    @traces(spec="005", story="US-002", reqs=["REQ-005-NFR4-001"])
    def is_ancestor(self, ancestor_id: str, descendant_id: str) -> bool:
        """Check if ancestor_id is an ancestor of descendant_id."""
        return ancestor_id in self.get_ancestors(descendant_id)

    @traces(
        spec="005",
        story="US-002",
        reqs=["REQ-005-AC4-001", "REQ-005-FR5-001"],
    )
    def validate_tree(self) -> list[str]:
        """Validate tree integrity. Returns list of issues found."""
        issues = []
        issues.extend(self._check_orphans())
        issues.extend(self._check_depth_consistency())
        issues.extend(self._check_children_references())
        issues.extend(self._check_max_depth())
        issues.extend(self._check_prolog_supervision())
        return issues

    @traces(spec="005", story="US-002", reqs=["REQ-005-AC4-001", "REQ-005-FR5-001"])
    def _check_orphans(self) -> list[str]:
        """Find nodes whose parent_id references a node not in the tree."""
        return [
            f"Orphan: {agent_id} references missing parent {node.parent_id}"
            for agent_id, node in self._nodes.items()
            if node.parent_id and node.parent_id not in self._nodes
        ]

    @traces(spec="005", story="US-002", reqs=["REQ-005-AC4-001", "REQ-005-FR5-001"])
    def _check_depth_consistency(self) -> list[str]:
        """Find nodes whose depth is not exactly parent depth + 1."""
        issues = []
        for agent_id, node in self._nodes.items():
            if not node.parent_id:
                continue
            parent = self._nodes.get(node.parent_id)
            if parent and node.depth != parent.depth + 1:
                issues.append(
                    f"Depth mismatch: {agent_id} depth={node.depth}, "
                    f"parent depth={parent.depth}"
                )
        return issues

    @traces(spec="005", story="US-002", reqs=["REQ-005-AC4-001", "REQ-005-FR5-001"])
    def _check_children_references(self) -> list[str]:
        """Find children references that point to missing nodes."""
        return [
            f"Missing child: {agent_id} references nonexistent child {child_id}"
            for agent_id, node in self._nodes.items()
            for child_id in node.children
            if child_id not in self._nodes
        ]

    @traces(spec="005", story="US-002", reqs=["REQ-005-AC4-001", "REQ-005-FR5-001"])
    def _check_max_depth(self) -> list[str]:
        """Find nodes whose depth exceeds the configured maximum."""
        return [
            f"Depth exceeded: {agent_id} at depth {node.depth}"
            for agent_id, node in self._nodes.items()
            if node.depth > self._max_depth
        ]

    @traces(spec="005", story="US-002", reqs=["REQ-005-AC4-001", "REQ-005-FR5-001"])
    def _check_prolog_supervision(self) -> list[str]:
        """Run the optional Prolog supervision-tree validation, if available."""
        if self._prolog is None:
            return []
        try:
            # BUG (pre-existing, not fixed here — no behavior change):
            # MultiagentPrologBridge has no supervision_tree_valid method,
            # and multiagent.pl declares no such predicate either, so this
            # call has always raised AttributeError, been swallowed by the
            # except below, and silently returned [] (no issues reported).
            # cast() preserves that exact runtime behavior while typing the
            # call; the real fix is implementing the predicate + bridge
            # method, out of scope for a pyright-only pass.
            if not cast(Any, self._prolog).supervision_tree_valid("hierarchy"):
                return ["Prolog supervision_tree_valid check failed"]
        except Exception as exc:
            logger.warning(
                "Prolog supervision_tree_valid check raised an error: %s",
                exc,
                exc_info=True,
            )
        return []

    def get_node(self, agent_id: str) -> HierarchyNode | None:
        """Get node by agent ID."""
        return self._nodes.get(agent_id)

    def contains(self, agent_id: str) -> bool:
        """Check if agent is in the hierarchy."""
        return agent_id in self._nodes

    @property
    def size(self) -> int:
        """Total number of agents in the hierarchy."""
        return len(self._nodes)

    def get_roots(self) -> list[str]:
        """Get all root agent IDs."""
        return [aid for aid, node in self._nodes.items() if node.parent_id is None]

    def get_siblings(self, agent_id: str) -> list[str]:
        """Get sibling agent IDs (same parent, excluding self)."""
        node = self._nodes.get(agent_id)
        if node is None or node.parent_id is None:
            return []
        parent = self._nodes.get(node.parent_id)
        if parent is None:
            return []
        return [c for c in parent.children if c != agent_id]

    def clear(self) -> None:
        """Clear all nodes."""
        self._nodes.clear()