yamlgraph 0.3.9__py3-none-any.whl

yamlgraph/executor_async.py

@@ -0,0 +1,288 @@
+"""Async Prompt Executor - Async interface for LLM calls.
+
+This module provides async versions of execute_prompt for use in
+async contexts like web servers or concurrent pipelines.
+
+Note: This is a foundation module. The underlying LLM calls still
+use sync HTTP clients wrapped with run_in_executor.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from collections.abc import AsyncIterator
+from typing import TYPE_CHECKING, TypeVar
+
+from langchain_core.messages import HumanMessage, SystemMessage
+from pydantic import BaseModel
+
+from yamlgraph.config import DEFAULT_TEMPERATURE
+from yamlgraph.executor import format_prompt, load_prompt
+from yamlgraph.utils.llm_factory import create_llm
+from yamlgraph.utils.llm_factory_async import invoke_async
+
+if TYPE_CHECKING:
+    from langgraph.graph.state import CompiledStateGraph
+from yamlgraph.utils.template import validate_variables
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T", bound=BaseModel)
+
+
+async def execute_prompt_async(
+    prompt_name: str,
+    variables: dict | None = None,
+    output_model: type[T] | None = None,
+    temperature: float = DEFAULT_TEMPERATURE,
+    provider: str | None = None,
+) -> T | str:
+    """Execute a YAML prompt asynchronously.
+
+    Async version of execute_prompt for use in async contexts.
+
+    Args:
+        prompt_name: Name of the prompt file (without .yaml)
+        variables: Variables to substitute in the template
+        output_model: Optional Pydantic model for structured output
+        temperature: LLM temperature setting
+        provider: LLM provider ("anthropic", "mistral", "openai")
+
+    Returns:
+        Parsed Pydantic model if output_model provided, else raw string
+
+    Example:
+        >>> result = await execute_prompt_async(
+        ...     "greet",
+        ...     variables={"name": "World"},
+        ...     output_model=GenericReport,
+        ... )
+    """
+    variables = variables or {}
+
+    # Load and validate prompt (sync - file I/O is fast)
+    prompt_config = load_prompt(prompt_name)
+
+    full_template = prompt_config.get("system", "") + prompt_config.get("user", "")
+    validate_variables(full_template, variables, prompt_name)
+
+    # Extract provider from YAML metadata if not provided
+    if provider is None and "provider" in prompt_config:
+        provider = prompt_config["provider"]
+        logger.debug(f"Using provider from YAML metadata: {provider}")
+
+    system_text = format_prompt(prompt_config.get("system", ""), variables)
+    user_text = format_prompt(prompt_config["user"], variables)
+
+    messages = []
+    if system_text:
+        messages.append(SystemMessage(content=system_text))
+    messages.append(HumanMessage(content=user_text))
+
+    # Create LLM (cached via factory)
+    llm = create_llm(temperature=temperature, provider=provider)
+
+    # Invoke asynchronously
+    return await invoke_async(llm, messages, output_model)
+
+
+async def execute_prompts_concurrent(
+    prompts: list[dict],
+) -> list[BaseModel | str]:
+    """Execute multiple prompts concurrently.
+
+    Useful for parallel LLM calls in pipelines.
+
+    Args:
+        prompts: List of dicts with keys:
+            - prompt_name: str (required)
+            - variables: dict (optional)
+            - output_model: Type[BaseModel] (optional)
+            - temperature: float (optional)
+            - provider: str (optional)
+
+    Returns:
+        List of results in same order as input prompts
+
+    Example:
+        >>> results = await execute_prompts_concurrent([
+        ...     {"prompt_name": "summarize", "variables": {"text": "..."}},
+        ...     {"prompt_name": "analyze", "variables": {"text": "..."}},
+        ... ])
+    """
+    tasks = []
+    for prompt_config in prompts:
+        task = execute_prompt_async(
+            prompt_name=prompt_config["prompt_name"],
+            variables=prompt_config.get("variables"),
+            output_model=prompt_config.get("output_model"),
+            temperature=prompt_config.get("temperature", DEFAULT_TEMPERATURE),
+            provider=prompt_config.get("provider"),
+        )
+        tasks.append(task)
+
+    return await asyncio.gather(*tasks)
+
+
+# ==============================================================================
+# Streaming Support (Phase 3 - Feature 004)
+# ==============================================================================
+
+
+async def execute_prompt_streaming(
+    prompt_name: str,
+    variables: dict | None = None,
+    temperature: float = DEFAULT_TEMPERATURE,
+    provider: str | None = None,
+) -> AsyncIterator[str]:
+    """Execute a YAML prompt with streaming token output.
+
+    Yields tokens as they are generated by the LLM. Does not support
+    structured output (output_model) - use execute_prompt_async for that.
+
+    Args:
+        prompt_name: Name of the prompt file (without .yaml)
+        variables: Variables to substitute in the template
+        temperature: LLM temperature setting
+        provider: LLM provider ("anthropic", "mistral", "openai")
+
+    Yields:
+        Token strings as they are generated
+
+    Example:
+        >>> async for token in execute_prompt_streaming("greet", {"name": "World"}):
+        ...     print(token, end="", flush=True)
+        Hello, World!
+    """
+    variables = variables or {}
+
+    # Load and validate prompt
+    prompt_config = load_prompt(prompt_name)
+
+    full_template = prompt_config.get("system", "") + prompt_config.get("user", "")
+    validate_variables(full_template, variables, prompt_name)
+
+    # Extract provider from YAML metadata if not provided
+    if provider is None and "provider" in prompt_config:
+        provider = prompt_config["provider"]
+        logger.debug(f"Using provider from YAML metadata: {provider}")
+
+    system_text = format_prompt(prompt_config.get("system", ""), variables)
+    user_text = format_prompt(prompt_config["user"], variables)
+
+    messages = []
+    if system_text:
+        messages.append(SystemMessage(content=system_text))
+    messages.append(HumanMessage(content=user_text))
+
+    # Create LLM (cached via factory)
+    llm = create_llm(temperature=temperature, provider=provider)
+
+    # Stream tokens
+    async for chunk in llm.astream(messages):
+        content = chunk.content
+        if content:  # Skip empty chunks
+            yield content
+
+
+# ==============================================================================
+# Async Graph Execution (Phase 2 - Feature 003)
+# ==============================================================================
+
+
+async def run_graph_async(
+    app,
+    initial_state: dict,
+    config: dict | None = None,
+) -> dict:
+    """Execute a compiled graph asynchronously.
+
+    Thin wrapper around LangGraph's ainvoke for consistent API.
+    Supports interrupt handling and Command resume.
+
+    Args:
+        app: Compiled LangGraph app (from graph.compile())
+        initial_state: Initial state dict or Command(resume=...) for resuming
+        config: LangGraph config with thread_id, e.g.
+                {"configurable": {"thread_id": "my-thread"}}
+
+    Returns:
+        Final state dict. If interrupted, contains "__interrupt__" key.
+
+    Example:
+        >>> app = load_and_compile_async("graphs/interview.yaml")
+        >>> result = await run_graph_async(
+        ...     app,
+        ...     {"query": "hello"},
+        ...     {"configurable": {"thread_id": "t1"}},
+        ... )
+        >>> if "__interrupt__" in result:
+        ...     # Handle interrupt - get user input
+        ...     result = await run_graph_async(
+        ...         app,
+        ...         Command(resume="user answer"),
+        ...         {"configurable": {"thread_id": "t1"}},
+        ...     )
+    """
+    config = config or {}
+    return await app.ainvoke(initial_state, config)
+
+
+def compile_graph_async(
+    graph,
+    config,
+) -> CompiledStateGraph:
+    """Compile a StateGraph with async-compatible checkpointer.
+
+    Uses async_mode=True when fetching checkpointer to get
+    AsyncRedisSaver instead of RedisSaver.
+
+    Args:
+        graph: StateGraph instance
+        config: GraphConfig with optional checkpointer field
+
+    Returns:
+        Compiled graph ready for ainvoke()
+    """
+    from yamlgraph.storage.checkpointer_factory import get_checkpointer
+
+    checkpointer_config = getattr(config, "checkpointer", None)
+    checkpointer = get_checkpointer(checkpointer_config, async_mode=True)
+
+    return graph.compile(checkpointer=checkpointer)
+
+
+async def load_and_compile_async(path: str) -> CompiledStateGraph:
+    """Load YAML and compile to async-ready graph.
+
+    Convenience function combining load_graph_config, compile_graph,
+    and compile_graph_async.
+
+    Args:
+        path: Path to YAML graph definition
+
+    Returns:
+        Compiled graph ready for ainvoke()
+
+    Example:
+        >>> app = await load_and_compile_async("graphs/interview.yaml")
+        >>> result = await run_graph_async(app, {"input": "hi"}, config)
+    """
+    from yamlgraph.graph_loader import compile_graph, load_graph_config
+
+    config = load_graph_config(path)
+    logger.info(f"Loaded graph config: {config.name} v{config.version}")
+
+    state_graph = compile_graph(config)
+    return compile_graph_async(state_graph, config)
+
+
+__all__ = [
+    "execute_prompt_async",
+    "execute_prompt_streaming",
+    "execute_prompts_concurrent",
+    "run_graph_async",
+    "compile_graph_async",
+    "load_and_compile_async",
+]