yamlgraph 0.3.9__py3-none-any.whl

yamlgraph/graph_loader.py

@@ -0,0 +1,451 @@
+"""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 collections.abc import Callable
+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_interrupt_node,
+    create_node_function,
+    create_passthrough_node,
+    create_subgraph_node,
+    create_tool_call_node,
+    resolve_class,
+)
+from yamlgraph.routing import make_expr_router_fn, make_router_fn
+from yamlgraph.storage.checkpointer_factory import get_checkpointer
+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,
+    load_python_function,
+    parse_python_tools,
+)
+from yamlgraph.tools.shell import parse_tools
+from yamlgraph.tools.websearch import parse_websearch_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, source_path: Path | None = None):
+        """Initialize from parsed YAML dict.
+
+        Args:
+            config: Parsed YAML configuration dictionary
+            source_path: Path to the source YAML file (for subgraph resolution)
+
+        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", {})
+        self.checkpointer = config.get("checkpointer")
+        # Store raw config for dynamic state building
+        self.raw_config = config
+        # Store source path for subgraph resolution
+        self.source_path = source_path
+        # Prompt resolution options (FR-A: graph-relative prompts)
+        self.prompts_relative = self.defaults.get("prompts_relative", False)
+        self.prompts_dir = self.defaults.get("prompts_dir")
+
+
+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, source_path=path.resolve())
+
+
+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], dict[str, Any], dict[str, Callable]]:
+    """Parse shell, Python, and websearch tools from config.
+
+    Args:
+        config: Graph configuration
+
+    Returns:
+        Tuple of (shell_tools, python_tools, websearch_tools, callable_registry)
+        callable_registry maps tool names to actual callable functions for tool_call nodes
+    """
+    tools = parse_tools(config.tools)
+    python_tools = parse_python_tools(config.tools)
+    websearch_tools = parse_websearch_tools(config.tools)
+
+    # Build callable registry for tool_call nodes
+    callable_registry: dict[str, Callable] = {}
+    for name, tool_config in python_tools.items():
+        try:
+            callable_registry[name] = load_python_function(tool_config)
+        except (ImportError, AttributeError) as e:
+            logger.warning(f"Failed to load tool '{name}': {e}")
+
+    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())}"
+        )
+    if websearch_tools:
+        logger.info(
+            f"Parsed {len(websearch_tools)} websearch tools: {', '.join(websearch_tools.keys())}"
+        )
+
+    return tools, python_tools, websearch_tools, callable_registry
+
+
+def _compile_node(
+    node_name: str,
+    node_config: dict[str, Any],
+    graph: StateGraph,
+    config: GraphConfig,
+    tools: dict[str, Any],
+    python_tools: dict[str, Any],
+    websearch_tools: dict[str, Any],
+    callable_registry: dict[str, Callable],
+) -> 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
+        websearch_tools: Web search tools registry (LangChain StructuredTool)
+        callable_registry: Loaded callable functions for tool_call nodes
+
+    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]
+
+    # Extract prompts path config from defaults (FR-A)
+    prompts_relative = config.defaults.get("prompts_relative", False)
+    prompts_dir = config.defaults.get("prompts_dir")
+    if prompts_dir:
+        prompts_dir = Path(prompts_dir)
+
+    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, websearch_tools, python_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, callable_registry
+        )
+        logger.info(f"Added node: {node_name} (type={node_type})")
+        return (node_name, (map_edge_fn, sub_node_name))
+    elif node_type == NodeType.TOOL_CALL:
+        # Dynamic tool call from state
+        node_fn = create_tool_call_node(node_name, enriched_config, callable_registry)
+        graph.add_node(node_name, node_fn)
+    elif node_type == NodeType.INTERRUPT:
+        # Human-in-the-loop interrupt node
+        node_fn = create_interrupt_node(
+            node_name,
+            enriched_config,
+            graph_path=config.source_path,
+            prompts_dir=prompts_dir,
+            prompts_relative=prompts_relative,
+        )
+        graph.add_node(node_name, node_fn)
+    elif node_type == NodeType.PASSTHROUGH:
+        # Simple state transformation node
+        node_fn = create_passthrough_node(node_name, enriched_config)
+        graph.add_node(node_name, node_fn)
+    elif node_type == NodeType.SUBGRAPH:
+        # Subgraph node - compose graphs from YAML
+        if not config.source_path:
+            raise ValueError(
+                f"Cannot resolve subgraph path for node '{node_name}': "
+                "parent graph has no source_path"
+            )
+        node_fn = create_subgraph_node(
+            node_name,
+            enriched_config,
+            parent_graph_path=config.source_path,
+        )
+        graph.add_node(node_name, node_fn)
+    else:
+        # LLM and router nodes
+        node_fn = create_node_function(
+            node_name,
+            enriched_config,
+            config.defaults,
+            graph_path=config.source_path,
+        )
+        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],
+    websearch_tools: dict[str, Any],
+    callable_registry: dict[str, Callable],
+) -> 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
+        websearch_tools: Web search tools registry
+        callable_registry: Loaded callable functions for tool_call nodes
+
+    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,
+            websearch_tools,
+            callable_registry,
+        )
+        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 from_node in map_nodes and to_node in map_nodes:
+        # Edge from map node TO another map node: sub_node → map_edge_fn
+        _, from_sub = map_nodes[from_node]
+        to_map_edge_fn, to_sub = map_nodes[to_node]
+        graph.add_conditional_edges(from_sub, to_map_edge_fn, [to_sub])
+    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, websearch_tools, callable_registry = _parse_all_tools(config)
+
+    # Compile all nodes
+    map_nodes = _compile_nodes(
+        config, graph, tools, python_tools, websearch_tools, callable_registry
+    )
+
+    # 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)
+
+
+def get_checkpointer_for_graph(config: GraphConfig, *, async_mode: bool = False):
+    """Get checkpointer from graph config.
+
+    Args:
+        config: Graph configuration
+        async_mode: If True, return async-compatible saver
+
+    Returns:
+        Configured checkpointer or None if not specified
+    """
+    return get_checkpointer(config.checkpointer, async_mode=async_mode)

yamlgraph/map_compiler.py

@@ -0,0 +1,150 @@
+"""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.constants import NodeType
+from yamlgraph.node_factory import create_node_function, create_tool_call_node
+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],
+    tools_registry: dict[str, Any] | None = None,
+) -> 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
+        tools_registry: Optional tools registry for tool_call sub-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")
+    sub_node_type = sub_node_config.get("type", "llm")
+
+    # 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 based on type
+    if sub_node_type == NodeType.TOOL_CALL:
+        if tools_registry is None:
+            raise ValueError(
+                f"Map node '{name}' has tool_call sub-node but no tools_registry"
+            )
+        sub_node = create_tool_call_node(sub_node_name, sub_node_config, tools_registry)
+    else:
+        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",
+]