yamlgraph 0.1.1__py3-none-any.whl

tests/unit/test_llm_factory_async.py

@@ -0,0 +1,118 @@
+"""Unit tests for async LLM factory module."""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from yamlgraph.utils.llm_factory_async import (
+    create_llm_async,
+    get_executor,
+    invoke_async,
+    shutdown_executor,
+)
+
+
+class TestGetExecutor:
+    """Tests for get_executor function."""
+
+    def teardown_method(self):
+        """Clean up executor after each test."""
+        shutdown_executor()
+
+    def test_creates_executor(self):
+        """Should create a ThreadPoolExecutor."""
+        executor = get_executor()
+        assert executor is not None
+
+    def test_returns_same_executor(self):
+        """Should return the same executor on subsequent calls."""
+        executor1 = get_executor()
+        executor2 = get_executor()
+        assert executor1 is executor2
+
+
+class TestShutdownExecutor:
+    """Tests for shutdown_executor function."""
+
+    def test_shutdown_cleans_up(self):
+        """Shutdown should clean up executor."""
+        # Create an executor
+        executor1 = get_executor()
+        assert executor1 is not None
+
+        # Shutdown
+        shutdown_executor()
+
+        # Next call should create a new executor
+        executor2 = get_executor()
+        assert executor2 is not executor1
+
+    def test_shutdown_when_none(self):
+        """Shutdown when no executor should not raise."""
+        shutdown_executor()  # Ensure clean state
+        shutdown_executor()  # Should not raise
+
+
+class TestCreateLLMAsync:
+    """Tests for create_llm_async function."""
+
+    def teardown_method(self):
+        """Clean up executor after each test."""
+        shutdown_executor()
+
+    @pytest.mark.asyncio
+    async def test_creates_llm(self):
+        """Should create an LLM instance."""
+        llm = await create_llm_async(provider="anthropic", temperature=0.5)
+        assert llm is not None
+        assert llm.temperature == 0.5
+
+    @pytest.mark.asyncio
+    async def test_uses_default_provider(self):
+        """Should use default provider when not specified."""
+        with patch.dict("os.environ", {"PROVIDER": ""}, clear=False):
+            llm = await create_llm_async(temperature=0.7)
+            # Default is anthropic
+            assert "anthropic" in llm.__class__.__name__.lower()
+
+
+class TestInvokeAsync:
+    """Tests for invoke_async function."""
+
+    def teardown_method(self):
+        """Clean up executor after each test."""
+        shutdown_executor()
+
+    @pytest.mark.asyncio
+    async def test_invoke_returns_string(self):
+        """Should return string content when no output model."""
+        mock_llm = MagicMock()
+        mock_response = MagicMock()
+        mock_response.content = "Hello, world!"
+        mock_llm.invoke.return_value = mock_response
+
+        messages = [MagicMock()]
+        result = await invoke_async(mock_llm, messages)
+
+        assert result == "Hello, world!"
+        mock_llm.invoke.assert_called_once_with(messages)
+
+    @pytest.mark.asyncio
+    async def test_invoke_with_output_model(self):
+        """Should use structured output when model provided."""
+        from pydantic import BaseModel
+
+        class TestOutput(BaseModel):
+            value: str
+
+        mock_llm = MagicMock()
+        mock_structured_llm = MagicMock()
+        mock_llm.with_structured_output.return_value = mock_structured_llm
+        mock_structured_llm.invoke.return_value = TestOutput(value="test")
+
+        messages = [MagicMock()]
+        result = await invoke_async(mock_llm, messages, output_model=TestOutput)
+
+        assert isinstance(result, TestOutput)
+        assert result.value == "test"
+        mock_llm.with_structured_output.assert_called_once_with(TestOutput)

tests/unit/test_loops.py

@@ -0,0 +1,403 @@
+"""Tests for Section 3: Self-Correction Loops (Reflexion).
+
+TDD tests for expression conditions, loop tracking, and cyclic graphs.
+"""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+# =============================================================================
+# Test: Expression Condition Parsing
+# =============================================================================
+
+
+class TestExpressionConditions:
+    """Tests for condition expression evaluation."""
+
+    def test_evaluate_condition_exists(self):
+        """evaluate_condition function should exist."""
+        from yamlgraph.utils.conditions import evaluate_condition
+
+        assert callable(evaluate_condition)
+
+    def test_less_than_comparison(self):
+        """Evaluates 'score < 0.8' correctly."""
+        from yamlgraph.utils.conditions import evaluate_condition
+
+        state = {"score": 0.5}
+        assert evaluate_condition("score < 0.8", state) is True
+
+        state = {"score": 0.9}
+        assert evaluate_condition("score < 0.8", state) is False
+
+    def test_greater_than_comparison(self):
+        """Evaluates 'score > 0.5' correctly."""
+        from yamlgraph.utils.conditions import evaluate_condition
+
+        state = {"score": 0.7}
+        assert evaluate_condition("score > 0.5", state) is True
+
+        state = {"score": 0.3}
+        assert evaluate_condition("score > 0.5", state) is False
+
+    def test_less_than_or_equal(self):
+        """Evaluates 'score <= 0.8' correctly."""
+        from yamlgraph.utils.conditions import evaluate_condition
+
+        state = {"score": 0.8}
+        assert evaluate_condition("score <= 0.8", state) is True
+
+        state = {"score": 0.9}
+        assert evaluate_condition("score <= 0.8", state) is False
+
+    def test_greater_than_or_equal(self):
+        """Evaluates 'score >= 0.8' correctly."""
+        from yamlgraph.utils.conditions import evaluate_condition
+
+        state = {"score": 0.8}
+        assert evaluate_condition("score >= 0.8", state) is True
+
+        state = {"score": 0.7}
+        assert evaluate_condition("score >= 0.8", state) is False
+
+    def test_equality_comparison(self):
+        """Evaluates 'status == \"approved\"' correctly."""
+        from yamlgraph.utils.conditions import evaluate_condition
+
+        state = {"status": "approved"}
+        assert evaluate_condition('status == "approved"', state) is True
+
+        state = {"status": "pending"}
+        assert evaluate_condition('status == "approved"', state) is False
+
+    def test_inequality_comparison(self):
+        """Evaluates 'error != null' correctly."""
+        from yamlgraph.utils.conditions import evaluate_condition
+
+        state = {"error": "something"}
+        assert evaluate_condition("error != null", state) is True
+
+        state = {"error": None}
+        assert evaluate_condition("error != null", state) is False
+
+    def test_nested_attribute_access(self):
+        """Evaluates 'critique.score >= 0.8' from state."""
+        from yamlgraph.utils.conditions import evaluate_condition
+
+        # Using object with attribute
+        critique = MagicMock()
+        critique.score = 0.85
+        state = {"critique": critique}
+        assert evaluate_condition("critique.score >= 0.8", state) is True
+
+        critique.score = 0.7
+        assert evaluate_condition("critique.score >= 0.8", state) is False
+
+    def test_compound_and_condition(self):
+        """Evaluates 'score < 0.8 and iteration < 3'."""
+        from yamlgraph.utils.conditions import evaluate_condition
+
+        state = {"score": 0.5, "iteration": 2}
+        assert evaluate_condition("score < 0.8 and iteration < 3", state) is True
+
+        state = {"score": 0.9, "iteration": 2}
+        assert evaluate_condition("score < 0.8 and iteration < 3", state) is False
+
+        state = {"score": 0.5, "iteration": 5}
+        assert evaluate_condition("score < 0.8 and iteration < 3", state) is False
+
+    def test_compound_or_condition(self):
+        """Evaluates 'approved == true or override == true'."""
+        from yamlgraph.utils.conditions import evaluate_condition
+
+        state = {"approved": True, "override": False}
+        assert evaluate_condition("approved == true or override == true", state) is True
+
+        state = {"approved": False, "override": True}
+        assert evaluate_condition("approved == true or override == true", state) is True
+
+        state = {"approved": False, "override": False}
+        assert (
+            evaluate_condition("approved == true or override == true", state) is False
+        )
+
+    def test_invalid_expression_raises(self):
+        """Malformed expression raises ValueError."""
+        from yamlgraph.utils.conditions import evaluate_condition
+
+        with pytest.raises(ValueError):
+            evaluate_condition("score <<< 0.8", {})
+
+    def test_missing_attribute_returns_false(self):
+        """Missing attribute in state returns False gracefully."""
+        from yamlgraph.utils.conditions import evaluate_condition
+
+        state = {}
+        # Should not raise, should return False for missing attribute
+        assert evaluate_condition("score < 0.8", state) is False
+
+
+# =============================================================================
+# Test: Loop Tracking
+# =============================================================================
+
+
+class TestLoopTracking:
+    """Tests for loop iteration tracking."""
+
+    def test_state_has_loop_counts_field(self):
+        """Dynamic state should have _loop_counts field."""
+        from yamlgraph.models.state_builder import build_state_class
+
+        State = build_state_class({"nodes": {}})
+        # Should have _loop_counts in annotations
+        assert "_loop_counts" in State.__annotations__
+
+        # And work at runtime
+        state = {"_loop_counts": {"critique": 2}}
+        assert state["_loop_counts"]["critique"] == 2
+
+    def test_node_increments_loop_counter(self):
+        """Each node execution increments its counter in _loop_counts."""
+        from yamlgraph.node_factory import create_node_function
+
+        node_config = {
+            "prompt": "test_prompt",
+            "state_key": "result",
+        }
+
+        with patch("yamlgraph.node_factory.execute_prompt") as mock_execute:
+            mock_execute.return_value = "test result"
+
+            node_fn = create_node_function("critique", node_config, {})
+
+            # First call - should initialize counter
+            state = {"message": "test"}
+            result = node_fn(state)
+            assert result.get("_loop_counts", {}).get("critique") == 1
+
+            # Second call - should increment
+            state = {"message": "test", "_loop_counts": {"critique": 1}}
+            result = node_fn(state)
+            assert result.get("_loop_counts", {}).get("critique") == 2
+
+
+# =============================================================================
+# Test: Loop Limits Configuration
+# =============================================================================
+
+
+class TestLoopLimits:
+    """Tests for loop_limits configuration."""
+
+    def test_parses_loop_limits_from_yaml(self):
+        """GraphConfig parses loop_limits section."""
+        from yamlgraph.graph_loader import GraphConfig
+
+        config_dict = {
+            "version": "1.0",
+            "name": "test",
+            "nodes": {
+                "draft": {"prompt": "draft"},
+                "critique": {"prompt": "critique"},
+            },
+            "edges": [
+                {"from": "START", "to": "draft"},
+                {"from": "draft", "to": "critique"},
+                {"from": "critique", "to": "END"},
+            ],
+            "loop_limits": {
+                "critique": 3,
+            },
+        }
+        config = GraphConfig(config_dict)
+        assert config.loop_limits == {"critique": 3}
+
+    def test_loop_limits_defaults_to_empty(self):
+        """Missing loop_limits defaults to empty dict."""
+        from yamlgraph.graph_loader import GraphConfig
+
+        config_dict = {
+            "version": "1.0",
+            "name": "test",
+            "nodes": {"node1": {"prompt": "p1"}},
+            "edges": [{"from": "START", "to": "node1"}, {"from": "node1", "to": "END"}],
+        }
+        config = GraphConfig(config_dict)
+        assert config.loop_limits == {}
+
+    def test_node_checks_loop_limit(self):
+        """Node execution checks loop limit before running."""
+        from yamlgraph.node_factory import create_node_function
+
+        node_config = {
+            "prompt": "test_prompt",
+            "state_key": "result",
+            "loop_limit": 3,  # Node-level limit
+        }
+
+        with patch("yamlgraph.node_factory.execute_prompt") as mock_execute:
+            mock_execute.return_value = "test result"
+
+            node_fn = create_node_function("critique", node_config, {})
+
+            # Under limit - should execute
+            state = {"_loop_counts": {"critique": 2}}
+            result = node_fn(state)
+            assert "result" in result
+
+            # At limit - should skip/terminate
+            state = {"_loop_counts": {"critique": 3}}
+            result = node_fn(state)
+            assert result.get("_loop_limit_reached") is True
+
+
+# =============================================================================
+# Test: Cyclic Edges
+# =============================================================================
+
+
+class TestCyclicEdges:
+    """Tests for cyclic graph support."""
+
+    def test_allows_backward_edges(self):
+        """Graph config allows edges pointing to earlier nodes."""
+        from yamlgraph.graph_loader import GraphConfig
+
+        config_dict = {
+            "version": "1.0",
+            "name": "test",
+            "nodes": {
+                "draft": {"prompt": "draft"},
+                "critique": {"prompt": "critique"},
+                "refine": {"prompt": "refine"},
+            },
+            "edges": [
+                {"from": "START", "to": "draft"},
+                {"from": "draft", "to": "critique"},
+                {
+                    "from": "critique",
+                    "to": "refine",
+                    "condition": "critique.score < 0.8",
+                },
+                {"from": "critique", "to": "END", "condition": "critique.score >= 0.8"},
+                {"from": "refine", "to": "critique"},  # Backward edge (cycle)
+            ],
+            "loop_limits": {"critique": 3},
+        }
+        # Should not raise
+        config = GraphConfig(config_dict)
+        assert config is not None
+
+    def test_compiles_cyclic_graph(self):
+        """Cyclic graph compiles to StateGraph."""
+        from yamlgraph.graph_loader import GraphConfig, compile_graph
+
+        config_dict = {
+            "version": "1.0",
+            "name": "test",
+            "nodes": {
+                "draft": {"prompt": "draft", "state_key": "current_draft"},
+                "critique": {"prompt": "critique", "state_key": "critique"},
+                "refine": {"prompt": "refine", "state_key": "current_draft"},
+            },
+            "edges": [
+                {"from": "START", "to": "draft"},
+                {"from": "draft", "to": "critique"},
+                {
+                    "from": "critique",
+                    "to": "refine",
+                    "condition": "critique.score < 0.8",
+                },
+                {"from": "critique", "to": "END", "condition": "critique.score >= 0.8"},
+                {"from": "refine", "to": "critique"},  # Cycle
+            ],
+            "loop_limits": {"critique": 3},
+        }
+        config = GraphConfig(config_dict)
+        graph = compile_graph(config)
+        assert graph is not None
+
+
+# =============================================================================
+# Test: Pydantic Models
+# =============================================================================
+
+
+class TestReflexionModels:
+    """Tests for DraftContent and Critique-like fixture models.
+
+    Note: Demo models were removed from yamlgraph.models in Section 10.
+    These tests use fixture models to prove the pattern still works.
+    """
+
+    def test_draft_content_model_exists(self):
+        """DraftContent-like fixture model can be created."""
+        from tests.conftest import FixtureDraftContent
+
+        assert FixtureDraftContent is not None
+
+    def test_draft_content_fields(self):
+        """DraftContent-like model has content and version fields."""
+        from tests.conftest import FixtureDraftContent
+
+        draft = FixtureDraftContent(content="Test essay", version=1)
+        assert draft.content == "Test essay"
+        assert draft.version == 1
+
+    def test_critique_model_exists(self):
+        """Critique-like fixture model can be created."""
+        from tests.conftest import FixtureCritique
+
+        assert FixtureCritique is not None
+
+    def test_critique_fields(self):
+        """Critique-like model has score, feedback, issues, should_refine fields."""
+        from tests.conftest import FixtureCritique
+
+        critique = FixtureCritique(
+            score=0.75,
+            feedback="Improve transitions",
+            issues=["Weak intro", "No conclusion"],
+            should_refine=True,
+        )
+        assert critique.score == 0.75
+        assert critique.feedback == "Improve transitions"
+        assert len(critique.issues) == 2
+        assert critique.should_refine is True
+
+
+# =============================================================================
+# Test: Reflexion Demo Graph
+# =============================================================================
+
+
+class TestReflexionDemoGraph:
+    """Tests for the reflexion-demo.yaml graph."""
+
+    def test_demo_graph_loads(self):
+        """reflexion-demo.yaml loads without error."""
+        from yamlgraph.graph_loader import load_graph_config
+
+        config = load_graph_config("graphs/reflexion-demo.yaml")
+        assert config.name == "reflexion-demo"
+        assert "draft" in config.nodes
+        assert "critique" in config.nodes
+        assert "refine" in config.nodes
+
+    def test_demo_graph_has_loop_limits(self):
+        """reflexion-demo.yaml has loop_limits configured."""
+        from yamlgraph.graph_loader import load_graph_config
+
+        config = load_graph_config("graphs/reflexion-demo.yaml")
+        assert "critique" in config.loop_limits
+        assert config.loop_limits["critique"] >= 3
+
+    def test_demo_graph_compiles(self):
+        """reflexion-demo.yaml compiles to StateGraph."""
+        from yamlgraph.graph_loader import compile_graph, load_graph_config
+
+        config = load_graph_config("graphs/reflexion-demo.yaml")
+        graph = compile_graph(config)
+        assert graph is not None

tests/unit/test_map_node.py

@@ -0,0 +1,144 @@
+"""Tests for type: map node functionality."""
+
+from unittest.mock import MagicMock
+
+import pytest
+
+from yamlgraph.map_compiler import compile_map_node, wrap_for_reducer
+
+
+class TestWrapForReducer:
+    """Tests for wrap_for_reducer helper."""
+
+    def test_wraps_result_in_list(self):
+        """Wrap node output for reducer aggregation."""
+
+        def simple_node(state: dict) -> dict:
+            return {"result": state["item"] * 2}
+
+        wrapped = wrap_for_reducer(simple_node, "collected", "result")
+        result = wrapped({"item": 5})
+
+        assert result == {"collected": [10]}
+
+    def test_preserves_map_index(self):
+        """Preserve _map_index in wrapped output."""
+
+        def node_fn(state: dict) -> dict:
+            return {"data": state["value"]}
+
+        wrapped = wrap_for_reducer(node_fn, "results", "data")
+        result = wrapped({"value": "test", "_map_index": 2})
+
+        assert result == {"results": [{"_map_index": 2, "value": "test"}]}
+
+    def test_extracts_state_key(self):
+        """Extract specific state_key from node result."""
+
+        def node_fn(state: dict) -> dict:
+            return {"frame_data": {"before": "a", "after": "b"}, "other": "ignore"}
+
+        wrapped = wrap_for_reducer(node_fn, "frames", "frame_data")
+        result = wrapped({})
+
+        assert result == {"frames": [{"before": "a", "after": "b"}]}
+
+
+class TestCompileMapNode:
+    """Tests for compile_map_node function."""
+
+    def test_creates_map_edge_function(self):
+        """compile_map_node returns a map edge function."""
+        config = {
+            "over": "{items}",
+            "as": "item",
+            "collect": "results",
+            "node": {"type": "llm", "prompt": "test", "state_key": "result"},
+        }
+        builder = MagicMock()
+        defaults = {}
+
+        map_edge, sub_node_name = compile_map_node("expand", config, builder, defaults)
+
+        # Should return callable and sub-node name
+        assert callable(map_edge)
+        assert sub_node_name == "_map_expand_sub"
+
+    def test_map_edge_returns_send_list(self):
+        """Map edge function returns list of Send objects."""
+        from langgraph.types import Send
+
+        config = {
+            "over": "{items}",
+            "as": "item",
+            "collect": "results",
+            "node": {"type": "llm", "prompt": "test", "state_key": "result"},
+        }
+        builder = MagicMock()
+        defaults = {}
+
+        map_edge, sub_node_name = compile_map_node("expand", config, builder, defaults)
+
+        state = {"items": ["a", "b", "c"]}
+        sends = map_edge(state)
+
+        assert len(sends) == 3
+        assert all(isinstance(s, Send) for s in sends)
+        assert sends[0].node == sub_node_name
+        assert sends[0].arg["item"] == "a"
+        assert sends[0].arg["_map_index"] == 0
+        assert sends[1].arg["item"] == "b"
+        assert sends[1].arg["_map_index"] == 1
+
+    def test_map_edge_empty_list(self):
+        """Empty list returns empty Send list."""
+        config = {
+            "over": "{items}",
+            "as": "item",
+            "collect": "results",
+            "node": {"type": "llm", "prompt": "test", "state_key": "result"},
+        }
+        builder = MagicMock()
+        defaults = {}
+
+        map_edge, _ = compile_map_node("expand", config, builder, defaults)
+
+        state = {"items": []}
+        sends = map_edge(state)
+
+        assert sends == []
+
+    def test_adds_wrapped_sub_node_to_builder(self):
+        """compile_map_node adds wrapped sub-node to builder."""
+        config = {
+            "over": "{items}",
+            "as": "item",
+            "collect": "results",
+            "node": {"type": "llm", "prompt": "test", "state_key": "result"},
+        }
+        builder = MagicMock()
+        defaults = {}
+
+        compile_map_node("expand", config, builder, defaults)
+
+        # Should call builder.add_node
+        builder.add_node.assert_called_once()
+        call_args = builder.add_node.call_args
+        assert call_args[0][0] == "_map_expand_sub"
+
+    def test_validates_over_is_list(self):
+        """Map edge validates that 'over' resolves to a list."""
+        config = {
+            "over": "{not_a_list}",
+            "as": "item",
+            "collect": "results",
+            "node": {"type": "llm", "prompt": "test", "state_key": "result"},
+        }
+        builder = MagicMock()
+        defaults = {}
+
+        map_edge, _ = compile_map_node("expand", config, builder, defaults)
+
+        state = {"not_a_list": "string"}
+        with pytest.raises(TypeError, match="must resolve to list"):
+            map_edge(state)