flyteplugins-agents-google 2.5.10__py3-none-any.whl

flyteplugins/agents/google/__init__.py

@@ -0,0 +1,23 @@
+"""Google ADK (Agent Development Kit) adapter for Flyte.
+
+Bring your own ``google-adk`` agent and run it durably on Flyte. ADK's ``Runner``
+owns the loop; Flyte is the runtime underneath: the tools you expose are Flyte tasks
+(durable child actions), each model turn is recorded for replay (``durable=True``),
+the run timeline renders into the task report, and ``memory_key`` gives cross-run
+conversation memory.
+
+- :func:`tool` — turn an ``@env.task`` into a Google ADK tool.
+- :func:`run_agent` — run the ADK agent loop inside your task and return the answer.
+- :func:`durable_model` — wrap a model so its turns are durable, for hand-built agent
+  trees (e.g. sub-agent transfers) passed to ``run_agent`` via ``agent=``.
+
+Set the model provider's API key in the environment (e.g. ``GOOGLE_API_KEY`` for
+Gemini) — wire it as a Flyte secret.
+"""
+
+from flyteplugins.agents.core import tool
+
+from ._durable import FlyteLlm, durable_model
+from ._run import run_agent, run_agent_sync
+
+__all__ = ["FlyteLlm", "durable_model", "run_agent", "run_agent_sync", "tool"]

flyteplugins/agents/google/_durable.py

@@ -0,0 +1,82 @@
+"""Durable model turns for Google ADK — trace the seam below the loop.
+
+ADK's ``Runner`` owns the loop, but every model turn flows through the agent's
+``BaseLlm.generate_content_async``. :class:`FlyteLlm` wraps that method so each
+(non-streaming) turn is recorded as a ``durable_step`` (a ``flyte.trace`` leaf):
+on a crash/retry the completed turns replay from their recorded ``LlmResponse``
+instead of re-calling (and re-billing) the model — while ADK still drives the loop.
+
+Streaming turns are passed through unmemoized; tool calls remain durable
+Flyte actions regardless.
+"""
+
+from __future__ import annotations
+
+import json
+import typing
+
+from flyteplugins.agents.core import durable_step, fingerprint, jsonable
+
+from google.adk.models import BaseLlm
+from google.adk.models.llm_response import LlmResponse
+
+
+class FlyteLlm(BaseLlm):
+    """A ``BaseLlm`` that records each model turn via ``durable_step`` for replay.
+
+    Wraps an inner ``BaseLlm`` (resolved from the agent's ``model``); ``model`` is set
+    to the inner model name so ADK behaves identically. Construct via
+    :func:`durable_model`.
+    """
+
+    inner: typing.Any = None
+
+    async def generate_content_async(
+        self, llm_request: typing.Any, stream: bool = False
+    ) -> typing.AsyncGenerator[typing.Any, None]:
+        if stream:
+            # Streamed turns are not memoized per-turn in this version.
+            async for response in self.inner.generate_content_async(llm_request, stream=True):
+                yield response
+            return
+
+        key = fingerprint({"model": self.model, "request": _request_key(llm_request)})
+
+        async def _run() -> list[dict]:
+            collected: list[dict] = []
+            async for response in self.inner.generate_content_async(llm_request, stream=False):
+                collected.append(response.model_dump(mode="json", exclude_none=True))
+            return collected
+
+        recorded = await durable_step(key, _run, name="adk_model_turn", dumps=json.dumps, loads=json.loads)
+        for payload in recorded:
+            yield LlmResponse.model_validate(payload)
+
+
+def _request_key(llm_request: typing.Any) -> typing.Any:
+    """A deterministic, JSON-able view of the request to key the durable turn."""
+    try:
+        return jsonable(llm_request.model_dump(mode="json", exclude_none=True))
+    except Exception:  # pragma: no cover - fall back to a coarse key
+        return jsonable(getattr(llm_request, "contents", None))
+
+
+def durable_model(model: typing.Any) -> typing.Any:
+    """Wrap ``model`` (a name string or ``BaseLlm``) so its turns are durable.
+
+    Returns a :class:`FlyteLlm` over the resolved inner model, or ``model`` unchanged
+    when it can't be wrapped (durability is best-effort, never fatal).
+    """
+    try:
+        from google.adk.models import BaseLlm as _BaseLlm
+        from google.adk.models import LLMRegistry
+
+        if isinstance(model, str):
+            inner = LLMRegistry.new_llm(model)
+        elif isinstance(model, _BaseLlm):
+            inner = model
+        else:
+            return model
+        return FlyteLlm(model=inner.model, inner=inner)
+    except Exception:  # pragma: no cover - never break a run over durability wiring
+        return model

flyteplugins/agents/google/_memory.py

@@ -0,0 +1,39 @@
+"""Cross-run memory for Google ADK — persist and restore the session transcript.
+
+ADK keeps the conversation as a list of ``Event``s on the session. For cross-run
+memory we persist those events to a keyed ``MemoryStore`` and replay them into a
+fresh session on the next run, so the agent continues the conversation. Keyed by a
+stable ``memory_key`` (a user/thread id); best-effort and never fatal.
+"""
+
+from __future__ import annotations
+
+import typing
+
+from flyteplugins.agents.core import resolve_memory
+
+# Path-addressed slot holding the thread's ADK event transcript in the MemoryStore.
+_EVENTS_PATH = "google/events.json"
+
+
+async def load_memory(memory_key: str | None) -> tuple[typing.Any, list[typing.Any]]:
+    """Return ``(store, prior_events)``; ``store`` is ``None`` when memory is off/unavailable."""
+    from google.adk.events import Event
+
+    store = await resolve_memory(memory_key)
+    if store is None:
+        return None, []
+
+    raw = await store.read_json.aio(_EVENTS_PATH)
+    events = [Event.model_validate(e) for e in (raw or [])]
+    return store, events
+
+
+async def save_memory(store: typing.Any, events: typing.Sequence[typing.Any]) -> None:
+    """Persist the session's events to the keyed store (no-op when ``store`` is ``None``)."""
+    if store is None:
+        return
+
+    payload = [e.model_dump(mode="json", exclude_none=True) for e in events]
+    await store.write_json.aio(_EVENTS_PATH, payload, actor="google-agent")
+    await store.save.aio()

flyteplugins/agents/google/_run.py

@@ -0,0 +1,205 @@
+"""``run_agent`` — run a Google ADK agent on Flyte using the SDK's own loop.
+
+ADK's ``Runner`` owns the agent loop (it drives the model + tools and yields
+``Event``s). ``run_agent`` runs that loop inside your ``@env.task``: it builds an
+``LlmAgent`` with Flyte-task tools, drives ``Runner.run_async``, renders the events
+into the Flyte report, and returns the final answer.
+
+Durability via the seam below the loop: with ``durable=True`` the agent's model is
+wrapped (:class:`FlyteLlm`) so each turn is recorded for replay. Cross-run memory
+via ``memory_key``: the session transcript is persisted to a keyed ``MemoryStore``
+and restored on the next run.
+
+API keys are read from the environment (e.g. ``GOOGLE_API_KEY``) — wire them as
+Flyte secrets.
+"""
+
+from __future__ import annotations
+
+import typing
+import uuid
+
+from flyte._task import AsyncFunctionTaskTemplate
+from flyteplugins.agents.core import ReportTimeline, abbrev, flush_report, sync_variant, tool
+
+from ._durable import durable_model
+from ._memory import load_memory, save_memory
+
+
+def _coerce_tool(t: typing.Any) -> typing.Any:
+    return tool(t) if isinstance(t, AsyncFunctionTaskTemplate) else t
+
+
+def _content_text(content: typing.Any) -> str:
+    parts = getattr(content, "parts", None) or []
+    return "".join(getattr(p, "text", "") or "" for p in parts)
+
+
+def _render(timeline: ReportTimeline, event: typing.Any) -> None:
+    content = getattr(event, "content", None)
+    for part in getattr(content, "parts", None) or []:
+        call = getattr(part, "function_call", None)
+        resp = getattr(part, "function_response", None)
+        text = getattr(part, "text", None)
+        if call is not None:
+            timeline.row(
+                icon="🛠️",
+                label=getattr(call, "name", ""),
+                meta="tool",
+                detail=abbrev(getattr(call, "args", ""), 160),
+            )
+        elif resp is not None:
+            timeline.row(
+                icon="🔧",
+                label=getattr(resp, "name", ""),
+                meta="tool result",
+                detail=abbrev(getattr(resp, "response", ""), 160),
+            )
+        elif text and text.strip():
+            timeline.row(icon="💬", label="assistant", detail=abbrev(text, 200))
+
+
+def _run_config(max_llm_calls: int | None) -> typing.Any:
+    """Build an ADK ``RunConfig`` that caps model calls; ``None`` → ADK's default (500)."""
+    if max_llm_calls is None:
+        return None
+    from google.adk.agents.run_config import RunConfig
+
+    return RunConfig(max_llm_calls=max_llm_calls)
+
+
+class _UsageSink:
+    """Tally model-turn count + token usage across ADK's event stream.
+
+    Each model-response ``Event`` carries genai ``usage_metadata`` (tool-result events
+    don't), so events with usage = model turns and we sum their token counts. Gemini's
+    ``cached_content_token_count`` is surfaced as ``cached`` (its context cache, like
+    Claude's cache-read tokens), ``thoughts_token_count`` as ``thinking`` for those models.
+    """
+
+    def __init__(self) -> None:
+        self.turns = self.prompt = self.completion = self.total = self.cached = self.thinking = 0
+
+    def add(self, event: typing.Any) -> None:
+        um = getattr(event, "usage_metadata", None)
+        if um is None:
+            return
+        self.turns += 1
+        self.prompt += getattr(um, "prompt_token_count", 0) or 0
+        self.completion += getattr(um, "candidates_token_count", 0) or 0
+        self.total += getattr(um, "total_token_count", 0) or 0
+        self.cached += getattr(um, "cached_content_token_count", 0) or 0
+        self.thinking += getattr(um, "thoughts_token_count", 0) or 0
+
+    def detail(self) -> str:
+        out = (
+            f"{self.turns} model turns · {self.prompt} prompt · "
+            f"{self.completion} completion · {self.total} total tokens"
+        )
+        if self.thinking:
+            out += f" · {self.thinking} thinking"
+        if self.cached:
+            out += f" · {self.cached} cached"
+        return out
+
+
+async def run_agent(
+    input: str,
+    *,
+    agent: typing.Any = None,
+    tools: typing.Sequence[typing.Any] = (),
+    model: str = "gemini-2.0-flash",
+    instructions: str | None = None,
+    name: str = "assistant",
+    max_llm_calls: int | None = None,
+    durable: bool = True,
+    observability: bool = True,
+    memory_key: str | None = None,
+    app_name: str = "flyte-agent",
+    user_id: str = "flyte-user",
+) -> str:
+    """Run a Google ADK agent with the given tools and prompt; return the final text.
+
+    Await this from an async task as ``await run_agent(...)``; from a sync task
+    use :func:`run_agent_sync` instead.
+
+    Call this from inside an ``@env.task`` — that task is the durable parent, and each
+    tool the agent calls runs as a durable Flyte child action. Provide either a
+    pre-built ``agent`` (an ADK ``LlmAgent``/``BaseAgent``) or ``tools`` + ``model`` +
+    ``instructions`` to have one built.
+
+    Args:
+        input: The user prompt.
+        agent: A pre-built ADK agent. Mutually exclusive with ``tools``.
+        tools: ``tool``-wrapped tools or bare ``@env.task`` templates.
+        model: Model name for the built agent (e.g. ``gemini-2.0-flash``).
+        instructions: System instruction for the built agent.
+        name: Agent name (a valid Python identifier). ADK injects this into the system
+            prompt as the model's "internal name", so it can surface in replies — keep it
+            natural (defaults to ``"assistant"``; avoid a brand-y/internal label).
+        max_llm_calls: Cap on model (LLM) calls before ADK raises
+            ``LlmCallsLimitExceededError`` (its runaway-loop guard, via
+            ``RunConfig.max_llm_calls``); ``None`` uses ADK's default of 500. Counts LLM
+            calls, not conversational turns (a tool round is ~2 calls). For a wall-clock
+            bound on the whole run, set ``timeout=`` on the enclosing ``@env.task``.
+        durable: Wrap the model so each turn is recorded/replayed via ``flyte.trace``.
+        observability: Render the run timeline into the Flyte task report.
+        memory_key: Stable id (user/thread) for cross-run memory. When set, the session
+            transcript is persisted and restored so a later run continues the conversation.
+        app_name: ADK app name (namespacing).
+        user_id: ADK user id.
+    """
+    from google.adk.agents import LlmAgent
+    from google.adk.runners import Runner
+    from google.adk.sessions.in_memory_session_service import InMemorySessionService
+    from google.genai import types
+
+    if agent is not None and tools:
+        raise ValueError("Pass either `agent` (with its own tools) or `tools`, not both.")
+
+    if agent is None:
+        agent = LlmAgent(
+            name=name,
+            model=durable_model(model) if durable else model,
+            instruction=instructions or "You are a helpful assistant.",
+            tools=[_coerce_tool(t) for t in tools],
+        )
+
+    session_service = InMemorySessionService()
+    store, prior_events = await load_memory(memory_key)
+    session_id = memory_key or uuid.uuid4().hex
+    session = await session_service.create_session(app_name=app_name, user_id=user_id, session_id=session_id)
+    for event in prior_events:
+        await session_service.append_event(session, event)
+
+    runner = Runner(agent=agent, app_name=app_name, session_service=session_service)
+    timeline = ReportTimeline() if observability else None
+    usage = _UsageSink() if observability else None
+    if timeline is not None:
+        timeline.heading("Google ADK agent")
+
+    final = ""
+    message = types.Content(role="user", parts=[types.Part.from_text(text=input)])
+    run_config = _run_config(max_llm_calls)
+    async for event in runner.run_async(
+        user_id=user_id, session_id=session_id, new_message=message, run_config=run_config
+    ):
+        if timeline is not None:
+            _render(timeline, event)
+        if usage is not None:
+            usage.add(event)
+        if event.is_final_response() and event.content is not None:
+            final = _content_text(event.content) or final
+
+    if memory_key:
+        latest = await session_service.get_session(app_name=app_name, user_id=user_id, session_id=session_id)
+        await save_memory(store, getattr(latest, "events", session.events))
+
+    if timeline is not None and usage is not None and usage.turns:
+        timeline.row(icon="📊", label="usage", meta="model", detail=usage.detail())
+    if observability:
+        await flush_report()
+    return final
+
+
+run_agent_sync = sync_variant(run_agent)

flyteplugins_agents_google-2.5.10.dist-info/METADATA

@@ -0,0 +1,96 @@
+Metadata-Version: 2.4
+Name: flyteplugins-agents-google
+Version: 2.5.10
+Summary: Run Google ADK (Agent Development Kit) agents on Flyte.
+Author-email: Samhita Alla <samhita@union.ai>
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: flyteplugins-agents-core
+Requires-Dist: google-adk>=2.0
+
+# flyteplugins-agents-google
+
+Run [Google ADK](https://github.com/google/adk-python) (Agent Development Kit)
+agents on Flyte. You keep writing ADK agents; Flyte is the runtime underneath.
+
+```bash
+pip install flyteplugins-agents-google
+```
+
+```python
+import flyte
+from flyteplugins.agents.google import tool, run_agent
+
+env = flyte.TaskEnvironment(
+    "google-agent",
+    secrets=[flyte.Secret(key="google_api_key", as_env_var="GOOGLE_API_KEY")],
+)
+
+@tool
+@env.task(cache="auto", retries=3)
+async def get_weather(city: str) -> str:
+    """Get the current weather for a city."""
+    return f"The weather in {city} is sunny, 22°C."
+
+@env.task(report=True, retries=3)
+async def city_agent(question: str) -> str:
+    return await run_agent(question, tools=[get_weather], model="gemini-2.0-flash")
+```
+
+## How it maps to Flyte
+
+- The SDK owns the loop — we don't reimplement it. ADK's `Runner` drives the
+  agent loop (model + tools, yielding `Event`s); `run_agent` builds an `LlmAgent`,
+  runs `Runner.run_async` inside your `@env.task`, and returns the final answer.
+- Tools as durable child actions. `tool` wraps an `@env.task` as the
+  Python function ADK calls; its body dispatches to `task.aio()`, so each tool call
+  runs as a durable Flyte child action. ADK derives the tool declaration from the
+  task signature.
+- Durable, replayable model turns. With `durable=True`, the agent's model is
+  wrapped (`FlyteLlm`) so each turn through `BaseLlm.generate_content_async` — the
+  seam below the loop — is recorded via `flyte.trace`. On a crash/retry, completed
+  turns replay from their recorded `LlmResponse` and tools are cache hits. (Same idea
+  as swapping OpenAI's `ModelProvider`: trace the model-call seam, not the loop.)
+- Observability: the turns and tool calls render into the task report.
+
+The API key is read from the environment (e.g. `GOOGLE_API_KEY` for Gemini, or your
+Vertex AI config), so it can't leak into task inputs — wire it as a Flyte secret.
+
+## Memory
+
+Pass `memory_key` (a user/thread id) for cross-run memory — the agent continues
+the same conversation across separate runs:
+
+```python
+await run_agent(message, model="gemini-2.0-flash", memory_key="user-alice")
+```
+
+ADK keeps the conversation as a list of `Event`s on the session; we persist those to a
+durable, keyed `MemoryStore` and restore them into a fresh session on the next run
+with the same key.
+
+## Examples
+
+See [`examples/`](examples/):
+
+- [`google_durable_agent.py`](examples/google_durable_agent.py) — a single durable
+  agent: tools as Flyte tasks, traced model turns, agent timeline in the report.
+- [`google_multi_agent.py`](examples/google_multi_agent.py) — multi-agent
+  orchestration: a planner agent decomposes a topic, researcher agents fan out in
+  parallel, an editor agent synthesizes — each agent its own durable action.
+- [`google_crash_resume.py`](examples/google_crash_resume.py) — crash & resume: the
+  task crashes on its first attempt; on retry the completed model turns replay from
+  their `flyte.trace` records and the tool calls are cache hits. Run on a backend.
+- [`google_memory.py`](examples/google_memory.py) — cross-run memory: two separate
+  runs share a `memory_key`; the agent learns a fact in run 1 and recalls it in run 2.
+- [`google_handoffs.py`](examples/google_handoffs.py) — native agent transfer: a triage
+  agent transfers to a billing or technical-support sub-agent, the whole agent tree
+  durable on Flyte. The specialist can pause on a Flyte condition (`flyte.new_condition`)
+  to have a human share details mid-conversation, then resume with them.
+
+## Conformance
+
+This adapter passes the shared `flyteplugins.agents.core.testing.assert_adapter_conforms`
+check — the same one every adapter runs — so it follows the common format
+(`tool` + `run_agent`, tool tasks wired to the resolver), shared with the
+OpenAI, Claude and Mistral adapters.

flyteplugins_agents_google-2.5.10.dist-info/RECORD

@@ -0,0 +1,8 @@
+flyteplugins/agents/google/__init__.py,sha256=GJ2Fmx2Gu5bKpZ8xikc8rqwehgMhFv6XT8O3OReTe_A,1067
+flyteplugins/agents/google/_durable.py,sha256=_Omx14XTEZePtv4KduxdfmrXhqFq5LsG38k5yNGkScs,3240
+flyteplugins/agents/google/_memory.py,sha256=J4nqiV8npRNyHjRc28H7kqtAOxfC-dv6dojTS-pdcQY,1484
+flyteplugins/agents/google/_run.py,sha256=9Ja9CfiqAhiZNaOpsjc_qXu-ypEPgPRdWca778jGMSs,8610
+flyteplugins_agents_google-2.5.10.dist-info/METADATA,sha256=MWz_OzW7s40bE964JEy85KGrBHS5ARJ4DxWmXuSrlqI,4219
+flyteplugins_agents_google-2.5.10.dist-info/WHEEL,sha256=K260EYznzXsJYBQGqmI8VTxEdiZYNvDZwW9cBh9-_MA,91
+flyteplugins_agents_google-2.5.10.dist-info/top_level.txt,sha256=cgd779rPu9EsvdtuYgUxNHHgElaQvPn74KhB5XSeMBE,13
+flyteplugins_agents_google-2.5.10.dist-info/RECORD,,

flyteplugins_agents_google-2.5.10.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_google-2.5.10.dist-info/top_level.txt

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