yamlgraph 0.3.9__py3-none-any.whl

tests/unit/test_agent_nodes.py

@@ -0,0 +1,355 @@
+"""Tests for agent nodes (type: agent).
+
+Agent nodes allow the LLM to autonomously decide which tools to call
+in a loop until it has enough information to respond.
+"""
+
+from unittest.mock import MagicMock, patch
+
+from yamlgraph.tools.agent import (
+    build_langchain_tool,
+    build_python_tool,
+    create_agent_node,
+)
+from yamlgraph.tools.python_tool import PythonToolConfig
+from yamlgraph.tools.shell import ShellToolConfig
+
+
+class TestBuildLangchainTool:
+    """Tests for build_langchain_tool function."""
+
+    def test_creates_tool_with_name(self):
+        """Tool has correct name."""
+        config = ShellToolConfig(
+            command="echo test",
+            description="Test tool",
+        )
+        tool = build_langchain_tool("my_tool", config)
+        assert tool.name == "my_tool"
+
+    def test_creates_tool_with_description(self):
+        """Tool has correct description."""
+        config = ShellToolConfig(
+            command="echo test",
+            description="A helpful test tool",
+        )
+        tool = build_langchain_tool("test", config)
+        assert tool.description == "A helpful test tool"
+
+    def test_tool_executes_command(self):
+        """Tool invocation runs shell command."""
+        config = ShellToolConfig(
+            command="echo {message}",
+            description="Echo a message",
+        )
+        tool = build_langchain_tool("echo", config)
+        result = tool.invoke({"message": "hello"})
+        assert "hello" in result
+
+
+class TestCreateAgentNode:
+    """Tests for create_agent_node function."""
+
+    @patch("yamlgraph.tools.agent.create_llm")
+    def test_agent_completes_without_tools(self, mock_create_llm):
+        """Agent can finish with no tool calls."""
+        # Mock LLM that returns a direct answer (no tool calls)
+        mock_llm = MagicMock()
+        mock_response = MagicMock()
+        mock_response.tool_calls = []
+        mock_response.content = "The answer is 42"
+        mock_llm.bind_tools.return_value = mock_llm
+        mock_llm.invoke.return_value = mock_response
+        mock_create_llm.return_value = mock_llm
+
+        tools = {
+            "search": ShellToolConfig(command="echo search", description="Search"),
+        }
+        node_config = {
+            "prompt": "agent",
+            "tools": ["search"],
+            "max_iterations": 5,
+            "state_key": "result",
+        }
+
+        node_fn = create_agent_node("agent", node_config, tools)
+        result = node_fn({"input": "What is the meaning of life?"})
+
+        assert result["result"] == "The answer is 42"
+        assert result["_agent_iterations"] == 1
+
+    @patch("yamlgraph.tools.agent.create_llm")
+    def test_agent_calls_tool(self, mock_create_llm):
+        """LLM tool call executes shell command."""
+        # Mock LLM that first calls a tool, then returns answer
+        mock_llm = MagicMock()
+
+        # First response: call a tool
+        first_response = MagicMock()
+        first_response.tool_calls = [
+            {"id": "call1", "name": "echo", "args": {"message": "test"}}
+        ]
+        first_response.content = ""
+
+        # Second response: final answer
+        second_response = MagicMock()
+        second_response.tool_calls = []
+        second_response.content = "I echoed: test"
+
+        mock_llm.bind_tools.return_value = mock_llm
+        mock_llm.invoke.side_effect = [first_response, second_response]
+        mock_create_llm.return_value = mock_llm
+
+        tools = {
+            "echo": ShellToolConfig(command="echo {message}", description="Echo"),
+        }
+        node_config = {
+            "prompt": "agent",
+            "tools": ["echo"],
+            "max_iterations": 5,
+            "state_key": "result",
+        }
+
+        node_fn = create_agent_node("agent", node_config, tools)
+        result = node_fn({"input": "Echo something"})
+
+        assert result["result"] == "I echoed: test"
+        assert result["_agent_iterations"] == 2
+
+    @patch("yamlgraph.tools.agent.create_llm")
+    def test_max_iterations_enforced(self, mock_create_llm):
+        """Stops after max_iterations reached."""
+        # Mock LLM that always calls a tool (never finishes)
+        mock_llm = MagicMock()
+        mock_response = MagicMock()
+        mock_response.tool_calls = [
+            {"id": "call1", "name": "search", "args": {"query": "more"}}
+        ]
+        mock_response.content = "Still searching..."
+        mock_llm.bind_tools.return_value = mock_llm
+        mock_llm.invoke.return_value = mock_response
+        mock_create_llm.return_value = mock_llm
+
+        tools = {
+            "search": ShellToolConfig(command="echo searching", description="Search"),
+        }
+        node_config = {
+            "prompt": "agent",
+            "tools": ["search"],
+            "max_iterations": 3,
+            "state_key": "result",
+        }
+
+        node_fn = create_agent_node("agent", node_config, tools)
+        result = node_fn({"input": "Search forever"})
+
+        # Should stop at max_iterations
+        assert result["_agent_limit_reached"] is True
+        assert mock_llm.invoke.call_count == 3
+
+    @patch("yamlgraph.tools.agent.create_llm")
+    def test_tool_result_returned_to_llm(self, mock_create_llm):
+        """LLM sees tool output in next turn."""
+        mock_llm = MagicMock()
+
+        # First: call tool
+        first_response = MagicMock()
+        first_response.tool_calls = [
+            {"id": "call1", "name": "calc", "args": {"expr": "2+2"}}
+        ]
+        first_response.content = ""
+
+        # Second: answer based on tool result
+        second_response = MagicMock()
+        second_response.tool_calls = []
+        second_response.content = "The result is 4"
+
+        mock_llm.bind_tools.return_value = mock_llm
+        mock_llm.invoke.side_effect = [first_response, second_response]
+        mock_create_llm.return_value = mock_llm
+
+        tools = {
+            "calc": ShellToolConfig(
+                command="echo 4",  # Simulates python calc
+                description="Calculate",
+            ),
+        }
+        node_config = {
+            "prompt": "agent",
+            "tools": ["calc"],
+            "max_iterations": 5,
+            "state_key": "answer",
+        }
+
+        node_fn = create_agent_node("agent", node_config, tools)
+        node_fn({"input": "What is 2+2?"})
+
+        # Check that second invoke received messages with tool result
+        second_call_messages = mock_llm.invoke.call_args_list[1][0][0]
+        # Should have: system, user, ai (with tool call), tool result
+        assert len(second_call_messages) >= 4
+
+    def test_default_max_iterations(self):
+        """Default max_iterations is 5."""
+        tools = {
+            "test": ShellToolConfig(command="echo test", description="Test"),
+        }
+        node_config = {
+            "prompt": "agent",
+            "tools": ["test"],
+            # No max_iterations specified
+        }
+
+        # Just verify it doesn't fail - actual behavior tested above
+        node_fn = create_agent_node("agent", node_config, tools)
+        assert callable(node_fn)
+
+
+class TestBuildPythonTool:
+    """Tests for build_python_tool function."""
+
+    def test_creates_tool_with_name(self):
+        """Tool has correct name."""
+        config = PythonToolConfig(
+            module="yamlgraph.utils.langsmith",
+            function="get_run_details",
+            description="Get run details",
+        )
+        tool = build_python_tool("get_run_details", config)
+        assert tool.name == "get_run_details"
+
+    def test_creates_tool_with_description(self):
+        """Tool has correct description."""
+        config = PythonToolConfig(
+            module="yamlgraph.utils.langsmith",
+            function="get_run_details",
+            description="Get details about a LangSmith run",
+        )
+        tool = build_python_tool("run_details", config)
+        assert tool.description == "Get details about a LangSmith run"
+
+    def test_tool_is_structured_tool(self):
+        """Tool is a LangChain StructuredTool."""
+        from langchain_core.tools import StructuredTool
+
+        config = PythonToolConfig(
+            module="yamlgraph.utils.langsmith",
+            function="get_run_details",
+            description="Get run details",
+        )
+        tool = build_python_tool("test_tool", config)
+        assert isinstance(tool, StructuredTool)
+
+    def test_tool_executes_function(self):
+        """Tool invocation calls the Python function."""
+        # Use a simple test function
+        config = PythonToolConfig(
+            module="os.path",
+            function="join",
+            description="Join paths",
+        )
+        tool = build_python_tool("path_join", config)
+        result = tool.invoke({"a": "/home", "p": "user"})
+        assert "/home" in result or "user" in result
+
+
+class TestAgentWithPythonTools:
+    """Tests for agent nodes using Python tools."""
+
+    @patch("yamlgraph.tools.agent.create_llm")
+    def test_agent_calls_python_tool(self, mock_create_llm):
+        """Agent can use Python tools."""
+        mock_llm = MagicMock()
+
+        # First response: call a python tool
+        first_response = MagicMock()
+        first_response.tool_calls = [
+            {
+                "id": "call1",
+                "name": "my_python_tool",
+                "args": {"a": "/home", "p": "user"},
+            }
+        ]
+        first_response.content = ""
+
+        # Second response: final answer
+        second_response = MagicMock()
+        second_response.tool_calls = []
+        second_response.content = "Path is /home/user"
+
+        mock_llm.bind_tools.return_value = mock_llm
+        mock_llm.invoke.side_effect = [first_response, second_response]
+        mock_create_llm.return_value = mock_llm
+
+        python_tools = {
+            "my_python_tool": PythonToolConfig(
+                module="os.path",
+                function="join",
+                description="Join paths",
+            ),
+        }
+        node_config = {
+            "prompt": "agent",
+            "tools": ["my_python_tool"],
+            "max_iterations": 5,
+            "state_key": "result",
+        }
+
+        node_fn = create_agent_node("agent", node_config, {}, python_tools=python_tools)
+        result = node_fn({"input": "Join home and user"})
+
+        assert result["result"] == "Path is /home/user"
+        assert result["_agent_iterations"] == 2
+
+    @patch("yamlgraph.tools.agent.create_llm")
+    def test_agent_mixes_shell_and_python_tools(self, mock_create_llm):
+        """Agent can use both shell and python tools."""
+        mock_llm = MagicMock()
+
+        # First: call shell tool
+        first_response = MagicMock()
+        first_response.tool_calls = [
+            {"id": "call1", "name": "echo_tool", "args": {"message": "hello"}}
+        ]
+        first_response.content = ""
+
+        # Second: call python tool
+        second_response = MagicMock()
+        second_response.tool_calls = [
+            {"id": "call2", "name": "path_tool", "args": {"a": "/", "p": "tmp"}}
+        ]
+        second_response.content = ""
+
+        # Third: final answer
+        third_response = MagicMock()
+        third_response.tool_calls = []
+        third_response.content = "Done with both tools"
+
+        mock_llm.bind_tools.return_value = mock_llm
+        mock_llm.invoke.side_effect = [first_response, second_response, third_response]
+        mock_create_llm.return_value = mock_llm
+
+        shell_tools = {
+            "echo_tool": ShellToolConfig(command="echo {message}", description="Echo"),
+        }
+        python_tools = {
+            "path_tool": PythonToolConfig(
+                module="os.path",
+                function="join",
+                description="Join paths",
+            ),
+        }
+        node_config = {
+            "prompt": "agent",
+            "tools": ["echo_tool", "path_tool"],
+            "max_iterations": 5,
+            "state_key": "result",
+        }
+
+        node_fn = create_agent_node(
+            "agent", node_config, shell_tools, python_tools=python_tools
+        )
+        result = node_fn({"input": "Use both tools"})
+
+        assert result["result"] == "Done with both tools"
+        assert result["_agent_iterations"] == 3

tests/unit/test_async_executor.py

@@ -0,0 +1,346 @@
+"""Tests for async executor - Phase 2 (003).
+
+TDD: RED phase - write tests first.
+"""
+
+import asyncio
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from pydantic import BaseModel
+
+from yamlgraph.executor_async import execute_prompt_async
+
+
+class MockResponse(BaseModel):
+    """Mock response model for testing."""
+
+    summary: str
+    score: int
+
+
+# ==============================================================================
+# execute_prompt_async tests (existing function - verify it works)
+# ==============================================================================
+
+
+@pytest.mark.asyncio
+async def test_execute_prompt_async_returns_string():
+    """execute_prompt_async returns string when no output_model."""
+    mock_llm = MagicMock()
+    mock_response = MagicMock()
+    mock_response.content = "Hello, World!"
+
+    with (
+        patch("yamlgraph.executor_async.create_llm", return_value=mock_llm),
+        patch(
+            "yamlgraph.executor_async.invoke_async", new_callable=AsyncMock
+        ) as mock_invoke,
+        patch("yamlgraph.executor_async.load_prompt") as mock_load,
+    ):
+        mock_load.return_value = {
+            "system": "You are helpful.",
+            "user": "Say hello to {name}",
+        }
+        mock_invoke.return_value = "Hello, World!"
+
+        result = await execute_prompt_async(
+            "greet",
+            variables={"name": "World"},
+        )
+
+        assert result == "Hello, World!"
+        mock_invoke.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_execute_prompt_async_with_output_model():
+    """execute_prompt_async returns parsed model when output_model provided."""
+    mock_llm = MagicMock()
+    expected = MockResponse(summary="Test", score=42)
+
+    with (
+        patch("yamlgraph.executor_async.create_llm", return_value=mock_llm),
+        patch(
+            "yamlgraph.executor_async.invoke_async", new_callable=AsyncMock
+        ) as mock_invoke,
+        patch("yamlgraph.executor_async.load_prompt") as mock_load,
+    ):
+        mock_load.return_value = {
+            "system": "Analyze this.",
+            "user": "Input: {text}",
+        }
+        mock_invoke.return_value = expected
+
+        result = await execute_prompt_async(
+            "analyze",
+            variables={"text": "test input"},
+            output_model=MockResponse,
+        )
+
+        assert isinstance(result, MockResponse)
+        assert result.summary == "Test"
+        assert result.score == 42
+
+
+@pytest.mark.asyncio
+async def test_execute_prompt_async_uses_provider_from_yaml():
+    """execute_prompt_async extracts provider from YAML metadata."""
+    mock_llm = MagicMock()
+
+    with (
+        patch(
+            "yamlgraph.executor_async.create_llm", return_value=mock_llm
+        ) as mock_create,
+        patch(
+            "yamlgraph.executor_async.invoke_async", new_callable=AsyncMock
+        ) as mock_invoke,
+        patch("yamlgraph.executor_async.load_prompt") as mock_load,
+    ):
+        mock_load.return_value = {
+            "system": "Hello",
+            "user": "{input}",
+            "provider": "openai",  # Provider in YAML
+        }
+        mock_invoke.return_value = "response"
+
+        await execute_prompt_async("test", variables={"input": "x"})
+
+        # Should use provider from YAML
+        mock_create.assert_called_once()
+        call_kwargs = mock_create.call_args.kwargs
+        assert call_kwargs.get("provider") == "openai"
+
+
+# ==============================================================================
+# run_graph_async tests (new function)
+# ==============================================================================
+
+
+@pytest.mark.asyncio
+async def test_run_graph_async_executes_graph():
+    """run_graph_async invokes graph asynchronously."""
+    from yamlgraph.executor_async import run_graph_async
+
+    # Mock compiled graph
+    mock_app = AsyncMock()
+    mock_app.ainvoke.return_value = {"output": "result", "current_step": "done"}
+
+    result = await run_graph_async(
+        mock_app,
+        initial_state={"input": "test"},
+        config={"configurable": {"thread_id": "t1"}},
+    )
+
+    assert result["output"] == "result"
+    mock_app.ainvoke.assert_called_once_with(
+        {"input": "test"},
+        {"configurable": {"thread_id": "t1"}},
+    )
+
+
+@pytest.mark.asyncio
+async def test_run_graph_async_with_checkpointer():
+    """run_graph_async works with checkpointer in config."""
+    from yamlgraph.executor_async import run_graph_async
+
+    mock_app = AsyncMock()
+    mock_app.ainvoke.return_value = {"result": "ok"}
+
+    result = await run_graph_async(
+        mock_app,
+        initial_state={"query": "hello"},
+        config={"configurable": {"thread_id": "test-thread"}},
+    )
+
+    assert result["result"] == "ok"
+
+
+@pytest.mark.asyncio
+async def test_run_graph_async_handles_interrupt():
+    """run_graph_async returns interrupt payload when graph pauses."""
+    from yamlgraph.executor_async import run_graph_async
+
+    mock_app = AsyncMock()
+    # Simulate interrupt response
+    interrupt_value = MagicMock()
+    interrupt_value.value = {"question": "What is your name?"}
+    mock_app.ainvoke.return_value = {"__interrupt__": (interrupt_value,)}
+
+    result = await run_graph_async(
+        mock_app,
+        initial_state={},
+        config={"configurable": {"thread_id": "t1"}},
+    )
+
+    assert "__interrupt__" in result
+    assert result["__interrupt__"][0].value == {"question": "What is your name?"}
+
+
+@pytest.mark.asyncio
+async def test_run_graph_async_resume_with_command():
+    """run_graph_async can resume with Command."""
+    from langgraph.types import Command
+
+    from yamlgraph.executor_async import run_graph_async
+
+    mock_app = AsyncMock()
+    mock_app.ainvoke.return_value = {"user_name": "Alice", "greeting": "Hello Alice!"}
+
+    result = await run_graph_async(
+        mock_app,
+        initial_state=Command(resume="Alice"),
+        config={"configurable": {"thread_id": "t1"}},
+    )
+
+    assert result["user_name"] == "Alice"
+    mock_app.ainvoke.assert_called_once()
+
+
+# ==============================================================================
+# compile_graph_async tests (new function)
+# ==============================================================================
+
+
+@pytest.mark.asyncio
+async def test_compile_graph_async_with_memory_checkpointer():
+    """compile_graph_async compiles graph with memory checkpointer."""
+    from yamlgraph.executor_async import compile_graph_async
+
+    # Use minimal test graph
+    with patch("yamlgraph.storage.checkpointer_factory.get_checkpointer") as mock_cp:
+        from langgraph.checkpoint.memory import MemorySaver
+
+        mock_cp.return_value = MemorySaver()
+
+        config = MagicMock()
+        config.checkpointer = {"type": "memory"}
+
+        mock_graph = MagicMock()
+        mock_compiled = MagicMock()
+        mock_graph.compile.return_value = mock_compiled
+
+        result = compile_graph_async(mock_graph, config)
+
+        mock_graph.compile.assert_called_once()
+        assert result == mock_compiled
+
+
+@pytest.mark.asyncio
+async def test_compile_graph_async_sets_async_mode():
+    """compile_graph_async passes async_mode=True to get_checkpointer."""
+    from yamlgraph.executor_async import compile_graph_async
+
+    with patch("yamlgraph.storage.checkpointer_factory.get_checkpointer") as mock_cp:
+        mock_cp.return_value = MagicMock()
+
+        mock_graph = MagicMock()
+        mock_graph.compile.return_value = MagicMock()
+
+        config = MagicMock()
+        config.checkpointer = {"type": "redis", "url": "redis://localhost"}
+
+        compile_graph_async(mock_graph, config)
+
+        mock_cp.assert_called_once_with(config.checkpointer, async_mode=True)
+
+
+# ==============================================================================
+# load_and_compile_async tests (convenience function)
+# ==============================================================================
+
+
+@pytest.mark.asyncio
+async def test_load_and_compile_async_returns_compiled_graph():
+    """load_and_compile_async loads YAML and returns compiled graph."""
+    from yamlgraph.executor_async import load_and_compile_async
+
+    with (
+        patch("yamlgraph.graph_loader.load_graph_config") as mock_load,
+        patch("yamlgraph.graph_loader.compile_graph") as mock_compile,
+        patch("yamlgraph.storage.checkpointer_factory.get_checkpointer") as mock_cp,
+    ):
+        mock_config = MagicMock()
+        mock_config.name = "test-graph"
+        mock_config.version = "1.0"
+        mock_config.checkpointer = {"type": "memory"}
+        mock_load.return_value = mock_config
+
+        mock_state_graph = MagicMock()
+        mock_compiled = MagicMock()
+        mock_state_graph.compile.return_value = mock_compiled
+        mock_compile.return_value = mock_state_graph
+
+        mock_cp.return_value = None
+
+        result = await load_and_compile_async("graphs/test.yaml")
+
+        assert result == mock_compiled
+        mock_load.assert_called_once_with("graphs/test.yaml")
+
+
+# ==============================================================================
+# Concurrent execution tests
+# ==============================================================================
+
+
+@pytest.mark.asyncio
+async def test_run_graphs_concurrent():
+    """Multiple graphs can run concurrently."""
+    from yamlgraph.executor_async import run_graph_async
+
+    mock_app1 = AsyncMock()
+    mock_app1.ainvoke.return_value = {"result": "first"}
+
+    mock_app2 = AsyncMock()
+    mock_app2.ainvoke.return_value = {"result": "second"}
+
+    results = await asyncio.gather(
+        run_graph_async(mock_app1, {}, {"configurable": {"thread_id": "t1"}}),
+        run_graph_async(mock_app2, {}, {"configurable": {"thread_id": "t2"}}),
+    )
+
+    assert results[0]["result"] == "first"
+    assert results[1]["result"] == "second"
+
+
+# ==============================================================================
+# Error handling tests
+# ==============================================================================
+
+
+@pytest.mark.asyncio
+async def test_run_graph_async_propagates_errors():
+    """run_graph_async propagates exceptions from graph execution."""
+    from yamlgraph.executor_async import run_graph_async
+
+    mock_app = AsyncMock()
+    mock_app.ainvoke.side_effect = ValueError("Graph execution failed")
+
+    with pytest.raises(ValueError, match="Graph execution failed"):
+        await run_graph_async(
+            mock_app,
+            initial_state={},
+            config={"configurable": {"thread_id": "t1"}},
+        )
+
+
+@pytest.mark.asyncio
+async def test_compile_graph_async_without_checkpointer():
+    """compile_graph_async works without checkpointer config."""
+    from yamlgraph.executor_async import compile_graph_async
+
+    mock_graph = MagicMock()
+    mock_compiled = MagicMock()
+    mock_graph.compile.return_value = mock_compiled
+
+    config = MagicMock()
+    config.checkpointer = None
+
+    with patch(
+        "yamlgraph.storage.checkpointer_factory.get_checkpointer", return_value=None
+    ):
+        result = compile_graph_async(mock_graph, config)
+
+    mock_graph.compile.assert_called_once_with(checkpointer=None)
+    assert result == mock_compiled