yamlgraph 0.3.9__py3-none-any.whl

tests/unit/test_conversation_memory.py

@@ -0,0 +1,276 @@
+"""Tests for Phase 6.3: Conversation Memory.
+
+Tests that agent nodes:
+1. Return messages to state for accumulation
+2. Store raw tool results in state
+3. Support multi-turn conversations via thread_id
+"""
+
+from unittest.mock import MagicMock, patch
+
+from langchain_core.messages import (
+    AIMessage,
+    HumanMessage,
+    SystemMessage,
+)
+
+
+class TestAgentReturnsMessages:
+    """Tests for message accumulation in agent state."""
+
+    def test_agent_returns_messages_in_state(self):
+        """Agent node should return messages for state accumulation."""
+        from yamlgraph.tools.agent import create_agent_node
+
+        # Setup mock LLM
+        mock_response = MagicMock()
+        mock_response.content = "Analysis complete"
+        mock_response.tool_calls = []
+
+        mock_llm = MagicMock()
+        mock_llm.bind_tools.return_value = mock_llm
+        mock_llm.invoke.return_value = mock_response
+
+        with patch("yamlgraph.tools.agent.create_llm", return_value=mock_llm):
+            node_fn = create_agent_node(
+                "agent",
+                {"tools": [], "state_key": "result"},
+                {},
+            )
+            result = node_fn({"input": "test"})
+
+        # Should include messages in output
+        assert "messages" in result, "Agent should return messages for accumulation"
+        assert (
+            len(result["messages"]) >= 2
+        ), "Should have at least system + user + AI messages"
+
+    def test_agent_messages_include_all_types(self):
+        """Agent should include system, user, AI, and tool messages."""
+        from yamlgraph.tools.agent import create_agent_node
+        from yamlgraph.tools.shell import ShellToolConfig
+
+        # Mock LLM that calls a tool then responds
+        # Use actual AIMessage for proper type checking
+        tool_response = AIMessage(
+            content="", tool_calls=[{"name": "test_tool", "args": {}, "id": "call_1"}]
+        )
+
+        final_response = AIMessage(content="Done")
+
+        mock_llm = MagicMock()
+        mock_llm.bind_tools.return_value = mock_llm
+        mock_llm.invoke.side_effect = [tool_response, final_response]
+
+        tool_config = ShellToolConfig(
+            command="echo test",
+            description="Test tool",
+        )
+
+        with (
+            patch("yamlgraph.tools.agent.create_llm", return_value=mock_llm),
+            patch("yamlgraph.tools.agent.execute_shell_tool") as mock_exec,
+        ):
+            mock_exec.return_value = MagicMock(success=True, output="tool output")
+
+            node_fn = create_agent_node(
+                "agent",
+                {"tools": ["test_tool"], "state_key": "result"},
+                {"test_tool": tool_config},
+            )
+            result = node_fn({"input": "test"})
+
+        messages = result["messages"]
+        types = [type(m).__name__ for m in messages]
+
+        assert "SystemMessage" in types, "Should include system message"
+        assert "HumanMessage" in types, "Should include human message"
+        assert "AIMessage" in types, "Should include AI message"
+        assert "ToolMessage" in types, "Should include tool message"
+
+
+class TestToolResultsPersistence:
+    """Tests for raw tool result storage."""
+
+    def test_tool_results_stored_in_state(self):
+        """Agent should store raw tool results in state."""
+        from yamlgraph.tools.agent import create_agent_node
+        from yamlgraph.tools.shell import ShellToolConfig
+
+        tool_response = MagicMock()
+        tool_response.content = ""
+        tool_response.tool_calls = [
+            {"name": "git_log", "args": {"count": "5"}, "id": "call_1"}
+        ]
+
+        final_response = MagicMock()
+        final_response.content = "Report ready"
+        final_response.tool_calls = []
+
+        mock_llm = MagicMock()
+        mock_llm.bind_tools.return_value = mock_llm
+        mock_llm.invoke.side_effect = [tool_response, final_response]
+
+        tool_config = ShellToolConfig(
+            command="git log -n {count}",
+            description="Get git log",
+        )
+
+        with (
+            patch("yamlgraph.tools.agent.create_llm", return_value=mock_llm),
+            patch("yamlgraph.tools.agent.execute_shell_tool") as mock_exec,
+        ):
+            mock_exec.return_value = MagicMock(
+                success=True, output="commit abc123\nAuthor: Test"
+            )
+
+            node_fn = create_agent_node(
+                "agent",
+                {
+                    "tools": ["git_log"],
+                    "state_key": "report",
+                    "tool_results_key": "_tool_results",
+                },
+                {"git_log": tool_config},
+            )
+            result = node_fn({"input": "analyze"})
+
+        assert "_tool_results" in result, "Should include tool_results in state"
+        assert len(result["_tool_results"]) == 1, "Should have one tool result"
+
+        tool_result = result["_tool_results"][0]
+        assert tool_result["tool"] == "git_log"
+        assert tool_result["args"] == {"count": "5"}
+        assert "commit abc123" in tool_result["output"]
+        assert tool_result["success"] is True
+
+    def test_tool_results_key_is_optional(self):
+        """Without tool_results_key, raw results are not stored."""
+        from yamlgraph.tools.agent import create_agent_node
+
+        mock_response = MagicMock()
+        mock_response.content = "Done"
+        mock_response.tool_calls = []
+
+        mock_llm = MagicMock()
+        mock_llm.bind_tools.return_value = mock_llm
+        mock_llm.invoke.return_value = mock_response
+
+        with patch("yamlgraph.tools.agent.create_llm", return_value=mock_llm):
+            node_fn = create_agent_node(
+                "agent",
+                {"tools": [], "state_key": "result"},  # No tool_results_key
+                {},
+            )
+            result = node_fn({"input": "test"})
+
+        # Should NOT have _tool_results if not configured
+        assert "_tool_results" not in result
+
+    def test_multiple_tool_calls_all_stored(self):
+        """Multiple tool calls should all be stored."""
+        from yamlgraph.tools.agent import create_agent_node
+        from yamlgraph.tools.shell import ShellToolConfig
+
+        tool_response = MagicMock()
+        tool_response.content = ""
+        tool_response.tool_calls = [
+            {"name": "tool_a", "args": {}, "id": "call_1"},
+            {"name": "tool_b", "args": {}, "id": "call_2"},
+        ]
+
+        final_response = MagicMock()
+        final_response.content = "Done"
+        final_response.tool_calls = []
+
+        mock_llm = MagicMock()
+        mock_llm.bind_tools.return_value = mock_llm
+        mock_llm.invoke.side_effect = [tool_response, final_response]
+
+        tools = {
+            "tool_a": ShellToolConfig(command="echo a", description="A"),
+            "tool_b": ShellToolConfig(command="echo b", description="B"),
+        }
+
+        with (
+            patch("yamlgraph.tools.agent.create_llm", return_value=mock_llm),
+            patch("yamlgraph.tools.agent.execute_shell_tool") as mock_exec,
+        ):
+            mock_exec.return_value = MagicMock(success=True, output="output")
+
+            node_fn = create_agent_node(
+                "agent",
+                {
+                    "tools": ["tool_a", "tool_b"],
+                    "state_key": "result",
+                    "tool_results_key": "_tool_results",
+                },
+                tools,
+            )
+            result = node_fn({"input": "test"})
+
+        assert len(result["_tool_results"]) == 2
+        tool_names = [r["tool"] for r in result["_tool_results"]]
+        assert "tool_a" in tool_names
+        assert "tool_b" in tool_names
+
+
+class TestMultiTurnConversation:
+    """Tests for multi-turn conversation support."""
+
+    def test_existing_messages_preserved(self):
+        """Agent should preserve existing messages from state."""
+        from yamlgraph.tools.agent import create_agent_node
+
+        mock_response = MagicMock()
+        mock_response.content = "Follow-up response"
+        mock_response.tool_calls = []
+
+        mock_llm = MagicMock()
+        mock_llm.bind_tools.return_value = mock_llm
+        mock_llm.invoke.return_value = mock_response
+
+        # Simulate state with existing messages
+        existing_messages = [
+            SystemMessage(content="You are helpful."),
+            HumanMessage(content="First question"),
+            AIMessage(content="First answer"),
+        ]
+
+        with patch("yamlgraph.tools.agent.create_llm", return_value=mock_llm):
+            node_fn = create_agent_node(
+                "agent",
+                {"tools": [], "state_key": "result"},
+                {},
+            )
+            result = node_fn(
+                {
+                    "input": "Follow-up question",
+                    "messages": existing_messages,
+                }
+            )
+
+        messages = result["messages"]
+        # Should include new messages (at minimum human + AI for this turn)
+        # The exact count depends on implementation - key is messages are returned
+        assert len(messages) >= 2, "Should return messages for accumulation"
+
+    def test_agent_state_message_reducer_works(self):
+        """Dynamic state's Annotated[list, add] should accumulate messages."""
+        from operator import add as add_op
+        from typing import get_type_hints
+
+        from yamlgraph.models.state_builder import build_state_class
+
+        State = build_state_class({"nodes": {"agent": {"type": "agent"}}})
+        hints = get_type_hints(State, include_extras=True)
+
+        # Check messages field has reducer annotation
+        messages_hint = hints.get("messages")
+        assert messages_hint is not None, "State should have messages field"
+
+        # The Annotated type should have add as metadata
+        if hasattr(messages_hint, "__metadata__"):
+            assert (
+                add_op in messages_hint.__metadata__
+            ), "messages should use add reducer"

tests/unit/test_database.py

@@ -0,0 +1,145 @@
+"""Tests for yamlgraph.storage.database module."""
+
+from yamlgraph.models import create_initial_state
+
+
+class TestYamlGraphDB:
+    """Tests for YamlGraphDB class."""
+
+    def test_db_initialization(self, temp_db):
+        """Database should initialize successfully."""
+        assert temp_db.db_path.exists()
+
+    def test_save_and_load_state(self, temp_db, sample_state):
+        """State should be saved and loaded correctly."""
+        thread_id = sample_state["thread_id"]
+        temp_db.save_state(thread_id, sample_state, status="completed")
+
+        loaded = temp_db.load_state(thread_id)
+        assert loaded is not None
+        assert loaded["topic"] == sample_state["topic"]
+        assert loaded["thread_id"] == thread_id
+
+    def test_load_nonexistent_state(self, temp_db):
+        """Loading nonexistent state should return None."""
+        result = temp_db.load_state("nonexistent")
+        assert result is None
+
+    def test_update_existing_state(self, temp_db, empty_state):
+        """Updating existing state should work."""
+        thread_id = empty_state["thread_id"]
+
+        # Save initial state
+        temp_db.save_state(thread_id, empty_state, status="running")
+
+        # Update state
+        empty_state["current_step"] = "generate"
+        temp_db.save_state(thread_id, empty_state, status="completed")
+
+        # Load and verify
+        loaded = temp_db.load_state(thread_id)
+        assert loaded["current_step"] == "generate"
+
+    def test_list_runs_empty(self, temp_db):
+        """List runs should return empty list when no runs."""
+        runs = temp_db.list_runs()
+        assert runs == []
+
+    def test_list_runs_with_data(self, temp_db):
+        """List runs should return saved runs."""
+        state1 = create_initial_state(topic="test1", thread_id="thread1")
+        state2 = create_initial_state(topic="test2", thread_id="thread2")
+
+        temp_db.save_state("thread1", state1, status="completed")
+        temp_db.save_state("thread2", state2, status="running")
+
+        runs = temp_db.list_runs()
+        assert len(runs) == 2
+        thread_ids = [r["thread_id"] for r in runs]
+        assert "thread1" in thread_ids
+        assert "thread2" in thread_ids
+
+    def test_list_runs_limit(self, temp_db):
+        """List runs should respect limit parameter."""
+        for i in range(5):
+            state = create_initial_state(topic=f"test{i}", thread_id=f"thread{i}")
+            temp_db.save_state(f"thread{i}", state)
+
+        runs = temp_db.list_runs(limit=3)
+        assert len(runs) == 3
+
+    def test_delete_run(self, temp_db, empty_state):
+        """Delete run should remove the state."""
+        thread_id = empty_state["thread_id"]
+        temp_db.save_state(thread_id, empty_state)
+
+        result = temp_db.delete_run(thread_id)
+        assert result is True
+
+        loaded = temp_db.load_state(thread_id)
+        assert loaded is None
+
+    def test_delete_nonexistent_run(self, temp_db):
+        """Deleting nonexistent run should return False."""
+        result = temp_db.delete_run("nonexistent")
+        assert result is False
+
+    def test_serialize_state_with_pydantic(self, temp_db, sample_state):
+        """State with Pydantic models should serialize correctly."""
+        thread_id = sample_state["thread_id"]
+        temp_db.save_state(thread_id, sample_state)
+
+        loaded = temp_db.load_state(thread_id)
+        # Pydantic models should be dicts after serialization
+        assert isinstance(loaded["generated"], dict)
+        assert loaded["generated"]["title"] == "Test Article"
+
+
+class TestConnectionPool:
+    """Tests for connection pooling."""
+
+    def test_pooled_mode_works(self, tmp_path):
+        """Database should work in pooled mode."""
+        from yamlgraph.storage.database import YamlGraphDB
+
+        db_path = tmp_path / "pooled_test.db"
+        db = YamlGraphDB(db_path=db_path, use_pool=True, pool_size=3)
+
+        try:
+            # Save and load should work
+            db.save_state("test-thread", {"topic": "test"})
+            loaded = db.load_state("test-thread")
+            assert loaded["topic"] == "test"
+        finally:
+            db.close()
+
+    def test_pool_reuses_connections(self, tmp_path):
+        """Pool should reuse connections."""
+        from yamlgraph.storage.database import ConnectionPool
+
+        db_path = tmp_path / "pool_test.db"
+        pool = ConnectionPool(db_path, pool_size=2)
+
+        # Get a connection, use it, return it
+        with pool.get_connection() as conn1:
+            conn1_id = id(conn1)
+
+        # Next connection should be the same one (reused from pool)
+        with pool.get_connection() as conn2:
+            conn2_id = id(conn2)
+
+        assert conn1_id == conn2_id
+        pool.close_all()
+
+    def test_close_method(self, tmp_path):
+        """Close should clean up connections."""
+        from yamlgraph.storage.database import YamlGraphDB
+
+        db_path = tmp_path / "close_test.db"
+        db = YamlGraphDB(db_path=db_path, use_pool=True, pool_size=2)
+
+        # Use the connection to create one
+        db.save_state("test", {"data": "value"})
+
+        # Close should not raise
+        db.close()

tests/unit/test_deprecation.py

@@ -0,0 +1,104 @@
+"""Tests for deprecation module.
+
+TDD tests for DeprecationError and deprecation utilities.
+"""
+
+import pytest
+
+
+class TestDeprecationError:
+    """Tests for DeprecationError exception."""
+
+    def test_deprecation_error_exists(self):
+        """DeprecationError should be importable."""
+        from yamlgraph.cli.deprecation import DeprecationError
+
+        assert issubclass(DeprecationError, Exception)
+
+    def test_deprecation_error_message(self):
+        """DeprecationError should include replacement command."""
+        from yamlgraph.cli.deprecation import DeprecationError
+
+        err = DeprecationError(
+            old_command="route",
+            new_command="graph run graphs/router-demo.yaml --var message=...",
+        )
+
+        assert "route" in str(err)
+        assert "graph run" in str(err)
+        assert "deprecated" in str(err).lower()
+
+    def test_deprecation_error_has_attributes(self):
+        """DeprecationError should expose old and new commands."""
+        from yamlgraph.cli.deprecation import DeprecationError
+
+        err = DeprecationError(
+            old_command="refine",
+            new_command="graph run graphs/reflexion-demo.yaml",
+        )
+
+        assert err.old_command == "refine"
+        assert err.new_command == "graph run graphs/reflexion-demo.yaml"
+
+
+class TestDeprecatedCommand:
+    """Tests for deprecated_command decorator/helper."""
+
+    def test_deprecated_command_exists(self):
+        """deprecated_command should be importable."""
+        from yamlgraph.cli.deprecation import deprecated_command
+
+        assert callable(deprecated_command)
+
+    def test_deprecated_command_raises(self):
+        """deprecated_command should raise DeprecationError."""
+        from yamlgraph.cli.deprecation import DeprecationError, deprecated_command
+
+        with pytest.raises(DeprecationError) as exc_info:
+            deprecated_command(
+                "route",
+                "graph run graphs/router-demo.yaml --var message=...",
+            )
+
+        assert "route" in str(exc_info.value)
+
+    def test_deprecated_command_with_mapping(self):
+        """deprecated_command should format with variable mapping."""
+        from yamlgraph.cli.deprecation import DeprecationError, deprecated_command
+
+        with pytest.raises(DeprecationError) as exc_info:
+            deprecated_command(
+                "refine --topic X",
+                "graph run graphs/reflexion-demo.yaml --var topic=X",
+            )
+
+        assert "topic=X" in str(exc_info.value)
+
+
+class TestCommandMappings:
+    """Tests for deprecated command mappings."""
+
+    def test_get_replacement_for_route(self):
+        """Should return replacement for route command."""
+        from yamlgraph.cli.deprecation import get_replacement_command
+
+        result = get_replacement_command("route", {"message": "hello"})
+        assert "graph run" in result
+        assert "router-demo.yaml" in result
+        assert "message=hello" in result
+
+    def test_get_replacement_for_refine(self):
+        """Should return replacement for refine command."""
+        from yamlgraph.cli.deprecation import get_replacement_command
+
+        result = get_replacement_command("refine", {"topic": "AI"})
+        assert "graph run" in result
+        assert "reflexion-demo.yaml" in result
+        assert "topic=AI" in result
+
+    def test_get_replacement_unknown_command(self):
+        """Unknown command returns None."""
+        from yamlgraph.cli.deprecation import get_replacement_command
+
+        result = get_replacement_command("unknown", {})
+        assert result is None

tests/unit/test_executor.py

@@ -0,0 +1,172 @@
+"""Tests for yamlgraph.executor module."""
+
+import pytest
+
+from yamlgraph.executor import format_prompt, load_prompt
+
+
+class TestLoadPrompt:
+    """Tests for load_prompt function."""
+
+    def test_load_existing_prompt(self):
+        """Should load an existing prompt file."""
+        prompt = load_prompt("generate")
+        assert "system" in prompt
+        assert "user" in prompt
+
+    def test_load_analyze_prompt(self):
+        """Should load analyze prompt."""
+        prompt = load_prompt("analyze")
+        assert "system" in prompt
+        assert "{content}" in prompt["user"]
+
+    def test_load_nonexistent_prompt(self):
+        """Should raise FileNotFoundError for missing prompt."""
+        with pytest.raises(FileNotFoundError):
+            load_prompt("nonexistent_prompt")
+
+
+class TestFormatPrompt:
+    """Tests for format_prompt function."""
+
+    def test_format_single_variable(self):
+        """Should format single variable."""
+        template = "Hello, {name}!"
+        result = format_prompt(template, {"name": "World"})
+        assert result == "Hello, World!"
+
+    def test_format_multiple_variables(self):
+        """Should format multiple variables."""
+        template = "Topic: {topic}, Style: {style}"
+        result = format_prompt(template, {"topic": "AI", "style": "casual"})
+        assert result == "Topic: AI, Style: casual"
+
+    def test_format_empty_variables(self):
+        """Should handle empty variables dict."""
+        template = "No variables here"
+        result = format_prompt(template, {})
+        assert result == "No variables here"
+
+    def test_format_missing_variable_raises(self):
+        """Should raise KeyError for missing variable."""
+        template = "Hello, {name}!"
+        with pytest.raises(KeyError):
+            format_prompt(template, {})
+
+    def test_format_with_numbers(self):
+        """Should handle numeric variables."""
+        template = "Count: {word_count}"
+        result = format_prompt(template, {"word_count": 300})
+        assert result == "Count: 300"
+
+
+class TestPromptExecutorGraphRelative:
+    """Tests for PromptExecutor with graph-relative prompts."""
+
+    def test_execute_with_graph_path_and_prompts_relative(self, tmp_path):
+        """Executor should resolve prompts relative to graph when configured."""
+        from unittest.mock import MagicMock, patch
+
+        from yamlgraph.executor import PromptExecutor
+
+        # Create graph-relative prompt structure
+        graph_dir = tmp_path / "questionnaires" / "audit"
+        prompts_dir = graph_dir / "prompts"
+        prompts_dir.mkdir(parents=True)
+
+        # Create colocated prompt
+        prompt_file = prompts_dir / "opening.yaml"
+        prompt_file.write_text(
+            """
+system: You are an audit assistant.
+user: Generate opening for {questionnaire_name}.
+"""
+        )
+
+        graph_path = graph_dir / "graph.yaml"
+        graph_path.touch()  # Just needs to exist for path resolution
+
+        # Mock LLM to avoid actual API calls
+        mock_llm = MagicMock()
+        mock_llm.invoke.return_value = MagicMock(content="Welcome to the audit.")
+
+        executor = PromptExecutor()
+
+        with patch.object(executor, "_get_llm", return_value=mock_llm):
+            # Should find prompts/opening.yaml relative to graph_path
+            result = executor.execute(
+                prompt_name="prompts/opening",
+                variables={"questionnaire_name": "Financial Audit"},
+                graph_path=graph_path,
+                prompts_relative=True,
+            )
+
+        assert result == "Welcome to the audit."
+        mock_llm.invoke.assert_called_once()
+
+    def test_execute_with_prompts_dir_override(self, tmp_path):
+        """Executor should use explicit prompts_dir when provided."""
+        from unittest.mock import MagicMock, patch
+
+        from yamlgraph.executor import PromptExecutor
+
+        # Create prompts in explicit directory
+        prompts_dir = tmp_path / "my_prompts"
+        prompts_dir.mkdir()
+
+        prompt_file = prompts_dir / "greeting.yaml"
+        prompt_file.write_text(
+            """
+system: You are helpful.
+user: Say hello to {name}.
+"""
+        )
+
+        mock_llm = MagicMock()
+        mock_llm.invoke.return_value = MagicMock(content="Hello!")
+
+        executor = PromptExecutor()
+
+        with patch.object(executor, "_get_llm", return_value=mock_llm):
+            result = executor.execute(
+                prompt_name="greeting",
+                variables={"name": "World"},
+                prompts_dir=prompts_dir,
+            )
+
+        assert result == "Hello!"
+
+    def test_execute_prompt_function_passes_path_params(self, tmp_path):
+        """execute_prompt() should accept and forward path params."""
+        from unittest.mock import MagicMock, patch
+
+        from yamlgraph.executor import execute_prompt
+
+        # Create test prompt
+        prompts_dir = tmp_path / "prompts"
+        prompts_dir.mkdir()
+        (prompts_dir / "test.yaml").write_text(
+            """
+system: Test system.
+user: Test {msg}.
+"""
+        )
+
+        mock_llm = MagicMock()
+        mock_llm.invoke.return_value = MagicMock(content="OK")
+
+        with patch("yamlgraph.executor.get_executor") as mock_get:
+            mock_executor = MagicMock()
+            mock_executor.execute.return_value = "OK"
+            mock_get.return_value = mock_executor
+
+            execute_prompt(
+                prompt_name="test",
+                variables={"msg": "hello"},
+                prompts_dir=prompts_dir,
+            )
+
+            # Verify path params were forwarded
+            mock_executor.execute.assert_called_once()
+            call_kwargs = mock_executor.execute.call_args.kwargs
+            assert call_kwargs["prompts_dir"] == prompts_dir