yamlgraph 0.3.9__py3-none-any.whl

yamlgraph/tools/agent.py

@@ -0,0 +1,320 @@
+"""Agent node factory for LLM-driven tool loops.
+
+This module provides the agent node type that allows the LLM to
+autonomously decide which tools to call until it has enough
+information to provide a final answer.
+"""
+
+from __future__ import annotations
+
+import inspect
+import logging
+from collections.abc import Callable
+from typing import Any
+
+from langchain_core.messages import HumanMessage, SystemMessage, ToolMessage
+
+from yamlgraph.tools.python_tool import PythonToolConfig, load_python_function
+from yamlgraph.tools.shell import ShellToolConfig, execute_shell_tool
+from yamlgraph.utils.llm_factory import create_llm
+from yamlgraph.utils.prompts import load_prompt
+
+logger = logging.getLogger(__name__)
+
+
+def build_langchain_tool(name: str, config: ShellToolConfig) -> Callable:
+    """Convert shell config to LangChain Tool.
+
+    Args:
+        name: Tool name for LLM to reference
+        config: Shell tool configuration
+
+    Returns:
+        LangChain-compatible tool function
+    """
+    import re
+
+    from langchain_core.tools import StructuredTool
+    from pydantic import Field, create_model
+
+    # Extract variable names from command template
+    var_names = re.findall(r"\{(\w+)\}", config.command)
+
+    # Create dynamic Pydantic model for tool args
+    if var_names:
+        fields = {
+            var: (str, Field(description=f"Value for {var}")) for var in var_names
+        }
+        ArgsModel = create_model(f"{name}_args", **fields)
+    else:
+        ArgsModel = None
+
+    def execute_tool_with_dict(**kwargs) -> str:
+        """Execute shell command with provided arguments."""
+        result = execute_shell_tool(config, kwargs)
+        if result.success:
+            return (
+                str(result.output).strip() if result.output is not None else "Success"
+            )
+        else:
+            return f"Error: {result.error}"
+
+    return StructuredTool.from_function(
+        func=execute_tool_with_dict,
+        name=name,
+        description=config.description,
+        args_schema=ArgsModel,
+    )
+
+
+def build_python_tool(name: str, config: PythonToolConfig) -> Any:
+    """Convert Python tool config to LangChain StructuredTool.
+
+    Args:
+        name: Tool name for LLM to reference
+        config: Python tool configuration
+
+    Returns:
+        LangChain StructuredTool
+    """
+    from langchain_core.tools import StructuredTool
+    from pydantic import Field, create_model
+
+    # Load the Python function
+    func = load_python_function(config)
+
+    # Build args schema from function signature
+    sig = inspect.signature(func)
+    fields = {}
+    for param_name, param in sig.parameters.items():
+        # Skip *args, **kwargs
+        if param.kind in (
+            inspect.Parameter.VAR_POSITIONAL,
+            inspect.Parameter.VAR_KEYWORD,
+        ):
+            continue
+
+        # Get type annotation or default to str
+        param_type = (
+            param.annotation if param.annotation != inspect.Parameter.empty else str
+        )
+
+        # Create field with description
+        fields[param_name] = (param_type, Field(description=f"Parameter: {param_name}"))
+
+    # Create dynamic Pydantic model
+    ArgsModel = create_model(f"{name}_args", **fields) if fields else None
+
+    def execute_python(**kwargs) -> str:
+        """Execute the Python function and return result as string."""
+        try:
+            result = func(**kwargs)
+            return str(result) if result is not None else "Success"
+        except Exception as e:
+            return f"Error: {e}"
+
+    return StructuredTool.from_function(
+        func=execute_python,
+        name=name,
+        description=config.description,
+        args_schema=ArgsModel,
+    )
+
+
+def create_agent_node(
+    node_name: str,
+    node_config: dict[str, Any],
+    tools: dict[str, ShellToolConfig],
+    websearch_tools: dict[str, Any] | None = None,
+    python_tools: dict[str, PythonToolConfig] | None = None,
+) -> Callable[[dict], dict]:
+    """Create an agent node that loops with tool calls.
+
+    The agent will:
+    1. Send the prompt to the LLM with available tools
+    2. If LLM returns tool calls, execute them and feed results back
+    3. Repeat until LLM returns without tool calls or max_iterations reached
+
+    Args:
+        node_name: Name of the node in the graph
+        node_config: Node configuration from YAML
+        tools: Registry of available shell tools
+        websearch_tools: Registry of web search tools (LangChain StructuredTool)
+        python_tools: Registry of Python tools (PythonToolConfig)
+
+    Returns:
+        Node function that runs the agent loop
+
+    Config options:
+        - tools: List of tool names to make available
+        - max_iterations: Max tool-call loops (default: 5)
+        - state_key: Key to store final answer (default: node_name)
+        - prompt: Prompt file name (default: "agent")
+        - tool_results_key: Optional key to store raw tool outputs
+    """
+    if websearch_tools is None:
+        websearch_tools = {}
+    if python_tools is None:
+        python_tools = {}
+
+    tool_names = node_config.get("tools", [])
+    max_iterations = node_config.get("max_iterations", 5)
+    state_key = node_config.get("state_key", node_name)
+    prompt_name = node_config.get("prompt", "agent")
+    tool_results_key = node_config.get("tool_results_key")
+
+    # Build LangChain tools from configs
+    lc_tools = []
+    tool_lookup = {}
+
+    for name in tool_names:
+        if name in tools:
+            # Shell tool - need to wrap
+            lc_tools.append(build_langchain_tool(name, tools[name]))
+            tool_lookup[name] = tools[name]
+        elif name in websearch_tools:
+            # Websearch tool - already a LangChain tool
+            lc_tools.append(websearch_tools[name])
+            tool_lookup[name] = websearch_tools[name]
+        elif name in python_tools:
+            # Python tool - wrap as LangChain tool
+            lc_tools.append(build_python_tool(name, python_tools[name]))
+            tool_lookup[name] = python_tools[name]
+        else:
+            logger.warning(
+                f"Tool '{name}' not found in shell, websearch, or python registries"
+            )
+
+    def node_fn(state: dict) -> dict:
+        """Execute the agent loop."""
+        # Load prompts - fail fast if missing
+        prompt_config = load_prompt(prompt_name)
+        system_prompt = prompt_config.get("system", "")
+        user_template = prompt_config.get("user", "{input}")
+
+        # Format user prompt with state - handle missing keys
+        import re
+
+        def replace_var(match):
+            key = match.group(1)
+            return str(state.get(key, f"{{{key}}}"))
+
+        user_prompt = re.sub(r"\{(\w+)\}", replace_var, user_template)
+
+        # Initialize messages - preserve existing if multi-turn
+        existing_messages = list(state.get("messages", []))
+        if existing_messages:
+            # Multi-turn: add new user message to existing conversation
+            messages = existing_messages + [HumanMessage(content=user_prompt)]
+        else:
+            # New conversation: start with system + user
+            messages = [
+                SystemMessage(content=system_prompt),
+                HumanMessage(content=user_prompt),
+            ]
+
+        # Track raw tool outputs for persistence
+        tool_results: list[dict] = []
+
+        # Get LLM with tools bound
+        llm = create_llm().bind_tools(lc_tools)
+
+        logger.info(
+            f"🤖 Starting agent loop: {node_name} (max {max_iterations} iterations)"
+        )
+        logger.debug(f"Tools available: {[t.name for t in lc_tools]}")
+        logger.debug(f"User prompt: {user_prompt[:100]}...")
+
+        for iteration in range(max_iterations):
+            logger.debug(f"Agent iteration {iteration + 1}/{max_iterations}")
+
+            # Get LLM response
+            response = llm.invoke(messages)
+            messages.append(response)
+
+            logger.debug(f"Response tool_calls: {response.tool_calls}")
+
+            # Check if LLM wants to call tools
+            if not response.tool_calls:
+                # Done - LLM finished reasoning
+                logger.info(f"✓ Agent completed after {iteration + 1} iterations")
+                result = {
+                    state_key: response.content,
+                    "current_step": node_name,
+                    "_agent_iterations": iteration + 1,
+                    "messages": messages,  # Return for accumulation
+                }
+                if tool_results_key and tool_results:
+                    result[tool_results_key] = tool_results
+                return result
+
+            # Execute tool calls
+            for tool_call in response.tool_calls:
+                tool_name = tool_call["name"]
+                tool_args = tool_call["args"]
+                tool_id = tool_call.get("id", f"call_{iteration}")
+
+                logger.info(f"🔧 Calling tool: {tool_name}({tool_args})")
+
+                # Execute the tool
+                tool_config = tool_lookup.get(tool_name)
+                if tool_config:
+                    # Check the type of tool config
+                    if isinstance(tool_config, ShellToolConfig):
+                        # Shell tool - use execute_shell_tool
+                        result = execute_shell_tool(tool_config, tool_args)
+                        output = (
+                            str(result.output)
+                            if result.success
+                            else f"Error: {result.error}"
+                        )
+                        success = result.success
+                    elif isinstance(tool_config, PythonToolConfig):
+                        # Python tool - load and execute function
+                        try:
+                            func = load_python_function(tool_config)
+                            output = str(func(**tool_args))
+                            success = True
+                        except Exception as e:
+                            output = f"Error: {e}"
+                            success = False
+                    else:
+                        # LangChain tool (websearch, etc) - invoke directly
+                        try:
+                            output = tool_config.invoke(tool_args)
+                            success = True
+                        except Exception as e:
+                            output = f"Error: {e}"
+                            success = False
+                else:
+                    output = f"Error: Unknown tool '{tool_name}'"
+                    success = False
+
+                # Store raw tool result for persistence
+                tool_results.append(
+                    {
+                        "tool": tool_name,
+                        "args": tool_args,
+                        "output": output,
+                        "success": success,
+                    }
+                )
+
+                # Add tool result to messages
+                messages.append(ToolMessage(content=output, tool_call_id=tool_id))
+
+        # Hit max iterations
+        logger.warning(f"Agent hit max iterations ({max_iterations})")
+        last_content = messages[-1].content if hasattr(messages[-1], "content") else ""
+        result = {
+            state_key: last_content,
+            "current_step": node_name,
+            "_agent_iterations": max_iterations,
+            "_agent_limit_reached": True,
+            "messages": messages,  # Return for accumulation
+        }
+        if tool_results_key and tool_results:
+            result[tool_results_key] = tool_results
+        return result
+
+    return node_fn

yamlgraph/tools/graph_linter.py

@@ -0,0 +1,388 @@
+"""Graph linter for validating YAML graph files.
+
+Checks for common issues:
+- Missing state declarations
+- Undefined tool references
+- Missing prompt files
+- Unreachable nodes
+- Invalid node types
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+from pathlib import Path
+from typing import Any
+
+import yaml
+from pydantic import BaseModel
+
+logger = logging.getLogger(__name__)
+
+# Valid node types
+VALID_NODE_TYPES = {"llm", "router", "agent", "map", "python"}
+
+# Built-in state fields that don't need declaration
+BUILTIN_STATE_FIELDS = {
+    "thread_id",
+    "current_step",
+    "error",
+    "errors",
+    "messages",
+    "_loop_counts",
+    "_loop_limit_reached",
+    "_agent_iterations",
+    "_agent_limit_reached",
+    "started_at",
+    "completed_at",
+}
+
+
+class LintIssue(BaseModel):
+    """A single lint issue found in the graph."""
+
+    severity: str  # "error", "warning", "info"
+    code: str  # e.g., "E001", "W002"
+    message: str
+    line: int | None = None
+    fix: str | None = None
+
+
+class LintResult(BaseModel):
+    """Result of linting a graph file."""
+
+    file: str
+    issues: list[LintIssue]
+    valid: bool
+
+
+def _load_graph(graph_path: Path) -> dict[str, Any]:
+    """Load and parse a YAML graph file."""
+    with open(graph_path) as f:
+        return yaml.safe_load(f) or {}
+
+
+def _extract_variables(text: str) -> set[str]:
+    """Extract {variable} placeholders from text.
+
+    Ignores escaped {{variable}} (doubled braces).
+    """
+    # Find all {word} patterns but not {{word}}
+    # First, temporarily replace {{ and }} to protect them
+    protected = text.replace("{{", "\x00").replace("}}", "\x01")
+    matches = re.findall(r"\{(\w+)\}", protected)
+    return set(matches)
+
+
+def _get_prompt_path(prompt_name: str, prompts_dir: Path) -> Path:
+    """Get the full path to a prompt file."""
+    return prompts_dir / f"{prompt_name}.yaml"
+
+
+def check_state_declarations(
+    graph_path: Path, project_root: Path | None = None
+) -> list[LintIssue]:
+    """Check if variables used in prompts/tools are declared in state.
+
+    Args:
+        graph_path: Path to the graph YAML file
+        project_root: Root directory containing prompts/ folder
+
+    Returns:
+        List of lint issues for missing state declarations
+    """
+    issues = []
+    graph = _load_graph(graph_path)
+
+    if project_root is None:
+        project_root = graph_path.parent
+
+    prompts_dir = project_root / "prompts"
+
+    # Get declared state variables
+    declared_state = set(graph.get("state", {}).keys())
+    declared_state.update(BUILTIN_STATE_FIELDS)
+
+    # Also include state_keys from nodes as they become available at runtime
+    for node_config in graph.get("nodes", {}).values():
+        if "state_key" in node_config:
+            declared_state.add(node_config["state_key"])
+
+    # Find tools used by agent nodes (their variables come from LLM, not state)
+    agent_tools: set[str] = set()
+    for node_config in graph.get("nodes", {}).values():
+        if node_config.get("type") == "agent":
+            agent_tools.update(node_config.get("tools", []))
+
+    # Check shell tool commands for variables (skip agent tools)
+    for tool_name, tool_config in graph.get("tools", {}).items():
+        if tool_config.get("type") == "shell":
+            # Skip tools used by agent nodes - their args come from LLM
+            if tool_name in agent_tools:
+                continue
+
+            command = tool_config.get("command", "")
+            variables = _extract_variables(command)
+            for var in variables:
+                if var not in declared_state:
+                    issues.append(
+                        LintIssue(
+                            severity="error",
+                            code="E001",
+                            message=f"Variable '{var}' used in tool '{tool_name}' "
+                            f"but not declared in state",
+                            fix=f"Add '{var}: str' to the state section",
+                        )
+                    )
+
+    # Check prompt files for variables
+    for _node_name, node_config in graph.get("nodes", {}).items():
+        prompt_name = node_config.get("prompt")
+        if prompt_name:
+            prompt_path = _get_prompt_path(prompt_name, prompts_dir)
+            if prompt_path.exists():
+                with open(prompt_path) as f:
+                    prompt_content = f.read()
+                variables = _extract_variables(prompt_content)
+
+                # Node-level variables provide values for prompt placeholders
+                node_variables = set(node_config.get("variables", {}).keys())
+
+                for var in variables:
+                    # Variable is valid if it's in state OR defined in node variables
+                    if var not in declared_state and var not in node_variables:
+                        issues.append(
+                            LintIssue(
+                                severity="error",
+                                code="E002",
+                                message=f"Variable '{var}' used in prompt "
+                                f"'{prompt_name}' but not declared in state",
+                                fix=f"Add '{var}: str' to the state section",
+                            )
+                        )
+
+    return issues
+
+
+def check_tool_references(graph_path: Path) -> list[LintIssue]:
+    """Check that all tool references in nodes are defined.
+
+    Args:
+        graph_path: Path to the graph YAML file
+
+    Returns:
+        List of lint issues for undefined/unused tools
+    """
+    issues = []
+    graph = _load_graph(graph_path)
+
+    defined_tools = set(graph.get("tools", {}).keys())
+    used_tools: set[str] = set()
+
+    # Find all tool references in nodes
+    for node_name, node_config in graph.get("nodes", {}).items():
+        node_tools = node_config.get("tools", [])
+        for tool in node_tools:
+            used_tools.add(tool)
+            if tool not in defined_tools:
+                issues.append(
+                    LintIssue(
+                        severity="error",
+                        code="E003",
+                        message=f"Tool '{tool}' referenced in node '{node_name}' "
+                        f"but not defined in tools section",
+                        fix=f"Add tool '{tool}' to the tools section or remove reference",
+                    )
+                )
+
+    # Check for unused tools
+    for tool in defined_tools - used_tools:
+        issues.append(
+            LintIssue(
+                severity="warning",
+                code="W001",
+                message=f"Tool '{tool}' is defined but never used",
+                fix=f"Remove unused tool '{tool}' from tools section",
+            )
+        )
+
+    return issues
+
+
+def check_prompt_files(
+    graph_path: Path, project_root: Path | None = None
+) -> list[LintIssue]:
+    """Check that all prompt files referenced by nodes exist.
+
+    Args:
+        graph_path: Path to the graph YAML file
+        project_root: Root directory containing prompts/ folder
+
+    Returns:
+        List of lint issues for missing prompt files
+    """
+    issues = []
+    graph = _load_graph(graph_path)
+
+    if project_root is None:
+        project_root = graph_path.parent
+
+    prompts_dir = project_root / "prompts"
+
+    for node_name, node_config in graph.get("nodes", {}).items():
+        prompt_name = node_config.get("prompt")
+        if prompt_name:
+            prompt_path = _get_prompt_path(prompt_name, prompts_dir)
+            if not prompt_path.exists():
+                issues.append(
+                    LintIssue(
+                        severity="error",
+                        code="E004",
+                        message=f"Prompt file '{prompt_name}.yaml' not found "
+                        f"for node '{node_name}'",
+                        fix=f"Create file: prompts/{prompt_name}.yaml",
+                    )
+                )
+
+    return issues
+
+
+def check_edge_coverage(graph_path: Path) -> list[LintIssue]:
+    """Check that all nodes are reachable and have paths to END.
+
+    Args:
+        graph_path: Path to the graph YAML file
+
+    Returns:
+        List of lint issues for unreachable/dead-end nodes
+    """
+    issues = []
+    graph = _load_graph(graph_path)
+
+    nodes = set(graph.get("nodes", {}).keys())
+    edges = graph.get("edges", [])
+
+    # Build adjacency lists
+    reachable_from_start: set[str] = set()
+    can_reach_end: set[str] = set()
+    nodes_in_edges: set[str] = set()
+
+    def normalize_targets(target) -> list[str]:
+        """Handle both single target and list of targets."""
+        if isinstance(target, list):
+            return target
+        return [target] if target else []
+
+    # Forward traversal from START
+    frontier = {"START"}
+    while frontier:
+        current = frontier.pop()
+        for edge in edges:
+            if edge.get("from") == current:
+                targets = normalize_targets(edge.get("to"))
+                for target in targets:
+                    nodes_in_edges.add(target)
+                    if target not in reachable_from_start and target != "END":
+                        reachable_from_start.add(target)
+                        frontier.add(target)
+
+    # Backward traversal from END
+    frontier = {"END"}
+    visited_backward: set[str] = set()
+    while frontier:
+        current = frontier.pop()
+        visited_backward.add(current)
+        for edge in edges:
+            targets = normalize_targets(edge.get("to"))
+            if current in targets:
+                source = edge.get("from")
+                nodes_in_edges.add(source)
+                if source not in can_reach_end and source != "START":
+                    can_reach_end.add(source)
+                    frontier.add(source)
+
+    # Check for orphaned nodes (not in any edge)
+    for node in nodes:
+        if node not in reachable_from_start:
+            issues.append(
+                LintIssue(
+                    severity="warning",
+                    code="W002",
+                    message=f"Node '{node}' is not reachable from START",
+                    fix=f"Add edge from START or another node to '{node}'",
+                )
+            )
+        elif node not in can_reach_end:
+            issues.append(
+                LintIssue(
+                    severity="warning",
+                    code="W003",
+                    message=f"Node '{node}' has no path to END",
+                    fix=f"Add edge from '{node}' to END or another node",
+                )
+            )
+
+    return issues
+
+
+def check_node_types(graph_path: Path) -> list[LintIssue]:
+    """Check that all node types are valid.
+
+    Args:
+        graph_path: Path to the graph YAML file
+
+    Returns:
+        List of lint issues for invalid node types
+    """
+    issues = []
+    graph = _load_graph(graph_path)
+
+    for node_name, node_config in graph.get("nodes", {}).items():
+        node_type = node_config.get("type")
+        if node_type and node_type not in VALID_NODE_TYPES:
+            issues.append(
+                LintIssue(
+                    severity="error",
+                    code="E005",
+                    message=f"Invalid node type '{node_type}' in node '{node_name}'",
+                    fix=f"Use one of: {', '.join(sorted(VALID_NODE_TYPES))}",
+                )
+            )
+
+    return issues
+
+
+def lint_graph(
+    graph_path: Path | str, project_root: Path | str | None = None
+) -> LintResult:
+    """Lint a YAML graph file for issues.
+
+    Args:
+        graph_path: Path to the graph YAML file
+        project_root: Root directory containing prompts/ folder
+
+    Returns:
+        LintResult with all issues found
+    """
+    graph_path = Path(graph_path)
+    if project_root:
+        project_root = Path(project_root)
+
+    all_issues: list[LintIssue] = []
+
+    # Run all checks
+    all_issues.extend(check_state_declarations(graph_path, project_root))
+    all_issues.extend(check_tool_references(graph_path))
+    all_issues.extend(check_prompt_files(graph_path, project_root))
+    all_issues.extend(check_edge_coverage(graph_path))
+    all_issues.extend(check_node_types(graph_path))
+
+    # Determine validity (no errors)
+    has_errors = any(issue.severity == "error" for issue in all_issues)
+
+    return LintResult(
+        file=str(graph_path),
+        issues=all_issues,
+        valid=not has_errors,
+    )