yamlgraph 0.1.1__py3-none-any.whl

yamlgraph/models/schemas.py

@@ -0,0 +1,124 @@
+"""Pydantic models for structured LLM outputs.
+
+This module contains FRAMEWORK models only - models used by the framework itself.
+Demo-specific output schemas are defined inline in graph YAML files.
+"""
+
+from datetime import datetime
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+# =============================================================================
+# Error Types
+# =============================================================================
+
+
+class ErrorType(str, Enum):
+    """Types of errors that can occur in the pipeline."""
+
+    LLM_ERROR = "llm_error"  # LLM API errors (rate limit, timeout, etc.)
+    VALIDATION_ERROR = "validation_error"  # Pydantic validation failures
+    PROMPT_ERROR = "prompt_error"  # Missing prompt, template errors
+    STATE_ERROR = "state_error"  # Missing required state data
+    UNKNOWN_ERROR = "unknown_error"  # Catch-all
+
+
+class PipelineError(BaseModel):
+    """Structured error information for pipeline failures."""
+
+    type: ErrorType = Field(description="Category of error")
+    message: str = Field(description="Human-readable error message")
+    node: str = Field(description="Node where error occurred")
+    timestamp: datetime = Field(default_factory=datetime.now)
+    retryable: bool = Field(
+        default=False, description="Whether this error can be retried"
+    )
+    details: dict[str, Any] = Field(
+        default_factory=dict, description="Additional error context"
+    )
+
+    @classmethod
+    def from_exception(
+        cls, e: Exception, node: str, error_type: ErrorType | None = None
+    ) -> "PipelineError":
+        """Create a PipelineError from an exception.
+
+        Args:
+            e: The exception that occurred
+            node: Name of the node where error occurred
+            error_type: Optional explicit error type
+
+        Returns:
+            PipelineError instance
+        """
+        # Infer error type from exception
+        if error_type is None:
+            exc_name = type(e).__name__.lower()
+            if "rate" in exc_name or "timeout" in exc_name or "api" in exc_name:
+                error_type = ErrorType.LLM_ERROR
+                retryable = True
+            elif "validation" in exc_name:
+                error_type = ErrorType.VALIDATION_ERROR
+                retryable = False
+            elif "file" in exc_name or "prompt" in exc_name:
+                error_type = ErrorType.PROMPT_ERROR
+                retryable = False
+            else:
+                error_type = ErrorType.UNKNOWN_ERROR
+                retryable = False
+        else:
+            retryable = error_type == ErrorType.LLM_ERROR
+
+        return cls(
+            type=error_type,
+            message=str(e),
+            node=node,
+            retryable=retryable,
+            details={"exception_type": type(e).__name__},
+        )
+
+
+# =============================================================================
+# Generic Report Model (Flexible for Any Use Case)
+# =============================================================================
+
+
+class GenericReport(BaseModel):
+    """Flexible report structure for any use case.
+
+    Use this when you don't need a custom schema - works for most
+    analysis and summary tasks. The LLM can populate any combination
+    of the optional fields as needed.
+
+    Example usage in graph YAML:
+        nodes:
+          analyze:
+            type: llm
+            prompt: my_analysis
+            output_model: yamlgraph.models.GenericReport
+
+    Example prompts can request specific sections:
+        "Analyze the repository and provide:
+         - A summary of findings
+         - Key findings as bullet points
+         - Recommendations for improvement"
+    """
+
+    title: str = Field(description="Report title")
+    summary: str = Field(description="Executive summary")
+    sections: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Named sections with any content (strings, dicts, lists)",
+    )
+    findings: list[str] = Field(
+        default_factory=list, description="Key findings or bullet points"
+    )
+    recommendations: list[str] = Field(
+        default_factory=list, description="Suggested actions or areas to focus on"
+    )
+    metadata: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Additional key-value data (author, version, tags, etc.)",
+    )

yamlgraph/models/state_builder.py

@@ -0,0 +1,236 @@
+"""Dynamic state class generation from graph configuration.
+
+Builds TypedDict programmatically from YAML graph config, eliminating
+the need for state_class coupling between YAML and Python.
+"""
+
+import logging
+from operator import add
+from typing import Annotated, Any, TypedDict
+
+logger = logging.getLogger(__name__)
+
+
+def sorted_add(existing: list, new: list) -> list:
+    """Reducer that adds items and sorts by _map_index if present.
+
+    Used for map node fan-in to guarantee order regardless of
+    parallel execution timing.
+
+    Args:
+        existing: Current list in state
+        new: New items to add
+
+    Returns:
+        Combined list sorted by _map_index (if items have it)
+    """
+    combined = (existing or []) + (new or [])
+
+    # Sort by _map_index if items have it
+    if combined and isinstance(combined[0], dict) and "_map_index" in combined[0]:
+        combined = sorted(combined, key=lambda x: x.get("_map_index", 0))
+
+    return combined
+
+
+# =============================================================================
+# Base Fields - Always included in generated state
+# =============================================================================
+
+# Infrastructure fields present in all graphs
+BASE_FIELDS: dict[str, type] = {
+    # Core tracking
+    "thread_id": str,
+    "current_step": str,
+    # Error handling - singular for current error
+    "error": Any,
+    # Error handling with reducer (accumulates)
+    "errors": Annotated[list, add],
+    # Messages with reducer (accumulates)
+    "messages": Annotated[list, add],
+    # Loop tracking
+    "_loop_counts": dict,
+    "_loop_limit_reached": bool,
+    "_agent_iterations": int,
+    "_agent_limit_reached": bool,
+    # Timestamps
+    "started_at": Any,
+    "completed_at": Any,
+}
+
+# Common input fields used across graph types
+# These are always included to support --var inputs
+COMMON_INPUT_FIELDS: dict[str, type] = {
+    "input": str,  # Agent prompt input
+    "topic": str,  # Content generation topic
+    "style": str,  # Writing style
+    "word_count": int,  # Target word count
+    "message": str,  # Router message input
+}
+
+# Type mapping for YAML state config
+TYPE_MAP: dict[str, type] = {
+    "str": str,
+    "string": str,
+    "int": int,
+    "integer": int,
+    "float": float,
+    "bool": bool,
+    "boolean": bool,
+    "list": list,
+    "dict": dict,
+    "any": Any,
+}
+
+
+def parse_state_config(state_config: dict) -> dict[str, type]:
+    """Parse YAML state section into field types.
+
+    Supports simple type strings:
+        state:
+          concept: str
+          count: int
+
+    Args:
+        state_config: Dict from YAML 'state' section
+
+    Returns:
+        Dict of field_name -> Python type
+    """
+    fields: dict[str, type] = {}
+
+    for field_name, type_spec in state_config.items():
+        if isinstance(type_spec, str):
+            # Simple type: "str", "int", etc.
+            normalized = type_spec.lower()
+            if normalized not in TYPE_MAP:
+                supported = ", ".join(sorted(set(TYPE_MAP.keys())))
+                logger.warning(
+                    f"Unknown type '{type_spec}' for state field '{field_name}'. "
+                    f"Supported types: {supported}. Defaulting to Any."
+                )
+            python_type = TYPE_MAP.get(normalized, Any)
+            fields[field_name] = python_type
+        else:
+            # Unknown format, use Any
+            logger.warning(
+                f"Invalid type specification for state field '{field_name}': "
+                f"expected string, got {type(type_spec).__name__}. Defaulting to Any."
+            )
+            fields[field_name] = Any
+
+    return fields
+
+
+def build_state_class(config: dict) -> type:
+    """Build TypedDict state class from graph configuration.
+
+    Dynamically generates a TypedDict with:
+    - Base infrastructure fields (errors, messages, thread_id, etc.)
+    - Common input fields (topic, style, input, message, etc.)
+    - Custom fields from YAML 'state' section
+    - Fields extracted from node state_key
+    - Special fields for agent/router node types
+
+    Args:
+        config: Parsed YAML graph configuration dict
+
+    Returns:
+        TypedDict class with total=False (all fields optional)
+    """
+    # Start with base and common fields
+    fields: dict[str, type] = {}
+    fields.update(BASE_FIELDS)
+    fields.update(COMMON_INPUT_FIELDS)
+
+    # Add custom state fields from YAML 'state' section
+    state_config = config.get("state", {})
+    custom_fields = parse_state_config(state_config)
+    fields.update(custom_fields)
+
+    # Extract fields from nodes
+    nodes = config.get("nodes", {})
+    node_fields = extract_node_fields(nodes)
+    fields.update(node_fields)
+
+    # Build TypedDict programmatically
+    return TypedDict("GraphState", fields, total=False)
+
+
+def extract_node_fields(nodes: dict) -> dict[str, type]:
+    """Extract state fields from node configurations.
+
+    Analyzes node configs to determine required state fields:
+    - state_key: Where node stores its output
+    - type: agent → adds input, _tool_results
+    - type: router → adds _route
+
+    Args:
+        nodes: Dict of node_name -> node_config
+
+    Returns:
+        Dict of field_name -> type for the state
+    """
+    fields: dict[str, type] = {}
+
+    for node_name, node_config in nodes.items():
+        if not isinstance(node_config, dict):
+            continue
+
+        # state_key → Any (accepts Pydantic models)
+        if state_key := node_config.get("state_key"):
+            fields[state_key] = Any
+
+        # Node type-specific fields
+        node_type = node_config.get("type", "llm")
+
+        if node_type == "agent":
+            fields["input"] = str
+            fields["_tool_results"] = list
+
+        elif node_type == "router":
+            fields["_route"] = str
+
+        elif node_type == "map":
+            # Map node collect field needs sorted reducer for ordered fan-in
+            if collect_key := node_config.get("collect"):
+                fields[collect_key] = Annotated[list, sorted_add]
+
+    return fields
+
+
+def create_initial_state(
+    topic: str = "",
+    style: str = "informative",
+    word_count: int = 300,
+    thread_id: str | None = None,
+    **kwargs: Any,
+) -> dict[str, Any]:
+    """Create an initial state for a new pipeline run.
+
+    Args:
+        topic: The topic to generate content about
+        style: Writing style (default: informative)
+        word_count: Target word count (default: 300)
+        thread_id: Optional thread ID (auto-generated if not provided)
+        **kwargs: Additional state fields (e.g., input for agents)
+
+    Returns:
+        Initialized state dictionary
+    """
+    import uuid
+    from datetime import datetime
+
+    return {
+        "thread_id": thread_id or uuid.uuid4().hex[:16],
+        "topic": topic,
+        "style": style,
+        "word_count": word_count,
+        "current_step": "init",
+        "error": None,
+        "errors": [],
+        "messages": [],
+        "started_at": datetime.now(),
+        "completed_at": None,
+        **kwargs,
+    }

yamlgraph/node_factory.py

@@ -0,0 +1,240 @@
+"""Node function factory for YAML-defined graphs.
+
+Creates LangGraph node functions from YAML configuration with support for:
+- Resume (skip if output exists)
+- Error handling (skip, retry, fail, fallback)
+- Router nodes with dynamic routing
+- Loop counting and limits
+"""
+
+import logging
+from typing import Any, Callable
+
+from yamlgraph.constants import ErrorHandler, NodeType
+from yamlgraph.error_handlers import (
+    check_loop_limit,
+    check_requirements,
+    handle_default,
+    handle_fail,
+    handle_fallback,
+    handle_retry,
+    handle_skip,
+)
+from yamlgraph.executor import execute_prompt
+from yamlgraph.utils.expressions import resolve_template
+from yamlgraph.utils.prompts import resolve_prompt_path
+
+# Type alias for dynamic state
+GraphState = dict[str, Any]
+
+logger = logging.getLogger(__name__)
+
+
+def resolve_class(class_path: str) -> type:
+    """Dynamically import and return a class from a module path.
+
+    Args:
+        class_path: Full path like "yamlgraph.models.GenericReport" or short name
+
+    Returns:
+        The class object
+    """
+    import importlib
+
+    parts = class_path.rsplit(".", 1)
+    if len(parts) != 2:
+        # Try to find in yamlgraph.models.schemas
+        try:
+            from yamlgraph.models import schemas
+
+            if hasattr(schemas, class_path):
+                return getattr(schemas, class_path)
+        except ImportError:
+            pass
+        raise ValueError(f"Invalid class path: {class_path}")
+
+    module_path, class_name = parts
+    module = importlib.import_module(module_path)
+    return getattr(module, class_name)
+
+
+def get_output_model_for_node(
+    node_config: dict[str, Any], prompts_dir: str | None = None
+) -> type | None:
+    """Get output model for a node, checking inline schema if no explicit model.
+
+    Priority:
+    1. Explicit output_model in node config (class path)
+    2. Inline schema in prompt YAML file
+    3. None (raw string output)
+
+    Args:
+        node_config: Node configuration from YAML
+        prompts_dir: Base prompts directory
+
+    Returns:
+        Pydantic model class or None
+    """
+    # 1. Check for explicit output_model
+    if model_path := node_config.get("output_model"):
+        return resolve_class(model_path)
+
+    # 2. Check for inline schema in prompt YAML
+    prompt_name = node_config.get("prompt")
+    if prompt_name:
+        try:
+            from yamlgraph.schema_loader import load_schema_from_yaml
+
+            yaml_path = resolve_prompt_path(prompt_name, prompts_dir)
+            return load_schema_from_yaml(yaml_path)
+        except FileNotFoundError:
+            # Prompt file doesn't exist yet - will fail later
+            pass
+
+    # 3. No output model
+    return None
+
+
+def create_node_function(
+    node_name: str,
+    node_config: dict,
+    defaults: dict,
+) -> Callable[[GraphState], dict]:
+    """Create a node function from YAML config.
+
+    Args:
+        node_name: Name of the node
+        node_config: Node configuration from YAML
+        defaults: Default configuration values
+
+    Returns:
+        Node function compatible with LangGraph
+    """
+    node_type = node_config.get("type", NodeType.LLM)
+    prompt_name = node_config.get("prompt")
+
+    # Resolve output model (explicit > inline schema > None)
+    output_model = get_output_model_for_node(node_config)
+
+    # Get config values (node > defaults)
+    temperature = node_config.get("temperature", defaults.get("temperature", 0.7))
+    provider = node_config.get("provider", defaults.get("provider"))
+    state_key = node_config.get("state_key", node_name)
+    variable_templates = node_config.get("variables", {})
+    requires = node_config.get("requires", [])
+
+    # Error handling
+    on_error = node_config.get("on_error")
+    max_retries = node_config.get("max_retries", 3)
+    fallback_config = node_config.get("fallback", {})
+    fallback_provider = fallback_config.get("provider") if fallback_config else None
+
+    # Router config
+    routes = node_config.get("routes", {})
+    default_route = node_config.get("default_route")
+
+    # Loop limit
+    loop_limit = node_config.get("loop_limit")
+
+    # Skip if exists (default true for resume support, false for loop nodes)
+    skip_if_exists = node_config.get("skip_if_exists", True)
+
+    def node_fn(state: dict) -> dict:
+        """Generated node function."""
+        loop_counts = dict(state.get("_loop_counts") or {})
+        current_count = loop_counts.get(node_name, 0)
+
+        # Check loop limit
+        if check_loop_limit(node_name, loop_limit, current_count):
+            return {"_loop_limit_reached": True, "current_step": node_name}
+
+        loop_counts[node_name] = current_count + 1
+
+        # Skip if output exists (resume support) - disabled for loop nodes
+        if skip_if_exists and state.get(state_key) is not None:
+            logger.info(f"Node {node_name} skipped - {state_key} already in state")
+            return {"current_step": node_name, "_loop_counts": loop_counts}
+
+        # Check requirements
+        if error := check_requirements(requires, state, node_name):
+            return {
+                "errors": [error],
+                "current_step": node_name,
+                "_loop_counts": loop_counts,
+            }
+
+        # Resolve variables
+        variables = {}
+        for key, template in variable_templates.items():
+            resolved = resolve_template(template, state)
+            if isinstance(resolved, list):
+                resolved = ", ".join(str(item) for item in resolved)
+            variables[key] = resolved
+
+        def attempt_execute(use_provider: str | None) -> tuple[Any, Exception | None]:
+            try:
+                result = execute_prompt(
+                    prompt_name=prompt_name,
+                    variables=variables,
+                    output_model=output_model,
+                    temperature=temperature,
+                    provider=use_provider,
+                )
+                return result, None
+            except Exception as e:
+                return None, e
+
+        result, error = attempt_execute(provider)
+
+        if error is None:
+            logger.info(f"Node {node_name} completed successfully")
+            update = {
+                state_key: result,
+                "current_step": node_name,
+                "_loop_counts": loop_counts,
+            }
+
+            # Router: add _route to state
+            if node_type == NodeType.ROUTER and routes:
+                route_key = getattr(result, "tone", None) or getattr(
+                    result, "intent", None
+                )
+                if route_key and route_key in routes:
+                    update["_route"] = routes[route_key]
+                elif default_route:
+                    update["_route"] = default_route
+                else:
+                    update["_route"] = list(routes.values())[0]
+                logger.info(f"Router {node_name} routing to: {update['_route']}")
+            return update
+
+        # Error handling - dispatch to strategy handlers
+        if on_error == ErrorHandler.SKIP:
+            handle_skip(node_name, error, loop_counts)
+            return {"current_step": node_name, "_loop_counts": loop_counts}
+
+        elif on_error == ErrorHandler.FAIL:
+            handle_fail(node_name, error)
+
+        elif on_error == ErrorHandler.RETRY:
+            result = handle_retry(
+                node_name,
+                lambda: attempt_execute(provider),
+                max_retries,
+            )
+            return result.to_state_update(state_key, node_name, loop_counts)
+
+        elif on_error == ErrorHandler.FALLBACK and fallback_provider:
+            result = handle_fallback(
+                node_name,
+                attempt_execute,
+                fallback_provider,
+            )
+            return result.to_state_update(state_key, node_name, loop_counts)
+
+        else:
+            result = handle_default(node_name, error)
+            return result.to_state_update(state_key, node_name, loop_counts)
+
+    node_fn.__name__ = f"{node_name}_node"
+    return node_fn

yamlgraph/routing.py

@@ -0,0 +1,87 @@
+"""Routing utilities for LangGraph edge conditions.
+
+Provides factory functions for creating router functions that determine
+which node to route to based on state values and expressions.
+"""
+
+import logging
+from collections.abc import Callable
+from typing import Any
+
+from langgraph.graph import END
+
+from yamlgraph.utils.conditions import evaluate_condition
+
+# Type alias for dynamic state
+GraphState = dict[str, Any]
+
+logger = logging.getLogger(__name__)
+
+
+def make_router_fn(targets: list[str]) -> Callable[[dict], str]:
+    """Create a router function that reads _route from state.
+
+    Used for type: router nodes with conditional edges to multiple targets.
+
+    NOTE: Use `state: dict` not `state: GraphState` - type hints cause
+    LangGraph to filter state fields. See docs/debug-router-type-hints.md
+
+    Args:
+        targets: List of valid target node names
+
+    Returns:
+        Router function that returns the target node name
+    """
+
+    def router_fn(state: dict) -> str:
+        route = state.get("_route")
+        logger.debug(f"Router: _route={route}, targets={targets}")
+        if route and route in targets:
+            logger.debug(f"Router: matched route {route}")
+            return route
+        # Default to first target
+        logger.debug(f"Router: defaulting to {targets[0]}")
+        return targets[0]
+
+    return router_fn
+
+
+def make_expr_router_fn(
+    edges: list[tuple[str, str]],
+    source_node: str,
+) -> Callable[[GraphState], str]:
+    """Create router that evaluates expression conditions.
+
+    Used for reflexion-style loops with expression-based conditions
+    like "critique.score < 0.8".
+
+    Args:
+        edges: List of (condition, target) tuples
+        source_node: Name of the source node (for logging)
+
+    Returns:
+        Router function that evaluates conditions and returns target
+    """
+
+    def expr_router_fn(state: GraphState) -> str:
+        # Check loop limit first
+        if state.get("_loop_limit_reached"):
+            return END
+
+        for condition, target in edges:
+            try:
+                if evaluate_condition(condition, state):
+                    logger.debug(
+                        f"Condition '{condition}' matched, routing to {target}"
+                    )
+                    return target
+            except ValueError as e:
+                logger.warning(f"Failed to evaluate condition '{condition}': {e}")
+        # No condition matched - this shouldn't happen with well-formed graphs
+        logger.warning(f"No condition matched for {source_node}, defaulting to END")
+        return END
+
+    return expr_router_fn
+
+
+__all__ = ["make_router_fn", "make_expr_router_fn"]