yamlgraph 0.1.1__py3-none-any.whl

tests/unit/test_executor_async.py

@@ -0,0 +1,179 @@
+"""Unit tests for async executor module."""
+
+import asyncio
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from yamlgraph.executor_async import execute_prompt_async, execute_prompts_concurrent
+from yamlgraph.utils.llm_factory_async import shutdown_executor
+
+
+class TestExecutePromptAsync:
+    """Tests for execute_prompt_async function."""
+
+    def teardown_method(self):
+        """Clean up executor after each test."""
+        shutdown_executor()
+
+    @pytest.mark.asyncio
+    async def test_executes_prompt(self):
+        """Should execute a prompt and return result."""
+        with patch("yamlgraph.executor_async.invoke_async") as mock_invoke:
+            mock_invoke.return_value = "Hello, World!"
+
+            result = await execute_prompt_async(
+                "greet",
+                variables={"name": "World", "style": "friendly"},
+            )
+
+            assert result == "Hello, World!"
+            mock_invoke.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_passes_output_model(self):
+        """Should pass output model to invoke_async."""
+        from pydantic import BaseModel
+
+        class TestModel(BaseModel):
+            greeting: str
+
+        with patch("yamlgraph.executor_async.invoke_async") as mock_invoke:
+            mock_invoke.return_value = TestModel(greeting="Hi")
+
+            result = await execute_prompt_async(
+                "greet",
+                variables={"name": "Test", "style": "casual"},
+                output_model=TestModel,
+            )
+
+            assert isinstance(result, TestModel)
+            # Check output_model was passed (positional arg)
+            call_args = mock_invoke.call_args
+            assert call_args[0][2] is TestModel  # 3rd positional arg
+
+    @pytest.mark.asyncio
+    async def test_validates_variables(self):
+        """Should raise error for missing variables."""
+        with pytest.raises(ValueError, match="Missing required variable"):
+            await execute_prompt_async("greet", variables={})
+
+    @pytest.mark.asyncio
+    async def test_uses_provider_from_yaml(self):
+        """Should extract provider from YAML metadata."""
+        with (
+            patch("yamlgraph.executor_async.load_prompt") as mock_load,
+            patch("yamlgraph.executor_async.invoke_async") as mock_invoke,
+            patch("yamlgraph.executor_async.create_llm") as mock_create_llm,
+        ):
+            mock_load.return_value = {
+                "system": "You are helpful.",
+                "user": "Hello {name}",
+                "provider": "mistral",
+            }
+            mock_invoke.return_value = "Response"
+            mock_create_llm.return_value = MagicMock()
+
+            await execute_prompt_async("test", variables={"name": "User"})
+
+            mock_create_llm.assert_called_once()
+            call_kwargs = mock_create_llm.call_args[1]
+            assert call_kwargs["provider"] == "mistral"
+
+
+class TestExecutePromptsConcurrent:
+    """Tests for execute_prompts_concurrent function."""
+
+    def teardown_method(self):
+        """Clean up executor after each test."""
+        shutdown_executor()
+
+    @pytest.mark.asyncio
+    async def test_executes_multiple_prompts(self):
+        """Should execute multiple prompts concurrently."""
+        with patch(
+            "yamlgraph.executor_async.execute_prompt_async", new_callable=AsyncMock
+        ) as mock_execute:
+            mock_execute.side_effect = ["Result 1", "Result 2", "Result 3"]
+
+            results = await execute_prompts_concurrent(
+                [
+                    {"prompt_name": "greet", "variables": {"name": "A", "style": "x"}},
+                    {"prompt_name": "greet", "variables": {"name": "B", "style": "y"}},
+                    {"prompt_name": "greet", "variables": {"name": "C", "style": "z"}},
+                ]
+            )
+
+            assert len(results) == 3
+            assert results == ["Result 1", "Result 2", "Result 3"]
+            assert mock_execute.call_count == 3
+
+    @pytest.mark.asyncio
+    async def test_preserves_order(self):
+        """Should return results in same order as input."""
+        with patch(
+            "yamlgraph.executor_async.execute_prompt_async", new_callable=AsyncMock
+        ) as mock_execute:
+            # Simulate varying response times
+            async def delayed_response(prompt_name, **kwargs):
+                name = kwargs.get("variables", {}).get("name", "")
+                if name == "slow":
+                    await asyncio.sleep(0.01)
+                return f"Response for {name}"
+
+            mock_execute.side_effect = delayed_response
+
+            results = await execute_prompts_concurrent(
+                [
+                    {
+                        "prompt_name": "greet",
+                        "variables": {"name": "slow", "style": "a"},
+                    },
+                    {
+                        "prompt_name": "greet",
+                        "variables": {"name": "fast", "style": "b"},
+                    },
+                ]
+            )
+
+            assert results[0] == "Response for slow"
+            assert results[1] == "Response for fast"
+
+    @pytest.mark.asyncio
+    async def test_empty_list(self):
+        """Should handle empty prompt list."""
+        results = await execute_prompts_concurrent([])
+        assert results == []
+
+    @pytest.mark.asyncio
+    async def test_passes_all_options(self):
+        """Should pass all options to execute_prompt_async."""
+        from pydantic import BaseModel
+
+        class TestModel(BaseModel):
+            value: str
+
+        with patch(
+            "yamlgraph.executor_async.execute_prompt_async", new_callable=AsyncMock
+        ) as mock_execute:
+            mock_execute.return_value = TestModel(value="test")
+
+            await execute_prompts_concurrent(
+                [
+                    {
+                        "prompt_name": "test",
+                        "variables": {"x": "y"},
+                        "output_model": TestModel,
+                        "temperature": 0.5,
+                        "provider": "openai",
+                    }
+                ]
+            )
+
+            mock_execute.assert_called_once_with(
+                prompt_name="test",
+                variables={"x": "y"},
+                output_model=TestModel,
+                temperature=0.5,
+                provider="openai",
+            )

tests/unit/test_export.py

@@ -0,0 +1,150 @@
+"""Tests for yamlgraph.storage.export module."""
+
+import json
+
+from tests.conftest import FixtureGeneratedContent
+from yamlgraph.storage.export import _serialize_state, export_state
+
+
+class TestExportState:
+    """Tests for export_state function."""
+
+    def test_export_creates_file(self, temp_output_dir, sample_state):
+        """Export should create a JSON file."""
+        filepath = export_state(sample_state, output_dir=temp_output_dir)
+        assert filepath.exists()
+        assert filepath.suffix == ".json"
+
+    def test_export_file_contains_valid_json(self, temp_output_dir, sample_state):
+        """Exported file should contain valid JSON."""
+        filepath = export_state(sample_state, output_dir=temp_output_dir)
+        with open(filepath) as f:
+            data = json.load(f)
+        assert "topic" in data
+        assert "thread_id" in data
+
+    def test_export_filename_format(self, temp_output_dir, sample_state):
+        """Filename should include prefix and thread_id."""
+        filepath = export_state(
+            sample_state,
+            output_dir=temp_output_dir,
+            prefix="test_export",
+        )
+        assert "test_export" in filepath.name
+        assert sample_state["thread_id"] in filepath.name
+
+    def test_export_creates_output_dir(self, tmp_path, sample_state):
+        """Export should create output directory if it doesn't exist."""
+        new_dir = tmp_path / "new_outputs"
+        filepath = export_state(sample_state, output_dir=new_dir)
+        assert new_dir.exists()
+        assert filepath.exists()
+
+
+class TestSerializeState:
+    """Tests for _serialize_state function."""
+
+    def test_serialize_simple_state(self, empty_state):
+        """Simple state should serialize unchanged."""
+        result = _serialize_state(empty_state)
+        assert result["topic"] == empty_state["topic"]
+        assert result["style"] == empty_state["style"]
+
+    def test_serialize_pydantic_models(self):
+        """Pydantic models should be converted to dicts."""
+        content = FixtureGeneratedContent(
+            title="Test",
+            content="Content",
+            word_count=1,
+            tags=["tag"],
+        )
+        state = {"generated": content}
+        result = _serialize_state(state)
+        assert isinstance(result["generated"], dict)
+        assert result["generated"]["title"] == "Test"
+
+    def test_serialize_preserves_none(self, empty_state):
+        """None values should be preserved."""
+        # Add a None field to test serialization
+        empty_state["generated"] = None
+        result = _serialize_state(empty_state)
+        assert result["generated"] is None
+        assert result["error"] is None
+
+
+class TestExportSummaryGeneric:
+    """Tests for generic export_summary behavior."""
+
+    def test_export_summary_with_any_pydantic_model(self):
+        """export_summary should work with any Pydantic model, not just demo-specific ones."""
+        from pydantic import BaseModel
+
+        from yamlgraph.storage.export import export_summary
+
+        class CustomModel(BaseModel):
+            name: str
+            value: int
+
+        state = {
+            "thread_id": "test-123",
+            "topic": "custom topic",
+            "custom_field": CustomModel(name="test", value=42),
+        }
+
+        summary = export_summary(state)
+
+        # Should include core fields
+        assert summary["thread_id"] == "test-123"
+        assert summary["topic"] == "custom topic"
+
+    def test_export_summary_extracts_scalar_fields(self):
+        """export_summary should extract key scalar fields from any model."""
+        from pydantic import BaseModel
+
+        from yamlgraph.storage.export import export_summary
+
+        class ReportContent(BaseModel):
+            headline: str
+            body: str
+            author: str
+
+        state = {
+            "thread_id": "report-1",
+            "topic": "report topic",
+            "report": ReportContent(
+                headline="Breaking News",
+                body="Content here...",
+                author="Alice",
+            ),
+        }
+
+        summary = export_summary(state)
+        # Should extract and include scalar fields
+        assert "report" in summary or any(k.startswith("report") for k in summary)
+
+    def test_export_summary_no_demo_model_dependencies(self):
+        """export_summary should not import demo-specific model types."""
+        import ast
+        import inspect
+
+        from yamlgraph.storage import export
+
+        source = inspect.getsource(export)
+        tree = ast.parse(source)
+
+        demo_models = {
+            "GeneratedContent",
+            "Analysis",
+            "ToneClassification",
+            "DraftContent",
+            "Critique",
+            "SearchResults",
+            "FinalReport",
+        }
+
+        for node in ast.walk(tree):
+            if isinstance(node, ast.ImportFrom):
+                if node.module and "schemas" in node.module:
+                    imported_names = {alias.name for alias in node.names}
+                    overlap = imported_names & demo_models
+                    assert not overlap, f"export.py imports demo models: {overlap}"

tests/unit/test_expressions.py

@@ -0,0 +1,178 @@
+"""Tests for state expression resolution."""
+
+import pytest
+from pydantic import BaseModel
+
+from yamlgraph.utils.expressions import (
+    resolve_state_expression,
+    resolve_state_path,
+    resolve_template,
+)
+
+
+class TestResolveStateExpression:
+    """Tests for resolve_state_expression function."""
+
+    def test_simple_key(self):
+        """Resolve simple state key."""
+        state = {"name": "test"}
+        result = resolve_state_expression("{name}", state)
+        assert result == "test"
+
+    def test_nested_path(self):
+        """Resolve nested path like {state.story.panels}."""
+        state = {"story": {"panels": ["a", "b", "c"]}}
+        result = resolve_state_expression("{state.story.panels}", state)
+        assert result == ["a", "b", "c"]
+
+    def test_state_prefix_stripped(self):
+        """The 'state.' prefix is optional and stripped."""
+        state = {"story": {"title": "My Story"}}
+        # With prefix
+        assert resolve_state_expression("{state.story.title}", state) == "My Story"
+        # Without prefix
+        assert resolve_state_expression("{story.title}", state) == "My Story"
+
+    def test_literal_passthrough(self):
+        """Non-expression strings pass through unchanged."""
+        result = resolve_state_expression("literal string", {})
+        assert result == "literal string"
+
+    def test_non_string_passthrough(self):
+        """Non-string values pass through unchanged."""
+        result = resolve_state_expression(42, {})
+        assert result == 42
+
+    def test_missing_key_raises(self):
+        """Missing key raises KeyError."""
+        state = {"foo": "bar"}
+        with pytest.raises(KeyError):
+            resolve_state_expression("{missing}", state)
+
+    def test_missing_nested_key_raises(self):
+        """Missing nested key raises KeyError."""
+        state = {"story": {"title": "test"}}
+        with pytest.raises(KeyError):
+            resolve_state_expression("{story.panels}", state)
+
+    def test_deeply_nested_path(self):
+        """Resolve deeply nested paths."""
+        state = {"a": {"b": {"c": {"d": "deep"}}}}
+        result = resolve_state_expression("{a.b.c.d}", state)
+        assert result == "deep"
+
+    def test_list_result(self):
+        """Can resolve to list values."""
+        state = {"items": [1, 2, 3]}
+        result = resolve_state_expression("{items}", state)
+        assert result == [1, 2, 3]
+
+    def test_dict_result(self):
+        """Can resolve to dict values."""
+        state = {"config": {"key": "value"}}
+        result = resolve_state_expression("{config}", state)
+        assert result == {"key": "value"}
+
+    def test_object_attribute_access(self):
+        """Can resolve object attributes (Pydantic models)."""
+
+        class MockModel:
+            def __init__(self):
+                self.title = "Test Title"
+                self.panels = ["panel 1", "panel 2"]
+
+        state = {"story": MockModel()}
+        result = resolve_state_expression("{state.story.panels}", state)
+        assert result == ["panel 1", "panel 2"]
+
+    def test_mixed_dict_and_object_access(self):
+        """Can resolve mixed dict and object paths."""
+
+        class Inner:
+            def __init__(self):
+                self.value = "found"
+
+        state = {"outer": {"middle": Inner()}}
+        result = resolve_state_expression("{outer.middle.value}", state)
+        assert result == "found"
+
+
+class TestResolveStatePath:
+    """Tests for resolve_state_path - the core resolution function."""
+
+    def test_simple_key(self):
+        """Should resolve simple key."""
+        state = {"score": 0.8}
+        assert resolve_state_path("score", state) == 0.8
+
+    def test_nested_dict_path(self):
+        """Should resolve nested dict path."""
+        state = {"critique": {"score": 0.9}}
+        assert resolve_state_path("critique.score", state) == 0.9
+
+    def test_deeply_nested(self):
+        """Should resolve deeply nested path."""
+        state = {"a": {"b": {"c": {"d": 42}}}}
+        assert resolve_state_path("a.b.c.d", state) == 42
+
+    def test_missing_key_returns_none(self):
+        """Should return None for missing key."""
+        state = {"a": 1}
+        assert resolve_state_path("b", state) is None
+
+    def test_missing_nested_returns_none(self):
+        """Should return None for missing nested path."""
+        state = {"a": {"b": 1}}
+        assert resolve_state_path("a.c", state) is None
+
+    def test_pydantic_model_attribute(self):
+        """Should resolve Pydantic model attribute."""
+
+        class Critique(BaseModel):
+            score: float
+            feedback: str
+
+        state = {"critique": Critique(score=0.75, feedback="Good")}
+        assert resolve_state_path("critique.score", state) == 0.75
+        assert resolve_state_path("critique.feedback", state) == "Good"
+
+    def test_empty_path_returns_none(self):
+        """Should return None for empty path."""
+        assert resolve_state_path("", {"a": 1}) is None
+
+
+class TestResolveTemplate:
+    """Tests for resolve_template - optional resolution returning None."""
+
+    def test_state_template(self):
+        """Should resolve {state.field} template."""
+        state = {"topic": "AI"}
+        assert resolve_template("{state.topic}", state) == "AI"
+
+    def test_nested_template(self):
+        """Should resolve nested path template."""
+        state = {"config": {"max_tokens": 100}}
+        assert resolve_template("{state.config.max_tokens}", state) == 100
+
+    def test_missing_returns_none(self):
+        """Should return None for missing path."""
+        state = {"a": 1}
+        assert resolve_template("{state.missing}", state) is None
+
+    def test_non_string_passthrough(self):
+        """Should pass through non-string values."""
+        assert resolve_template(123, {}) == 123
+
+    def test_non_state_template_passthrough(self):
+        """Should pass through non-state templates."""
+        assert resolve_template("{other.field}", {}) == "{other.field}"
+        assert resolve_template("plain text", {}) == "plain text"
+
+    def test_pydantic_model(self):
+        """Should resolve Pydantic model attribute."""
+
+        class Draft(BaseModel):
+            text: str
+
+        state = {"draft": Draft(text="Content")}
+        assert resolve_template("{state.draft.text}", state) == "Content"

tests/unit/test_format_prompt.py

@@ -0,0 +1,145 @@
+"""Tests for prompt formatting with Jinja2 support."""
+
+import pytest
+
+from yamlgraph.executor import format_prompt
+
+
+class TestFormatPrompt:
+    """Test the format_prompt function with both simple and Jinja2 templates."""
+
+    def test_simple_format_basic(self):
+        """Test basic string formatting with {variable} syntax."""
+        template = "Hello {name}!"
+        variables = {"name": "World"}
+        result = format_prompt(template, variables)
+        assert result == "Hello World!"
+
+    def test_simple_format_multiple_variables(self):
+        """Test formatting with multiple variables."""
+        template = "Topic: {topic}, Style: {style}, Words: {word_count}"
+        variables = {"topic": "AI", "style": "casual", "word_count": 500}
+        result = format_prompt(template, variables)
+        assert result == "Topic: AI, Style: casual, Words: 500"
+
+    def test_simple_format_missing_variable(self):
+        """Test that missing variables raise KeyError."""
+        template = "Hello {name}!"
+        variables = {}
+        with pytest.raises(KeyError):
+            format_prompt(template, variables)
+
+    def test_jinja2_basic_variable(self):
+        """Test Jinja2 template with basic {{ variable }} syntax."""
+        template = "Hello {{ name }}!"
+        variables = {"name": "World"}
+        result = format_prompt(template, variables)
+        assert result == "Hello World!"
+
+    def test_jinja2_for_loop(self):
+        """Test Jinja2 template with for loop."""
+        template = """{% for item in items %}
+- {{ item }}
+{% endfor %}"""
+        variables = {"items": ["apple", "banana", "cherry"]}
+        result = format_prompt(template, variables)
+        # Jinja2 preserves whitespace from template
+        assert "- apple" in result
+        assert "- banana" in result
+        assert "- cherry" in result
+
+    def test_jinja2_conditional(self):
+        """Test Jinja2 template with if/else."""
+        template = """{% if premium %}Premium User{% else %}Regular User{% endif %}"""
+
+        result_premium = format_prompt(template, {"premium": True})
+        assert result_premium == "Premium User"
+
+        result_regular = format_prompt(template, {"premium": False})
+        assert result_regular == "Regular User"
+
+    def test_jinja2_filter_slice(self):
+        """Test Jinja2 template with slice filter."""
+        template = "Summary: {{ text[:50] }}..."
+        variables = {
+            "text": "This is a very long text that should be truncated to show only first fifty characters"
+        }
+        result = format_prompt(template, variables)
+        # Check that the text is sliced to 50 characters
+        assert result.startswith(
+            "Summary: This is a very long text that should be truncated"
+        )
+        assert result.endswith("...")
+        assert len(result) < len(variables["text"]) + len("Summary: ...")
+
+    def test_jinja2_filter_upper(self):
+        """Test Jinja2 template with upper filter."""
+        template = "{{ name | upper }}"
+        variables = {"name": "world"}
+        result = format_prompt(template, variables)
+        assert result == "WORLD"
+
+    def test_jinja2_complex_template(self):
+        """Test complex Jinja2 template with loops and conditionals."""
+        template = """Items in {{ category }}:
+{% for item in items %}
+  {% if item.available %}
+  - {{ item.name }}: ${{ item.price }}
+  {% endif %}
+{% endfor %}"""
+        variables = {
+            "category": "Fruits",
+            "items": [
+                {"name": "Apple", "price": 1.50, "available": True},
+                {"name": "Banana", "price": 0.75, "available": False},
+                {"name": "Cherry", "price": 2.00, "available": True},
+            ],
+        }
+        result = format_prompt(template, variables)
+        assert "Apple: $1.5" in result
+        assert "Cherry: $2.0" in result
+        assert "Banana" not in result
+
+    def test_jinja2_missing_variable_graceful(self):
+        """Test that Jinja2 missing variables are handled (rendered as empty by default)."""
+        template = "Hello {{ name }}!"
+        variables = {}
+        result = format_prompt(template, variables)
+        # Jinja2 by default renders undefined variables as empty strings
+        assert result == "Hello !"
+
+    def test_detection_uses_jinja2_for_double_braces(self):
+        """Test that {{ triggers Jinja2 mode."""
+        template = "Value: {{ x }}"
+        variables = {"x": 42}
+        result = format_prompt(template, variables)
+        assert result == "Value: 42"
+
+    def test_detection_uses_jinja2_for_statements(self):
+        """Test that {% triggers Jinja2 mode."""
+        template = "{% if true %}Yes{% endif %}"
+        variables = {}
+        result = format_prompt(template, variables)
+        assert result == "Yes"
+
+    def test_backward_compatibility_no_jinja2_syntax(self):
+        """Test that templates without Jinja2 syntax still use simple format."""
+        # This ensures backward compatibility
+        template = "Simple {var} template"
+        variables = {"var": "test"}
+        result = format_prompt(template, variables)
+        assert result == "Simple test template"
+
+    def test_empty_template(self):
+        """Test formatting empty template."""
+        template = ""
+        variables = {}
+        result = format_prompt(template, variables)
+        assert result == ""
+
+    def test_template_with_no_placeholders(self):
+        """Test template with no variables."""
+        template = "Just plain text"
+        variables = {"unused": "value"}
+        result = format_prompt(template, variables)
+        assert result == "Just plain text"