yamlgraph 0.3.9__py3-none-any.whl

yamlgraph/utils/llm_factory_async.py

@@ -0,0 +1,105 @@
+"""Async LLM Factory - Async versions of LLM creation.
+
+This module provides async-compatible LLM creation with support for
+non-blocking I/O operations in async contexts.
+
+Note: This module is a foundation for future async support. Currently,
+LangChain's LLM implementations use sync HTTP clients internally, so
+this wraps them for use in async contexts via run_in_executor.
+"""
+
+import asyncio
+import logging
+from concurrent.futures import ThreadPoolExecutor
+from functools import partial
+from typing import TypeVar
+
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.messages import BaseMessage
+from pydantic import BaseModel
+
+from yamlgraph.utils.llm_factory import ProviderType, create_llm
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T", bound=BaseModel)
+
+# Shared executor for running sync LLM calls
+_executor: ThreadPoolExecutor | None = None
+
+
+def get_executor() -> ThreadPoolExecutor:
+    """Get or create the shared thread pool executor."""
+    global _executor
+    if _executor is None:
+        _executor = ThreadPoolExecutor(max_workers=4)
+    return _executor
+
+
+async def create_llm_async(
+    provider: ProviderType | None = None,
+    model: str | None = None,
+    temperature: float = 0.7,
+) -> BaseChatModel:
+    """Create an LLM instance asynchronously.
+
+    Currently wraps the sync create_llm. Future versions may use
+    native async LLM implementations.
+
+    Args:
+        provider: LLM provider ("anthropic", "mistral", "openai")
+        model: Model name
+        temperature: Temperature for generation
+
+    Returns:
+        Configured LLM instance
+    """
+    loop = asyncio.get_event_loop()
+    return await loop.run_in_executor(
+        get_executor(),
+        partial(create_llm, provider=provider, model=model, temperature=temperature),
+    )
+
+
+async def invoke_async(
+    llm: BaseChatModel,
+    messages: list[BaseMessage],
+    output_model: type[T] | None = None,
+) -> T | str:
+    """Invoke LLM asynchronously.
+
+    Runs the sync invoke in a thread pool to avoid blocking.
+
+    Args:
+        llm: The LLM instance
+        messages: Messages to send
+        output_model: Optional Pydantic model for structured output
+
+    Returns:
+        LLM response (parsed model or string)
+    """
+    loop = asyncio.get_event_loop()
+
+    def sync_invoke() -> T | str:
+        if output_model:
+            structured_llm = llm.with_structured_output(output_model)
+            return structured_llm.invoke(messages)
+        else:
+            response = llm.invoke(messages)
+            return response.content
+
+    return await loop.run_in_executor(get_executor(), sync_invoke)
+
+
+def shutdown_executor() -> None:
+    """Shutdown the thread pool executor.
+
+    Call this during application shutdown to clean up resources.
+    """
+    global _executor
+    if _executor is not None:
+        _executor.shutdown(wait=True)
+        _executor = None
+
+
+__all__ = ["create_llm_async", "invoke_async", "shutdown_executor"]

yamlgraph/utils/logging.py

@@ -0,0 +1,104 @@
+"""Structured logging configuration for yamlgraph.
+
+Provides consistent logging across all modules with JSON-formatted
+output for production environments.
+"""
+
+import logging
+import os
+import sys
+
+
+class StructuredFormatter(logging.Formatter):
+    """Formatter that outputs structured log messages.
+
+    In production (LOG_FORMAT=json), outputs JSON lines.
+    In development, outputs human-readable format.
+    """
+
+    def __init__(self, use_json: bool = False):
+        super().__init__()
+        self.use_json = use_json
+
+    def format(self, record: logging.LogRecord) -> str:
+        """Format a log record."""
+        if self.use_json:
+            import json
+
+            log_data = {
+                "timestamp": self.formatTime(record),
+                "level": record.levelname,
+                "logger": record.name,
+                "message": record.getMessage(),
+            }
+            # Add extra fields if present
+            if hasattr(record, "extra"):
+                log_data.update(record.extra)
+            if record.exc_info:
+                log_data["exception"] = self.formatException(record.exc_info)
+            return json.dumps(log_data)
+        else:
+            # Human-readable format
+            base = f"{self.formatTime(record)} [{record.levelname}] {record.name}: {record.getMessage()}"
+            if record.exc_info:
+                base += f"\n{self.formatException(record.exc_info)}"
+            return base
+
+
+def setup_logging(
+    level: str | None = None,
+    use_json: bool | None = None,
+) -> logging.Logger:
+    """Configure logging for yamlgraph.
+
+    Args:
+        level: Log level (DEBUG, INFO, WARNING, ERROR).
+               Defaults to LOG_LEVEL env var or INFO.
+        use_json: If True, output JSON lines.
+                  Defaults to LOG_FORMAT=json env var.
+
+    Returns:
+        Root logger for the yamlgraph package
+    """
+    if level is None:
+        level = os.getenv("LOG_LEVEL", "INFO")
+
+    if use_json is None:
+        use_json = os.getenv("LOG_FORMAT", "").lower() == "json"
+
+    # Get the yamlgraph logger
+    logger = logging.getLogger("yamlgraph")
+    logger.setLevel(getattr(logging, level.upper()))
+
+    # Remove existing handlers
+    logger.handlers.clear()
+
+    # Add handler with formatter
+    handler = logging.StreamHandler(sys.stderr)
+    handler.setFormatter(StructuredFormatter(use_json=use_json))
+    logger.addHandler(handler)
+
+    # Don't propagate to root logger
+    logger.propagate = False
+
+    return logger
+
+
+def get_logger(name: str) -> logging.Logger:
+    """Get a logger for a specific module.
+
+    Args:
+        name: Module name (typically __name__)
+
+    Returns:
+        Logger instance
+
+    Example:
+        >>> logger = get_logger(__name__)
+        >>> logger.info("Processing started", extra={"topic": "AI"})
+    """
+    return logging.getLogger(name)
+
+
+# Initialize logging on import
+_root_logger = setup_logging()

yamlgraph/utils/prompts.py

@@ -0,0 +1,171 @@
+"""Unified prompt loading and path resolution.
+
+This module consolidates prompt loading logic used by executor.py
+and node_factory.py into a single, testable module.
+
+Search order for prompts:
+1. If prompts_relative + prompts_dir + graph_path: graph_path.parent/prompts_dir/{prompt_name}.yaml
+2. If prompts_dir specified: prompts_dir/{prompt_name}.yaml
+3. If prompts_relative + graph_path: graph_path.parent/{prompt_name}.yaml
+4. Default: PROMPTS_DIR/{prompt_name}.yaml
+5. Fallback: {parent}/prompts/{basename}.yaml (external examples)
+"""
+
+from pathlib import Path
+
+import yaml
+
+from yamlgraph.config import PROMPTS_DIR
+
+
+def resolve_prompt_path(
+    prompt_name: str,
+    prompts_dir: Path | None = None,
+    graph_path: Path | None = None,
+    prompts_relative: bool = False,
+) -> Path:
+    """Resolve a prompt name to its full YAML file path.
+
+    Resolution order:
+    1. If prompts_relative + prompts_dir + graph_path: graph_path.parent/prompts_dir/{prompt_name}.yaml
+    2. If prompts_dir specified: prompts_dir/{prompt_name}.yaml
+    3. If prompts_relative + graph_path: graph_path.parent/{prompt_name}.yaml
+    4. Default: PROMPTS_DIR/{prompt_name}.yaml
+    5. Fallback: {parent}/prompts/{basename}.yaml (external examples)
+
+    Args:
+        prompt_name: Prompt name like "greet" or "prompts/opening"
+        prompts_dir: Explicit prompts directory (combined with graph_path if prompts_relative=True)
+        graph_path: Path to the graph YAML file (for relative resolution)
+        prompts_relative: If True, resolve relative to graph_path.parent
+
+    Returns:
+        Path to the YAML file
+
+    Raises:
+        FileNotFoundError: If prompt file doesn't exist
+        ValueError: If prompts_relative=True but graph_path not provided
+
+    Examples:
+        >>> resolve_prompt_path("greet")
+        PosixPath('/path/to/prompts/greet.yaml')
+
+        >>> resolve_prompt_path("prompts/opening", graph_path=Path("graphs/demo.yaml"), prompts_relative=True)
+        PosixPath('/path/to/graphs/prompts/opening.yaml')
+
+        >>> resolve_prompt_path("opening", prompts_dir="prompts", graph_path=Path("graphs/demo.yaml"), prompts_relative=True)
+        PosixPath('/path/to/graphs/prompts/opening.yaml')
+    """
+    # Validate prompts_relative requires graph_path
+    if prompts_relative and graph_path is None and prompts_dir is None:
+        raise ValueError("graph_path required when prompts_relative=True")
+
+    # 1. Graph-relative with explicit prompts_dir (combine them)
+    if prompts_relative and prompts_dir is not None and graph_path is not None:
+        graph_dir = Path(graph_path).parent
+        yaml_path = graph_dir / prompts_dir / f"{prompt_name}.yaml"
+        if yaml_path.exists():
+            return yaml_path
+        # Fall through if not found
+
+    # 2. Explicit prompts_dir (absolute path or CWD-relative)
+    if prompts_dir is not None:
+        prompts_dir = Path(prompts_dir)
+        yaml_path = prompts_dir / f"{prompt_name}.yaml"
+        if yaml_path.exists():
+            return yaml_path
+        # Fall through to other resolution methods
+
+    # 3. Graph-relative resolution (without explicit prompts_dir)
+    if prompts_relative and graph_path is not None:
+        graph_dir = Path(graph_path).parent
+        yaml_path = graph_dir / f"{prompt_name}.yaml"
+        if yaml_path.exists():
+            return yaml_path
+        # Fall through to default
+
+    # 4. Default: use global PROMPTS_DIR
+    default_dir = PROMPTS_DIR if prompts_dir is None else prompts_dir
+    yaml_path = Path(default_dir) / f"{prompt_name}.yaml"
+    if yaml_path.exists():
+        return yaml_path
+
+    # 5. Fallback: external example location {parent}/prompts/{basename}.yaml
+    parts = prompt_name.rsplit("/", 1)
+    if len(parts) == 2:
+        parent_dir, basename = parts
+        alt_path = Path(parent_dir) / "prompts" / f"{basename}.yaml"
+        if alt_path.exists():
+            return alt_path
+
+    raise FileNotFoundError(f"Prompt not found: {yaml_path}")
+
+
+def load_prompt(
+    prompt_name: str,
+    prompts_dir: Path | None = None,
+    graph_path: Path | None = None,
+    prompts_relative: bool = False,
+) -> dict:
+    """Load a YAML prompt template.
+
+    Args:
+        prompt_name: Name of the prompt file (without .yaml extension)
+        prompts_dir: Optional prompts directory override
+        graph_path: Path to the graph YAML file (for relative resolution)
+        prompts_relative: If True, resolve relative to graph_path.parent
+
+    Returns:
+        Dictionary with prompt content (typically 'system' and 'user' keys)
+
+    Raises:
+        FileNotFoundError: If prompt file doesn't exist
+    """
+    path = resolve_prompt_path(
+        prompt_name,
+        prompts_dir=prompts_dir,
+        graph_path=graph_path,
+        prompts_relative=prompts_relative,
+    )
+
+    with open(path) as f:
+        return yaml.safe_load(f)
+
+
+def load_prompt_path(
+    prompt_name: str,
+    prompts_dir: Path | None = None,
+    graph_path: Path | None = None,
+    prompts_relative: bool = False,
+) -> tuple[Path, dict]:
+    """Load a prompt and return both path and content.
+
+    Useful when you need both the file path (for schema loading)
+    and the content (for prompt execution).
+
+    Args:
+        prompt_name: Name of the prompt file (without .yaml extension)
+        prompts_dir: Optional prompts directory override
+        graph_path: Path to the graph YAML file (for relative resolution)
+        prompts_relative: If True, resolve relative to graph_path.parent
+
+    Returns:
+        Tuple of (path, content_dict)
+
+    Raises:
+        FileNotFoundError: If prompt file doesn't exist
+    """
+    path = resolve_prompt_path(
+        prompt_name,
+        prompts_dir=prompts_dir,
+        graph_path=graph_path,
+        prompts_relative=prompts_relative,
+    )
+
+    with open(path) as f:
+        content = yaml.safe_load(f)
+
+    return path, content
+
+
+__all__ = ["resolve_prompt_path", "load_prompt", "load_prompt_path"]

yamlgraph/utils/sanitize.py

@@ -0,0 +1,98 @@
+"""Input sanitization utilities.
+
+Provides functions for validating and sanitizing user input
+to prevent prompt injection and other security issues.
+"""
+
+import re
+from typing import NamedTuple
+
+from yamlgraph.config import DANGEROUS_PATTERNS, MAX_TOPIC_LENGTH
+
+
+class SanitizationResult(NamedTuple):
+    """Result of input sanitization."""
+
+    value: str
+    is_safe: bool
+    warnings: list[str]
+
+
+def sanitize_topic(topic: str) -> SanitizationResult:
+    """Sanitize a topic string for use in prompts.
+
+    Checks for:
+    - Length limits
+    - Potential prompt injection patterns
+    - Control characters
+
+    Args:
+        topic: The raw topic string
+
+    Returns:
+        SanitizationResult with cleaned value and safety status
+
+    Example:
+        >>> result = sanitize_topic("machine learning")
+        >>> result.is_safe
+        True
+        >>> result = sanitize_topic("ignore previous instructions")
+        >>> result.is_safe
+        False
+    """
+    warnings = []
+    cleaned = topic.strip()
+
+    # Check length
+    if len(cleaned) > MAX_TOPIC_LENGTH:
+        cleaned = cleaned[:MAX_TOPIC_LENGTH]
+        warnings.append(f"Topic truncated to {MAX_TOPIC_LENGTH} characters")
+
+    # Check for empty
+    if not cleaned:
+        return SanitizationResult(
+            value="",
+            is_safe=False,
+            warnings=["Topic cannot be empty"],
+        )
+
+    # Remove control characters (except newlines)
+    cleaned = re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]", "", cleaned)
+
+    # Check for dangerous patterns (case-insensitive)
+    topic_lower = cleaned.lower()
+    for pattern in DANGEROUS_PATTERNS:
+        if pattern.lower() in topic_lower:
+            return SanitizationResult(
+                value=cleaned,
+                is_safe=False,
+                warnings=[f"Topic contains potentially unsafe pattern: '{pattern}'"],
+            )
+
+    return SanitizationResult(
+        value=cleaned,
+        is_safe=True,
+        warnings=warnings,
+    )
+
+
+def sanitize_variables(variables: dict) -> dict:
+    """Sanitize a dictionary of template variables.
+
+    Args:
+        variables: Dictionary of variable name -> value
+
+    Returns:
+        Sanitized dictionary with cleaned values
+    """
+    sanitized = {}
+
+    for key, value in variables.items():
+        if isinstance(value, str):
+            # Remove control characters but preserve newlines
+            cleaned = re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]", "", value)
+            sanitized[key] = cleaned
+        else:
+            sanitized[key] = value
+
+    return sanitized

yamlgraph/utils/template.py

@@ -0,0 +1,102 @@
+"""Template utilities - Variable extraction and validation.
+
+This module provides functions to extract required variables from
+prompt templates and validate that all required variables are provided
+before execution.
+
+Supports both simple {variable} placeholders and Jinja2 templates.
+"""
+
+import logging
+import re
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+def extract_variables(template: str) -> set[str]:
+    """Extract all variable names required by a template.
+
+    Handles both simple {var} and Jinja2 {{ var }}, {% for x in var %} syntax.
+
+    Args:
+        template: Template string with placeholders
+
+    Returns:
+        Set of variable names required by the template
+
+    Examples:
+        >>> extract_variables("Hello {name}")
+        {'name'}
+
+        >>> extract_variables("{% for item in items %}{{ item }}{% endfor %}")
+        {'items'}
+    """
+    variables: set[str] = set()
+
+    # Simple format: {var} - but NOT {{ (Jinja2)
+    # Match {word} but not {{word}} - use negative lookbehind/lookahead
+    simple_pattern = r"(?<!\{)\{(\w+)\}(?!\})"
+    variables.update(re.findall(simple_pattern, template))
+
+    # Jinja2 variable: {{ var }} or {{ var.field }}
+    jinja_var_pattern = r"\{\{\s*(\w+)"
+    variables.update(re.findall(jinja_var_pattern, template))
+
+    # Jinja2 loop: {% for x in var %}
+    jinja_loop_pattern = r"\{%\s*for\s+\w+\s+in\s+(\w+)"
+    variables.update(re.findall(jinja_loop_pattern, template))
+
+    # Jinja2 condition: {% if var %} or {% if var.field %}
+    jinja_if_pattern = r"\{%\s*if\s+(\w+)"
+    variables.update(re.findall(jinja_if_pattern, template))
+
+    # Remove loop iteration variables (they're not inputs)
+    # e.g., in "{% for item in items %}", "item" is not required
+    loop_iter_pattern = r"\{%\s*for\s+(\w+)\s+in"
+    loop_vars = set(re.findall(loop_iter_pattern, template))
+    variables -= loop_vars
+
+    # Remove common non-input variables
+    # - state: injected by node_factory
+    # - loop: Jinja2 loop context
+    # - range: Jinja2 builtin function
+    excluded = {"state", "loop", "range", "true", "false", "none"}
+    variables -= excluded
+
+    return variables
+
+
+def validate_variables(
+    template: str,
+    provided: dict[str, Any],
+    prompt_name: str,
+) -> None:
+    """Validate that all required template variables are provided.
+
+    Raises ValueError with helpful message listing all missing variables.
+
+    Args:
+        template: Template string with placeholders
+        provided: Dictionary of provided variable values
+        prompt_name: Name of the prompt (for error messages)
+
+    Raises:
+        ValueError: If any required variables are missing
+
+    Examples:
+        >>> validate_variables("Hello {name}", {"name": "World"}, "greet")
+        # No error
+
+        >>> validate_variables("Hello {name}", {}, "greet")
+        ValueError: Missing required variable(s) for prompt 'greet': name
+    """
+    required = extract_variables(template)
+    provided_keys = set(provided.keys())
+    missing = required - provided_keys
+
+    if missing:
+        raise ValueError(
+            f"Missing required variable(s) for prompt '{prompt_name}': "
+            f"{', '.join(sorted(missing))}"
+        )