yamlgraph 0.1.1__py3-none-any.whl

yamlgraph/tools/python_tool.py

@@ -0,0 +1,178 @@
+"""Python function loader for type: python nodes.
+
+This module enables YAML graphs to call arbitrary Python functions
+by specifying the module path and function name.
+"""
+
+from __future__ import annotations
+
+import importlib
+import logging
+import os
+import sys
+from dataclasses import dataclass
+from typing import Any, Callable
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class PythonToolConfig:
+    """Configuration for a Python tool.
+
+    Attributes:
+        module: Full module path (e.g., "examples.storyboard.nodes.image_node")
+        function: Function name within the module
+        description: Human-readable description
+    """
+
+    module: str
+    function: str
+    description: str = ""
+
+
+def load_python_function(config: PythonToolConfig) -> Callable:
+    """Load a Python function from module path.
+
+    Args:
+        config: Python tool configuration
+
+    Returns:
+        The loaded function
+
+    Raises:
+        ImportError: If module cannot be imported
+        AttributeError: If function not found in module
+    """
+    # Ensure current working directory is in path for project imports
+    cwd = os.getcwd()
+    if cwd not in sys.path:
+        sys.path.insert(0, cwd)
+
+    try:
+        module = importlib.import_module(config.module)
+    except ImportError as e:
+        logger.error(f"Failed to import module: {config.module}")
+        raise ImportError(f"Cannot import module '{config.module}': {e}") from e
+
+    try:
+        func = getattr(module, config.function)
+    except AttributeError as e:
+        logger.error(f"Function not found: {config.function} in {config.module}")
+        raise AttributeError(
+            f"Function '{config.function}' not found in module '{config.module}'"
+        ) from e
+
+    if not callable(func):
+        raise TypeError(f"'{config.function}' in '{config.module}' is not callable")
+
+    logger.debug(f"Loaded Python function: {config.module}.{config.function}")
+    return func
+
+
+def parse_python_tools(tools_config: dict[str, Any]) -> dict[str, PythonToolConfig]:
+    """Parse Python tools from YAML tools section.
+
+    Only extracts tools with type: python.
+
+    Args:
+        tools_config: Dict from YAML tools: section
+
+    Returns:
+        Registry mapping tool names to PythonToolConfig objects
+    """
+    registry: dict[str, PythonToolConfig] = {}
+
+    for name, config in tools_config.items():
+        if config.get("type") != "python":
+            continue
+
+        if "module" not in config or "function" not in config:
+            logger.warning(
+                f"Python tool '{name}' missing 'module' or 'function', skipping"
+            )
+            continue
+
+        registry[name] = PythonToolConfig(
+            module=config["module"],
+            function=config["function"],
+            description=config.get("description", ""),
+        )
+
+    return registry
+
+
+def create_python_node(
+    node_name: str,
+    node_config: dict[str, Any],
+    python_tools: dict[str, PythonToolConfig],
+) -> Callable[[dict[str, Any]], dict]:
+    """Create a node that executes a Python function.
+
+    The function receives the full state dict and should return
+    a partial state update dict.
+
+    Args:
+        node_name: Name of the node in the graph
+        node_config: Node configuration from YAML
+        python_tools: Registry of available Python tools
+
+    Returns:
+        Node function that executes the Python function
+    """
+    tool_name = node_config.get("tool") or node_config.get("function")
+    if not tool_name:
+        raise ValueError(f"Python node '{node_name}' must specify 'tool' or 'function'")
+
+    if tool_name not in python_tools:
+        raise KeyError(f"Python tool '{tool_name}' not found in tools registry")
+
+    tool_config = python_tools[tool_name]
+    state_key = node_config.get("state_key", node_name)
+    on_error = node_config.get("on_error", "fail")
+
+    # Load the function at node creation time
+    func = load_python_function(tool_config)
+
+    def node_fn(state: dict[str, Any]) -> dict:
+        """Execute the Python function and return state update."""
+        logger.info(f"🐍 Executing Python node: {node_name} -> {tool_name}")
+
+        try:
+            result = func(state)
+
+            # If function returns a dict, merge with node metadata
+            if isinstance(result, dict):
+                result["current_step"] = node_name
+                return result
+            else:
+                # Function returned a single value, store in state_key
+                return {
+                    state_key: result,
+                    "current_step": node_name,
+                }
+
+        except Exception as e:
+            logger.error(f"Python node {node_name} failed: {e}")
+
+            if on_error == "skip":
+                from yamlgraph.models import ErrorType, PipelineError
+
+                errors = list(state.get("errors") or [])
+                errors.append(
+                    PipelineError(
+                        node=node_name,
+                        type=ErrorType.UNKNOWN_ERROR,
+                        message=str(e),
+                    )
+                )
+                return {
+                    state_key: None,
+                    "current_step": node_name,
+                    "errors": errors,
+                }
+            else:
+                raise
+
+    node_fn.__name__ = f"{node_name}_python_node"
+    return node_fn

yamlgraph/tools/shell.py

@@ -0,0 +1,205 @@
+"""Shell tool executor for running commands with variable substitution.
+
+This module provides the core shell tool execution functionality,
+allowing YAML-defined tools to run shell commands with parsed output.
+
+Security Model:
+    Variables are sanitized with shlex.quote() to prevent shell injection.
+    The command template itself is trusted (from YAML config), but all
+    user-provided variable values are escaped before substitution.
+
+    Example:
+        command: "git log --author={author}"
+        variables: {"author": "$(rm -rf /)"}
+        → Executed: git log --author='$(rm -rf /)'  (safely quoted)
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import shlex
+import subprocess
+from dataclasses import dataclass, field
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ShellToolConfig:
+    """Configuration for a shell tool.
+
+    Attributes:
+        command: Shell command with {variable} placeholders
+        description: Human-readable description for LLM tool selection
+        parse: Output parsing mode - 'text', 'json', or 'none'
+        timeout: Max seconds before command is killed
+        working_dir: Directory to run command in
+        env: Additional environment variables
+        success_codes: Exit codes considered successful
+    """
+
+    command: str
+    description: str = ""
+    parse: str = "text"  # text | json | none
+    timeout: int = 30
+    working_dir: str = "."
+    env: dict[str, str] = field(default_factory=dict)
+    success_codes: list[int] = field(default_factory=lambda: [0])
+
+
+@dataclass
+class ToolResult:
+    """Result from executing a shell tool.
+
+    Attributes:
+        success: Whether the command succeeded
+        output: Parsed output (str, dict, or None based on parse mode)
+        error: Error message if failed
+    """
+
+    success: bool
+    output: Any = None
+    error: str | None = None
+
+
+def sanitize_variables(variables: dict[str, Any]) -> dict[str, str]:
+    """Sanitize variable values for safe shell substitution.
+
+    Uses shlex.quote() to escape values, preventing shell injection.
+
+    Args:
+        variables: Raw variable values
+
+    Returns:
+        Sanitized variables safe for shell interpolation
+    """
+    sanitized = {}
+    for key, value in variables.items():
+        if value is None:
+            sanitized[key] = ""
+        elif isinstance(value, (list, dict)):
+            # Convert complex types to JSON, then quote
+            sanitized[key] = shlex.quote(json.dumps(value))
+        else:
+            sanitized[key] = shlex.quote(str(value))
+    return sanitized
+
+
+def execute_shell_tool(
+    config: ShellToolConfig,
+    variables: dict[str, Any],
+    sanitize: bool = True,
+) -> ToolResult:
+    """Execute shell command with variable substitution.
+
+    Args:
+        config: Tool configuration with command template
+        variables: Values to substitute into command placeholders
+        sanitize: Whether to sanitize variables with shlex.quote (default True)
+
+    Returns:
+        ToolResult with success status and parsed output or error
+    """
+    # Sanitize variables to prevent shell injection
+    safe_vars = sanitize_variables(variables) if sanitize else variables
+
+    # Substitute variables into command
+    try:
+        command = config.command.format(**safe_vars)
+    except KeyError as e:
+        return ToolResult(
+            success=False,
+            error=f"Missing variable: {e}",
+        )
+
+    logger.debug(f"Executing: {command}")
+
+    # Build environment
+    env = os.environ.copy()
+    env.update(config.env)
+
+    try:
+        # Security: shell=True is required for command templates with pipes/redirects.
+        # All user-provided variables are sanitized via shlex.quote() in sanitize_variables()
+        # before substitution, preventing shell injection attacks. The command template
+        # itself comes from trusted YAML configuration, not user input.
+        result = subprocess.run(
+            command,
+            shell=True,  # nosec B602
+            capture_output=True,
+            text=True,
+            timeout=config.timeout,
+            cwd=config.working_dir,
+            env=env,
+        )
+    except subprocess.TimeoutExpired:
+        return ToolResult(
+            success=False,
+            error=f"Command timed out after {config.timeout} seconds",
+        )
+    except Exception as e:
+        return ToolResult(
+            success=False,
+            error=f"Execution error: {e}",
+        )
+
+    # Check exit code
+    if result.returncode not in config.success_codes:
+        return ToolResult(
+            success=False,
+            output=result.stdout,
+            error=result.stderr or f"Exit code {result.returncode}",
+        )
+
+    # Parse output
+    output = result.stdout
+    if config.parse == "json":
+        try:
+            output = json.loads(output)
+        except json.JSONDecodeError as e:
+            return ToolResult(
+                success=False,
+                error=f"JSON parse error: {e}",
+            )
+    elif config.parse == "none":
+        output = None
+
+    return ToolResult(success=True, output=output)
+
+
+def parse_tools(tools_config: dict[str, Any]) -> dict[str, ShellToolConfig]:
+    """Parse tools: section from YAML into ShellToolConfig registry.
+
+    Only parses shell tools (type: shell or no type specified with command).
+    Skips Python tools (type: python).
+
+    Args:
+        tools_config: Dict from YAML tools: section
+
+    Returns:
+        Registry mapping tool names to ShellToolConfig objects
+    """
+    registry: dict[str, ShellToolConfig] = {}
+
+    for name, config in tools_config.items():
+        # Skip Python tools
+        if config.get("type") == "python":
+            continue
+        # Skip tools without command (invalid shell tools)
+        if "command" not in config:
+            continue
+
+        registry[name] = ShellToolConfig(
+            command=config["command"],
+            description=config.get("description", ""),
+            parse=config.get("parse", "text"),
+            timeout=config.get("timeout", 30),
+            working_dir=config.get("working_dir", "."),
+            env=config.get("env", {}),
+            success_codes=config.get("success_codes", [0]),
+        )
+
+    return registry

yamlgraph/utils/__init__.py

@@ -0,0 +1,47 @@
+"""Utility functions for observability and logging."""
+
+from yamlgraph.utils.conditions import evaluate_condition
+from yamlgraph.utils.expressions import (
+    resolve_state_expression,
+    resolve_state_path,
+    resolve_template,
+)
+from yamlgraph.utils.langsmith import (
+    get_client,
+    get_latest_run_id,
+    get_project_name,
+    get_run_url,
+    is_tracing_enabled,
+    log_execution,
+    print_run_tree,
+)
+from yamlgraph.utils.logging import get_logger, setup_logging
+from yamlgraph.utils.prompts import load_prompt, load_prompt_path, resolve_prompt_path
+from yamlgraph.utils.template import extract_variables, validate_variables
+
+__all__ = [
+    # Conditions
+    "evaluate_condition",
+    # Expression resolution (consolidated)
+    "resolve_state_path",
+    "resolve_state_expression",
+    "resolve_template",
+    # LangSmith
+    "get_client",
+    "get_project_name",
+    "is_tracing_enabled",
+    "get_latest_run_id",
+    "print_run_tree",
+    "log_execution",
+    "get_run_url",
+    # Logging
+    "get_logger",
+    "setup_logging",
+    # Prompts
+    "resolve_prompt_path",
+    "load_prompt",
+    "load_prompt_path",
+    # Template utilities
+    "extract_variables",
+    "validate_variables",
+]

yamlgraph/utils/conditions.py

@@ -0,0 +1,157 @@
+"""Condition expression evaluation for graph routing.
+
+Provides safe evaluation of condition expressions without using eval().
+Supports comparisons and compound boolean expressions.
+
+Examples:
+    >>> evaluate_condition("score < 0.8", {"score": 0.5})
+    True
+    >>> evaluate_condition("a > 1 and b < 2", {"a": 2, "b": 1})
+    True
+"""
+
+import re
+from typing import Any
+
+from yamlgraph.utils.expressions import resolve_state_path
+
+# Regex patterns for expression parsing
+# Valid operators: <=, >=, ==, !=, <, > (strict matching)
+COMPARISON_PATTERN = re.compile(
+    r"^\s*([a-zA-Z_][\w.]*)\s*(<=|>=|==|!=|<(?!<)|>(?!>))\s*(.+?)\s*$"
+)
+COMPOUND_AND_PATTERN = re.compile(r"\s+and\s+", re.IGNORECASE)
+COMPOUND_OR_PATTERN = re.compile(r"\s+or\s+", re.IGNORECASE)
+
+
+def resolve_value(path: str, state: dict) -> Any:
+    """Resolve a dotted path to a value from state.
+
+    Delegates to consolidated resolve_state_path in expressions module.
+
+    Args:
+        path: Dotted path like "critique.score"
+        state: State dictionary
+
+    Returns:
+        Resolved value or None if not found
+    """
+    return resolve_state_path(path, state)
+
+
+def parse_literal(value_str: str) -> Any:
+    """Parse a literal value from expression.
+
+    Args:
+        value_str: String representation of value
+
+    Returns:
+        Parsed Python value
+    """
+    value_str = value_str.strip()
+
+    # Handle quoted strings
+    if (value_str.startswith('"') and value_str.endswith('"')) or (
+        value_str.startswith("'") and value_str.endswith("'")
+    ):
+        return value_str[1:-1]
+
+    # Handle special keywords
+    if value_str.lower() == "true":
+        return True
+    if value_str.lower() == "false":
+        return False
+    if value_str.lower() == "null" or value_str.lower() == "none":
+        return None
+
+    # Handle numbers
+    try:
+        if "." in value_str:
+            return float(value_str)
+        return int(value_str)
+    except ValueError:
+        # Return as string if not a number
+        return value_str
+
+
+def evaluate_comparison(
+    left_path: str, operator: str, right_str: str, state: dict[str, Any]
+) -> bool:
+    """Evaluate a single comparison expression.
+
+    Args:
+        left_path: Dotted path to left value
+        operator: Comparison operator
+        right_str: String representation of right value
+        state: State dictionary
+
+    Returns:
+        Boolean result of comparison
+    """
+    left_value = resolve_value(left_path, state)
+    right_value = parse_literal(right_str)
+
+    # Handle missing left value
+    if left_value is None and operator not in ("==", "!="):
+        return False
+
+    try:
+        if operator == "<":
+            return left_value < right_value
+        elif operator == ">":
+            return left_value > right_value
+        elif operator == "<=":
+            return left_value <= right_value
+        elif operator == ">=":
+            return left_value >= right_value
+        elif operator == "==":
+            return left_value == right_value
+        elif operator == "!=":
+            return left_value != right_value
+        else:
+            raise ValueError(f"Unknown operator: {operator}")
+    except TypeError:
+        # Comparison failed (e.g., comparing None with number)
+        return False
+
+
+def evaluate_condition(expr: str, state: dict) -> bool:
+    """Safely evaluate a condition expression against state.
+
+    Uses pattern matching - no eval() for security.
+
+    Args:
+        expr: Condition expression like "score < 0.8" or "a > 1 and b < 2"
+        state: State dictionary to evaluate against
+
+    Returns:
+        Boolean result of evaluation
+
+    Raises:
+        ValueError: If expression is malformed
+
+    Examples:
+        >>> evaluate_condition("score < 0.8", {"score": 0.5})
+        True
+        >>> evaluate_condition("critique.score >= 0.8", {"critique": obj})
+        True
+    """
+    expr = expr.strip()
+
+    # Handle compound OR (lower precedence)
+    if COMPOUND_OR_PATTERN.search(expr):
+        parts = COMPOUND_OR_PATTERN.split(expr)
+        return any(evaluate_condition(part, state) for part in parts)
+
+    # Handle compound AND
+    if COMPOUND_AND_PATTERN.search(expr):
+        parts = COMPOUND_AND_PATTERN.split(expr)
+        return all(evaluate_condition(part, state) for part in parts)
+
+    # Parse single comparison
+    match = COMPARISON_PATTERN.match(expr)
+    if not match:
+        raise ValueError(f"Invalid condition expression: {expr}")
+
+    left_path, operator, right_str = match.groups()
+    return evaluate_comparison(left_path, operator, right_str, state)

yamlgraph/utils/expressions.py

@@ -0,0 +1,111 @@
+"""Expression resolution utilities for YAML graphs.
+
+Consolidated module for all state path/expression resolution.
+Use these functions instead of duplicating resolution logic elsewhere.
+"""
+
+from typing import Any
+
+
+def resolve_state_path(path: str, state: dict[str, Any]) -> Any:
+    """Resolve a dotted path to a value from state.
+
+    Core resolution function - handles nested dict access and object attributes.
+    This is the single source of truth for path resolution.
+
+    Args:
+        path: Dotted path like "critique.score" or "story.panels"
+        state: State dictionary
+
+    Returns:
+        Resolved value or None if not found
+    """
+    if not path:
+        return None
+
+    parts = path.split(".")
+    value = state
+
+    for part in parts:
+        if value is None:
+            return None
+        if isinstance(value, dict):
+            value = value.get(part)
+        else:
+            # Try attribute access for objects (Pydantic models, etc.)
+            value = getattr(value, part, None)
+
+    return value
+
+
+def resolve_state_expression(expr: str | Any, state: dict[str, Any]) -> Any:
+    """Resolve {state.path.to.value} expressions.
+
+    Supports expressions like:
+        - "{name}" -> state["name"]
+        - "{state.story.panels}" -> state["story"]["panels"]
+        - "{story.title}" -> state["story"]["title"]
+
+    Non-expression values (no braces) pass through unchanged.
+
+    Args:
+        expr: Expression string like "{state.story.panels}" or any value
+        state: Current graph state dict
+
+    Returns:
+        Resolved value from state, or original value if not an expression
+
+    Raises:
+        KeyError: If path cannot be resolved in state
+    """
+    if not isinstance(expr, str):
+        return expr
+
+    if not (expr.startswith("{") and expr.endswith("}")):
+        return expr
+
+    path = expr[1:-1]  # Remove braces
+
+    # Handle "state." prefix (optional)
+    if path.startswith("state."):
+        path = path[6:]  # Remove "state."
+
+    # Navigate nested path
+    value = state
+    for key in path.split("."):
+        if isinstance(value, dict) and key in value:
+            value = value[key]
+        elif hasattr(value, key):
+            # Support object attribute access (Pydantic models, etc.)
+            value = getattr(value, key)
+        else:
+            raise KeyError(f"Cannot resolve '{key}' in path '{expr}'")
+
+    return value
+
+
+def resolve_template(template: str | Any, state: dict[str, Any]) -> Any:
+    """Resolve a {state.field} template to its value.
+
+    Unlike resolve_state_expression, returns None instead of raising
+    when path not found. Used for optional variable resolution.
+
+    Args:
+        template: Template string like "{state.field}" or "{state.obj.attr}"
+        state: Current pipeline state
+
+    Returns:
+        Resolved value or None if not found
+    """
+    STATE_PREFIX = "{state."
+    STATE_SUFFIX = "}"
+
+    if not isinstance(template, str):
+        return template
+
+    if not (template.startswith(STATE_PREFIX) and template.endswith(STATE_SUFFIX)):
+        return template
+
+    # Extract path: "{state.foo.bar}" -> "foo.bar"
+    path = template[len(STATE_PREFIX) : -len(STATE_SUFFIX)]
+    return resolve_state_path(path, state)