yamlgraph 0.1.1__py3-none-any.whl

tests/unit/test_state_config.py

@@ -0,0 +1,197 @@
+"""Tests for YAML state configuration parsing.
+
+Tests the parse_state_config function and state: section handling
+in build_state_class.
+"""
+
+from typing import Any
+
+from yamlgraph.models.state_builder import (
+    TYPE_MAP,
+    build_state_class,
+    parse_state_config,
+)
+
+
+class TestParseStateConfig:
+    """Tests for parse_state_config function."""
+
+    def test_empty_config(self):
+        """Empty state config returns empty dict."""
+        result = parse_state_config({})
+        assert result == {}
+
+    def test_simple_str_type(self):
+        """Parse 'str' type."""
+        result = parse_state_config({"concept": "str"})
+        assert result == {"concept": str}
+
+    def test_simple_int_type(self):
+        """Parse 'int' type."""
+        result = parse_state_config({"count": "int"})
+        assert result == {"count": int}
+
+    def test_simple_float_type(self):
+        """Parse 'float' type."""
+        result = parse_state_config({"score": "float"})
+        assert result == {"score": float}
+
+    def test_simple_bool_type(self):
+        """Parse 'bool' type."""
+        result = parse_state_config({"enabled": "bool"})
+        assert result == {"enabled": bool}
+
+    def test_simple_list_type(self):
+        """Parse 'list' type."""
+        result = parse_state_config({"items": "list"})
+        assert result == {"items": list}
+
+    def test_simple_dict_type(self):
+        """Parse 'dict' type."""
+        result = parse_state_config({"metadata": "dict"})
+        assert result == {"metadata": dict}
+
+    def test_any_type(self):
+        """Parse 'any' type."""
+        result = parse_state_config({"data": "any"})
+        assert result == {"data": Any}
+
+    def test_type_aliases(self):
+        """Type aliases like 'string', 'integer', 'boolean' work."""
+        result = parse_state_config(
+            {
+                "name": "string",
+                "age": "integer",
+                "active": "boolean",
+            }
+        )
+        assert result == {"name": str, "age": int, "active": bool}
+
+    def test_case_insensitive(self):
+        """Type names are case-insensitive."""
+        result = parse_state_config(
+            {
+                "a": "STR",
+                "b": "Int",
+                "c": "FLOAT",
+            }
+        )
+        assert result == {"a": str, "b": int, "c": float}
+
+    def test_multiple_fields(self):
+        """Parse multiple fields."""
+        result = parse_state_config(
+            {
+                "concept": "str",
+                "count": "int",
+                "score": "float",
+            }
+        )
+        assert result == {"concept": str, "count": int, "score": float}
+
+    def test_unknown_type_defaults_to_any(self):
+        """Unknown type strings default to Any."""
+        result = parse_state_config({"custom": "unknown_type"})
+        assert result == {"custom": Any}
+
+    def test_non_string_value_defaults_to_any(self):
+        """Non-string values default to Any."""
+        result = parse_state_config(
+            {
+                "nested": {"type": "str"},  # Dict value, not string
+                "number": 123,  # Int value, not string
+            }
+        )
+        assert result == {"nested": Any, "number": Any}
+
+
+class TestTypeMap:
+    """Tests for TYPE_MAP constant."""
+
+    def test_all_basic_types_present(self):
+        """TYPE_MAP contains all basic Python types."""
+        assert "str" in TYPE_MAP
+        assert "int" in TYPE_MAP
+        assert "float" in TYPE_MAP
+        assert "bool" in TYPE_MAP
+        assert "list" in TYPE_MAP
+        assert "dict" in TYPE_MAP
+        assert "any" in TYPE_MAP
+
+    def test_aliases_present(self):
+        """TYPE_MAP contains common aliases."""
+        assert "string" in TYPE_MAP
+        assert "integer" in TYPE_MAP
+        assert "boolean" in TYPE_MAP
+
+
+class TestBuildStateClassWithStateConfig:
+    """Tests for build_state_class with state: section."""
+
+    def test_state_section_adds_fields(self):
+        """State section fields are included in generated class."""
+        config = {
+            "state": {"concept": "str", "count": "int"},
+            "nodes": {},
+            "edges": [],
+        }
+        state_class = build_state_class(config)
+        annotations = state_class.__annotations__
+
+        assert "concept" in annotations
+        assert "count" in annotations
+
+    def test_state_section_empty(self):
+        """Empty state section doesn't break build."""
+        config = {
+            "state": {},
+            "nodes": {},
+            "edges": [],
+        }
+        state_class = build_state_class(config)
+        # Should still have base fields
+        assert "thread_id" in state_class.__annotations__
+
+    def test_state_section_missing(self):
+        """Missing state section is handled."""
+        config = {
+            "nodes": {},
+            "edges": [],
+        }
+        state_class = build_state_class(config)
+        # Should still have base fields
+        assert "thread_id" in state_class.__annotations__
+
+    def test_custom_field_overrides_common(self):
+        """Custom state field can override common field type."""
+        config = {
+            "state": {"topic": "int"},  # Override str default
+            "nodes": {},
+            "edges": [],
+        }
+        state_class = build_state_class(config)
+        # The custom field should be present
+        assert "topic" in state_class.__annotations__
+
+    def test_storyboard_example(self):
+        """Test storyboard-style config with concept field."""
+        config = {
+            "state": {"concept": "str"},
+            "nodes": {
+                "expand_story": {
+                    "type": "llm",
+                    "state_key": "story",
+                }
+            },
+            "edges": [],
+        }
+        state_class = build_state_class(config)
+        annotations = state_class.__annotations__
+
+        # Custom field from state section
+        assert "concept" in annotations
+        # Output field from node
+        assert "story" in annotations
+        # Base infrastructure fields
+        assert "thread_id" in annotations
+        assert "errors" in annotations

tests/unit/test_template.py

@@ -0,0 +1,190 @@
+"""Tests for yamlgraph.utils.template module - Variable extraction and validation."""
+
+import pytest
+
+
+class TestExtractVariables:
+    """Tests for extract_variables function."""
+
+    def test_extract_simple_variables(self):
+        """Should extract {var} placeholders."""
+        from yamlgraph.utils.template import extract_variables
+
+        template = "Hello {name}, your style is {style}."
+        variables = extract_variables(template)
+        assert variables == {"name", "style"}
+
+    def test_extract_single_variable(self):
+        """Should extract a single variable."""
+        from yamlgraph.utils.template import extract_variables
+
+        template = "Welcome {user}!"
+        variables = extract_variables(template)
+        assert variables == {"user"}
+
+    def test_extract_no_variables(self):
+        """Should return empty set when no variables."""
+        from yamlgraph.utils.template import extract_variables
+
+        template = "No variables here"
+        variables = extract_variables(template)
+        assert variables == set()
+
+    def test_extract_duplicate_variables(self):
+        """Should deduplicate variables."""
+        from yamlgraph.utils.template import extract_variables
+
+        template = "{name} and {name} again"
+        variables = extract_variables(template)
+        assert variables == {"name"}
+
+    def test_extract_jinja2_variable(self):
+        """Should extract {{ var }} Jinja2 variables."""
+        from yamlgraph.utils.template import extract_variables
+
+        template = "Hello {{ name }}!"
+        variables = extract_variables(template)
+        assert "name" in variables
+
+    def test_extract_jinja2_variable_with_field_access(self):
+        """Should extract base variable from {{ var.field }}."""
+        from yamlgraph.utils.template import extract_variables
+
+        template = "User: {{ user.name }}"
+        variables = extract_variables(template)
+        assert "user" in variables
+
+    def test_extract_jinja2_loop_variable(self):
+        """Should extract iterable from {% for x in items %}."""
+        from yamlgraph.utils.template import extract_variables
+
+        template = "{% for item in items %}{{ item.name }}{% endfor %}"
+        variables = extract_variables(template)
+        assert "items" in variables
+        # 'item' is a loop variable, not a required input
+        assert "item" not in variables
+
+    def test_extract_jinja2_if_variable(self):
+        """Should extract variable from {% if condition %}."""
+        from yamlgraph.utils.template import extract_variables
+
+        template = "{% if show_details %}Details here{% endif %}"
+        variables = extract_variables(template)
+        assert "show_details" in variables
+
+    def test_exclude_state_variable(self):
+        """State is injected by framework, not a required input."""
+        from yamlgraph.utils.template import extract_variables
+
+        template = "{{ state.topic }}"
+        variables = extract_variables(template)
+        # state is excluded - it's injected by node_factory
+        assert "state" not in variables
+
+    def test_exclude_jinja2_builtins(self):
+        """Should exclude Jinja2 builtins like loop, range."""
+        from yamlgraph.utils.template import extract_variables
+
+        template = "{% for i in range(10) %}{{ loop.index }}{% endfor %}"
+        variables = extract_variables(template)
+        assert "range" not in variables
+        assert "loop" not in variables
+
+    def test_mixed_simple_and_jinja2(self):
+        """Should handle templates mixing {var} and {{ var }}."""
+        from yamlgraph.utils.template import extract_variables
+
+        template = "Simple {name} and Jinja2 {{ topic }}"
+        variables = extract_variables(template)
+        assert "name" in variables
+        assert "topic" in variables
+
+
+class TestValidateVariables:
+    """Tests for validate_variables function."""
+
+    def test_validate_all_provided(self):
+        """Should not raise when all variables provided."""
+        from yamlgraph.utils.template import validate_variables
+
+        template = "Hello {name}, style: {style}"
+        # Should not raise
+        validate_variables(template, {"name": "World", "style": "formal"}, "greet")
+
+    def test_validate_missing_single_variable(self):
+        """Should raise ValueError for single missing variable."""
+        from yamlgraph.utils.template import validate_variables
+
+        template = "Hello {name}, style: {style}"
+        with pytest.raises(ValueError, match="Missing required variable.*name"):
+            validate_variables(template, {"style": "formal"}, "greet")
+
+    def test_validate_missing_multiple_variables(self):
+        """Should list ALL missing variables in error."""
+        from yamlgraph.utils.template import validate_variables
+
+        template = "Hello {name}, style: {style}"
+        with pytest.raises(ValueError) as exc_info:
+            validate_variables(template, {}, "greet")
+        error_msg = str(exc_info.value)
+        assert "name" in error_msg
+        assert "style" in error_msg
+
+    def test_validate_extra_variables_ok(self):
+        """Should not raise when extra variables provided."""
+        from yamlgraph.utils.template import validate_variables
+
+        template = "Hello {name}"
+        # Should not raise - extra vars are fine
+        validate_variables(template, {"name": "World", "extra": "ignored"}, "greet")
+
+    def test_validate_prompt_name_in_error(self):
+        """Error message should include prompt name."""
+        from yamlgraph.utils.template import validate_variables
+
+        template = "Hello {name}"
+        with pytest.raises(ValueError, match="greet"):
+            validate_variables(template, {}, "greet")
+
+    def test_validate_empty_template(self):
+        """Should not raise for template without variables."""
+        from yamlgraph.utils.template import validate_variables
+
+        template = "No variables here"
+        # Should not raise
+        validate_variables(template, {}, "static")
+
+    def test_validate_jinja2_template(self):
+        """Should validate Jinja2 templates correctly."""
+        from yamlgraph.utils.template import validate_variables
+
+        template = "{% for item in items %}{{ item }}{% endfor %}"
+        with pytest.raises(ValueError, match="items"):
+            validate_variables(template, {}, "list_template")
+
+
+class TestExecutePromptValidation:
+    """Integration tests for validation in execute_prompt."""
+
+    def test_execute_prompt_raises_on_missing_variable(self):
+        """Should raise clear error when required variable is missing."""
+        from yamlgraph.executor import execute_prompt
+
+        with pytest.raises(ValueError, match="Missing required variable.*name"):
+            execute_prompt(
+                prompt_name="greet",
+                variables={"style": "formal"},  # Missing 'name'
+            )
+
+    def test_execute_prompt_lists_all_missing_variables(self):
+        """Error should list ALL missing variables, not just first."""
+        from yamlgraph.executor import execute_prompt
+
+        with pytest.raises(ValueError) as exc_info:
+            execute_prompt(
+                prompt_name="greet",
+                variables={},  # Missing both 'name' and 'style'
+            )
+        error_msg = str(exc_info.value)
+        assert "name" in error_msg
+        assert "style" in error_msg

tests/unit/test_tool_nodes.py

@@ -0,0 +1,129 @@
+"""Tests for tool nodes (type: tool)."""
+
+import pytest
+
+from yamlgraph.tools.nodes import create_tool_node
+from yamlgraph.tools.shell import ShellToolConfig
+
+
+class TestCreateToolNode:
+    """Tests for create_tool_node function."""
+
+    def test_executes_named_tool(self):
+        """Node runs correct tool from registry."""
+        tools = {
+            "echo_tool": ShellToolConfig(command="echo hello"),
+            "other_tool": ShellToolConfig(command="echo other"),
+        }
+        node_config = {"tool": "echo_tool"}
+
+        node_fn = create_tool_node("test_node", node_config, tools)
+        result = node_fn({})
+
+        assert result["test_node"].strip() == "hello"
+        assert result["current_step"] == "test_node"
+
+    def test_resolves_variables_from_state(self):
+        """State values passed to tool."""
+        tools = {
+            "greet": ShellToolConfig(command="echo Hello {name}"),
+        }
+        node_config = {
+            "tool": "greet",
+            "variables": {"name": "{state.user_name}"},
+        }
+
+        node_fn = create_tool_node("greet_node", node_config, tools)
+        result = node_fn({"user_name": "Alice"})
+
+        assert "Alice" in result["greet_node"]
+
+    def test_stores_result_in_state_key(self):
+        """Tool output saved to custom state_key."""
+        tools = {
+            "data_tool": ShellToolConfig(command="echo data_value"),
+        }
+        node_config = {
+            "tool": "data_tool",
+            "state_key": "my_data",
+        }
+
+        node_fn = create_tool_node("fetch_node", node_config, tools)
+        result = node_fn({})
+
+        assert "my_data" in result
+        assert result["my_data"].strip() == "data_value"
+
+    def test_on_error_skip(self):
+        """Failed tool skipped when on_error: skip."""
+        tools = {
+            "fail_tool": ShellToolConfig(command="exit 1"),
+        }
+        node_config = {
+            "tool": "fail_tool",
+            "on_error": "skip",
+        }
+
+        node_fn = create_tool_node("fail_node", node_config, tools)
+        result = node_fn({})
+
+        # Should not raise, should return with error info
+        assert result["current_step"] == "fail_node"
+        assert "errors" in result or result.get("fail_node") is None
+
+    def test_on_error_fail_raises(self):
+        """Failed tool raises when on_error: fail."""
+        tools = {
+            "fail_tool": ShellToolConfig(command="exit 1"),
+        }
+        node_config = {
+            "tool": "fail_tool",
+            "on_error": "fail",
+        }
+
+        node_fn = create_tool_node("fail_node", node_config, tools)
+
+        with pytest.raises(RuntimeError):
+            node_fn({})
+
+    def test_nested_state_variable(self):
+        """Nested state values like {state.location.lat} resolved."""
+        tools = {
+            "geo": ShellToolConfig(command="echo {lat},{lon}"),
+        }
+        node_config = {
+            "tool": "geo",
+            "variables": {
+                "lat": "{state.location.lat}",
+                "lon": "{state.location.lon}",
+            },
+        }
+
+        node_fn = create_tool_node("geo_node", node_config, tools)
+        result = node_fn({"location": {"lat": "37.7749", "lon": "-122.4194"}})
+
+        assert "37.7749" in result["geo_node"]
+        assert "-122.4194" in result["geo_node"]
+
+    def test_missing_tool_raises(self):
+        """Unknown tool name raises error."""
+        tools = {}
+        node_config = {"tool": "nonexistent"}
+
+        with pytest.raises(KeyError):
+            create_tool_node("bad_node", node_config, tools)
+
+    def test_json_parse_tool(self):
+        """Tool with parse: json returns dict."""
+        tools = {
+            "json_tool": ShellToolConfig(
+                command="echo '{{\"count\": 42}}'",
+                parse="json",
+            ),
+        }
+        node_config = {"tool": "json_tool"}
+
+        node_fn = create_tool_node("json_node", node_config, tools)
+        result = node_fn({})
+
+        assert result["json_node"] == {"count": 42}

yamlgraph/__init__.py

@@ -0,0 +1,35 @@
+"""YamlGraph - YAML-first LLM pipeline framework.
+
+Framework for building LLM pipelines with YAML configuration.
+State is generated dynamically from graph config.
+"""
+
+from yamlgraph.builder import build_graph, build_resume_graph, run_pipeline
+from yamlgraph.executor import execute_prompt, get_executor
+from yamlgraph.models import (
+    ErrorType,
+    GenericReport,
+    PipelineError,
+    build_state_class,
+    create_initial_state,
+)
+from yamlgraph.storage import YamlGraphDB
+
+__all__ = [
+    # Builder
+    "build_graph",
+    "build_resume_graph",
+    "run_pipeline",
+    # Executor
+    "execute_prompt",
+    "get_executor",
+    # Framework models
+    "ErrorType",
+    "PipelineError",
+    "GenericReport",
+    # Dynamic state
+    "build_state_class",
+    "create_initial_state",
+    # Storage
+    "YamlGraphDB",
+]

yamlgraph/builder.py

@@ -0,0 +1,110 @@
+"""Graph builders for yamlgraph pipelines.
+
+Provides functions to build pipeline graphs from YAML configuration.
+
+Pipeline Architecture
+=====================
+
+The main pipeline follows this flow:
+
+```mermaid
+graph LR
+    A[generate] -->|content| B{should_continue}
+    B -->|continue| C[analyze]
+    B -->|end| E[END]
+    C -->|analysis| D[summarize]
+    D --> E[END]
+```
+
+State Flow:
+- generate: Creates structured content from topic
+- analyze: Produces analysis from generated content
+- summarize: Combines all outputs into final_summary
+
+Graph Definition:
+- Pipelines are defined in graphs/*.yaml
+- Loaded and compiled via graph_loader module
+"""
+
+from pathlib import Path
+from typing import Any
+
+from langgraph.graph import StateGraph
+
+from yamlgraph.config import DEFAULT_GRAPH
+from yamlgraph.graph_loader import load_and_compile
+from yamlgraph.models import create_initial_state
+
+# Type alias for dynamic state
+GraphState = dict[str, Any]
+
+
+def build_graph(
+    graph_path: Path | str | None = None,
+    checkpointer: Any | None = None,
+) -> StateGraph:
+    """Build a pipeline graph from YAML with optional checkpointer.
+
+    Args:
+        graph_path: Path to YAML graph definition.
+                   Defaults to graphs/yamlgraph.yaml
+        checkpointer: Optional LangGraph checkpointer for state persistence.
+                     Use get_checkpointer() from storage.checkpointer.
+
+    Returns:
+        StateGraph ready for compilation
+    """
+    path = Path(graph_path) if graph_path else DEFAULT_GRAPH
+    graph = load_and_compile(path)
+
+    # Checkpointer is applied at compile time
+    if checkpointer is not None:
+        # Store reference for compile() to use
+        graph._checkpointer = checkpointer
+
+    return graph
+
+
+def build_resume_graph() -> StateGraph:
+    """Build a graph for resuming an interrupted pipeline.
+
+    This is an alias for build_graph(). Resume works automatically
+    because nodes skip execution if their output already exists in state
+    (skip_if_exists behavior).
+
+    To resume:
+    1. Load saved state from database
+    2. Invoke graph with that state
+    3. Nodes with existing outputs are skipped
+
+    Returns:
+        StateGraph for resume (same as main pipeline)
+    """
+    return build_graph()
+
+
+def run_pipeline(
+    topic: str,
+    style: str = "informative",
+    word_count: int = 300,
+    graph_path: Path | str | None = None,
+) -> GraphState:
+    """Run the complete pipeline with given inputs.
+
+    Args:
+        topic: Topic to generate content about
+        style: Writing style
+        word_count: Target word count
+        graph_path: Optional path to graph YAML
+
+    Returns:
+        Final state with all outputs
+    """
+    graph = build_graph(graph_path).compile()
+    initial_state = create_initial_state(
+        topic=topic,
+        style=style,
+        word_count=word_count,
+    )
+
+    return graph.invoke(initial_state)