flyteplugins-agents-core 2.5.9__py3-none-any.whl

flyteplugins/agents/core/__init__.py

@@ -0,0 +1,42 @@
+"""flyteplugins-agents-core — the shared contract every agent-SDK adapter implements.
+
+This package holds the SDK-agnostic machinery that makes Flyte the durable
+runtime under any agent framework, so each ``flyteplugins-agents-<sdk>`` adapter
+stays thin and consistent:
+
+- :func:`durable_step` — record a call as a durable, replayable ``flyte.trace``
+  leaf (the model-turn durability mechanism), keyed by a fingerprint.
+- :func:`fingerprint` — a deterministic memo key from a request payload.
+- :class:`ToolTaskResolver` / :func:`attach_tool_resolver` — make a tool-backing
+  task resolve to itself on the worker.
+- :class:`ReportTimeline` — render agent events into the Flyte task report.
+- :mod:`flyteplugins.agents.core.testing` — :func:`assert_adapter_conforms`, the
+  CI-enforced conformance check every adapter runs.
+
+The division of labor every adapter follows: the agent run is a Flyte
+``@env.task`` (durable parent), each model turn is a ``flyte.trace`` (a
+memoized, replayable leaf), and each tool is a Flyte task invoked as a
+durable child action.
+"""
+
+from ._durable import durable_step
+from ._fingerprint import fingerprint, jsonable
+from ._memory import resolve_memory
+from ._observability import ReportTimeline, abbrev, duration_ms, flush_report
+from ._tools import ToolTaskResolver, attach_tool_resolver, coerce_tool_args, task_json_schema, tool
+
+__all__ = [
+    "ReportTimeline",
+    "ToolTaskResolver",
+    "abbrev",
+    "attach_tool_resolver",
+    "coerce_tool_args",
+    "durable_step",
+    "duration_ms",
+    "fingerprint",
+    "flush_report",
+    "jsonable",
+    "resolve_memory",
+    "task_json_schema",
+    "tool",
+]

flyteplugins/agents/core/_durable.py

@@ -0,0 +1,54 @@
+"""Durable, replayable steps via ``flyte.trace`` — the shared mechanism.
+
+A model turn (or any expensive, replay-worthy call) is recorded as a
+``flyte.trace`` leaf so that, inside a Flyte task, a crash/retry replays the
+recorded result instead of re-running it. The real call is usually
+non-serializable (closures, SDK objects), but ``flyte.trace`` keys its memo on
+the decorated function's serializable arguments. :func:`durable_step` resolves
+that by capturing the real call in a closure and feeding the trace only a
+deterministic ``request_key``.
+"""
+
+from __future__ import annotations
+
+import typing
+
+import flyte
+
+T = typing.TypeVar("T")
+
+
+async def durable_step(
+    request_key: str,
+    run: typing.Callable[[], typing.Awaitable[typing.Any]],
+    *,
+    name: str = "durable_step",
+    dumps: typing.Callable[[typing.Any], str] = lambda v: v,
+    loads: typing.Callable[[str], typing.Any] = lambda v: v,
+) -> typing.Any:
+    """Run ``run()`` once as a durable, replayable trace step keyed by ``request_key``.
+
+    The real (possibly non-serializable) work is captured in the ``run`` closure,
+    so the traced function only ever sees the serializable ``request_key``. The
+    result is serialized with ``dumps`` for the trace record (a ``str`` is stored
+    inline and is human-readable in the Flyte UI) and rebuilt with ``loads`` on
+    the way out and on replay.
+
+    Args:
+        request_key: A deterministic fingerprint of the call; the trace memo key.
+        run: Zero-arg async callable performing the real work.
+        name: Label for the trace action in the Flyte UI.
+        dumps: Serialize the result to a ``str`` for the trace record.
+        loads: Rebuild the result from the recorded ``str``.
+
+    Outside a task context ``flyte.trace`` is a transparent pass-through, so this
+    also works unchanged for local runs and unit tests.
+    """
+
+    async def _step(key: str) -> str:
+        return dumps(await run())
+
+    # Name the closure so the trace action reads meaningfully in the UI.
+    _step.__name__ = name
+    _step.__qualname__ = name
+    return loads(await flyte.trace(_step)(request_key))

flyteplugins/agents/core/_fingerprint.py

@@ -0,0 +1,33 @@
+"""Deterministic request fingerprints for keying durable steps."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import typing
+
+
+def fingerprint(payload: typing.Mapping[str, typing.Any]) -> str:
+    """A deterministic sha256 hex of a request ``payload``.
+
+    Must be a pure function of the semantic request so replays line up. Pass a
+    mapping of the request's identifying fields (e.g. system prompt, input
+    items, model, tool names) — not callables or live objects. Anything not
+    natively JSON-serializable is coerced with ``str``.
+    """
+    blob = json.dumps(payload, sort_keys=True, default=str)
+    return hashlib.sha256(blob.encode("utf-8")).hexdigest()
+
+
+def jsonable(obj: typing.Any) -> typing.Any:
+    """Best-effort conversion of an SDK object to something JSON-dumpable."""
+    if obj is None or isinstance(obj, (str, int, float, bool)):
+        return obj
+    for attr in ("to_json_dict", "model_dump"):
+        method = getattr(obj, attr, None)
+        if callable(method):
+            try:
+                return method()
+            except Exception:  # pragma: no cover - defensive
+                pass
+    return str(obj)

flyteplugins/agents/core/_memory.py

@@ -0,0 +1,54 @@
+"""Cross-run agent memory — a thin handle over Flyte's keyed ``MemoryStore``.
+
+"Agent memory" = durable state addressable by a stable ``memory_key`` (a user or
+thread id), shared across runs and workers. We don't reinvent it: Flyte already
+ships a keyed, blob-store-backed store (``flyte.ai.agents.memory.MemoryStore``)
+that resolves a deterministic remote path from the key and the active run context
+— stripping the per-run scratch so two runs with the same key share one store,
+and encapsulating all the storage details.
+
+It is decoupled from the agent loop - importing it does not pull in any agent
+harness — so adapters use it purely as a durable store. Each adapter then maps its
+SDK's own conversation state onto the returned store.
+
+The store carries two complementary surfaces, so one ``memory_key`` covers both
+kinds of memory:
+
+- ``messages`` (``append``/``extend``) — the conversation transcript;
+- path-addressed ``read_json``/``write_json``/``list_paths`` — durable named facts
+  (the ``remember`` / ``recall`` substrate), with audit + version history.
+"""
+
+from __future__ import annotations
+
+import typing
+
+from flyte._logging import logger
+
+if typing.TYPE_CHECKING:
+    from flyte.ai.agents.memory import MemoryStore
+
+
+async def resolve_memory(memory_key: str | None, *, audit: bool = True) -> "MemoryStore | None":
+    """Open (or create) the keyed agent-memory store, or ``None`` if unavailable.
+
+    Best-effort: returns ``None`` (and logs) when ``memory_key`` is falsy or no
+    durable store can be resolved — e.g. no Flyte context/org, or an invalid key —
+    so memory never breaks a run. Call from inside an ``@env.task``; the store's
+    remote path is derived from the run context and is stable across runs for the
+    same key.
+
+    Args:
+        memory_key: A stable single-segment id (a user/thread id). ``None`` or
+            empty disables memory.
+        audit: Keep the store's append-only audit log (the ``MemoryStore`` default).
+    """
+    if not memory_key:
+        return None
+    try:
+        from flyte.ai.agents.memory import MemoryStore
+
+        return await MemoryStore.get_or_create.aio(key=memory_key, audit=audit)
+    except Exception:  # pragma: no cover - memory is best-effort, never fatal
+        logger.warning("Could not open agent memory for key %r; continuing without memory.", memory_key)
+        return None

flyteplugins/agents/core/_observability.py

@@ -0,0 +1,35 @@
+"""Shared report timeline — render agent events into the Flyte task report.
+
+Each adapter maps its own SDK's trace/span/event shape onto :class:`ReportTimeline`
+rows; the row formatting, abbreviation, and best-effort error handling now live in
+``flyte.report`` (:class:`flyte.report.Timeline`), so the native ``flyte.ai.agents``
+loop and every adapter render through one implementation. This module keeps the
+adapter-facing names (``ReportTimeline`` defaulting to the ``Agent`` tab, plus the
+``abbrev``/``duration_ms``/``flush_report`` helpers) stable.
+"""
+
+from __future__ import annotations
+
+import flyte.report
+from flyte.report import Timeline, duration_ms
+from flyte.report import abbreviate as abbrev
+
+__all__ = ["ReportTimeline", "abbrev", "duration_ms", "flush_report"]
+
+
+class ReportTimeline(Timeline):
+    """A :class:`flyte.report.Timeline` that defaults to the ``Agent`` report tab."""
+
+    def __init__(self, tab_name: str = "Agent"):
+        super().__init__(tab_name)
+
+
+async def flush_report() -> None:
+    """Flush the active Flyte report — a best-effort no-op when there is none.
+
+    Adapters call this once after a run so the rendered timeline is published.
+    """
+    try:
+        await flyte.report.flush()
+    except Exception:  # pragma: no cover - no active report / local run
+        pass

flyteplugins/agents/core/_tools.py

@@ -0,0 +1,181 @@
+"""Make a tool-backing Flyte task resolve to itself on the worker.
+
+Every adapter stacks its ``tool`` on top of ``@env.task``, which rebinds
+the module attribute to the tool and shadows the task. Without a guard, the
+worker's default resolver loads the tool, the task runner calls the tool's
+``execute``, and the task re-dispatches itself — recursing without end.
+
+:class:`ToolTaskResolver` recovers the real task via the tool's
+``__wrapped_task__`` hook; :func:`attach_tool_resolver` wires it onto the task.
+Each adapter's tool object only has to expose ``__wrapped_task__`` returning its
+backing :class:`~flyte._task.TaskTemplate`.
+"""
+
+from __future__ import annotations
+
+import functools
+import typing
+from functools import partial
+
+from flyte._internal.resolvers.default import DefaultTaskResolver
+from flyte._task import AsyncFunctionTaskTemplate, TaskTemplate
+
+
+class ToolTaskResolver(DefaultTaskResolver):
+    """Resolver for a task shadowed at module scope by a tool wrapper.
+
+    Recovers the underlying task via the wrapper's ``__wrapped_task__`` hook so
+    the worker runs the task's own body instead of re-dispatching the tool.
+    """
+
+    @property
+    def import_path(self) -> str:
+        return "flyteplugins.agents.core._tools.ToolTaskResolver"
+
+    def load_task(self, loader_args):  # type: ignore[override]
+        task_def = super().load_task(loader_args)
+        if isinstance(task_def, TaskTemplate):
+            return task_def
+        unwrapped = getattr(task_def, "__wrapped_task__", None)
+        if isinstance(unwrapped, TaskTemplate):
+            return unwrapped
+        return task_def
+
+
+def attach_tool_resolver(task: typing.Any) -> None:
+    """Point a tool-backing ``@env.task`` at :class:`ToolTaskResolver`.
+
+    No-op for anything that isn't an async-function task or that already declares
+    a custom resolver, so the default resolver is left untouched elsewhere.
+    """
+    if isinstance(task, AsyncFunctionTaskTemplate) and task.task_resolver is None:
+        task.task_resolver = ToolTaskResolver()
+
+
+def coerce_tool_args(task: AsyncFunctionTaskTemplate, kwargs: dict[str, typing.Any]) -> dict[str, typing.Any]:
+    """Coerce LLM-supplied tool args to the task's declared parameter types.
+
+    LLMs emit JSON numbers without a fractional part as ``int`` (e.g. ``42``), but
+    Flyte's type engine rejects an ``int`` where a ``float`` is declared — so the tool
+    call fails during input conversion, before the child action is even submitted (the
+    action never appears in the UI, and the task body never runs). This converts ``int``
+    -> ``float`` for float-annotated params (leaving ``bool`` alone). Best-effort: returns
+    the kwargs unchanged if the task's annotations can't be resolved.
+    """
+    try:
+        hints = typing.get_type_hints(task.func)
+    except Exception:  # pragma: no cover - unresolved annotations; pass through
+        return kwargs
+    coerced = dict(kwargs)
+    for name, value in kwargs.items():
+        if hints.get(name) is float and isinstance(value, int) and not isinstance(value, bool):
+            coerced[name] = float(value)
+    return coerced
+
+
+def tool(
+    func: AsyncFunctionTaskTemplate | typing.Callable | None = None,
+    *,
+    name: str | None = None,
+    description: str | None = None,
+) -> typing.Callable:
+    """Wrap a Flyte ``@env.task`` as a plain async tool function — the generic default.
+
+    For SDKs that accept plain Python callables as tools (deriving the schema from the
+    signature + docstring), this is the whole adapter ``tool``: the returned
+    function carries the task's signature (``functools.wraps``), dispatches to
+    ``task.aio()`` (so each call is a durable Flyte child action), exposes
+    ``__wrapped_task__``, and wires the backing task to :class:`ToolTaskResolver`.
+    Adapters whose SDK needs a native tool type (e.g. OpenAI's
+    ``FunctionTool``, Claude's MCP ``SdkMcpTool``) provide their own instead.
+
+    Also accepts any other callable — a plain function or an instance of a callable
+    class defining ``__call__`` — and returns it usable as a tool as-is, since the
+    plain-callable SDKs derive the schema by inspecting the callable (a class instance
+    is inspected through its ``__call__``). A ``name`` or ``description`` override is
+    applied to the callable best-effort.
+
+    Usable bare, parametrized or as a direct call::
+
+        @tool
+        @env.task
+        async def get_weather(city: str) -> str: ...
+    """
+
+    if func is None:
+        return partial(tool, name=name, description=description)
+    if isinstance(func, AsyncFunctionTaskTemplate):
+        return _task_to_tool(func, name=name, description=description)
+    if not callable(func):
+        raise TypeError(f"tool() expects a Flyte @env.task or a callable, got {type(func).__name__!r}.")
+    # A plain function or a callable class instance is already usable as a tool as-is.
+    return _relabel_callable(func, name=name, description=description)
+
+
+def _relabel_callable(
+    func: typing.Callable,
+    *,
+    name: str | None = None,
+    description: str | None = None,
+) -> typing.Callable:
+    """Apply ``name`` / ``description`` overrides to a callable in place, best-effort.
+
+    Functions accept ``__name__`` / ``__doc__`` assignment; a callable class instance
+    may reject it (e.g. ``__slots__``), in which case the override is silently skipped —
+    the callable is still returned and usable as a tool.
+    """
+    if name:
+        try:
+            func.__name__ = name  # type: ignore[attr-defined]
+        except (
+            AttributeError,
+            TypeError,
+        ):  # pragma: no cover - slotted/immutable callable
+            pass
+    if description:
+        try:
+            func.__doc__ = description
+        except (
+            AttributeError,
+            TypeError,
+        ):  # pragma: no cover - slotted/immutable callable
+            pass
+    return func
+
+
+def _task_to_tool(
+    task: AsyncFunctionTaskTemplate,
+    *,
+    name: str | None = None,
+    description: str | None = None,
+) -> typing.Callable:
+    inner = task.func
+
+    @functools.wraps(inner)
+    async def tool(*args: typing.Any, **kwargs: typing.Any) -> typing.Any:
+        # In a Flyte task context this submits a durable child action; locally it runs
+        # inline. ``functools.wraps`` keeps ``inner``'s signature so the SDK builds the
+        # right tool declaration; ``coerce_tool_args`` relaxes LLM int->float args.
+        return await task.aio(*args, **coerce_tool_args(task, kwargs))
+
+    if name:
+        tool.__name__ = name
+    if description:
+        tool.__doc__ = description
+
+    # The wrapper shadows the task at module scope; expose the real task and wire the
+    # shared resolver so it resolves to itself on the worker (no recursion).
+    tool.__wrapped_task__ = task  # type: ignore[attr-defined]
+    tool.task = task  # type: ignore[attr-defined]
+    attach_tool_resolver(task)
+    return tool
+
+
+def task_json_schema(task: AsyncFunctionTaskTemplate) -> dict[str, typing.Any]:
+    """The JSON schema of a Flyte task's inputs, via the Flyte type engine.
+
+    Useful for adapters whose SDK wants a JSON-schema tool definition and that
+    prefer Flyte's type-engine schema (correct ``Literal`` enums, ``File``/``Dir``
+    /``DataFrame``, dataclasses) over the SDK's own signature inspection.
+    """
+    return task.json_schema

flyteplugins/agents/core/testing.py

@@ -0,0 +1,65 @@
+"""Conformance harness — enforce the common adapter format.
+
+Every ``flyteplugins.agents.<sdk>`` adapter ships a one-line test::
+
+    from flyteplugins.agents.core.testing import assert_adapter_conforms
+    import flyteplugins.agents.openai as adapter
+
+    def test_conformance():
+        assert_adapter_conforms(adapter)
+
+CI then fails if an adapter drifts from the shared format.
+"""
+
+from __future__ import annotations
+
+import inspect
+import typing
+
+# The keyword surface every ``run_agent`` must accept, so the call shape is the
+# same across SDKs.
+REQUIRED_RUN_AGENT_PARAMS = ("tools", "model", "instructions", "durable", "observability", "memory_key")
+
+
+def assert_adapter_conforms(adapter: typing.Any) -> None:
+    """Assert that an adapter module implements the common agent-adapter contract.
+
+    The contract:
+
+    1. exports a callable ``tool`` that turns an ``@env.task`` into the
+       SDK's tool type, attaching :class:`~flyteplugins.agents.core.ToolTaskResolver`
+       and exposing ``__wrapped_task__`` (so the task does not self-recurse on the
+       worker);
+    2. exports an async ``run_agent`` accepting the standard keyword surface.
+
+    Raises ``AssertionError`` with a specific message on any deviation.
+    """
+    import flyte
+
+    from flyteplugins.agents.core import ToolTaskResolver
+
+    name = getattr(adapter, "__name__", repr(adapter))
+
+    tool = getattr(adapter, "tool", None)
+    assert callable(tool), f"{name}: must export a callable `tool`"
+
+    run_agent = getattr(adapter, "run_agent", None)
+    assert inspect.iscoroutinefunction(run_agent), f"{name}: must export an async `run_agent`"
+    params = inspect.signature(run_agent).parameters
+    for required in REQUIRED_RUN_AGENT_PARAMS:
+        assert required in params, f"{name}: `run_agent` must accept `{required}`"
+
+    # tool on a task must expose the real task and wire the resolver.
+    env = flyte.TaskEnvironment("agents-core-conformance")
+
+    @tool
+    @env.task
+    def _sample(city: str) -> str:
+        """A sample tool task."""
+        return city
+
+    wrapped = getattr(_sample, "__wrapped_task__", None)
+    assert wrapped is not None, f"{name}: `tool` result must expose `__wrapped_task__`"
+    assert isinstance(getattr(wrapped, "task_resolver", None), ToolTaskResolver), (
+        f"{name}: the tool's backing task must use ToolTaskResolver (or it self-recurses on the worker)"
+    )

flyteplugins_agents_core-2.5.9.dist-info/METADATA

@@ -0,0 +1,44 @@
+Metadata-Version: 2.4
+Name: flyteplugins-agents-core
+Version: 2.5.9
+Summary: Shared contract for Flyte agent-SDK adapters.
+Author-email: Samhita Alla <samhita@union.ai>
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: flyte
+
+# flyteplugins-agents-core
+
+The shared contract every Flyte agent-SDK adapter is built on, so they all
+follow one common format.
+
+It is SDK-agnostic — no agent SDK is a dependency here. Adapters depend on it and
+implement the contract; `assert_adapter_conforms` enforces that in CI.
+
+## What's in it
+
+- `durable_step(request_key, run, *, name, dumps, loads)` — record a call as a
+  durable, replayable `flyte.trace` leaf (the model-turn durability mechanism).
+  The real, non-serializable call is captured in the `run` closure, so the trace
+  only ever sees the serializable `request_key`.
+- `fingerprint(payload)` — a deterministic memo key from a request payload.
+- `ToolTaskResolver` / `attach_tool_resolver(task)` — make a tool-backing task
+  resolve to itself on the worker (via the tool's `__wrapped_task__` hook)
+  instead of re-dispatching, which would otherwise recurse indefinitely.
+- `ReportTimeline` — render agent events into a tab of the Flyte task report.
+- `flyteplugins.agents.core.testing.assert_adapter_conforms(module)` — the
+  conformance check every adapter runs as a one-line test.
+
+## The rule every adapter follows
+
+The agent run is a Flyte `@env.task` (the durable parent); each model turn
+is a `flyte.trace` (a memoized, replayable leaf); each tool is a Flyte task
+invoked as a durable child action.
+
+## Writing an adapter
+
+A `flyteplugins-agents-<sdk>` package depends on this core and exposes, at
+minimum, `tool` and `run_agent` (see the conformance contract). Its tool
+object exposes `__wrapped_task__`; `run_agent` runs the SDK's loop inside the
+calling `@env.task`, wraps the model in a durable provider built on `durable_step`,
+and renders the trace via `ReportTimeline`.

flyteplugins_agents_core-2.5.9.dist-info/RECORD

@@ -0,0 +1,11 @@
+flyteplugins/agents/core/__init__.py,sha256=5ZPSn5d2i4L66u8UjcMeOjwX9WZu95Gg7CYoEMcums4,1654
+flyteplugins/agents/core/_durable.py,sha256=Pxh0S1uju4EhfbO7KwMwZmdr-tsRpF3Uhrj_MfmO-fo,2128
+flyteplugins/agents/core/_fingerprint.py,sha256=Rg3f9YkPLYaaeGeVyRN6-8IsRwkplqsFAjTqxF1ERYA,1169
+flyteplugins/agents/core/_memory.py,sha256=9xSerEnMhFRGZyC8Hu4JYvU9hdyacj6iWxU-lOlighs,2380
+flyteplugins/agents/core/_observability.py,sha256=xFZgdww7RjnnIkclAxASrCfwrd6AVt70CFavj7XlsDg,1320
+flyteplugins/agents/core/_tools.py,sha256=oKN2msr5hGpvUW0Vzy6EI4nbhkgIJyjv8GO4gnc2rRk,7383
+flyteplugins/agents/core/testing.py,sha256=k3cROYEGlS_SmO_jz3EI5kTLOx7ZYsgd2mZkFTR34ic,2398
+flyteplugins_agents_core-2.5.9.dist-info/METADATA,sha256=Z0MEvUwrwFJ1yfmFhNidHeojmi7BU8yWffumjaTD080,1972
+flyteplugins_agents_core-2.5.9.dist-info/WHEEL,sha256=K260EYznzXsJYBQGqmI8VTxEdiZYNvDZwW9cBh9-_MA,91
+flyteplugins_agents_core-2.5.9.dist-info/top_level.txt,sha256=cgd779rPu9EsvdtuYgUxNHHgElaQvPn74KhB5XSeMBE,13
+flyteplugins_agents_core-2.5.9.dist-info/RECORD,,

flyteplugins_agents_core-2.5.9.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (83.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

flyteplugins_agents_core-2.5.9.dist-info/top_level.txt

@@ -0,0 +1 @@
+flyteplugins