metaflow-durability 0.1.0__py3-none-any.whl

metaflow_durability/__init__.py

@@ -0,0 +1,21 @@
+from typing import TYPE_CHECKING
+from typing import Any
+
+try:
+    from metaflow_durability._version import version as __version__
+except ImportError:
+    __version__ = "unknown"
+
+if TYPE_CHECKING:
+    from metaflow_durability.durability import MetaflowDurability
+
+__all__ = ["MetaflowDurability"]
+
+
+def __getattr__(name: str) -> Any:
+    # Lazy so that importing metaflow_durability does not require pydantic-ai.
+    if name == "MetaflowDurability":
+        from metaflow_durability.durability import MetaflowDurability
+
+        return MetaflowDurability
+    raise AttributeError(f"module 'metaflow_durability' has no attribute {name!r}")

metaflow_durability/_version.py

@@ -0,0 +1,24 @@
+# file generated by vcs-versioning
+# don't change, don't track in version control
+from __future__ import annotations
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+version: str
+__version__: str
+__version_tuple__: tuple[int | str, ...]
+version_tuple: tuple[int | str, ...]
+commit_id: str | None
+__commit_id__: str | None
+
+__version__ = version = '0.1.0'
+__version_tuple__ = version_tuple = (0, 1, 0)
+
+__commit_id__ = commit_id = None

metaflow_durability/agent_decorator.py

@@ -0,0 +1,346 @@
+"""AgentDecorator — Metaflow StepDecorator for the @agent pattern.
+
+Runs the full LLM loop within one Metaflow step. Tools are user-provided
+callables; schemas are auto-generated via openai-function-calling.
+
+LLM resolution (in priority order):
+1. ``client=`` parameter — user provides an OpenAI-compatible client
+2. Co-stacked inference decorator (``@llamacpp``, ``@vllm``) — detected
+   at runtime via the step's decorator list
+3. ``model=`` parameter — passed to litellm with any extra kwargs
+
+This module is imported lazily by Metaflow's extension loader. Users
+access the decorator via ``from metaflow import agent``.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Callable
+from typing import Any
+
+from metaflow.decorators import StepDecorator
+
+
+def _build_tool_schemas(
+    tools: list[Callable[..., Any]],
+) -> list[dict[str, Any]]:
+    """Generate OpenAI function-calling schemas from typed callables."""
+    from openai_function_calling import FunctionInferrer
+
+    schemas: list[dict[str, Any]] = []
+    for tool in tools:
+        if hasattr(tool, "schema"):
+            schemas.append(tool.schema)
+        else:
+            inferred = FunctionInferrer.infer_from_function_reference(tool)
+            schemas.append({"type": "function", "function": inferred.to_json_schema()})
+    return schemas
+
+
+def _build_tool_registry(
+    tools: list[Callable[..., Any]],
+) -> dict[str, Callable[..., Any]]:
+    """Map function names to callables for dispatch."""
+    return {tool.__name__: tool for tool in tools}
+
+
+def _resolve_completion_fn(
+    attributes: dict[str, Any],
+    step_decorators: list[Any],
+) -> Callable[..., Any]:
+    """Resolve the LLM completion function from available sources.
+
+    Priority:
+    1. Explicit client object → client.chat.completions.create
+    2. Co-stacked inference decorator → decorator's OpenAI client
+    3. litellm model string → litellm.completion with kwargs
+    """
+    client = attributes.get("client")
+    if client is not None:
+
+        def _client_completion(
+            messages: list[dict[str, Any]],
+            tools: list[dict[str, Any]] | None = None,
+            **kwargs: Any,
+        ) -> Any:
+            call_kwargs: dict[str, Any] = {"messages": messages}
+            if tools:
+                call_kwargs["tools"] = tools
+            return client.chat.completions.create(**call_kwargs)
+
+        return _client_completion
+
+    # Check for co-stacked inference decorators (@llamacpp, @vllm)
+    for deco in step_decorators:
+        if hasattr(deco, "get_openai_client"):
+            infer_client = deco.get_openai_client()
+
+            def _infer_completion(
+                messages: list[dict[str, Any]],
+                tools: list[dict[str, Any]] | None = None,
+                _c: Any = infer_client,
+                **kwargs: Any,
+            ) -> Any:
+                call_kwargs: dict[str, Any] = {"messages": messages}
+                if tools:
+                    call_kwargs["tools"] = tools
+                return _c.chat.completions.create(**call_kwargs)
+
+            return _infer_completion
+
+    model = attributes.get("model")
+    if model is not None:
+        import litellm
+
+        litellm.suppress_debug_info = True
+        litellm_kwargs = dict(attributes.get("litellm_kwargs") or {})
+
+        def _litellm_completion(
+            messages: list[dict[str, Any]],
+            tools: list[dict[str, Any]] | None = None,
+            **kwargs: Any,
+        ) -> Any:
+            call_kwargs: dict[str, Any] = {
+                "model": model,
+                "messages": messages,
+                **litellm_kwargs,
+            }
+            if tools:
+                call_kwargs["tools"] = tools
+            return litellm.completion(**call_kwargs)
+
+        return _litellm_completion
+
+    raise ValueError(
+        "@agent requires one of: model= (litellm string), client= (OpenAI-compatible), "
+        "or a co-stacked inference decorator (@llamacpp, @vllm)."
+    )
+
+
+def _run_agent_loop(
+    flow: Any,
+    completion_fn: Callable[..., Any],
+    system_prompt: str,
+    tool_schemas: list[dict[str, Any]],
+    tool_registry: dict[str, Callable[..., Any]],
+    max_iterations: int,
+) -> None:
+    """Run the full agent loop: LLM → tools → LLM → ... → done."""
+    from metaflow_durability.dispatch import InProcessExecutor
+
+    messages: list[dict[str, Any]] = list(flow.messages)
+    all_dispatched: list[dict[str, Any]] = list(getattr(flow, "dispatched_tools", []))
+
+    if not messages or messages[0].get("role") != "system":
+        messages = [{"role": "system", "content": system_prompt}, *messages]
+
+    executor = InProcessExecutor(tool_registry)
+
+    for iteration in range(max_iterations):
+        response = completion_fn(
+            messages=messages,
+            tools=tool_schemas if tool_schemas else None,
+        )
+
+        choice = response.choices[0]
+        message = choice.message
+
+        assistant_msg: dict[str, Any] = {"role": "assistant"}
+        if message.content:
+            assistant_msg["content"] = message.content
+        if message.tool_calls:
+            assistant_msg["tool_calls"] = [
+                {
+                    "id": tc.id,
+                    "type": "function",
+                    "function": {
+                        "name": tc.function.name,
+                        "arguments": tc.function.arguments,
+                    },
+                }
+                for tc in message.tool_calls
+            ]
+
+        messages = [*messages, assistant_msg]
+
+        if not message.tool_calls:
+            flow.messages = messages
+            flow.final_response = message.content
+            flow.iteration = iteration
+            flow.dispatched_tools = all_dispatched
+            return
+
+        parsed_calls: list[dict[str, Any]] = []
+        for tc in message.tool_calls:
+            try:
+                args = json.loads(tc.function.arguments)
+            except (json.JSONDecodeError, TypeError):
+                args = {}
+
+            parsed_calls.append(
+                {
+                    "id": tc.id,
+                    "name": tc.function.name,
+                    "arguments": args,
+                    "tier": "default",
+                }
+            )
+
+        dispatch_result = executor.execute(parsed_calls)
+        messages = [*messages, *dispatch_result.tool_messages]
+        all_dispatched.extend(dispatch_result.to_dicts())
+
+        tc_names = [d["tool_name"] for d in dispatch_result.to_dicts()]
+        print(f"  iter {iteration}: [LLM -> {', '.join(tc_names)}]")
+
+    flow.messages = messages
+    flow.final_response = None
+    flow.iteration = max_iterations
+    flow.dispatched_tools = all_dispatched
+
+
+class AgentDecorator(StepDecorator):
+    """Metaflow step decorator that runs the full agent loop within one step.
+
+    LLM source (pick one):
+    - ``model`` — litellm model string (e.g. "anthropic/claude-sonnet-4")
+    - ``client`` — pre-configured OpenAI-compatible client object
+    - neither — auto-detect from co-stacked @llamacpp / @vllm decorator
+
+    For litellm, pass provider-specific kwargs via ``litellm_kwargs``::
+
+        @agent(
+            model="bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0",
+            litellm_kwargs={"aws_region_name": "us-west-2"},
+            tools=[search],
+        )
+    """
+
+    name = "agent"
+
+    defaults: dict[str, Any] = {
+        "model": None,
+        "client": None,
+        "system_prompt": "You are a helpful assistant with access to tools.",
+        "tools": None,
+        "tool_schemas": None,
+        "litellm_kwargs": None,
+        "max_iterations": 10,
+    }  # type: ignore[assignment]
+
+    _tool_schemas: list[dict[str, Any]] = []
+    _tool_registry: dict[str, Callable[..., Any]] = {}
+    _step_decorators: list[Any] = []
+
+    def step_init(
+        self,
+        flow,
+        graph,
+        step_name,
+        decorators,
+        environment,
+        flow_datastore,
+        logger,
+    ):
+        """Resolve tool schemas/registry and capture sibling decorators."""
+        tools = self.attributes.get("tools")
+        tool_schemas = self.attributes.get("tool_schemas")
+
+        if tools:
+            self._tool_schemas = (
+                tool_schemas if tool_schemas else _build_tool_schemas(tools)
+            )
+            self._tool_registry = _build_tool_registry(tools)
+        elif tool_schemas:
+            self._tool_schemas = tool_schemas
+            self._tool_registry = {}
+        else:
+            self._tool_schemas = []
+            self._tool_registry = {}
+
+        self._step_decorators = list(decorators)
+
+        # Auto-attach @card(type="agent") if not already present
+        try:
+            from metaflow.plugins.cards.card_decorator import CardDecorator
+
+            has_agent_card = any(
+                isinstance(d, CardDecorator) and d.attributes.get("type") == "agent"
+                for d in decorators
+            )
+            if not has_agent_card:
+                card_attrs = CardDecorator.defaults.copy()
+                card_attrs["type"] = "agent"
+                card_attrs["id"] = "agent"
+                card_deco = CardDecorator(attributes=card_attrs)
+                decorators.append(card_deco)
+                card_deco.step_init(
+                    flow,
+                    graph,
+                    step_name,
+                    decorators,
+                    environment,
+                    flow_datastore,
+                    logger,
+                )
+        except ImportError:
+            pass
+
+    def task_pre_step(
+        self,
+        step_name,
+        task_datastore,
+        metadata,
+        run_id,
+        task_id,
+        flow,
+        graph,
+        retry_count,
+        max_user_code_retries,
+        ubf_context,
+        inputs,
+    ):
+        """Initialize artifacts before the step runs."""
+        if not hasattr(flow, "messages") or flow.messages is None:
+            flow.messages = []
+        if not hasattr(flow, "iteration"):
+            flow.iteration = 0
+        if not hasattr(flow, "dispatched_tools"):
+            flow.dispatched_tools = []
+        if not hasattr(flow, "final_response"):
+            flow.final_response = None
+
+    def task_decorate(
+        self,
+        step_func,
+        flow,
+        graph,
+        retry_count,
+        max_user_code_retries,
+        ubf_context,
+    ):
+        """Wrap step to run the full agent loop, then the user's step body."""
+
+        def _wrapped_step():
+            step_func()
+            completion_fn = _resolve_completion_fn(
+                self.attributes, self._step_decorators
+            )
+            _run_agent_loop(
+                flow=flow,
+                completion_fn=completion_fn,
+                system_prompt=self.attributes["system_prompt"],
+                tool_schemas=self._tool_schemas,
+                tool_registry=self._tool_registry,
+                max_iterations=self.attributes["max_iterations"],
+            )
+
+        return _wrapped_step
+
+    def task_post_step(
+        self, step_name, flow, graph, retry_count, max_user_code_retries
+    ):
+        """Ensure dispatched_tools artifact exists."""
+        if not hasattr(flow, "dispatched_tools"):
+            flow.dispatched_tools = []

metaflow_durability/dispatch.py

@@ -0,0 +1,155 @@
+"""Dispatch records — structured artifacts for tool execution within agent steps.
+
+Each tool call dispatched by the @agent decorator produces a DispatchRecord.
+Records are collected per-iteration as self.dispatched_tools on the agent step's
+task, making them queryable via the Client API without adding graph nodes.
+
+The in-process executor runs tools sequentially in the same Python process.
+The interface is designed so that swapping in Runner-based dispatch (child flow
+runs) or native self.dispatch() (graph sub-tasks) requires changing only the
+executor, not the record format.
+"""
+
+from __future__ import annotations
+
+import json
+import time
+from collections.abc import Callable
+from dataclasses import asdict
+from dataclasses import dataclass
+from dataclasses import field
+from typing import Any
+from typing import Protocol
+
+
+@dataclass(frozen=True)
+class DispatchRecord:
+    """One tool execution's full record.
+
+    Stored as a list on self.dispatched_tools per iteration. Designed to be
+    stable across executor backends: in-process, Runner, or graph-node.
+    """
+
+    tool_call_id: str
+    tool_name: str
+    arguments: dict[str, Any]
+    tier: str
+    result: str
+    started_at: float
+    finished_at: float
+    error: str | None = None
+
+    @property
+    def duration_ms(self) -> float:
+        return (self.finished_at - self.started_at) * 1000
+
+    def to_dict(self) -> dict[str, Any]:
+        d = asdict(self)
+        d["duration_ms"] = self.duration_ms
+        return d
+
+    @classmethod
+    def from_dict(cls, d: dict[str, Any]) -> DispatchRecord:
+        d = dict(d)
+        d.pop("duration_ms", None)
+        return cls(**d)
+
+
+@dataclass
+class DispatchResult:
+    """Aggregated result of dispatching all tool calls in one iteration.
+
+    Contains the individual records plus the tool messages ready to append
+    to the conversation history.
+    """
+
+    records: list[DispatchRecord] = field(default_factory=list)
+    tool_messages: list[dict[str, Any]] = field(default_factory=list)
+
+    def to_dicts(self) -> list[dict[str, Any]]:
+        return [r.to_dict() for r in self.records]
+
+    @classmethod
+    def from_dicts(cls, dicts: list[dict[str, Any]]) -> DispatchResult:
+        records = [DispatchRecord.from_dict(d) for d in dicts]
+        return cls(records=records)
+
+
+class ToolExecutor(Protocol):
+    """Protocol for tool execution backends.
+
+    In-process, Runner, and eventual graph-node executors all implement this.
+    """
+
+    def execute(
+        self,
+        tool_calls: list[dict[str, Any]],
+    ) -> DispatchResult: ...
+
+
+class InProcessExecutor:
+    """Execute tool calls sequentially in the current process.
+
+    Tools are plain callables. The executor calls each with ``**arguments``
+    from the LLM's tool_call and converts the return value to a string
+    for the tool message.
+    """
+
+    def __init__(self, tool_registry: dict[str, Callable[..., Any]]) -> None:
+        self._registry = tool_registry
+
+    def execute(
+        self,
+        tool_calls: list[dict[str, Any]],
+    ) -> DispatchResult:
+        records: list[DispatchRecord] = []
+        tool_messages: list[dict[str, Any]] = []
+
+        for tc in tool_calls:
+            tool_name = tc["name"]
+            arguments = tc["arguments"]
+            tool_call_id = tc["id"]
+            tier = tc.get("tier", "default")
+
+            started_at = time.monotonic()
+            error: str | None = None
+            result_str: str
+
+            fn = self._registry.get(tool_name)
+            if fn is None:
+                result_str = f"Unknown tool: {tool_name}"
+                error = result_str
+            else:
+                try:
+                    result = fn(**arguments)
+                    result_str = (
+                        result if isinstance(result, str) else json.dumps(result)
+                    )
+                except Exception as exc:
+                    result_str = f"Error executing {tool_name}: {exc}"
+                    error = str(exc)
+
+            finished_at = time.monotonic()
+
+            records.append(
+                DispatchRecord(
+                    tool_call_id=tool_call_id,
+                    tool_name=tool_name,
+                    arguments=arguments,
+                    tier=tier,
+                    result=result_str,
+                    started_at=started_at,
+                    finished_at=finished_at,
+                    error=error,
+                )
+            )
+
+            tool_messages.append(
+                {
+                    "role": "tool",
+                    "tool_call_id": tool_call_id,
+                    "content": result_str,
+                }
+            )
+
+        return DispatchResult(records=records, tool_messages=tool_messages)

metaflow_durability/durability.py

@@ -0,0 +1,115 @@
+"""pydantic-ai durable execution backend for Metaflow.
+
+Attach `MetaflowDurability()` to a pydantic-ai `Agent`. Inside a running Metaflow task,
+every durable operation (model request, tool call, tool discovery, ...) is journaled to
+the task's `agent_journal` artifact. On `@retry`, operations recorded by the previous
+attempt are replayed from the journal instead of re-executed. Outside a Metaflow task
+the capability is transparent.
+"""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Awaitable
+from collections.abc import Callable
+from typing import Any
+
+from pydantic_ai.durable_exec import JSON_CODEC
+from pydantic_ai.durable_exec import BaseDurabilityCapability
+from pydantic_ai.durable_exec import DurabilityEngineSpec
+from pydantic_ai.durable_exec import DurableOperationId
+from pydantic_ai.durable_exec import JournalCallableOperationBackend
+from pydantic_ai.durable_exec import RoleBasedOperationConfig
+
+from metaflow_durability import journal
+
+
+class MetaflowSerializationError(Exception):
+    """A durable operation's result could not be encoded for the journal."""
+
+
+class _MetaflowBackend(JournalCallableOperationBackend[None]):
+    def __init__(self, agent_name: str, default_model_id: str | None) -> None:
+        super().__init__(
+            agent_name=agent_name,
+            default_model_id=default_model_id,
+            config=RoleBasedOperationConfig(
+                model=None, event=None, capability=None, tool=None
+            ),
+        )
+
+    async def execute(
+        self,
+        *,
+        operation_id: DurableOperationId,
+        name: str,
+        body: Callable[[], Awaitable[object]],
+        cache_key: tuple[object, ...],
+        config: None,
+    ) -> object:
+        run = journal.current_run.get()
+        if run is None:
+            return await body()
+        kind = type(operation_id).__name__
+        input_hash = journal.input_hash(cache_key)
+        prior = run.lookup(name, input_hash)
+        if prior is not None:
+            run.record(
+                journal.JournalEntry(
+                    name=name,
+                    kind=kind,
+                    input_hash=input_hash,
+                    payload=prior["payload"],
+                    duration_ms=prior["duration_ms"],
+                    replayed=True,
+                )
+            )
+            return prior["payload"]
+        started = time.monotonic()
+        payload = await body()
+        run.record(
+            journal.JournalEntry(
+                name=name,
+                kind=kind,
+                input_hash=input_hash,
+                payload=payload,
+                duration_ms=(time.monotonic() - started) * 1000,
+                replayed=False,
+            )
+        )
+        return payload
+
+
+class MetaflowDurability(BaseDurabilityCapability[Any]):
+    """Journal and replay a pydantic-ai agent's operations inside a Metaflow step."""
+
+    engine_spec = DurabilityEngineSpec(
+        engine_name="Metaflow",
+        durable_unit_noun="operation",
+        durable_container_noun="task",
+        codec=JSON_CODEC,
+        serialization_failure=lambda exc: MetaflowSerializationError(str(exc)),
+        toolset_lifecycles={
+            "function": "enter-always",
+            "mcp": "enter-in-durable-unit",
+            "dynamic": "enter-never",
+        },
+        sequential_tools_in_durable_context=True,
+    )
+
+    @property
+    def in_durable_context(self) -> bool:
+        return journal.active() is not None
+
+    def get_durable_operation_backend(self) -> _MetaflowBackend:
+        return _MetaflowBackend(self.name, self.default_model_id)
+
+    async def wrap_run(self, ctx: Any, *, handler: Any) -> Any:
+        task_journal = journal.active()
+        if task_journal is None:
+            return await super().wrap_run(ctx, handler=handler)
+        token = journal.current_run.set(task_journal.start_run(self.name, self.agent))
+        try:
+            return await super().wrap_run(ctx, handler=handler)
+        finally:
+            journal.current_run.reset(token)