markitect-main/markitect/prompts/dependencies/models.py

"""
Dependency tracking models for prompt artifact graphs.

Implements FR-6: Dependency Tracking
- FR-6.1: Directed dependency edges between artifacts
- FR-6.2: Cross-space dependency graph
- FR-6.3: Circular dependency detection
"""

from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Dict, Any, List, Optional, Set, Tuple
import uuid


class EdgeType(Enum):
    """Type of dependency relationship between artifacts."""

    REQUIRES = "requires"
    GENERATES = "generates"
    INCLUDES = "includes"


class CircularDependencyError(Exception):
    """
    Raised when a circular dependency is detected in the graph.

    Attributes:
        cycle: List of node IDs forming the cycle
    """

    def __init__(self, message: str, cycle: Optional[List[str]] = None):
        super().__init__(message)
        self.cycle = cycle or []


@dataclass
class DependencyEdge:
    """
    Persistent dependency edge between artifacts.

    Extends the ephemeral DependencyEdge from execution/manifest.py
    with persistence fields (id, run_id, created_at).

    Attributes:
        id: Unique edge identifier
        source_artifact_id: Source artifact ID
        target_artifact_id: Target artifact ID
        run_id: ID of the run that created this edge
        edge_type: Type of dependency relationship
        created_at: When this edge was created
    """

    id: str
    source_artifact_id: str
    target_artifact_id: str
    run_id: str
    edge_type: EdgeType
    created_at: datetime = field(default_factory=datetime.utcnow)

    @classmethod
    def create(
        cls,
        source_artifact_id: str,
        target_artifact_id: str,
        run_id: str,
        edge_type: EdgeType,
    ) -> "DependencyEdge":
        """
        Create a new dependency edge.

        Args:
            source_artifact_id: Source artifact ID
            target_artifact_id: Target artifact ID
            run_id: Run that established this dependency
            edge_type: Type of dependency

        Returns:
            New DependencyEdge instance
        """
        return cls(
            id=str(uuid.uuid4()),
            source_artifact_id=source_artifact_id,
            target_artifact_id=target_artifact_id,
            run_id=run_id,
            edge_type=edge_type,
        )

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            "id": self.id,
            "source_artifact_id": self.source_artifact_id,
            "target_artifact_id": self.target_artifact_id,
            "run_id": self.run_id,
            "edge_type": self.edge_type.value,
            "created_at": self.created_at.isoformat(),
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "DependencyEdge":
        """Create from dictionary."""
        return cls(
            id=data["id"],
            source_artifact_id=data["source_artifact_id"],
            target_artifact_id=data["target_artifact_id"],
            run_id=data["run_id"],
            edge_type=EdgeType(data["edge_type"]),
            created_at=datetime.fromisoformat(data["created_at"]),
        )


# DFS node coloring for cycle detection
_WHITE = 0  # Unvisited
_GRAY = 1   # In current DFS path
_BLACK = 2  # Fully processed


class DependencyGraph:
    """
    In-memory directed graph for dependency analysis.

    Maintains forward and reverse adjacency lists for efficient
    traversal in both directions. Supports cycle detection via
    3-color DFS and topological sorting.

    Implements FR-6.2: Cross-space dependency graph
    Implements FR-6.3: Circular dependency detection
    """

    def __init__(self) -> None:
        self._forward: Dict[str, Set[str]] = {}
        self._reverse: Dict[str, Set[str]] = {}
        self._edge_types: Dict[Tuple[str, str], EdgeType] = {}

    @property
    def nodes(self) -> Set[str]:
        """All nodes in the graph."""
        return set(self._forward.keys()) | set(self._reverse.keys())

    @property
    def edge_count(self) -> int:
        """Total number of edges."""
        return sum(len(targets) for targets in self._forward.values())

    def add_edge(
        self,
        source: str,
        target: str,
        edge_type: EdgeType = EdgeType.REQUIRES,
    ) -> None:
        """
        Add a directed edge from source to target.

        Args:
            source: Source node ID
            target: Target node ID
            edge_type: Type of dependency
        """
        if source not in self._forward:
            self._forward[source] = set()
        if target not in self._forward:
            self._forward[target] = set()
        if source not in self._reverse:
            self._reverse[source] = set()
        if target not in self._reverse:
            self._reverse[target] = set()

        self._forward[source].add(target)
        self._reverse[target].add(source)
        self._edge_types[(source, target)] = edge_type

    def get_successors(self, node: str) -> Set[str]:
        """Get direct successors (dependencies) of a node."""
        return set(self._forward.get(node, set()))

    def get_predecessors(self, node: str) -> Set[str]:
        """Get direct predecessors (dependents) of a node."""
        return set(self._reverse.get(node, set()))

    def get_edge_type(self, source: str, target: str) -> Optional[EdgeType]:
        """Get the edge type between two nodes."""
        return self._edge_types.get((source, target))

    def has_edge(self, source: str, target: str) -> bool:
        """Check if an edge exists between source and target."""
        return target in self._forward.get(source, set())

    def detect_cycles(self) -> List[List[str]]:
        """
        Detect all cycles in the graph using 3-color DFS.

        Returns:
            List of cycles, where each cycle is a list of node IDs
        """
        color: Dict[str, int] = {node: _WHITE for node in self.nodes}
        parent: Dict[str, Optional[str]] = {node: None for node in self.nodes}
        cycles: List[List[str]] = []

        def _dfs(node: str) -> None:
            color[node] = _GRAY
            for neighbor in self._forward.get(node, set()):
                if color.get(neighbor, _WHITE) == _GRAY:
                    # Back edge found - extract cycle
                    cycle = [neighbor]
                    current = node
                    while current != neighbor:
                        cycle.append(current)
                        current = parent[current]
                    cycle.append(neighbor)
                    cycle.reverse()
                    cycles.append(cycle)
                elif color.get(neighbor, _WHITE) == _WHITE:
                    parent[neighbor] = node
                    _dfs(neighbor)
            color[node] = _BLACK

        for node in self.nodes:
            if color.get(node, _WHITE) == _WHITE:
                _dfs(node)

        return cycles

    def has_cycle(self) -> bool:
        """Check if the graph contains any cycles."""
        color: Dict[str, int] = {node: _WHITE for node in self.nodes}

        def _dfs(node: str) -> bool:
            color[node] = _GRAY
            for neighbor in self._forward.get(node, set()):
                if color.get(neighbor, _WHITE) == _GRAY:
                    return True
                if color.get(neighbor, _WHITE) == _WHITE and _dfs(neighbor):
                    return True
            color[node] = _BLACK
            return False

        for node in self.nodes:
            if color.get(node, _WHITE) == _WHITE:
                if _dfs(node):
                    return True
        return False

    def topological_sort(self) -> List[str]:
        """
        Topologically sort the graph nodes.

        Returns:
            Nodes in topological order (dependencies before dependents)

        Raises:
            CircularDependencyError: If graph contains cycles
        """
        color: Dict[str, int] = {node: _WHITE for node in self.nodes}
        result: List[str] = []
        cycle_path: List[str] = []

        def _dfs(node: str) -> bool:
            color[node] = _GRAY
            cycle_path.append(node)
            for neighbor in sorted(self._forward.get(node, set())):
                if color.get(neighbor, _WHITE) == _GRAY:
                    # Extract cycle for error reporting
                    idx = cycle_path.index(neighbor)
                    cycle = cycle_path[idx:] + [neighbor]
                    raise CircularDependencyError(
                        f"Circular dependency detected: {' -> '.join(cycle)}",
                        cycle=cycle,
                    )
                if color.get(neighbor, _WHITE) == _WHITE:
                    _dfs(neighbor)
            cycle_path.pop()
            color[node] = _BLACK
            result.append(node)

        for node in sorted(self.nodes):
            if color.get(node, _WHITE) == _WHITE:
                _dfs(node)

        result.reverse()
        return result

    def get_subgraph(self, node_ids: Set[str]) -> "DependencyGraph":
        """
        Extract a subgraph containing only the specified nodes.

        Args:
            node_ids: Set of node IDs to include

        Returns:
            New DependencyGraph with only the specified nodes and their edges
        """
        subgraph = DependencyGraph()
        for source in node_ids:
            for target in self._forward.get(source, set()):
                if target in node_ids:
                    edge_type = self._edge_types.get(
                        (source, target), EdgeType.REQUIRES
                    )
                    subgraph.add_edge(source, target, edge_type)
        # Ensure isolated nodes are present
        for node_id in node_ids:
            if node_id not in subgraph._forward:
                subgraph._forward[node_id] = set()
            if node_id not in subgraph._reverse:
                subgraph._reverse[node_id] = set()
        return subgraph