yamlgraph 0.3.9__py3-none-any.whl

examples/__init__.py

@@ -0,0 +1 @@
+"""Examples package - sample projects demonstrating framework features."""

examples/codegen/__init__.py

@@ -0,0 +1,5 @@
+"""Codegen Example - Implementation Agent.
+
+A standalone example demonstrating YAMLGraph for code analysis
+and implementation planning.
+"""

examples/codegen/models/__init__.py

@@ -0,0 +1,13 @@
+"""Codegen models for structured LLM outputs."""
+
+from examples.codegen.models.schemas import (
+    DiscoveryFindings,
+    DiscoveryPlan,
+    DiscoveryTask,
+)
+
+__all__ = [
+    "DiscoveryTask",
+    "DiscoveryPlan",
+    "DiscoveryFindings",
+]

examples/codegen/models/schemas.py

@@ -0,0 +1,76 @@
+"""Schemas for impl-agent discovery and planning.
+
+These models define the structure for todo-driven discovery workflow.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field
+
+
+class DiscoveryTask(BaseModel):
+    """A single discovery task for impl-agent.
+
+    Represents one investigation step in the discovery plan.
+    """
+
+    id: int = Field(description="Unique task identifier")
+    task: str = Field(description="Human-readable description of what to discover")
+    tool: str = Field(description="Name of the tool to call")
+    args: dict[str, Any] = Field(description="Arguments to pass to the tool")
+    rationale: str = Field(description="Why this discovery task matters")
+    status: Literal["pending", "done", "failed"] = Field(
+        default="pending",
+        description="Current status of the task",
+    )
+    priority: int = Field(
+        default=1,
+        description="Priority (1=high, 2=medium, 3=low)",
+        ge=1,
+        le=3,
+    )
+
+
+class DiscoveryResult(BaseModel):
+    """Result of executing a discovery task.
+
+    Captures both successful results and error information.
+    """
+
+    task_id: int = Field(description="ID of the task that was executed")
+    tool: str = Field(description="Name of the tool that was called")
+    success: bool = Field(description="Whether the tool call succeeded")
+    result: Any = Field(
+        default=None,
+        description="Result from the tool (if successful)",
+    )
+    error: str | None = Field(
+        default=None,
+        description="Error message (if failed)",
+    )
+
+
+class DiscoveryPlan(BaseModel):
+    """Complete discovery plan with ordered tasks.
+
+    Generated by the plan_discovery LLM node.
+    """
+
+    tasks: list[DiscoveryTask] = Field(
+        default_factory=list,
+        description="Ordered list of discovery tasks to execute",
+    )
+
+
+class DiscoveryFindings(BaseModel):
+    """Accumulated discovery results.
+
+    Collected by the execute_discovery map node.
+    """
+
+    results: list[DiscoveryResult] = Field(
+        default_factory=list,
+        description="Results from all executed discovery tasks",
+    )

examples/codegen/tests/__init__.py

@@ -0,0 +1 @@
+"""Tests for codegen example."""

examples/codegen/tests/test_ai_helpers.py

@@ -0,0 +1,235 @@
+"""Tests for AI helper tools (context compression, diff preview, similarity)."""
+
+from examples.codegen.tools.ai_helpers import (
+    diff_preview,
+    find_similar_code,
+    summarize_module,
+)
+
+
+class TestSummarizeModule:
+    """Tests for summarize_module function."""
+
+    def test_summarizes_real_module(self):
+        """Summarizes a real module with classes and functions."""
+        result = summarize_module("yamlgraph/executor.py")
+
+        assert "error" not in result
+        assert "summary" in result
+        assert isinstance(result["summary"], str)
+        # Should be significantly shorter than original file
+        assert len(result["summary"]) < 2000
+
+    def test_includes_class_names(self):
+        """Summary includes class names."""
+        result = summarize_module("yamlgraph/models/schemas.py")
+
+        assert "error" not in result
+        # Should mention key classes
+        assert "class" in result["summary"].lower() or "Class" in result["summary"]
+
+    def test_includes_function_signatures(self):
+        """Summary includes function signatures (not bodies)."""
+        result = summarize_module("yamlgraph/executor.py")
+
+        assert "error" not in result
+        # Should have function defs
+        assert "def " in result["summary"]
+        # Bodies are stripped - should be much shorter than reading whole file
+
+    def test_includes_docstring(self):
+        """Summary includes module docstring."""
+        result = summarize_module("yamlgraph/executor.py")
+
+        assert "error" not in result
+        # Module docstring should be preserved
+        assert len(result["summary"]) > 50  # Not empty
+
+    def test_returns_error_for_invalid_file(self):
+        """Returns error for non-existent file."""
+        result = summarize_module("nonexistent/file.py")
+
+        assert "error" in result
+
+    def test_max_length_parameter(self):
+        """Respects max_length parameter."""
+        result = summarize_module("yamlgraph/executor.py", max_length=500)
+
+        assert "error" not in result
+        assert len(result["summary"]) <= 550  # Some tolerance
+
+    def test_returns_line_count(self):
+        """Returns original line count for context."""
+        result = summarize_module("yamlgraph/executor.py")
+
+        assert "error" not in result
+        assert "original_lines" in result
+        assert result["original_lines"] > 50
+
+
+class TestDiffPreview:
+    """Tests for diff_preview function."""
+
+    def test_shows_add_diff(self):
+        """Shows diff for adding a line."""
+        result = diff_preview(
+            file_path="yamlgraph/__init__.py",
+            line=1,
+            action="ADD",
+            new_code="# New comment",
+        )
+
+        assert "error" not in result
+        assert "diff" in result
+        assert "+" in result["diff"]  # Added line marker
+
+    def test_shows_modify_diff(self):
+        """Shows diff for modifying a line."""
+        # Read first line to modify it
+        result = diff_preview(
+            file_path="yamlgraph/__init__.py",
+            line=1,
+            action="MODIFY",
+            new_code='"""Modified docstring."""',
+        )
+
+        assert "error" not in result
+        assert "diff" in result
+        # Should show both old (-) and new (+)
+
+    def test_shows_delete_diff(self):
+        """Shows diff for deleting a line."""
+        result = diff_preview(
+            file_path="yamlgraph/__init__.py",
+            line=1,
+            action="DELETE",
+            new_code="",
+        )
+
+        assert "error" not in result
+        assert "diff" in result
+        assert "-" in result["diff"]  # Deleted line marker
+
+    def test_returns_error_for_invalid_file(self):
+        """Returns error for non-existent file."""
+        result = diff_preview(
+            file_path="nonexistent/file.py",
+            line=1,
+            action="ADD",
+            new_code="test",
+        )
+
+        assert "error" in result
+
+    def test_returns_error_for_invalid_line(self):
+        """Returns error for line beyond file length."""
+        result = diff_preview(
+            file_path="yamlgraph/__init__.py",
+            line=99999,
+            action="MODIFY",
+            new_code="test",
+        )
+
+        assert "error" in result
+
+    def test_validates_syntax_of_result(self):
+        """Optionally validates syntax of resulting code."""
+        result = diff_preview(
+            file_path="yamlgraph/__init__.py",
+            line=1,
+            action="MODIFY",
+            new_code="def broken(:",  # Invalid syntax - colon after open paren
+            validate_syntax=True,
+        )
+
+        # Should still return diff but flag syntax issue
+        assert "syntax_valid" in result
+        assert result["syntax_valid"] is False
+
+
+class TestFindSimilarCode:
+    """Tests for find_similar_code function."""
+
+    def test_finds_similar_functions(self):
+        """Finds functions with similar structure."""
+        # Look for functions similar to git_blame (simple tool pattern)
+        result = find_similar_code(
+            file_path="examples/codegen/tools/git_tools.py",
+            symbol_name="git_blame",
+            project_path="examples/codegen/tools",
+        )
+
+        assert "error" not in result
+        assert "similar" in result
+        assert isinstance(result["similar"], list)
+
+    def test_returns_file_and_line(self):
+        """Each result includes file and line."""
+        result = find_similar_code(
+            file_path="examples/codegen/tools/git_tools.py",
+            symbol_name="git_blame",
+            project_path="examples/codegen/tools",
+        )
+
+        assert "error" not in result
+        for item in result["similar"]:
+            assert "file" in item
+            assert "name" in item
+            assert "line" in item
+
+    def test_includes_similarity_reason(self):
+        """Results explain why they're similar."""
+        result = find_similar_code(
+            file_path="examples/codegen/tools/git_tools.py",
+            symbol_name="git_blame",
+            project_path="examples/codegen/tools",
+        )
+
+        assert "error" not in result
+        for item in result["similar"]:
+            assert "reason" in item
+
+    def test_returns_error_for_invalid_file(self):
+        """Returns error for non-existent file."""
+        result = find_similar_code(
+            file_path="nonexistent/file.py",
+            symbol_name="foo",
+            project_path="yamlgraph",
+        )
+
+        assert "error" in result
+
+    def test_returns_error_for_invalid_symbol(self):
+        """Returns error for non-existent symbol."""
+        result = find_similar_code(
+            file_path="yamlgraph/executor.py",
+            symbol_name="nonexistent_function_xyz",
+            project_path="yamlgraph",
+        )
+
+        assert "error" in result
+
+    def test_max_results_parameter(self):
+        """Respects max_results parameter."""
+        result = find_similar_code(
+            file_path="examples/codegen/tools/git_tools.py",
+            symbol_name="git_blame",
+            project_path="examples/codegen/tools",
+            max_results=2,
+        )
+
+        assert "error" not in result
+        assert len(result["similar"]) <= 2
+
+    def test_includes_code_snippet(self):
+        """Results include code snippet."""
+        result = find_similar_code(
+            file_path="examples/codegen/tools/git_tools.py",
+            symbol_name="git_blame",
+            project_path="examples/codegen/tools",
+        )
+
+        assert "error" not in result
+        for item in result["similar"]:
+            assert "snippet" in item
+            assert "def " in item["snippet"]

examples/codegen/tests/test_ast_analysis.py

@@ -0,0 +1,174 @@
+"""Tests for AST-based code analysis tools."""
+
+import tempfile
+from pathlib import Path
+
+from examples.codegen.tools.ast_analysis import get_module_structure
+
+
+class TestGetModuleStructure:
+    """Tests for get_module_structure function."""
+
+    def test_extracts_classes_with_line_numbers(self):
+        """Classes include name, bases, methods, line numbers."""
+        # Use a known file in the project
+        result = get_module_structure("yamlgraph/models/schemas.py")
+
+        assert "error" not in result
+        assert "classes" in result
+        assert len(result["classes"]) > 0
+
+        # Each class should have required fields
+        for cls in result["classes"]:
+            assert "name" in cls
+            assert "line" in cls
+            assert "end_line" in cls
+            assert isinstance(cls["line"], int)
+            assert isinstance(cls["end_line"], int)
+            assert cls["end_line"] >= cls["line"]
+
+    def test_extracts_functions_with_signature(self):
+        """Functions include name, args, returns, decorators."""
+        result = get_module_structure("yamlgraph/executor.py")
+
+        assert "error" not in result
+        assert "functions" in result
+        assert len(result["functions"]) > 0
+
+        # Each function should have required fields
+        for func in result["functions"]:
+            assert "name" in func
+            assert "args" in func
+            assert "line" in func
+            assert "end_line" in func
+            assert isinstance(func["args"], list)
+
+    def test_extracts_imports(self):
+        """Import statements are extracted."""
+        result = get_module_structure("yamlgraph/executor.py")
+
+        assert "error" not in result
+        assert "imports" in result
+        assert len(result["imports"]) > 0
+
+    def test_extracts_module_docstring(self):
+        """Module-level docstring is extracted."""
+        result = get_module_structure("yamlgraph/executor.py")
+
+        assert "error" not in result
+        assert "docstring" in result
+        # executor.py should have a docstring
+        assert result["docstring"] is not None
+
+    def test_handles_missing_file(self):
+        """Returns error for non-existent file."""
+        result = get_module_structure("nonexistent_file_12345.py")
+
+        assert "error" in result
+        assert "not found" in result["error"].lower()
+
+    def test_handles_syntax_error(self):
+        """Returns error for file with syntax errors."""
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
+            f.write("def broken(\n")  # Invalid syntax
+            temp_path = f.name
+
+        try:
+            result = get_module_structure(temp_path)
+            assert "error" in result
+            assert "syntax" in result["error"].lower()
+        finally:
+            Path(temp_path).unlink()
+
+    def test_includes_file_path_in_result(self):
+        """Result includes the file path."""
+        result = get_module_structure("yamlgraph/executor.py")
+
+        assert "error" not in result
+        assert "file" in result
+        assert "executor.py" in result["file"]
+
+    def test_extracts_class_methods(self):
+        """Class methods are listed."""
+        result = get_module_structure("yamlgraph/models/schemas.py")
+
+        assert "error" not in result
+        # At least some classes should have methods
+        # Note: Pydantic models may not have explicit methods
+        assert "methods" in result["classes"][0]
+
+    def test_extracts_class_bases(self):
+        """Class base classes are extracted."""
+        result = get_module_structure("yamlgraph/models/schemas.py")
+
+        assert "error" not in result
+        # Pydantic models inherit from BaseModel
+        pydantic_classes = [
+            c for c in result["classes"] if "BaseModel" in c.get("bases", [])
+        ]
+        assert len(pydantic_classes) > 0
+
+    def test_extracts_function_decorators(self):
+        """Function decorators are extracted."""
+        # Create a temp file with decorators
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
+            f.write("""
+def simple_decorator(func):
+    return func
+
+@simple_decorator
+def decorated_function():
+    pass
+""")
+            temp_path = f.name
+
+        try:
+            result = get_module_structure(temp_path)
+            assert "error" not in result
+            decorated = [
+                f for f in result["functions"] if f["name"] == "decorated_function"
+            ]
+            assert len(decorated) == 1
+            assert len(decorated[0]["decorators"]) > 0
+        finally:
+            Path(temp_path).unlink()
+
+    def test_extracts_function_return_type(self):
+        """Function return type annotations are extracted."""
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
+            f.write("""
+def typed_function(x: int, y: str) -> bool:
+    return True
+""")
+            temp_path = f.name
+
+        try:
+            result = get_module_structure(temp_path)
+            assert "error" not in result
+            typed = [f for f in result["functions"] if f["name"] == "typed_function"]
+            assert len(typed) == 1
+            assert typed[0]["returns"] == "bool"
+            assert typed[0]["args"] == ["x", "y"]
+        finally:
+            Path(temp_path).unlink()
+
+    def test_extracts_function_docstring(self):
+        """Function docstrings are extracted."""
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
+            f.write('''
+def documented_function():
+    """This is the docstring."""
+    pass
+''')
+            temp_path = f.name
+
+        try:
+            result = get_module_structure(temp_path)
+            assert "error" not in result
+            documented = [
+                f for f in result["functions"] if f["name"] == "documented_function"
+            ]
+            assert len(documented) == 1
+            assert documented[0]["docstring"] == "This is the docstring."
+        finally:
+            Path(temp_path).unlink()

examples/codegen/tests/test_code_analysis.py

@@ -0,0 +1,134 @@
+"""Tests for code-analysis graph."""
+
+from pathlib import Path
+
+
+class TestCodeAnalysisGraphStructure:
+    """Tests for code-analysis.yaml graph structure."""
+
+    def test_graph_file_exists(self):
+        """Graph file should exist."""
+        graph_path = Path("graphs/code-analysis.yaml")
+        assert graph_path.exists(), "graphs/code-analysis.yaml not found"
+
+    def test_graph_loads_successfully(self):
+        """Graph should load without errors."""
+        from yamlgraph.graph_loader import load_and_compile
+
+        graph = load_and_compile("graphs/code-analysis.yaml")
+        assert graph is not None
+
+    def test_graph_has_required_nodes(self):
+        """Graph should have run_analysis and generate_recommendations nodes."""
+        from yamlgraph.graph_loader import load_graph_config
+
+        config = load_graph_config("graphs/code-analysis.yaml")
+        assert "run_analysis" in config.nodes
+        assert "generate_recommendations" in config.nodes
+
+    def test_graph_has_analysis_tools(self):
+        """Graph should define analysis tools."""
+        from yamlgraph.graph_loader import load_graph_config
+
+        config = load_graph_config("graphs/code-analysis.yaml")
+        tools = config.tools
+
+        # Should have at least these tools
+        expected_tools = ["run_ruff", "run_tests", "run_bandit"]
+        for tool in expected_tools:
+            assert tool in tools, f"Missing tool: {tool}"
+
+
+class TestCodeAnalysisPrompts:
+    """Tests for code-analysis prompts."""
+
+    def test_analyzer_prompt_exists(self):
+        """Analyzer prompt should exist."""
+        prompt_path = Path("prompts/code-analysis/analyzer.yaml")
+        assert prompt_path.exists(), "prompts/code-analysis/analyzer.yaml not found"
+
+    def test_recommend_prompt_exists(self):
+        """Recommend prompt should exist."""
+        prompt_path = Path("prompts/code-analysis/recommend.yaml")
+        assert prompt_path.exists(), "prompts/code-analysis/recommend.yaml not found"
+
+    def test_analyzer_prompt_has_system_and_user(self):
+        """Analyzer prompt should have system and user sections."""
+        import yaml
+
+        with open("prompts/code-analysis/analyzer.yaml") as f:
+            prompt = yaml.safe_load(f)
+
+        assert "system" in prompt, "Missing system prompt"
+        assert "user" in prompt, "Missing user prompt"
+
+    def test_recommend_prompt_has_system_and_user(self):
+        """Recommend prompt should have system and user sections."""
+        import yaml
+
+        with open("prompts/code-analysis/recommend.yaml") as f:
+            prompt = yaml.safe_load(f)
+
+        assert "system" in prompt, "Missing system prompt"
+        assert "user" in prompt, "Missing user prompt"
+
+
+class TestCodeAnalysisTools:
+    """Tests for shell tool commands."""
+
+    def test_ruff_tool_command_valid(self):
+        """Ruff tool should have valid command structure."""
+        from yamlgraph.graph_loader import load_graph_config
+
+        config = load_graph_config("graphs/code-analysis.yaml")
+        ruff_tool = config.tools.get("run_ruff", {})
+
+        assert "command" in ruff_tool
+        assert "ruff" in ruff_tool["command"]
+
+    def test_tests_tool_command_valid(self):
+        """Tests tool should have valid command structure."""
+        from yamlgraph.graph_loader import load_graph_config
+
+        config = load_graph_config("graphs/code-analysis.yaml")
+        tests_tool = config.tools.get("run_tests", {})
+
+        assert "command" in tests_tool
+        assert "pytest" in tests_tool["command"]
+
+    def test_bandit_tool_command_valid(self):
+        """Bandit tool should have valid command structure."""
+        from yamlgraph.graph_loader import load_graph_config
+
+        config = load_graph_config("graphs/code-analysis.yaml")
+        bandit_tool = config.tools.get("run_bandit", {})
+
+        assert "command" in bandit_tool
+        assert "bandit" in bandit_tool["command"]
+
+
+class TestCodeAnalysisCompilation:
+    """Tests for graph compilation."""
+
+    def test_graph_compiles_with_checkpointer(self):
+        """Graph should compile with SQLite checkpointer."""
+        from langgraph.checkpoint.sqlite import SqliteSaver
+
+        from yamlgraph.graph_loader import load_and_compile
+
+        graph = load_and_compile("graphs/code-analysis.yaml")
+
+        with SqliteSaver.from_conn_string(":memory:") as checkpointer:
+            compiled = graph.compile(checkpointer=checkpointer)
+            assert compiled is not None
+
+    def test_graph_has_entry_point(self):
+        """Graph should have START -> run_analysis edge."""
+        from yamlgraph.graph_loader import load_graph_config
+
+        config = load_graph_config("graphs/code-analysis.yaml")
+
+        # Find edge from START
+        start_edges = [e for e in config.edges if e.get("from") == "START"]
+        assert len(start_edges) > 0, "No edge from START"
+        assert start_edges[0]["to"] == "run_analysis"