yamlgraph 0.1.1__py3-none-any.whl

yamlgraph/storage/export.py

@@ -0,0 +1,269 @@
+"""JSON Export - Serialize pipeline results.
+
+Provides functions to export pipeline state and results
+to JSON format for sharing and archival.
+"""
+
+import json
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+from pydantic import BaseModel
+
+from yamlgraph.config import OUTPUTS_DIR
+
+
+def export_state(
+    state: dict,
+    output_dir: str | Path | None = None,
+    prefix: str = "export",
+) -> Path:
+    """Export pipeline state to JSON file.
+
+    Args:
+        state: State dictionary to export
+        output_dir: Directory for output files (default: outputs/)
+        prefix: Filename prefix
+
+    Returns:
+        Path to the created file
+    """
+    if output_dir is None:
+        output_dir = OUTPUTS_DIR
+    output_path = Path(output_dir)
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    thread_id = state.get("thread_id", "unknown")
+    filename = f"{prefix}_{thread_id}_{timestamp}.json"
+
+    filepath = output_path / filename
+
+    # Convert state to JSON-serializable format
+    export_data = _serialize_state(state)
+
+    with open(filepath, "w") as f:
+        json.dump(export_data, f, indent=2, default=str)
+
+    return filepath
+
+
+def _serialize_state(state: dict) -> dict:
+    """Convert state to JSON-serializable format.
+
+    Handles Pydantic models and other complex types.
+
+    Args:
+        state: State dictionary
+
+    Returns:
+        JSON-serializable dictionary
+    """
+    result = {}
+
+    for key, value in state.items():
+        if isinstance(value, BaseModel):
+            result[key] = value.model_dump()
+        elif hasattr(value, "__dict__"):
+            result[key] = _serialize_object(value)
+        else:
+            result[key] = value
+
+    return result
+
+
+def _serialize_object(obj: Any) -> Any:
+    """Recursively serialize an object.
+
+    Args:
+        obj: Object to serialize
+
+    Returns:
+        JSON-serializable representation
+    """
+    if isinstance(obj, BaseModel):
+        return obj.model_dump()
+    elif isinstance(obj, dict):
+        return {k: _serialize_object(v) for k, v in obj.items()}
+    elif isinstance(obj, (list, tuple)):
+        return [_serialize_object(item) for item in obj]
+    elif hasattr(obj, "isoformat"):
+        return obj.isoformat()
+    else:
+        return obj
+
+
+def load_export(filepath: str | Path) -> dict:
+    """Load an exported JSON file.
+
+    Args:
+        filepath: Path to JSON file
+
+    Returns:
+        Loaded dictionary
+    """
+    with open(filepath) as f:
+        return json.load(f)
+
+
+def list_exports(
+    output_dir: str | Path = "outputs", prefix: str = "export"
+) -> list[Path]:
+    """List all export files in a directory.
+
+    Args:
+        output_dir: Directory to search
+        prefix: Filename prefix filter
+
+    Returns:
+        List of matching file paths, sorted by modification time
+    """
+    output_path = Path(output_dir)
+    if not output_path.exists():
+        return []
+
+    files = list(output_path.glob(f"{prefix}_*.json"))
+    return sorted(files, key=lambda f: f.stat().st_mtime, reverse=True)
+
+
+def export_summary(state: dict) -> dict:
+    """Create a summary export (without full content).
+
+    Useful for quick review of pipeline results.
+    Works generically with any Pydantic models in state.
+
+    Args:
+        state: Full state dictionary
+
+    Returns:
+        Summary dictionary with key information only
+    """
+    # Internal keys to skip
+    internal_keys = frozenset(
+        {"_route", "_loop_counts", "thread_id", "topic", "current_step", "error"}
+    )
+
+    summary = {
+        "thread_id": state.get("thread_id"),
+        "topic": state.get("topic"),
+        "current_step": state.get("current_step"),
+        "error": state.get("error"),
+    }
+
+    # Process all non-internal fields generically
+    for key, value in state.items():
+        if key in internal_keys or value is None:
+            continue
+
+        if isinstance(value, BaseModel):
+            # Extract scalar fields from any Pydantic model
+            summary[key] = _extract_scalar_summary(value)
+        elif isinstance(value, str):
+            # For strings, include presence only
+            summary[f"has_{key}"] = bool(value)
+
+    return summary
+
+
+def _extract_scalar_summary(model: BaseModel) -> dict[str, Any]:
+    """Extract scalar fields from a Pydantic model for summary.
+
+    Args:
+        model: Any Pydantic model
+
+    Returns:
+        Dict with scalar field names and values (strings truncated)
+    """
+    result = {}
+    for field_name, field_value in model.model_dump().items():
+        if isinstance(field_value, str):
+            # Truncate long strings
+            result[field_name] = (
+                field_value[:100] + "..." if len(field_value) > 100 else field_value
+            )
+        elif isinstance(field_value, (int, float, bool)):
+            result[field_name] = field_value
+        elif isinstance(field_value, list):
+            result[f"{field_name}_count"] = len(field_value)
+    return result
+
+
+def export_result(
+    state: dict,
+    export_config: dict,
+    base_path: str | Path = "outputs",
+) -> list[Path]:
+    """Export state fields to files.
+
+    Args:
+        state: Final graph state
+        export_config: Mapping of field -> export settings
+        base_path: Base directory for exports
+
+    Returns:
+        List of paths to exported files
+
+    Example config:
+        {
+            "final_summary": {"format": "markdown", "filename": "summary.md"},
+            "generated": {"format": "json", "filename": "content.json"},
+        }
+    """
+    base_path = Path(base_path)
+    thread_id = state.get("thread_id", "unknown")
+    output_dir = base_path / thread_id
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    exported = []
+
+    for field, settings in export_config.items():
+        if field not in state or state[field] is None:
+            continue
+
+        value = state[field]
+        filename = settings.get("filename", f"{field}.txt")
+        format_type = settings.get("format", "text")
+
+        file_path = output_dir / filename
+
+        if format_type == "json":
+            content = _serialize_to_json(value)
+            file_path.write_text(content)
+        elif format_type == "markdown":
+            content = _serialize_to_markdown(value)
+            file_path.write_text(content)
+        else:
+            file_path.write_text(str(value))
+
+        exported.append(file_path)
+
+    return exported
+
+
+def _serialize_to_json(value: Any) -> str:
+    """Serialize value to JSON string."""
+    if isinstance(value, BaseModel):
+        return value.model_dump_json(indent=2)
+    return json.dumps(value, default=str, indent=2)
+
+
+def _serialize_to_markdown(value: Any) -> str:
+    """Serialize value to Markdown string."""
+    if isinstance(value, BaseModel):
+        return _pydantic_to_markdown(value)
+    return str(value)
+
+
+def _pydantic_to_markdown(model: BaseModel) -> str:
+    """Convert Pydantic model to Markdown."""
+    lines = [f"# {model.__class__.__name__}", ""]
+    for field, value in model.model_dump().items():
+        if isinstance(value, list):
+            lines.append(f"## {field.replace('_', ' ').title()}")
+            for item in value:
+                lines.append(f"- {item}")
+            lines.append("")
+        else:
+            lines.append(f"**{field.replace('_', ' ').title()}**: {value}")
+    return "\n".join(lines)

yamlgraph/tools/__init__.py

@@ -0,0 +1 @@
+"""Shell tool execution utilities."""

yamlgraph/tools/agent.py

@@ -0,0 +1,235 @@
+"""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 logging
+from collections.abc import Callable
+from typing import Any
+
+import yaml
+from langchain_core.messages import HumanMessage, SystemMessage, ToolMessage
+
+from yamlgraph.config import PROMPTS_DIR
+from yamlgraph.tools.shell import ShellToolConfig, execute_shell_tool
+from yamlgraph.utils.llm_factory import create_llm
+
+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 _load_prompt(prompt_name: str) -> tuple[str, str]:
+    """Load system and user prompts from YAML file.
+
+    Args:
+        prompt_name: Name of prompt file (without .yaml)
+
+    Returns:
+        Tuple of (system_prompt, user_template)
+    """
+    prompt_path = PROMPTS_DIR / f"{prompt_name}.yaml"
+    if not prompt_path.exists():
+        raise FileNotFoundError(f"Prompt file not found: {prompt_path}")
+
+    with open(prompt_path) as f:
+        prompt_config = yaml.safe_load(f)
+
+    return prompt_config.get("system", ""), prompt_config.get("user", "{input}")
+
+
+def create_agent_node(
+    node_name: str,
+    node_config: dict[str, Any],
+    tools: dict[str, ShellToolConfig],
+) -> 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 tools
+
+    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
+    """
+    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 shell configs
+    lc_tools = [build_langchain_tool(name, tools[name]) for name in tool_names]
+    tool_lookup = {name: tools[name] for name in tool_names}
+
+    def node_fn(state: dict) -> dict:
+        """Execute the agent loop."""
+        # Load prompts - fail fast if missing
+        system_prompt, user_template = _load_prompt(prompt_name)
+
+        # 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:
+                    result = execute_shell_tool(tool_config, tool_args)
+                    output = (
+                        str(result.output)
+                        if result.success
+                        else f"Error: {result.error}"
+                    )
+                    success = result.success
+                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/nodes.py

@@ -0,0 +1,124 @@
+"""Node factories for tool and agent nodes.
+
+This module provides functions to create graph nodes that execute
+shell tools, either deterministically (tool nodes) or via LLM
+decision-making (agent nodes).
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from typing import Any
+
+from yamlgraph.models.schemas import ErrorType, PipelineError
+from yamlgraph.tools.shell import ShellToolConfig, execute_shell_tool
+from yamlgraph.utils.expressions import resolve_template
+
+# Type alias for state - dynamic TypedDict at runtime
+GraphState = dict[str, Any]
+
+logger = logging.getLogger(__name__)
+
+
+def resolve_state_variable(template: str, state: dict[str, Any]) -> str:
+    """Resolve {state.path.to.value} to actual state value.
+
+    Note: Uses consolidated resolve_template from expressions module.
+
+    Args:
+        template: String with {state.key} or {state.nested.key} placeholders
+        state: Current graph state
+
+    Returns:
+        Resolved string value
+    """
+    value = resolve_template(template, state)
+    # resolve_template returns the template unchanged if not a state expression
+    if value is template:
+        return template
+    return str(value) if value is not None else ""
+
+
+def resolve_variables(
+    variables_config: dict[str, str],
+    state: dict[str, Any],
+) -> dict[str, Any]:
+    """Resolve all variable templates against state.
+
+    Args:
+        variables_config: Dict of {var_name: template_string}
+        state: Current graph state
+
+    Returns:
+        Dict of {var_name: resolved_value}
+    """
+    resolved = {}
+    for key, template in variables_config.items():
+        resolved[key] = resolve_state_variable(template, state)
+    return resolved
+
+
+def create_tool_node(
+    node_name: str,
+    node_config: dict[str, Any],
+    tools: dict[str, ShellToolConfig],
+) -> Callable[[GraphState], dict]:
+    """Create a node that executes a shell tool.
+
+    Args:
+        node_name: Name of the node in the graph
+        node_config: Node configuration from YAML
+        tools: Registry of available tools
+
+    Returns:
+        Node function that executes the tool
+
+    Raises:
+        KeyError: If tool name not in registry
+    """
+    tool_name = node_config["tool"]
+    tool_config = tools[tool_name]  # Raise KeyError if not found
+    state_key = node_config.get("state_key", node_name)
+    on_error = node_config.get("on_error", "fail")
+    variables_template = node_config.get("variables", {})
+
+    def node_fn(state: GraphState) -> dict:
+        """Execute the shell tool and return state update."""
+        # Resolve variables from state
+        variables = resolve_variables(variables_template, state)
+
+        logger.info(f"🔧 Executing tool: {tool_name}")
+        result = execute_shell_tool(tool_config, variables)
+
+        if not result.success:
+            logger.warning(f"Tool {tool_name} failed: {result.error}")
+
+            if on_error == "skip":
+                # Return with error tracked but don't raise
+                errors = list(state.get("errors") or [])
+                errors.append(
+                    PipelineError(
+                        node=node_name,
+                        type=ErrorType.UNKNOWN_ERROR,
+                        message=result.error or "Tool execution failed",
+                    )
+                )
+                return {
+                    state_key: None,
+                    "current_step": node_name,
+                    "errors": errors,
+                }
+            else:
+                # on_error == "fail" - raise exception
+                raise RuntimeError(
+                    f"Tool '{tool_name}' failed in node '{node_name}': {result.error}"
+                )
+
+        logger.info(f"✓ Tool {tool_name} completed")
+        return {
+            state_key: result.output,
+            "current_step": node_name,
+        }
+
+    return node_fn