yamlgraph 0.3.9__py3-none-any.whl

tests/unit/test_websearch.py

@@ -0,0 +1,234 @@
+"""Tests for yamlgraph.tools.websearch module."""
+
+from unittest.mock import MagicMock, patch
+
+
+class TestWebSearchToolConfig:
+    """Tests for WebSearchToolConfig dataclass."""
+
+    def test_config_defaults(self):
+        """Config should have sensible defaults."""
+        from yamlgraph.tools.websearch import WebSearchToolConfig
+
+        config = WebSearchToolConfig()
+        assert config.provider == "duckduckgo"
+        assert config.max_results == 5
+        assert config.description == ""
+        assert config.timeout == 30
+
+    def test_config_custom_values(self):
+        """Config should accept custom values."""
+        from yamlgraph.tools.websearch import WebSearchToolConfig
+
+        config = WebSearchToolConfig(
+            provider="duckduckgo",
+            max_results=10,
+            description="Search the web",
+            timeout=60,
+        )
+        assert config.max_results == 10
+        assert config.description == "Search the web"
+        assert config.timeout == 60
+
+
+class TestWebSearchResult:
+    """Tests for WebSearchResult dataclass."""
+
+    def test_result_success(self):
+        """Result should store successful search data."""
+        from yamlgraph.tools.websearch import WebSearchResult
+
+        result = WebSearchResult(
+            success=True,
+            results=[{"title": "Test", "url": "https://test.com", "body": "Content"}],
+            query="test query",
+        )
+        assert result.success is True
+        assert len(result.results) == 1
+        assert result.error is None
+
+    def test_result_failure(self):
+        """Result should store error information."""
+        from yamlgraph.tools.websearch import WebSearchResult
+
+        result = WebSearchResult(
+            success=False,
+            results=[],
+            query="test query",
+            error="Search failed",
+        )
+        assert result.success is False
+        assert result.error == "Search failed"
+
+
+class TestExecuteWebSearch:
+    """Tests for execute_web_search function."""
+
+    @patch("yamlgraph.tools.websearch.DUCKDUCKGO_AVAILABLE", True)
+    @patch("yamlgraph.tools.websearch.DDGS")
+    def test_successful_search(self, mock_ddgs_class):
+        """Should return results from DuckDuckGo."""
+        from yamlgraph.tools.websearch import WebSearchToolConfig, execute_web_search
+
+        # Mock the DDGS context manager
+        mock_ddgs = MagicMock()
+        mock_ddgs.__enter__ = MagicMock(return_value=mock_ddgs)
+        mock_ddgs.__exit__ = MagicMock(return_value=False)
+        mock_ddgs.text.return_value = [
+            {"title": "Result 1", "href": "https://example.com", "body": "Content 1"},
+            {"title": "Result 2", "href": "https://example2.com", "body": "Content 2"},
+        ]
+        mock_ddgs_class.return_value = mock_ddgs
+
+        config = WebSearchToolConfig(max_results=5)
+        result = execute_web_search("python programming", config)
+
+        assert result.success is True
+        assert len(result.results) == 2
+        assert result.results[0]["title"] == "Result 1"
+        mock_ddgs.text.assert_called_once_with("python programming", max_results=5)
+
+    @patch("yamlgraph.tools.websearch.DUCKDUCKGO_AVAILABLE", True)
+    @patch("yamlgraph.tools.websearch.DDGS")
+    def test_empty_results(self, mock_ddgs_class):
+        """Should handle empty search results."""
+        from yamlgraph.tools.websearch import WebSearchToolConfig, execute_web_search
+
+        mock_ddgs = MagicMock()
+        mock_ddgs.__enter__ = MagicMock(return_value=mock_ddgs)
+        mock_ddgs.__exit__ = MagicMock(return_value=False)
+        mock_ddgs.text.return_value = []
+        mock_ddgs_class.return_value = mock_ddgs
+
+        config = WebSearchToolConfig()
+        result = execute_web_search("obscure query xyz123", config)
+
+        assert result.success is True
+        assert result.results == []
+
+    @patch("yamlgraph.tools.websearch.DUCKDUCKGO_AVAILABLE", True)
+    @patch("yamlgraph.tools.websearch.DDGS")
+    def test_search_error(self, mock_ddgs_class):
+        """Should handle search errors gracefully."""
+        from yamlgraph.tools.websearch import WebSearchToolConfig, execute_web_search
+
+        mock_ddgs = MagicMock()
+        mock_ddgs.__enter__ = MagicMock(return_value=mock_ddgs)
+        mock_ddgs.__exit__ = MagicMock(return_value=False)
+        mock_ddgs.text.side_effect = Exception("Network error")
+        mock_ddgs_class.return_value = mock_ddgs
+
+        config = WebSearchToolConfig()
+        result = execute_web_search("test query", config)
+
+        assert result.success is False
+        assert "Network error" in result.error
+
+    def test_empty_query(self):
+        """Should handle empty query."""
+        from yamlgraph.tools.websearch import WebSearchToolConfig, execute_web_search
+
+        config = WebSearchToolConfig()
+        result = execute_web_search("", config)
+
+        assert result.success is False
+        assert "empty" in result.error.lower()
+
+
+class TestFormatSearchResults:
+    """Tests for format_search_results function."""
+
+    def test_format_results_as_text(self):
+        """Should format results as readable text."""
+        from yamlgraph.tools.websearch import WebSearchResult, format_search_results
+
+        result = WebSearchResult(
+            success=True,
+            query="python",
+            results=[
+                {
+                    "title": "Python.org",
+                    "href": "https://python.org",
+                    "body": "Official site",
+                },
+                {
+                    "title": "Python Wiki",
+                    "href": "https://wiki.python.org",
+                    "body": "Wiki",
+                },
+            ],
+        )
+
+        formatted = format_search_results(result)
+
+        assert "Python.org" in formatted
+        assert "https://python.org" in formatted
+        assert "Official site" in formatted
+
+    def test_format_empty_results(self):
+        """Should handle empty results."""
+        from yamlgraph.tools.websearch import WebSearchResult, format_search_results
+
+        result = WebSearchResult(success=True, query="xyz", results=[])
+
+        formatted = format_search_results(result)
+
+        assert "no results" in formatted.lower()
+
+    def test_format_error_result(self):
+        """Should format error message."""
+        from yamlgraph.tools.websearch import WebSearchResult, format_search_results
+
+        result = WebSearchResult(
+            success=False,
+            query="test",
+            results=[],
+            error="Search failed",
+        )
+
+        formatted = format_search_results(result)
+
+        assert "error" in formatted.lower()
+        assert "Search failed" in formatted
+
+
+class TestCreateWebSearchTool:
+    """Tests for create_web_search_tool function."""
+
+    def test_create_langchain_tool(self):
+        """Should create a LangChain-compatible tool."""
+        from yamlgraph.tools.websearch import (
+            WebSearchToolConfig,
+            create_web_search_tool,
+        )
+
+        config = WebSearchToolConfig(description="Search the web for information")
+        tool = create_web_search_tool("search_web", config)
+
+        assert tool.name == "search_web"
+        assert tool.description == "Search the web for information"
+
+    @patch("yamlgraph.tools.websearch.DUCKDUCKGO_AVAILABLE", True)
+    @patch("yamlgraph.tools.websearch.DDGS")
+    def test_tool_invocation(self, mock_ddgs_class):
+        """Created tool should be invocable."""
+        from yamlgraph.tools.websearch import (
+            WebSearchToolConfig,
+            create_web_search_tool,
+        )
+
+        mock_ddgs = MagicMock()
+        mock_ddgs.__enter__ = MagicMock(return_value=mock_ddgs)
+        mock_ddgs.__exit__ = MagicMock(return_value=False)
+        mock_ddgs.text.return_value = [
+            {"title": "Test", "href": "https://test.com", "body": "Test content"},
+        ]
+        mock_ddgs_class.return_value = mock_ddgs
+
+        config = WebSearchToolConfig()
+        tool = create_web_search_tool("search_web", config)
+
+        result = tool.invoke({"query": "test"})
+
+        assert "Test" in result
+        assert "https://test.com" in result

yamlgraph/__init__.py

@@ -0,0 +1,35 @@
+"""YamlGraph - YAML-first LLM pipeline framework.
+
+Framework for building LLM pipelines with YAML configuration.
+State is generated dynamically from graph config.
+"""
+
+from yamlgraph.builder import build_graph, build_resume_graph, run_pipeline
+from yamlgraph.executor import execute_prompt, get_executor
+from yamlgraph.models import (
+    ErrorType,
+    GenericReport,
+    PipelineError,
+    build_state_class,
+    create_initial_state,
+)
+from yamlgraph.storage import YamlGraphDB
+
+__all__ = [
+    # Builder
+    "build_graph",
+    "build_resume_graph",
+    "run_pipeline",
+    # Executor
+    "execute_prompt",
+    "get_executor",
+    # Framework models
+    "ErrorType",
+    "PipelineError",
+    "GenericReport",
+    # Dynamic state
+    "build_state_class",
+    "create_initial_state",
+    # Storage
+    "YamlGraphDB",
+]

yamlgraph/builder.py

@@ -0,0 +1,110 @@
+"""Graph builders for yamlgraph pipelines.
+
+Provides functions to build pipeline graphs from YAML configuration.
+
+Pipeline Architecture
+=====================
+
+The main pipeline follows this flow:
+
+```mermaid
+graph LR
+    A[generate] -->|content| B{should_continue}
+    B -->|continue| C[analyze]
+    B -->|end| E[END]
+    C -->|analysis| D[summarize]
+    D --> E[END]
+```
+
+State Flow:
+- generate: Creates structured content from topic
+- analyze: Produces analysis from generated content
+- summarize: Combines all outputs into final_summary
+
+Graph Definition:
+- Pipelines are defined in graphs/*.yaml
+- Loaded and compiled via graph_loader module
+"""
+
+from pathlib import Path
+from typing import Any
+
+from langgraph.graph import StateGraph
+
+from yamlgraph.config import DEFAULT_GRAPH
+from yamlgraph.graph_loader import load_and_compile
+from yamlgraph.models import create_initial_state
+
+# Type alias for dynamic state
+GraphState = dict[str, Any]
+
+
+def build_graph(
+    graph_path: Path | str | None = None,
+    checkpointer: Any | None = None,
+) -> StateGraph:
+    """Build a pipeline graph from YAML with optional checkpointer.
+
+    Args:
+        graph_path: Path to YAML graph definition.
+                   Defaults to graphs/yamlgraph.yaml
+        checkpointer: Optional LangGraph checkpointer for state persistence.
+                     Use get_checkpointer() from storage.checkpointer.
+
+    Returns:
+        StateGraph ready for compilation
+    """
+    path = Path(graph_path) if graph_path else DEFAULT_GRAPH
+    graph = load_and_compile(path)
+
+    # Checkpointer is applied at compile time
+    if checkpointer is not None:
+        # Store reference for compile() to use
+        graph._checkpointer = checkpointer
+
+    return graph
+
+
+def build_resume_graph() -> StateGraph:
+    """Build a graph for resuming an interrupted pipeline.
+
+    This is an alias for build_graph(). Resume works automatically
+    because nodes skip execution if their output already exists in state
+    (skip_if_exists behavior).
+
+    To resume:
+    1. Load saved state from database
+    2. Invoke graph with that state
+    3. Nodes with existing outputs are skipped
+
+    Returns:
+        StateGraph for resume (same as main pipeline)
+    """
+    return build_graph()
+
+
+def run_pipeline(
+    topic: str,
+    style: str = "informative",
+    word_count: int = 300,
+    graph_path: Path | str | None = None,
+) -> GraphState:
+    """Run the complete pipeline with given inputs.
+
+    Args:
+        topic: Topic to generate content about
+        style: Writing style
+        word_count: Target word count
+        graph_path: Optional path to graph YAML
+
+    Returns:
+        Final state with all outputs
+    """
+    graph = build_graph(graph_path).compile()
+    initial_state = create_initial_state(
+        topic=topic,
+        style=style,
+        word_count=word_count,
+    )
+
+    return graph.invoke(initial_state)

yamlgraph/cli/__init__.py

@@ -0,0 +1,159 @@
+"""YamlGraph CLI - Command-line interface for yamlgraph.
+
+This package provides the CLI entry point and command implementations.
+
+Usage:
+    yamlgraph graph run graphs/yamlgraph.yaml --var topic="AI" --var style=casual
+    yamlgraph graph run graphs/router-demo.yaml --var message="hello"
+    yamlgraph graph list
+    yamlgraph list-runs
+    yamlgraph resume --thread-id abc123
+    yamlgraph trace --run-id <run-id>
+"""
+
+import argparse
+
+# Import submodules for package access
+from yamlgraph.cli import commands, validators
+from yamlgraph.cli.commands import (
+    cmd_export,
+    cmd_list_runs,
+    cmd_resume,
+    cmd_trace,
+)
+
+__all__ = [
+    # Submodules
+    "commands",
+    "validators",
+    # Entry points
+    "main",
+    "create_parser",
+]
+
+
+def create_parser() -> argparse.ArgumentParser:
+    """Create and configure the CLI argument parser.
+
+    Returns:
+        Configured ArgumentParser for testing and main().
+    """
+    parser = argparse.ArgumentParser(
+        description="YAMLGraph - YAML-first LLM Pipeline Framework",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+
+    subparsers = parser.add_subparsers(dest="command", help="Available commands")
+
+    # List runs command
+    list_parser = subparsers.add_parser("list-runs", help="List recent runs")
+    list_parser.add_argument(
+        "--limit", "-l", type=int, default=10, help="Maximum runs to show"
+    )
+    list_parser.set_defaults(func=cmd_list_runs)
+
+    # Resume command
+    resume_parser = subparsers.add_parser("resume", help="Resume a pipeline")
+    resume_parser.add_argument(
+        "--thread-id", "-i", required=True, help="Thread ID to resume"
+    )
+    resume_parser.set_defaults(func=cmd_resume)
+
+    # Trace command
+    trace_parser = subparsers.add_parser("trace", help="Show execution trace")
+    trace_parser.add_argument(
+        "--run-id", "-r", help="Run ID (uses latest if not provided)"
+    )
+    trace_parser.add_argument(
+        "--verbose", "-v", action="store_true", help="Include timing details"
+    )
+    trace_parser.set_defaults(func=cmd_trace)
+
+    # Export command
+    export_parser = subparsers.add_parser("export", help="Export a run to JSON")
+    export_parser.add_argument(
+        "--thread-id", "-i", required=True, help="Thread ID to export"
+    )
+    export_parser.set_defaults(func=cmd_export)
+
+    # Graph command group (universal runner)
+    from yamlgraph.cli.graph_commands import cmd_graph_dispatch
+
+    graph_parser = subparsers.add_parser(
+        "graph", help="Universal graph runner and utilities"
+    )
+    graph_subparsers = graph_parser.add_subparsers(
+        dest="graph_command", help="Graph commands"
+    )
+
+    # graph run
+    graph_run_parser = graph_subparsers.add_parser("run", help="Run any graph")
+    graph_run_parser.add_argument("graph_path", help="Path to graph YAML file")
+    graph_run_parser.add_argument(
+        "--var",
+        "-v",
+        action="append",
+        default=[],
+        help="Set state variable (key=value), can repeat",
+    )
+    graph_run_parser.add_argument(
+        "--thread", "-t", type=str, default=None, help="Thread ID for persistence"
+    )
+    graph_run_parser.add_argument(
+        "--export", "-e", action="store_true", help="Export results to files"
+    )
+    graph_run_parser.add_argument(
+        "--full", "-f", action="store_true", help="Show full output without truncation"
+    )
+
+    # graph list
+    graph_subparsers.add_parser("list", help="List available graphs")
+
+    # graph info
+    graph_info_parser = graph_subparsers.add_parser(
+        "info", help="Show graph information"
+    )
+    graph_info_parser.add_argument("graph_path", help="Path to graph YAML file")
+
+    # graph validate
+    graph_validate_parser = graph_subparsers.add_parser(
+        "validate", help="Validate graph YAML schema"
+    )
+    graph_validate_parser.add_argument("graph_path", help="Path to graph YAML file")
+
+    # graph lint
+    graph_lint_parser = graph_subparsers.add_parser(
+        "lint", help="Lint graph for issues (missing state, unused tools, etc.)"
+    )
+    graph_lint_parser.add_argument(
+        "graph_path", nargs="+", help="Path(s) to graph YAML file(s)"
+    )
+
+    # graph mermaid
+    graph_mermaid_parser = graph_subparsers.add_parser(
+        "mermaid", help="Generate Mermaid diagram from graph"
+    )
+    graph_mermaid_parser.add_argument("graph_path", help="Path to graph YAML file")
+    graph_mermaid_parser.add_argument(
+        "--output", "-o", type=str, help="Output file (default: stdout)"
+    )
+
+    graph_parser.set_defaults(func=cmd_graph_dispatch)
+
+    return parser
+
+
+def main():
+    """Main CLI entry point."""
+    parser = create_parser()
+    args = parser.parse_args()
+
+    if not args.command:
+        parser.print_help()
+        return
+
+    args.func(args)
+
+
+if __name__ == "__main__":
+    main()

yamlgraph/cli/__main__.py

@@ -0,0 +1,6 @@
+"""Allow running yamlgraph.cli as a module: python -m yamlgraph.cli"""
+
+from yamlgraph.cli import main
+
+if __name__ == "__main__":
+    main()