yamlgraph 0.3.9__py3-none-any.whl

yamlgraph/models/graph_schema.py

@@ -0,0 +1,181 @@
+"""Pydantic schemas for YAML graph configuration validation.
+
+Provides structured validation for graph YAML files with clear error messages.
+"""
+
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field, field_validator, model_validator
+
+from yamlgraph.constants import ErrorHandler, NodeType
+
+
+class SubgraphNodeConfig(BaseModel):
+    """Configuration for a subgraph node."""
+
+    type: Literal["subgraph"]
+    graph: str = Field(
+        ..., description="Path to subgraph YAML file (relative to parent)"
+    )
+    mode: Literal["invoke", "direct"] = Field(
+        default="invoke",
+        description="invoke: explicit state mapping; direct: shared schema",
+    )
+    input_mapping: dict[str, str] | Literal["auto", "*"] = Field(
+        default_factory=dict,
+        description="Map parent state fields to child input (mode=invoke only)",
+    )
+    output_mapping: dict[str, str] | Literal["auto", "*"] = Field(
+        default_factory=dict,
+        description="Map child output fields to parent state (mode=invoke only)",
+    )
+    interrupt_output_mapping: dict[str, str] | Literal["auto", "*"] = Field(
+        default_factory=dict,
+        description="Map child state to parent when subgraph interrupts (FR-006)",
+    )
+    checkpointer: str | None = Field(
+        default=None,
+        description="Override parent checkpointer",
+    )
+
+    model_config = {"extra": "allow"}
+
+    @model_validator(mode="after")
+    def validate_config(self) -> "SubgraphNodeConfig":
+        """Validate subgraph configuration."""
+        if not self.graph.endswith((".yaml", ".yml")):
+            raise ValueError(f"Subgraph must be a YAML file: {self.graph}")
+        if self.mode == "direct" and (self.input_mapping or self.output_mapping):
+            raise ValueError("mode=direct does not support input/output mappings")
+        return self
+
+
+class NodeConfig(BaseModel):
+    """Configuration for a single graph node."""
+
+    type: str = Field(default=NodeType.LLM, description="Node type")
+    prompt: str | None = Field(default=None, description="Prompt template name")
+    state_key: str | None = Field(default=None, description="State key for output")
+    temperature: float | None = Field(default=None, ge=0, le=2)
+    provider: str | None = Field(default=None)
+    on_error: str | None = Field(default=None)
+    fallback: dict[str, Any] | None = Field(default=None)
+    variables: dict[str, str] = Field(default_factory=dict)
+    requires: list[str] = Field(default_factory=list)
+    routes: dict[str, str] | None = Field(default=None, description="Router routes")
+
+    # Map node fields
+    over: str | None = Field(default=None, description="Map over expression")
+    # 'as' is reserved in Python, handled specially
+    item_var: str | None = Field(default=None, alias="as")
+    node: dict[str, Any] | None = Field(default=None, description="Map sub-node")
+    collect: str | None = Field(default=None, description="Map collect key")
+
+    # Tool/Agent fields
+    tools: list[str] = Field(default_factory=list)
+    max_iterations: int = Field(default=10, ge=1)
+
+    model_config = {"extra": "allow", "populate_by_name": True}
+
+    @field_validator("on_error")
+    @classmethod
+    def validate_on_error(cls, v: str | None) -> str | None:
+        """Validate on_error is a known handler."""
+        if v is not None and v not in ErrorHandler.all_values():
+            valid = ", ".join(ErrorHandler.all_values())
+            raise ValueError(f"Invalid on_error '{v}'. Valid: {valid}")
+        return v
+
+    @model_validator(mode="after")
+    def validate_node_requirements(self) -> "NodeConfig":
+        """Validate node has required fields based on type."""
+        if NodeType.requires_prompt(self.type) and not self.prompt:
+            raise ValueError(f"Node type '{self.type}' requires 'prompt' field")
+
+        if self.type == NodeType.ROUTER and not self.routes:
+            raise ValueError("Router node requires 'routes' field")
+
+        if self.type == NodeType.MAP:
+            if not self.over:
+                raise ValueError("Map node requires 'over' field")
+            if not self.item_var:
+                raise ValueError("Map node requires 'as' field")
+            if not self.node:
+                raise ValueError("Map node requires 'node' field")
+            if not self.collect:
+                raise ValueError("Map node requires 'collect' field")
+
+        return self
+
+
+class EdgeConfig(BaseModel):
+    """Configuration for a graph edge."""
+
+    from_node: str = Field(..., alias="from", description="Source node")
+    to: str | list[str] = Field(..., description="Target node(s)")
+    condition: str | None = Field(default=None, description="Condition expression")
+
+    model_config = {"populate_by_name": True}
+
+
+class GraphConfigSchema(BaseModel):
+    """Full YAML graph configuration schema.
+
+    Use this for validating graph YAML files with Pydantic.
+    """
+
+    version: str = Field(default="1.0")
+    name: str = Field(default="unnamed")
+    description: str = Field(default="")
+    defaults: dict[str, Any] = Field(default_factory=dict)
+    nodes: dict[str, NodeConfig] = Field(...)
+    edges: list[EdgeConfig] = Field(...)
+    tools: dict[str, Any] = Field(default_factory=dict)
+    state_class: str = Field(default="")
+    loop_limits: dict[str, int] = Field(default_factory=dict)
+
+    model_config = {"extra": "allow"}
+
+    @model_validator(mode="after")
+    def validate_router_targets(self) -> "GraphConfigSchema":
+        """Validate router routes point to existing nodes."""
+        for node_name, node in self.nodes.items():
+            if node.type == NodeType.ROUTER and node.routes:
+                for route_key, target in node.routes.items():
+                    if target not in self.nodes:
+                        raise ValueError(
+                            f"Router '{node_name}' route '{route_key}' "
+                            f"targets nonexistent node '{target}'"
+                        )
+        return self
+
+    @model_validator(mode="after")
+    def validate_edge_nodes(self) -> "GraphConfigSchema":
+        """Validate edge sources and targets exist."""
+        valid_nodes = set(self.nodes.keys()) | {"START", "END"}
+
+        for edge in self.edges:
+            if edge.from_node not in valid_nodes:
+                raise ValueError(f"Edge 'from' node '{edge.from_node}' not found")
+
+            targets = edge.to if isinstance(edge.to, list) else [edge.to]
+            for target in targets:
+                if target not in valid_nodes:
+                    raise ValueError(f"Edge 'to' node '{target}' not found")
+
+        return self
+
+
+def validate_graph_schema(config: dict[str, Any]) -> GraphConfigSchema:
+    """Validate a graph configuration dict using Pydantic.
+
+    Args:
+        config: Raw parsed YAML configuration
+
+    Returns:
+        Validated GraphConfigSchema
+
+    Raises:
+        pydantic.ValidationError: If validation fails
+    """
+    return GraphConfigSchema.model_validate(config)

yamlgraph/models/schemas.py

@@ -0,0 +1,124 @@
+"""Pydantic models for structured LLM outputs.
+
+This module contains FRAMEWORK models only - models used by the framework itself.
+Demo-specific output schemas are defined inline in graph YAML files.
+"""
+
+from datetime import datetime
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+# =============================================================================
+# Error Types
+# =============================================================================
+
+
+class ErrorType(str, Enum):
+    """Types of errors that can occur in the pipeline."""
+
+    LLM_ERROR = "llm_error"  # LLM API errors (rate limit, timeout, etc.)
+    VALIDATION_ERROR = "validation_error"  # Pydantic validation failures
+    PROMPT_ERROR = "prompt_error"  # Missing prompt, template errors
+    STATE_ERROR = "state_error"  # Missing required state data
+    UNKNOWN_ERROR = "unknown_error"  # Catch-all
+
+
+class PipelineError(BaseModel):
+    """Structured error information for pipeline failures."""
+
+    type: ErrorType = Field(description="Category of error")
+    message: str = Field(description="Human-readable error message")
+    node: str = Field(description="Node where error occurred")
+    timestamp: datetime = Field(default_factory=datetime.now)
+    retryable: bool = Field(
+        default=False, description="Whether this error can be retried"
+    )
+    details: dict[str, Any] = Field(
+        default_factory=dict, description="Additional error context"
+    )
+
+    @classmethod
+    def from_exception(
+        cls, e: Exception, node: str, error_type: ErrorType | None = None
+    ) -> "PipelineError":
+        """Create a PipelineError from an exception.
+
+        Args:
+            e: The exception that occurred
+            node: Name of the node where error occurred
+            error_type: Optional explicit error type
+
+        Returns:
+            PipelineError instance
+        """
+        # Infer error type from exception
+        if error_type is None:
+            exc_name = type(e).__name__.lower()
+            if "rate" in exc_name or "timeout" in exc_name or "api" in exc_name:
+                error_type = ErrorType.LLM_ERROR
+                retryable = True
+            elif "validation" in exc_name:
+                error_type = ErrorType.VALIDATION_ERROR
+                retryable = False
+            elif "file" in exc_name or "prompt" in exc_name:
+                error_type = ErrorType.PROMPT_ERROR
+                retryable = False
+            else:
+                error_type = ErrorType.UNKNOWN_ERROR
+                retryable = False
+        else:
+            retryable = error_type == ErrorType.LLM_ERROR
+
+        return cls(
+            type=error_type,
+            message=str(e),
+            node=node,
+            retryable=retryable,
+            details={"exception_type": type(e).__name__},
+        )
+
+
+# =============================================================================
+# Generic Report Model (Flexible for Any Use Case)
+# =============================================================================
+
+
+class GenericReport(BaseModel):
+    """Flexible report structure for any use case.
+
+    Use this when you don't need a custom schema - works for most
+    analysis and summary tasks. The LLM can populate any combination
+    of the optional fields as needed.
+
+    Example usage in graph YAML:
+        nodes:
+          analyze:
+            type: llm
+            prompt: my_analysis
+            output_model: yamlgraph.models.GenericReport
+
+    Example prompts can request specific sections:
+        "Analyze the repository and provide:
+         - A summary of findings
+         - Key findings as bullet points
+         - Recommendations for improvement"
+    """
+
+    title: str = Field(description="Report title")
+    summary: str = Field(description="Executive summary")
+    sections: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Named sections with any content (strings, dicts, lists)",
+    )
+    findings: list[str] = Field(
+        default_factory=list, description="Key findings or bullet points"
+    )
+    recommendations: list[str] = Field(
+        default_factory=list, description="Suggested actions or areas to focus on"
+    )
+    metadata: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Additional key-value data (author, version, tags, etc.)",
+    )

yamlgraph/models/state_builder.py

@@ -0,0 +1,236 @@
+"""Dynamic state class generation from graph configuration.
+
+Builds TypedDict programmatically from YAML graph config, eliminating
+the need for state_class coupling between YAML and Python.
+"""
+
+import logging
+from operator import add
+from typing import Annotated, Any, TypedDict
+
+logger = logging.getLogger(__name__)
+
+
+def sorted_add(existing: list, new: list) -> list:
+    """Reducer that adds items and sorts by _map_index if present.
+
+    Used for map node fan-in to guarantee order regardless of
+    parallel execution timing.
+
+    Args:
+        existing: Current list in state
+        new: New items to add
+
+    Returns:
+        Combined list sorted by _map_index (if items have it)
+    """
+    combined = (existing or []) + (new or [])
+
+    # Sort by _map_index if items have it
+    if combined and isinstance(combined[0], dict) and "_map_index" in combined[0]:
+        combined = sorted(combined, key=lambda x: x.get("_map_index", 0))
+
+    return combined
+
+
+# =============================================================================
+# Base Fields - Always included in generated state
+# =============================================================================
+
+# Infrastructure fields present in all graphs
+BASE_FIELDS: dict[str, type] = {
+    # Core tracking
+    "thread_id": str,
+    "current_step": str,
+    # Error handling - singular for current error
+    "error": Any,
+    # Error handling with reducer (accumulates)
+    "errors": Annotated[list, add],
+    # Messages with reducer (accumulates)
+    "messages": Annotated[list, add],
+    # Loop tracking
+    "_loop_counts": dict,
+    "_loop_limit_reached": bool,
+    "_agent_iterations": int,
+    "_agent_limit_reached": bool,
+    # Timestamps
+    "started_at": Any,
+    "completed_at": Any,
+}
+
+# Common input fields used across graph types
+# These are always included to support --var inputs
+COMMON_INPUT_FIELDS: dict[str, type] = {
+    "input": str,  # Agent prompt input
+    "topic": str,  # Content generation topic
+    "style": str,  # Writing style
+    "word_count": int,  # Target word count
+    "message": str,  # Router message input
+}
+
+# Type mapping for YAML state config
+TYPE_MAP: dict[str, type] = {
+    "str": str,
+    "string": str,
+    "int": int,
+    "integer": int,
+    "float": float,
+    "bool": bool,
+    "boolean": bool,
+    "list": list,
+    "dict": dict,
+    "any": Any,
+}
+
+
+def parse_state_config(state_config: dict) -> dict[str, type]:
+    """Parse YAML state section into field types.
+
+    Supports simple type strings:
+        state:
+          concept: str
+          count: int
+
+    Args:
+        state_config: Dict from YAML 'state' section
+
+    Returns:
+        Dict of field_name -> Python type
+    """
+    fields: dict[str, type] = {}
+
+    for field_name, type_spec in state_config.items():
+        if isinstance(type_spec, str):
+            # Simple type: "str", "int", etc.
+            normalized = type_spec.lower()
+            if normalized not in TYPE_MAP:
+                supported = ", ".join(sorted(set(TYPE_MAP.keys())))
+                logger.warning(
+                    f"Unknown type '{type_spec}' for state field '{field_name}'. "
+                    f"Supported types: {supported}. Defaulting to Any."
+                )
+            python_type = TYPE_MAP.get(normalized, Any)
+            fields[field_name] = python_type
+        else:
+            # Unknown format, use Any
+            logger.warning(
+                f"Invalid type specification for state field '{field_name}': "
+                f"expected string, got {type(type_spec).__name__}. Defaulting to Any."
+            )
+            fields[field_name] = Any
+
+    return fields
+
+
+def build_state_class(config: dict) -> type:
+    """Build TypedDict state class from graph configuration.
+
+    Dynamically generates a TypedDict with:
+    - Base infrastructure fields (errors, messages, thread_id, etc.)
+    - Common input fields (topic, style, input, message, etc.)
+    - Custom fields from YAML 'state' section
+    - Fields extracted from node state_key
+    - Special fields for agent/router node types
+
+    Args:
+        config: Parsed YAML graph configuration dict
+
+    Returns:
+        TypedDict class with total=False (all fields optional)
+    """
+    # Start with base and common fields
+    fields: dict[str, type] = {}
+    fields.update(BASE_FIELDS)
+    fields.update(COMMON_INPUT_FIELDS)
+
+    # Add custom state fields from YAML 'state' section
+    state_config = config.get("state", {})
+    custom_fields = parse_state_config(state_config)
+    fields.update(custom_fields)
+
+    # Extract fields from nodes
+    nodes = config.get("nodes", {})
+    node_fields = extract_node_fields(nodes)
+    fields.update(node_fields)
+
+    # Build TypedDict programmatically
+    return TypedDict("GraphState", fields, total=False)
+
+
+def extract_node_fields(nodes: dict) -> dict[str, type]:
+    """Extract state fields from node configurations.
+
+    Analyzes node configs to determine required state fields:
+    - state_key: Where node stores its output
+    - type: agent → adds input, _tool_results
+    - type: router → adds _route
+
+    Args:
+        nodes: Dict of node_name -> node_config
+
+    Returns:
+        Dict of field_name -> type for the state
+    """
+    fields: dict[str, type] = {}
+
+    for _node_name, node_config in nodes.items():
+        if not isinstance(node_config, dict):
+            continue
+
+        # state_key → Any (accepts Pydantic models)
+        if state_key := node_config.get("state_key"):
+            fields[state_key] = Any
+
+        # Node type-specific fields
+        node_type = node_config.get("type", "llm")
+
+        if node_type == "agent":
+            fields["input"] = str
+            fields["_tool_results"] = list
+
+        elif node_type == "router":
+            fields["_route"] = str
+
+        elif node_type == "map":
+            # Map node collect field needs sorted reducer for ordered fan-in
+            if collect_key := node_config.get("collect"):
+                fields[collect_key] = Annotated[list, sorted_add]
+
+    return fields
+
+
+def create_initial_state(
+    topic: str = "",
+    style: str = "informative",
+    word_count: int = 300,
+    thread_id: str | None = None,
+    **kwargs: Any,
+) -> dict[str, Any]:
+    """Create an initial state for a new pipeline run.
+
+    Args:
+        topic: The topic to generate content about
+        style: Writing style (default: informative)
+        word_count: Target word count (default: 300)
+        thread_id: Optional thread ID (auto-generated if not provided)
+        **kwargs: Additional state fields (e.g., input for agents)
+
+    Returns:
+        Initialized state dictionary
+    """
+    import uuid
+    from datetime import datetime
+
+    return {
+        "thread_id": thread_id or uuid.uuid4().hex[:16],
+        "topic": topic,
+        "style": style,
+        "word_count": word_count,
+        "current_step": "init",
+        "error": None,
+        "errors": [],
+        "messages": [],
+        "started_at": datetime.now(),
+        "completed_at": None,
+        **kwargs,
+    }