yamlgraph 0.1.1__py3-none-any.whl

yamlgraph/utils/prompts.py

@@ -0,0 +1,116 @@
+"""Unified prompt loading and path resolution.
+
+This module consolidates prompt loading logic used by executor.py
+and node_factory.py into a single, testable module.
+
+Search order for prompts:
+1. {prompts_dir}/{prompt_name}.yaml (standard location)
+2. {parent}/prompts/{basename}.yaml (external examples like examples/storyboard/...)
+"""
+
+from pathlib import Path
+
+import yaml
+
+from yamlgraph.config import PROMPTS_DIR
+
+
+def resolve_prompt_path(
+    prompt_name: str,
+    prompts_dir: Path | None = None,
+) -> Path:
+    """Resolve a prompt name to its full YAML file path.
+
+    Search order:
+    1. prompts_dir/{prompt_name}.yaml (default: prompts/)
+    2. {parent}/prompts/{basename}.yaml (for external examples)
+
+    Args:
+        prompt_name: Prompt name like "greet" or "examples/storyboard/expand_story"
+        prompts_dir: Base prompts directory (defaults to PROMPTS_DIR from config)
+
+    Returns:
+        Path to the YAML file
+
+    Raises:
+        FileNotFoundError: If prompt file doesn't exist
+
+    Examples:
+        >>> resolve_prompt_path("greet")
+        PosixPath('/path/to/prompts/greet.yaml')
+
+        >>> resolve_prompt_path("map-demo/generate_ideas")
+        PosixPath('/path/to/prompts/map-demo/generate_ideas.yaml')
+    """
+    if prompts_dir is None:
+        prompts_dir = PROMPTS_DIR
+
+    prompts_dir = Path(prompts_dir)
+
+    # Try standard location first: prompts_dir/{prompt_name}.yaml
+    yaml_path = prompts_dir / f"{prompt_name}.yaml"
+    if yaml_path.exists():
+        return yaml_path
+
+    # Try external example location: {parent}/prompts/{basename}.yaml
+    # e.g., "examples/storyboard/expand_story" -> "examples/storyboard/prompts/expand_story.yaml"
+    parts = prompt_name.rsplit("/", 1)
+    if len(parts) == 2:
+        parent_dir, basename = parts
+        alt_path = Path(parent_dir) / "prompts" / f"{basename}.yaml"
+        if alt_path.exists():
+            return alt_path
+
+    raise FileNotFoundError(f"Prompt not found: {yaml_path}")
+
+
+def load_prompt(
+    prompt_name: str,
+    prompts_dir: Path | None = None,
+) -> dict:
+    """Load a YAML prompt template.
+
+    Args:
+        prompt_name: Name of the prompt file (without .yaml extension)
+        prompts_dir: Optional prompts directory override
+
+    Returns:
+        Dictionary with prompt content (typically 'system' and 'user' keys)
+
+    Raises:
+        FileNotFoundError: If prompt file doesn't exist
+    """
+    path = resolve_prompt_path(prompt_name, prompts_dir)
+
+    with open(path) as f:
+        return yaml.safe_load(f)
+
+
+def load_prompt_path(
+    prompt_name: str,
+    prompts_dir: Path | None = None,
+) -> tuple[Path, dict]:
+    """Load a prompt and return both path and content.
+
+    Useful when you need both the file path (for schema loading)
+    and the content (for prompt execution).
+
+    Args:
+        prompt_name: Name of the prompt file (without .yaml extension)
+        prompts_dir: Optional prompts directory override
+
+    Returns:
+        Tuple of (path, content_dict)
+
+    Raises:
+        FileNotFoundError: If prompt file doesn't exist
+    """
+    path = resolve_prompt_path(prompt_name, prompts_dir)
+
+    with open(path) as f:
+        content = yaml.safe_load(f)
+
+    return path, content
+
+
+__all__ = ["resolve_prompt_path", "load_prompt", "load_prompt_path"]

yamlgraph/utils/sanitize.py

@@ -0,0 +1,98 @@
+"""Input sanitization utilities.
+
+Provides functions for validating and sanitizing user input
+to prevent prompt injection and other security issues.
+"""
+
+import re
+from typing import NamedTuple
+
+from yamlgraph.config import DANGEROUS_PATTERNS, MAX_TOPIC_LENGTH
+
+
+class SanitizationResult(NamedTuple):
+    """Result of input sanitization."""
+
+    value: str
+    is_safe: bool
+    warnings: list[str]
+
+
+def sanitize_topic(topic: str) -> SanitizationResult:
+    """Sanitize a topic string for use in prompts.
+
+    Checks for:
+    - Length limits
+    - Potential prompt injection patterns
+    - Control characters
+
+    Args:
+        topic: The raw topic string
+
+    Returns:
+        SanitizationResult with cleaned value and safety status
+
+    Example:
+        >>> result = sanitize_topic("machine learning")
+        >>> result.is_safe
+        True
+        >>> result = sanitize_topic("ignore previous instructions")
+        >>> result.is_safe
+        False
+    """
+    warnings = []
+    cleaned = topic.strip()
+
+    # Check length
+    if len(cleaned) > MAX_TOPIC_LENGTH:
+        cleaned = cleaned[:MAX_TOPIC_LENGTH]
+        warnings.append(f"Topic truncated to {MAX_TOPIC_LENGTH} characters")
+
+    # Check for empty
+    if not cleaned:
+        return SanitizationResult(
+            value="",
+            is_safe=False,
+            warnings=["Topic cannot be empty"],
+        )
+
+    # Remove control characters (except newlines)
+    cleaned = re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]", "", cleaned)
+
+    # Check for dangerous patterns (case-insensitive)
+    topic_lower = cleaned.lower()
+    for pattern in DANGEROUS_PATTERNS:
+        if pattern.lower() in topic_lower:
+            return SanitizationResult(
+                value=cleaned,
+                is_safe=False,
+                warnings=[f"Topic contains potentially unsafe pattern: '{pattern}'"],
+            )
+
+    return SanitizationResult(
+        value=cleaned,
+        is_safe=True,
+        warnings=warnings,
+    )
+
+
+def sanitize_variables(variables: dict) -> dict:
+    """Sanitize a dictionary of template variables.
+
+    Args:
+        variables: Dictionary of variable name -> value
+
+    Returns:
+        Sanitized dictionary with cleaned values
+    """
+    sanitized = {}
+
+    for key, value in variables.items():
+        if isinstance(value, str):
+            # Remove control characters but preserve newlines
+            cleaned = re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]", "", value)
+            sanitized[key] = cleaned
+        else:
+            sanitized[key] = value
+
+    return sanitized

yamlgraph/utils/template.py

@@ -0,0 +1,102 @@
+"""Template utilities - Variable extraction and validation.
+
+This module provides functions to extract required variables from
+prompt templates and validate that all required variables are provided
+before execution.
+
+Supports both simple {variable} placeholders and Jinja2 templates.
+"""
+
+import logging
+import re
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+def extract_variables(template: str) -> set[str]:
+    """Extract all variable names required by a template.
+
+    Handles both simple {var} and Jinja2 {{ var }}, {% for x in var %} syntax.
+
+    Args:
+        template: Template string with placeholders
+
+    Returns:
+        Set of variable names required by the template
+
+    Examples:
+        >>> extract_variables("Hello {name}")
+        {'name'}
+
+        >>> extract_variables("{% for item in items %}{{ item }}{% endfor %}")
+        {'items'}
+    """
+    variables: set[str] = set()
+
+    # Simple format: {var} - but NOT {{ (Jinja2)
+    # Match {word} but not {{word}} - use negative lookbehind/lookahead
+    simple_pattern = r"(?<!\{)\{(\w+)\}(?!\})"
+    variables.update(re.findall(simple_pattern, template))
+
+    # Jinja2 variable: {{ var }} or {{ var.field }}
+    jinja_var_pattern = r"\{\{\s*(\w+)"
+    variables.update(re.findall(jinja_var_pattern, template))
+
+    # Jinja2 loop: {% for x in var %}
+    jinja_loop_pattern = r"\{%\s*for\s+\w+\s+in\s+(\w+)"
+    variables.update(re.findall(jinja_loop_pattern, template))
+
+    # Jinja2 condition: {% if var %} or {% if var.field %}
+    jinja_if_pattern = r"\{%\s*if\s+(\w+)"
+    variables.update(re.findall(jinja_if_pattern, template))
+
+    # Remove loop iteration variables (they're not inputs)
+    # e.g., in "{% for item in items %}", "item" is not required
+    loop_iter_pattern = r"\{%\s*for\s+(\w+)\s+in"
+    loop_vars = set(re.findall(loop_iter_pattern, template))
+    variables -= loop_vars
+
+    # Remove common non-input variables
+    # - state: injected by node_factory
+    # - loop: Jinja2 loop context
+    # - range: Jinja2 builtin function
+    excluded = {"state", "loop", "range", "true", "false", "none"}
+    variables -= excluded
+
+    return variables
+
+
+def validate_variables(
+    template: str,
+    provided: dict[str, Any],
+    prompt_name: str,
+) -> None:
+    """Validate that all required template variables are provided.
+
+    Raises ValueError with helpful message listing all missing variables.
+
+    Args:
+        template: Template string with placeholders
+        provided: Dictionary of provided variable values
+        prompt_name: Name of the prompt (for error messages)
+
+    Raises:
+        ValueError: If any required variables are missing
+
+    Examples:
+        >>> validate_variables("Hello {name}", {"name": "World"}, "greet")
+        # No error
+
+        >>> validate_variables("Hello {name}", {}, "greet")
+        ValueError: Missing required variable(s) for prompt 'greet': name
+    """
+    required = extract_variables(template)
+    provided_keys = set(provided.keys())
+    missing = required - provided_keys
+
+    if missing:
+        raise ValueError(
+            f"Missing required variable(s) for prompt '{prompt_name}': "
+            f"{', '.join(sorted(missing))}"
+        )

yamlgraph/utils/validators.py

@@ -0,0 +1,181 @@
+"""Graph configuration validators.
+
+Validation functions for YAML graph configuration structures.
+"""
+
+from typing import Any
+
+from yamlgraph.constants import ErrorHandler, NodeType
+
+
+def validate_required_sections(config: dict[str, Any]) -> None:
+    """Validate required top-level sections exist.
+
+    Args:
+        config: Parsed YAML configuration dictionary
+
+    Raises:
+        ValueError: If required sections are missing
+    """
+    if not config.get("nodes"):
+        raise ValueError("Graph config missing required 'nodes' section")
+    if not config.get("edges"):
+        raise ValueError("Graph config missing required 'edges' section")
+
+
+def validate_node_prompt(node_name: str, node_config: dict[str, Any]) -> None:
+    """Validate node has required prompt if applicable.
+
+    Args:
+        node_name: Name of the node
+        node_config: Node configuration dictionary
+
+    Raises:
+        ValueError: If prompt is required but missing
+    """
+    node_type = node_config.get("type", NodeType.LLM)
+    # Only llm and router nodes require prompts
+    # tool, python, agent, and map nodes don't require prompts
+    if NodeType.requires_prompt(node_type) and not node_config.get("prompt"):
+        raise ValueError(f"Node '{node_name}' missing required 'prompt' field")
+
+
+def validate_router_node(
+    node_name: str, node_config: dict[str, Any], all_nodes: dict[str, Any]
+) -> None:
+    """Validate router node has routes pointing to valid nodes.
+
+    Args:
+        node_name: Name of the node
+        node_config: Node configuration dictionary
+        all_nodes: All nodes in the graph for target validation
+
+    Raises:
+        ValueError: If router configuration is invalid
+    """
+    if node_config.get("type") != NodeType.ROUTER:
+        return
+
+    if not node_config.get("routes"):
+        raise ValueError(f"Router node '{node_name}' missing required 'routes' field")
+
+    for route_key, target_node in node_config["routes"].items():
+        if target_node not in all_nodes:
+            raise ValueError(
+                f"Router node '{node_name}' route '{route_key}' points to "
+                f"nonexistent node '{target_node}'"
+            )
+
+
+def validate_edges(edges: list[dict[str, Any]]) -> None:
+    """Validate each edge has required from/to fields.
+
+    Args:
+        edges: List of edge configurations
+
+    Raises:
+        ValueError: If edge is missing required fields
+    """
+    for i, edge in enumerate(edges):
+        if "from" not in edge:
+            raise ValueError(f"Edge {i} missing required 'from' field")
+        if "to" not in edge:
+            raise ValueError(f"Edge {i} missing required 'to' field")
+
+        # Validate condition expressions at load time
+        if "condition" in edge:
+            validate_condition_expression(edge["condition"], i)
+
+
+def validate_condition_expression(condition: str, edge_index: int) -> None:
+    """Validate a condition expression has valid syntax.
+
+    Performs compile-time validation of condition expressions to catch
+    syntax errors early rather than at runtime.
+
+    Supports expression conditions like "score < 0.8", "a.b >= 1 and c == 'done'"
+
+    Args:
+        condition: Condition expression like "score < 0.8"
+        edge_index: Edge index for error messages
+
+    Raises:
+        ValueError: If condition has invalid syntax
+    """
+    import re
+
+    # Expression syntax check - must match comparison pattern
+    # Valid: "score < 0.8", "a.b >= 1", "x == 'done'"
+    # Also valid: compound expressions "a > 1 and b < 2"
+    comparison_pattern = r"[a-zA-Z_][\w.]*\s*(<=|>=|==|!=|<|>)\s*.+"
+    compound_pattern = r"\s+(and|or)\s+"
+
+    # Split by and/or and validate each part
+    parts = re.split(compound_pattern, condition, flags=re.IGNORECASE)
+    # parts includes the 'and'/'or' tokens, so filter to just comparisons
+    comparisons = [p.strip() for p in parts if p.strip().lower() not in ("and", "or")]
+
+    for part in comparisons:
+        if not re.match(comparison_pattern, part.strip()):
+            raise ValueError(
+                f"Edge {edge_index} has invalid condition syntax: '{condition}'. "
+                f"Expected format: 'field <op> value' (e.g., 'score < 0.8')"
+            )
+
+
+def validate_on_error(node_name: str, node_config: dict[str, Any]) -> None:
+    """Validate on_error value is valid.
+
+    Args:
+        node_name: Name of the node
+        node_config: Node configuration dictionary
+
+    Raises:
+        ValueError: If on_error value is invalid
+    """
+    on_error = node_config.get("on_error")
+    if on_error and on_error not in ErrorHandler.all_values():
+        raise ValueError(
+            f"Node '{node_name}' has invalid on_error value '{on_error}'. "
+            f"Valid values: {', '.join(ErrorHandler.all_values())}"
+        )
+
+
+def validate_map_node(node_name: str, node_config: dict[str, Any]) -> None:
+    """Validate map node has required fields.
+
+    Args:
+        node_name: Name of the node
+        node_config: Node configuration dictionary
+
+    Raises:
+        ValueError: If map node configuration is invalid
+    """
+    if node_config.get("type") != NodeType.MAP:
+        return
+
+    required_fields = ["over", "as", "node", "collect"]
+    for field in required_fields:
+        if field not in node_config:
+            raise ValueError(f"Map node '{node_name}' missing required '{field}' field")
+
+
+def validate_config(config: dict[str, Any]) -> None:
+    """Validate YAML configuration structure.
+
+    Args:
+        config: Parsed YAML dictionary
+
+    Raises:
+        ValueError: If required fields are missing or invalid
+    """
+    validate_required_sections(config)
+
+    nodes = config["nodes"]
+    for node_name, node_config in nodes.items():
+        validate_node_prompt(node_name, node_config)
+        validate_router_node(node_name, node_config, nodes)
+        validate_on_error(node_name, node_config)
+        validate_map_node(node_name, node_config)
+
+    validate_edges(config["edges"])