yamlgraph 0.3.9__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_streaming.py

@@ -0,0 +1,307 @@
+"""Tests for streaming support - Phase 3 (004).
+
+TDD: RED phase - write tests first.
+"""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+# ==============================================================================
+# execute_prompt_streaming tests
+# ==============================================================================
+
+
+@pytest.mark.asyncio
+async def test_execute_prompt_streaming_yields_tokens():
+    """execute_prompt_streaming yields tokens from LLM stream."""
+    from yamlgraph.executor_async import execute_prompt_streaming
+
+    # Mock LLM with astream method
+    mock_chunk1 = MagicMock()
+    mock_chunk1.content = "Hello"
+    mock_chunk2 = MagicMock()
+    mock_chunk2.content = " World"
+    mock_chunk3 = MagicMock()
+    mock_chunk3.content = "!"
+
+    async def mock_astream(*args, **kwargs):
+        for chunk in [mock_chunk1, mock_chunk2, mock_chunk3]:
+            yield chunk
+
+    mock_llm = MagicMock()
+    mock_llm.astream = mock_astream
+
+    with (
+        patch("yamlgraph.executor_async.create_llm", return_value=mock_llm),
+        patch("yamlgraph.executor_async.load_prompt") as mock_load,
+    ):
+        mock_load.return_value = {
+            "system": "You are helpful.",
+            "user": "Say hello",
+        }
+
+        tokens = []
+        async for token in execute_prompt_streaming("greet", variables={}):
+            tokens.append(token)
+
+        assert tokens == ["Hello", " World", "!"]
+
+
+@pytest.mark.asyncio
+async def test_execute_prompt_streaming_with_variables():
+    """execute_prompt_streaming formats template with variables."""
+    from yamlgraph.executor_async import execute_prompt_streaming
+
+    mock_chunk = MagicMock()
+    mock_chunk.content = "Hi Alice!"
+
+    async def mock_astream(*args, **kwargs):
+        yield mock_chunk
+
+    mock_llm = MagicMock()
+    mock_llm.astream = mock_astream
+
+    with (
+        patch("yamlgraph.executor_async.create_llm", return_value=mock_llm),
+        patch("yamlgraph.executor_async.load_prompt") as mock_load,
+        patch("yamlgraph.executor_async.format_prompt", return_value="Say hello to Alice") as mock_format,
+    ):
+        mock_load.return_value = {
+            "system": "",
+            "user": "Say hello to {name}",
+        }
+
+        tokens = []
+        async for token in execute_prompt_streaming("greet", variables={"name": "Alice"}):
+            tokens.append(token)
+
+        assert tokens == ["Hi Alice!"]
+        mock_format.assert_called()
+
+
+@pytest.mark.asyncio
+async def test_execute_prompt_streaming_uses_provider():
+    """execute_prompt_streaming passes provider to create_llm."""
+    from yamlgraph.executor_async import execute_prompt_streaming
+
+    mock_chunk = MagicMock()
+    mock_chunk.content = "test"
+
+    async def mock_astream(*args, **kwargs):
+        yield mock_chunk
+
+    mock_llm = MagicMock()
+    mock_llm.astream = mock_astream
+
+    with (
+        patch("yamlgraph.executor_async.create_llm", return_value=mock_llm) as mock_create,
+        patch("yamlgraph.executor_async.load_prompt") as mock_load,
+    ):
+        mock_load.return_value = {"system": "", "user": "test"}
+
+        async for _ in execute_prompt_streaming("test", variables={}, provider="openai"):
+            pass
+
+        mock_create.assert_called_once()
+        call_kwargs = mock_create.call_args.kwargs
+        assert call_kwargs.get("provider") == "openai"
+
+
+@pytest.mark.asyncio
+async def test_execute_prompt_streaming_handles_empty_chunks():
+    """execute_prompt_streaming skips empty chunks."""
+    from yamlgraph.executor_async import execute_prompt_streaming
+
+    mock_chunk1 = MagicMock()
+    mock_chunk1.content = "Hello"
+    mock_chunk2 = MagicMock()
+    mock_chunk2.content = ""  # Empty
+    mock_chunk3 = MagicMock()
+    mock_chunk3.content = None  # None
+    mock_chunk4 = MagicMock()
+    mock_chunk4.content = "World"
+
+    async def mock_astream(*args, **kwargs):
+        for chunk in [mock_chunk1, mock_chunk2, mock_chunk3, mock_chunk4]:
+            yield chunk
+
+    mock_llm = MagicMock()
+    mock_llm.astream = mock_astream
+
+    with (
+        patch("yamlgraph.executor_async.create_llm", return_value=mock_llm),
+        patch("yamlgraph.executor_async.load_prompt") as mock_load,
+    ):
+        mock_load.return_value = {"system": "", "user": "test"}
+
+        tokens = []
+        async for token in execute_prompt_streaming("test", variables={}):
+            tokens.append(token)
+
+        # Only non-empty tokens
+        assert tokens == ["Hello", "World"]
+
+
+@pytest.mark.asyncio
+async def test_execute_prompt_streaming_propagates_errors():
+    """execute_prompt_streaming propagates LLM errors."""
+    from yamlgraph.executor_async import execute_prompt_streaming
+
+    async def mock_astream(*args, **kwargs):
+        yield MagicMock(content="start")
+        raise ValueError("LLM error")
+
+    mock_llm = MagicMock()
+    mock_llm.astream = mock_astream
+
+    with (
+        patch("yamlgraph.executor_async.create_llm", return_value=mock_llm),
+        patch("yamlgraph.executor_async.load_prompt") as mock_load,
+    ):
+        mock_load.return_value = {"system": "", "user": "test"}
+
+        tokens = []
+        with pytest.raises(ValueError, match="LLM error"):
+            async for token in execute_prompt_streaming("test", variables={}):
+                tokens.append(token)
+
+        # Should have received first token before error
+        assert tokens == ["start"]
+
+
+# ==============================================================================
+# Streaming with output collection
+# ==============================================================================
+
+
+@pytest.mark.asyncio
+async def test_execute_prompt_streaming_collect():
+    """execute_prompt_streaming can collect all tokens into string."""
+    from yamlgraph.executor_async import execute_prompt_streaming
+
+    mock_chunks = [MagicMock(content=c) for c in ["The ", "quick ", "brown ", "fox"]]
+
+    async def mock_astream(*args, **kwargs):
+        for chunk in mock_chunks:
+            yield chunk
+
+    mock_llm = MagicMock()
+    mock_llm.astream = mock_astream
+
+    with (
+        patch("yamlgraph.executor_async.create_llm", return_value=mock_llm),
+        patch("yamlgraph.executor_async.load_prompt") as mock_load,
+    ):
+        mock_load.return_value = {"system": "", "user": "test"}
+
+        # Collect all tokens
+        result = "".join([token async for token in execute_prompt_streaming("test", {})])
+
+        assert result == "The quick brown fox"
+
+
+# ==============================================================================
+# Streaming node factory tests
+# ==============================================================================
+
+
+@pytest.mark.asyncio
+async def test_create_streaming_node_yields_tokens():
+    """create_streaming_node creates node that yields tokens."""
+    from yamlgraph.node_factory import create_streaming_node
+
+    mock_chunks = [MagicMock(content=c) for c in ["Hello", " ", "World"]]
+
+    async def mock_streaming(*args, **kwargs):
+        for chunk in mock_chunks:
+            yield chunk.content
+
+    with patch("yamlgraph.executor_async.execute_prompt_streaming", mock_streaming):
+        node_config = {
+            "prompt": "greet",
+            "state_key": "response",
+        }
+        streaming_node = create_streaming_node("generate", node_config)
+
+        state = {"input": "test"}
+        tokens = []
+        async for token in streaming_node(state):
+            tokens.append(token)
+
+        assert tokens == ["Hello", " ", "World"]
+
+
+@pytest.mark.asyncio
+async def test_streaming_node_with_callback():
+    """Streaming node can use callback for each token."""
+    from yamlgraph.node_factory import create_streaming_node
+
+    async def mock_streaming(*args, **kwargs):
+        for token in ["A", "B", "C"]:
+            yield token
+
+    collected = []
+
+    def token_callback(token: str):
+        collected.append(token)
+
+    with patch("yamlgraph.executor_async.execute_prompt_streaming", mock_streaming):
+        node_config = {
+            "prompt": "test",
+            "state_key": "output",
+            "on_token": token_callback,
+        }
+        streaming_node = create_streaming_node("stream_node", node_config)
+
+        # Consume the generator
+        async for _ in streaming_node({}):
+            pass
+
+        assert collected == ["A", "B", "C"]
+
+
+# ==============================================================================
+# YAML config tests
+# ==============================================================================
+
+
+def test_node_config_stream_true_creates_streaming_node():
+    """Node with stream: true creates streaming node."""
+    from yamlgraph.node_factory import create_node_function
+
+    node_config = {
+        "prompt": "greet",
+        "state_key": "response",
+        "stream": True,
+    }
+
+    with patch("yamlgraph.node_factory.create_streaming_node") as mock_create:
+        mock_create.return_value = MagicMock()
+
+        # This should detect stream: true and use create_streaming_node
+        _result = create_node_function("generate", node_config, defaults={})
+
+        mock_create.assert_called_once_with("generate", node_config)
+
+
+def test_node_config_stream_false_creates_regular_node():
+    """Node with stream: false creates regular node."""
+    from yamlgraph.node_factory import create_node_function
+
+    node_config = {
+        "prompt": "greet",
+        "state_key": "response",
+        "stream": False,
+    }
+
+    with (
+        patch("yamlgraph.node_factory.create_streaming_node") as mock_streaming,
+        patch("yamlgraph.node_factory.execute_prompt") as mock_execute,
+    ):
+        mock_execute.return_value = "result"
+
+        _result = create_node_function("generate", node_config, defaults={})
+
+        # Should NOT call create_streaming_node
+        mock_streaming.assert_not_called()