yamlgraph 0.1.1__py3-none-any.whl

yamlgraph/graph_loader.py

@@ -0,0 +1,337 @@
+"""YAML Graph Loader - Compile YAML to LangGraph.
+
+This module provides functionality to load graph definitions from YAML files
+and compile them into LangGraph StateGraph instances.
+"""
+
+import logging
+from pathlib import Path
+from typing import Any
+
+import yaml
+from langgraph.graph import END, StateGraph
+
+from yamlgraph.constants import NodeType
+from yamlgraph.map_compiler import compile_map_node
+from yamlgraph.models.state_builder import build_state_class
+from yamlgraph.node_factory import create_node_function, resolve_class
+from yamlgraph.routing import make_expr_router_fn, make_router_fn
+from yamlgraph.tools.agent import create_agent_node
+from yamlgraph.tools.nodes import create_tool_node
+from yamlgraph.tools.python_tool import create_python_node, parse_python_tools
+from yamlgraph.tools.shell import parse_tools
+from yamlgraph.utils.validators import validate_config
+
+# Type alias for dynamic state
+GraphState = dict[str, Any]
+
+logger = logging.getLogger(__name__)
+
+
+class GraphConfig:
+    """Parsed graph configuration from YAML."""
+
+    def __init__(self, config: dict):
+        """Initialize from parsed YAML dict.
+
+        Args:
+            config: Parsed YAML configuration dictionary
+
+        Raises:
+            ValueError: If config is invalid
+        """
+        # Validate before storing
+        validate_config(config)
+
+        self.version = config.get("version", "1.0")
+        self.name = config.get("name", "unnamed")
+        self.description = config.get("description", "")
+        self.defaults = config.get("defaults", {})
+        self.nodes = config.get("nodes", {})
+        self.edges = config.get("edges", [])
+        self.tools = config.get("tools", {})
+        self.state_class = config.get("state_class", "")
+        self.loop_limits = config.get("loop_limits", {})
+        # Store raw config for dynamic state building
+        self.raw_config = config
+
+
+def load_graph_config(path: str | Path) -> GraphConfig:
+    """Load and parse a YAML graph definition.
+
+    Args:
+        path: Path to the YAML file
+
+    Returns:
+        GraphConfig instance
+
+    Raises:
+        FileNotFoundError: If the file doesn't exist
+        ValueError: If the YAML is invalid or missing required fields
+    """
+    path = Path(path)
+    if not path.exists():
+        raise FileNotFoundError(f"Graph config not found: {path}")
+
+    with open(path) as f:
+        config = yaml.safe_load(f)
+
+    return GraphConfig(config)
+
+
+def _resolve_state_class(config: GraphConfig) -> type:
+    """Resolve the state class for the graph.
+
+    Uses dynamic state generation unless explicit state_class is set
+    (deprecated).
+
+    Args:
+        config: Graph configuration
+
+    Returns:
+        TypedDict class for graph state
+    """
+    if config.state_class and config.state_class != "yamlgraph.models.GraphState":
+        import warnings
+
+        warnings.warn(
+            f"state_class '{config.state_class}' is deprecated. "
+            "State is now auto-generated from graph config.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return resolve_class(config.state_class)
+    return build_state_class(config.raw_config)
+
+
+def _parse_all_tools(
+    config: GraphConfig,
+) -> tuple[dict[str, Any], dict[str, Any]]:
+    """Parse shell and Python tools from config.
+
+    Args:
+        config: Graph configuration
+
+    Returns:
+        Tuple of (shell_tools, python_tools) dictionaries
+    """
+    tools = parse_tools(config.tools)
+    python_tools = parse_python_tools(config.tools)
+
+    if tools:
+        logger.info(f"Parsed {len(tools)} shell tools: {', '.join(tools.keys())}")
+    if python_tools:
+        logger.info(
+            f"Parsed {len(python_tools)} Python tools: {', '.join(python_tools.keys())}"
+        )
+
+    return tools, python_tools
+
+
+def _compile_node(
+    node_name: str,
+    node_config: dict[str, Any],
+    graph: StateGraph,
+    config: GraphConfig,
+    tools: dict[str, Any],
+    python_tools: dict[str, Any],
+) -> tuple[str, Any] | None:
+    """Compile a single node and add to graph.
+
+    Args:
+        node_name: Name of the node
+        node_config: Node configuration dict
+        graph: StateGraph to add node to
+        config: Full graph config for defaults
+        tools: Shell tools registry
+        python_tools: Python tools registry
+
+    Returns:
+        Tuple of (node_name, map_info) for map nodes, None otherwise
+    """
+    # Copy node config and add loop_limit if specified
+    enriched_config = dict(node_config)
+    if node_name in config.loop_limits:
+        enriched_config["loop_limit"] = config.loop_limits[node_name]
+
+    node_type = node_config.get("type", NodeType.LLM)
+
+    if node_type == NodeType.TOOL:
+        node_fn = create_tool_node(node_name, enriched_config, tools)
+        graph.add_node(node_name, node_fn)
+    elif node_type == NodeType.PYTHON:
+        node_fn = create_python_node(node_name, enriched_config, python_tools)
+        graph.add_node(node_name, node_fn)
+    elif node_type == NodeType.AGENT:
+        node_fn = create_agent_node(node_name, enriched_config, tools)
+        graph.add_node(node_name, node_fn)
+    elif node_type == NodeType.MAP:
+        map_edge_fn, sub_node_name = compile_map_node(
+            node_name, enriched_config, graph, config.defaults
+        )
+        logger.info(f"Added node: {node_name} (type={node_type})")
+        return (node_name, (map_edge_fn, sub_node_name))
+    else:
+        # LLM and router nodes
+        node_fn = create_node_function(node_name, enriched_config, config.defaults)
+        graph.add_node(node_name, node_fn)
+
+    logger.info(f"Added node: {node_name} (type={node_type})")
+    return None
+
+
+def _compile_nodes(
+    config: GraphConfig,
+    graph: StateGraph,
+    tools: dict[str, Any],
+    python_tools: dict[str, Any],
+) -> dict[str, tuple]:
+    """Compile all nodes and add to graph.
+
+    Args:
+        config: Graph configuration
+        graph: StateGraph to add nodes to
+        tools: Shell tools registry
+        python_tools: Python tools registry
+
+    Returns:
+        Dict of map_nodes: name -> (map_edge_fn, sub_node_name)
+    """
+    map_nodes: dict[str, tuple] = {}
+
+    for node_name, node_config in config.nodes.items():
+        result = _compile_node(
+            node_name, node_config, graph, config, tools, python_tools
+        )
+        if result:
+            map_nodes[result[0]] = result[1]
+
+    return map_nodes
+
+
+def _process_edge(
+    edge: dict[str, Any],
+    graph: StateGraph,
+    map_nodes: dict[str, tuple],
+    router_edges: dict[str, list],
+    expression_edges: dict[str, list[tuple[str, str]]],
+) -> None:
+    """Process a single edge and add to graph or edge tracking dicts.
+
+    Args:
+        edge: Edge configuration dict
+        graph: StateGraph to add edges to
+        map_nodes: Map node tracking dict
+        router_edges: Dict to collect router edges
+        expression_edges: Dict to collect expression-based edges
+    """
+    from_node = edge["from"]
+    to_node = edge["to"]
+    condition = edge.get("condition")
+    edge_type = edge.get("type")
+
+    if from_node == "START":
+        graph.set_entry_point(to_node)
+    elif isinstance(to_node, str) and to_node in map_nodes:
+        # Edge TO a map node: use conditional edge with Send function
+        map_edge_fn, sub_node_name = map_nodes[to_node]
+        graph.add_conditional_edges(from_node, map_edge_fn, [sub_node_name])
+    elif from_node in map_nodes:
+        # Edge FROM a map node: wire sub_node to next_node for fan-in
+        _, sub_node_name = map_nodes[from_node]
+        target = END if to_node == "END" else to_node
+        graph.add_edge(sub_node_name, target)
+    elif edge_type == "conditional" and isinstance(to_node, list):
+        # Router-style conditional edge: store for later processing
+        router_edges[from_node] = to_node
+    elif condition:
+        # Expression-based condition (e.g., "critique.score < 0.8")
+        if from_node not in expression_edges:
+            expression_edges[from_node] = []
+        target = END if to_node == "END" else to_node
+        expression_edges[from_node].append((condition, target))
+    elif to_node == "END":
+        graph.add_edge(from_node, END)
+    else:
+        graph.add_edge(from_node, to_node)
+
+
+def _add_conditional_edges(
+    graph: StateGraph,
+    router_edges: dict[str, list],
+    expression_edges: dict[str, list[tuple[str, str]]],
+) -> None:
+    """Add router and expression conditional edges to graph.
+
+    Args:
+        graph: StateGraph to add edges to
+        router_edges: Router-style conditional edges
+        expression_edges: Expression-based conditional edges
+    """
+    # Add router conditional edges
+    for source_node, target_nodes in router_edges.items():
+        route_mapping = {target: target for target in target_nodes}
+        graph.add_conditional_edges(
+            source_node,
+            make_router_fn(target_nodes),
+            route_mapping,
+        )
+
+    # Add expression-based conditional edges
+    for source_node, expr_edges in expression_edges.items():
+        targets = {target for _, target in expr_edges}
+        targets.add(END)  # Always include END as fallback
+        route_mapping = {t: (END if t == END else t) for t in targets}
+        graph.add_conditional_edges(
+            source_node,
+            make_expr_router_fn(expr_edges, source_node),
+            route_mapping,
+        )
+
+
+def compile_graph(config: GraphConfig) -> StateGraph:
+    """Compile a GraphConfig to a LangGraph StateGraph.
+
+    Args:
+        config: Parsed graph configuration
+
+    Returns:
+        StateGraph ready for compilation
+    """
+    # Build state class and create graph
+    state_class = _resolve_state_class(config)
+    graph = StateGraph(state_class)
+
+    # Parse all tools
+    tools, python_tools = _parse_all_tools(config)
+
+    # Compile all nodes
+    map_nodes = _compile_nodes(config, graph, tools, python_tools)
+
+    # Process edges
+    router_edges: dict[str, list] = {}
+    expression_edges: dict[str, list[tuple[str, str]]] = {}
+
+    for edge in config.edges:
+        _process_edge(edge, graph, map_nodes, router_edges, expression_edges)
+
+    # Add conditional edges
+    _add_conditional_edges(graph, router_edges, expression_edges)
+
+    return graph
+
+
+def load_and_compile(path: str | Path) -> StateGraph:
+    """Load YAML and compile to StateGraph.
+
+    Convenience function combining load_graph_config and compile_graph.
+
+    Args:
+        path: Path to YAML graph definition
+
+    Returns:
+        StateGraph ready for compilation
+    """
+    config = load_graph_config(path)
+    logger.info(f"Loaded graph config: {config.name} v{config.version}")
+    return compile_graph(config)

yamlgraph/map_compiler.py

@@ -0,0 +1,138 @@
+"""Map node compiler - Handles type: map node compilation.
+
+This module provides functionality to compile map nodes that fan out
+to sub-nodes for parallel processing using LangGraph's Send mechanism.
+"""
+
+import logging
+from collections.abc import Callable
+from typing import Any
+
+from langgraph.graph import StateGraph
+from langgraph.types import Send
+
+from yamlgraph.node_factory import create_node_function
+from yamlgraph.utils.expressions import resolve_state_expression
+
+logger = logging.getLogger(__name__)
+
+
+def wrap_for_reducer(
+    node_fn: Callable[[dict], dict],
+    collect_key: str,
+    state_key: str,
+) -> Callable[[dict], dict]:
+    """Wrap sub-node output for Annotated reducer aggregation.
+
+    Handles error propagation: if a map branch fails, the error is
+    included in the result with the _map_index for tracking.
+
+    Args:
+        node_fn: The original node function
+        collect_key: State key where results are collected
+        state_key: Key to extract from node result
+
+    Returns:
+        Wrapped function that outputs in reducer-compatible format
+    """
+
+    def wrapped(state: dict) -> dict:
+        try:
+            result = node_fn(state)
+        except Exception as e:
+            # Propagate error with map index
+            from yamlgraph.models import PipelineError
+
+            error_result = {
+                "_map_index": state.get("_map_index", 0),
+                "_error": str(e),
+                "_error_type": type(e).__name__,
+            }
+            return {
+                collect_key: [error_result],
+                "errors": [PipelineError.from_exception(e, node="map_subnode")],
+            }
+
+        # Check if result contains an error
+        if "errors" in result or "error" in result:
+            error_result = {
+                "_map_index": state.get("_map_index", 0),
+                "_error": str(result.get("errors") or result.get("error")),
+            }
+            # Preserve errors in output
+            output = {collect_key: [error_result]}
+            if "errors" in result:
+                output["errors"] = result["errors"]
+            return output
+
+        extracted = result.get(state_key, result)
+
+        # Convert Pydantic models to dicts
+        if hasattr(extracted, "model_dump"):
+            extracted = extracted.model_dump()
+
+        # Include _map_index if present for ordering
+        if "_map_index" in state:
+            if isinstance(extracted, dict):
+                extracted = {"_map_index": state["_map_index"], **extracted}
+            else:
+                extracted = {"_map_index": state["_map_index"], "value": extracted}
+
+        return {collect_key: [extracted]}
+
+    return wrapped
+
+
+def compile_map_node(
+    name: str,
+    config: dict[str, Any],
+    builder: StateGraph,
+    defaults: dict[str, Any],
+) -> tuple[Callable[[dict], list[Send]], str]:
+    """Compile type: map node using LangGraph Send.
+
+    Creates a sub-node and returns a map edge function that fans out
+    to the sub-node for each item in the list.
+
+    Args:
+        name: Name of the map node
+        config: Map node configuration with 'over', 'as', 'node', 'collect'
+        builder: StateGraph builder to add sub-node to
+        defaults: Default configuration for nodes
+
+    Returns:
+        Tuple of (map_edge_function, sub_node_name)
+    """
+    over_expr = config["over"]
+    item_var = config["as"]
+    sub_node_name = f"_map_{name}_sub"
+    collect_key = config["collect"]
+    sub_node_config = dict(config["node"])  # Copy to avoid mutating original
+    state_key = sub_node_config.get("state_key", "result")
+
+    # Auto-inject the 'as' variable into sub-node's variables
+    # So the prompt can access it as {item_var}
+    sub_variables = dict(sub_node_config.get("variables", {}))
+    sub_variables[item_var] = f"{{state.{item_var}}}"
+    sub_node_config["variables"] = sub_variables
+
+    # Create sub-node from config
+    sub_node = create_node_function(sub_node_name, sub_node_config, defaults)
+    wrapped_node = wrap_for_reducer(sub_node, collect_key, state_key)
+    builder.add_node(sub_node_name, wrapped_node)
+
+    # Create fan-out edge function using Send
+    def map_edge(state: dict) -> list[Send]:
+        items = resolve_state_expression(over_expr, state)
+
+        if not isinstance(items, list):
+            raise TypeError(
+                f"Map 'over' must resolve to list, got {type(items).__name__}"
+            )
+
+        return [
+            Send(sub_node_name, {**state, item_var: item, "_map_index": i})
+            for i, item in enumerate(items)
+        ]
+
+    return map_edge, sub_node_name

yamlgraph/models/__init__.py

@@ -0,0 +1,36 @@
+"""Pydantic models and state definitions.
+
+Framework models for error handling and generic reports.
+State is now generated dynamically by state_builder.py.
+"""
+
+from yamlgraph.models.graph_schema import (
+    EdgeConfig,
+    GraphConfigSchema,
+    NodeConfig,
+    validate_graph_schema,
+)
+from yamlgraph.models.schemas import (
+    ErrorType,
+    GenericReport,
+    PipelineError,
+)
+from yamlgraph.models.state_builder import (
+    build_state_class,
+    create_initial_state,
+)
+
+__all__ = [
+    # Framework models
+    "ErrorType",
+    "PipelineError",
+    "GenericReport",
+    # Graph config schema
+    "GraphConfigSchema",
+    "NodeConfig",
+    "EdgeConfig",
+    "validate_graph_schema",
+    # Dynamic state generation
+    "build_state_class",
+    "create_initial_state",
+]

yamlgraph/models/graph_schema.py

@@ -0,0 +1,141 @@
+"""Pydantic schemas for YAML graph configuration validation.
+
+Provides structured validation for graph YAML files with clear error messages.
+"""
+
+from typing import Any
+
+from pydantic import BaseModel, Field, field_validator, model_validator
+
+from yamlgraph.constants import ErrorHandler, NodeType
+
+
+class NodeConfig(BaseModel):
+    """Configuration for a single graph node."""
+
+    type: str = Field(default=NodeType.LLM, description="Node type")
+    prompt: str | None = Field(default=None, description="Prompt template name")
+    state_key: str | None = Field(default=None, description="State key for output")
+    temperature: float | None = Field(default=None, ge=0, le=2)
+    provider: str | None = Field(default=None)
+    on_error: str | None = Field(default=None)
+    fallback: dict[str, Any] | None = Field(default=None)
+    variables: dict[str, str] = Field(default_factory=dict)
+    requires: list[str] = Field(default_factory=list)
+    routes: dict[str, str] | None = Field(default=None, description="Router routes")
+
+    # Map node fields
+    over: str | None = Field(default=None, description="Map over expression")
+    # 'as' is reserved in Python, handled specially
+    item_var: str | None = Field(default=None, alias="as")
+    node: dict[str, Any] | None = Field(default=None, description="Map sub-node")
+    collect: str | None = Field(default=None, description="Map collect key")
+
+    # Tool/Agent fields
+    tools: list[str] = Field(default_factory=list)
+    max_iterations: int = Field(default=10, ge=1)
+
+    model_config = {"extra": "allow", "populate_by_name": True}
+
+    @field_validator("on_error")
+    @classmethod
+    def validate_on_error(cls, v: str | None) -> str | None:
+        """Validate on_error is a known handler."""
+        if v is not None and v not in ErrorHandler.all_values():
+            valid = ", ".join(ErrorHandler.all_values())
+            raise ValueError(f"Invalid on_error '{v}'. Valid: {valid}")
+        return v
+
+    @model_validator(mode="after")
+    def validate_node_requirements(self) -> "NodeConfig":
+        """Validate node has required fields based on type."""
+        if NodeType.requires_prompt(self.type) and not self.prompt:
+            raise ValueError(f"Node type '{self.type}' requires 'prompt' field")
+
+        if self.type == NodeType.ROUTER and not self.routes:
+            raise ValueError("Router node requires 'routes' field")
+
+        if self.type == NodeType.MAP:
+            if not self.over:
+                raise ValueError("Map node requires 'over' field")
+            if not self.item_var:
+                raise ValueError("Map node requires 'as' field")
+            if not self.node:
+                raise ValueError("Map node requires 'node' field")
+            if not self.collect:
+                raise ValueError("Map node requires 'collect' field")
+
+        return self
+
+
+class EdgeConfig(BaseModel):
+    """Configuration for a graph edge."""
+
+    from_node: str = Field(..., alias="from", description="Source node")
+    to: str | list[str] = Field(..., description="Target node(s)")
+    condition: str | None = Field(default=None, description="Condition expression")
+
+    model_config = {"populate_by_name": True}
+
+
+class GraphConfigSchema(BaseModel):
+    """Full YAML graph configuration schema.
+
+    Use this for validating graph YAML files with Pydantic.
+    """
+
+    version: str = Field(default="1.0")
+    name: str = Field(default="unnamed")
+    description: str = Field(default="")
+    defaults: dict[str, Any] = Field(default_factory=dict)
+    nodes: dict[str, NodeConfig] = Field(...)
+    edges: list[EdgeConfig] = Field(...)
+    tools: dict[str, Any] = Field(default_factory=dict)
+    state_class: str = Field(default="")
+    loop_limits: dict[str, int] = Field(default_factory=dict)
+
+    model_config = {"extra": "allow"}
+
+    @model_validator(mode="after")
+    def validate_router_targets(self) -> "GraphConfigSchema":
+        """Validate router routes point to existing nodes."""
+        for node_name, node in self.nodes.items():
+            if node.type == NodeType.ROUTER and node.routes:
+                for route_key, target in node.routes.items():
+                    if target not in self.nodes:
+                        raise ValueError(
+                            f"Router '{node_name}' route '{route_key}' "
+                            f"targets nonexistent node '{target}'"
+                        )
+        return self
+
+    @model_validator(mode="after")
+    def validate_edge_nodes(self) -> "GraphConfigSchema":
+        """Validate edge sources and targets exist."""
+        valid_nodes = set(self.nodes.keys()) | {"START", "END"}
+
+        for edge in self.edges:
+            if edge.from_node not in valid_nodes:
+                raise ValueError(f"Edge 'from' node '{edge.from_node}' not found")
+
+            targets = edge.to if isinstance(edge.to, list) else [edge.to]
+            for target in targets:
+                if target not in valid_nodes:
+                    raise ValueError(f"Edge 'to' node '{target}' not found")
+
+        return self
+
+
+def validate_graph_schema(config: dict[str, Any]) -> GraphConfigSchema:
+    """Validate a graph configuration dict using Pydantic.
+
+    Args:
+        config: Raw parsed YAML configuration
+
+    Returns:
+        Validated GraphConfigSchema
+
+    Raises:
+        pydantic.ValidationError: If validation fails
+    """
+    return GraphConfigSchema.model_validate(config)