temporal-agent-harness 0.3.0__py3-none-any.whl

temporal_agent_harness/__init__.py

@@ -0,0 +1,15 @@
+"""Temporal-native agent harness (experimental).
+
+A Temporal-native outer agent harness that gives you the full power of Temporal's
+durable-execution primitives while letting you keep using the AI SDKs you already know
+(the inner harness), via first-class integrations under :mod:`temporal_agent_harness.ai_sdks`.
+
+Subpackages:
+  * :mod:`temporal_agent_harness.harness`  — the core harness: the agent workflow runner,
+    the agent/subagent protocol, tool definitions, and human-in-the-loop tool approvals.
+  * :mod:`temporal_agent_harness.ai_sdks`  — integrations that make AI SDK calls durable
+    Temporal activities (currently the Google Gemini SDK).
+  * :mod:`temporal_agent_harness.plugin`   — ``AgentHarnessPlugin``, the one Temporal plugin
+    that wires a client + worker for all of the above: it supplies the harness's data
+    converter and registers its activities, so a worker declares only its workflows.
+"""

temporal_agent_harness/ai_sdks/__init__.py

@@ -0,0 +1,11 @@
+"""First-class integrations between the Temporal-native harness and AI SDKs.
+
+Each integration turns an AI SDK's calls into durable Temporal activities, so workflow code
+can use the SDK it already knows while inheriting Temporal's durability, retries, and
+observability. One subpackage per SDK; more are expected over time.
+
+  * :mod:`temporal_agent_harness.ai_sdks.google_genai_plugin` — the Google Gemini SDK.
+  * :mod:`temporal_agent_harness.ai_sdks.openai_agents` — the OpenAI Agents SDK.
+  * :mod:`temporal_agent_harness.ai_sdks.pydantic_ai_harness` — Pydantic AI (glue for the
+    unmodified upstream ``pydantic_ai.durable_exec.temporal`` plugin).
+"""

temporal_agent_harness/ai_sdks/google_genai_plugin/__init__.py

@@ -0,0 +1,67 @@
+"""First-class Temporal integration for the Google Gemini SDK.
+
+.. warning::
+    This module is experimental and may change in future versions.
+    Use with caution in production environments.
+
+This integration lets you use the Gemini SDK's async client with full
+automatic function calling (AFC) support, plus a Temporal-aware shim of
+the Interactions API. Every API call becomes a **durable Temporal
+activity**. Define tools with :func:`harness.agent.tool_defn` (inline,
+deterministic in-workflow) or :func:`harness.agent.activity_tool_defn`
+(durable, activity-backed); for the latter, :func:`harness.agent.tool_activity`
+returns the activity to register on the worker.
+
+No credentials are fetched in the workflow, and no auth material appears in
+Temporal's event history.
+
+- :class:`GoogleGenAIPlugin` — registers all of the Gemini activities
+  using a caller-provided ``genai.Client`` on the worker side.
+- :func:`google_genai_client` — call from a workflow to get an
+  ``AsyncClient`` whose ``models.generate_content*``, ``files``,
+  ``file_search_stores``, and ``interactions`` modules all route through
+  Temporal activities.
+- :func:`function_param` — derive an Interactions-API ``FunctionParam``
+  tool declaration from a tool's model-facing signature. The Interactions
+  API has no AFC, so the workflow drives the tool-calling loop itself;
+  pair this helper with the ``gemini.interactions.create(...)`` call on
+  the client returned by :func:`google_genai_client`.
+
+Quickstart::
+
+    # ---- worker setup (outside the Temporal Python Sandbox) ----
+    client = genai.Client(api_key=os.environ["GOOGLE_API_KEY"])
+    plugin = GoogleGenAIPlugin(client)
+
+    @agent.activity_tool_defn()
+    async def get_weather(state: str) -> str: ...
+
+    Worker(..., activities=[agent.tool_activity(get_weather), ...])
+
+    # ---- workflow (inside the Temporal Python Sandbox) ----
+    @workflow.defn
+    class AgentWorkflow:
+        @workflow.run
+        async def run(self, query: str) -> str:
+            client = google_genai_client()
+            response = await client.models.generate_content(
+                model="gemini-2.5-flash",
+                contents=query,
+                config=types.GenerateContentConfig(tools=[get_weather]),
+            )
+            return response.text
+"""
+
+from __future__ import annotations
+
+from ._google_genai_plugin import GoogleGenAIPlugin
+from ._interactions_workflow import function_param
+from .workflow import (
+    google_genai_client,
+)
+
+__all__ = [
+    "GoogleGenAIPlugin",
+    "function_param",
+    "google_genai_client",
+]

temporal_agent_harness/ai_sdks/google_genai_plugin/_gemini_activity.py

@@ -0,0 +1,235 @@
+"""Temporal activity that executes Gemini SDK API calls with real credentials.
+
+The ``TemporalApiClient`` in the workflow dispatches calls here. This
+activity holds a user-provided ``genai.Client`` and forwards structured
+requests. Credentials are fetched/refreshed only within the activity —
+they never appear in workflow event history.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Sequence
+from contextlib import AsyncExitStack
+from typing import Any, Callable
+
+import google.auth.credentials
+from google.genai import Client as GeminiClient
+from google.genai import types
+from google.genai.types import HttpOptions
+from google.genai.types import HttpResponse as SdkHttpResponse
+
+from temporalio import activity
+
+from temporal_agent_harness.harness.agent_protocol import ReplyDelta
+from temporal_agent_harness.harness.agent_workflow import AgentWorkflowRunner
+from ._interactions_activity import make_gemini_interactions_create_streamed
+from ._models import (
+    _GeminiApiRequest,
+    _GeminiApiResponse,
+    _GeminiApiStreamedResponse,
+    _GeminiDownloadFileRequest,
+    _GeminiRegisterFilesRequest,
+    _GeminiUploadFileRequest,
+    _GeminiUploadToFileSearchStoreRequest,
+)
+
+
+def _extract_text_delta(body: str) -> str:
+    """Pull the prose text out of a single streamed Gemini response chunk.
+
+    Chunks come in as raw HTTP response bodies — usually a JSON object
+    representing one partial ``GenerateContentResponse``, sometimes
+    prefixed with an SSE ``data:`` field. Tolerant of both shapes;
+    returns ``""`` for non-text chunks (function-call parts, malformed
+    bodies, etc.).
+    """
+    body = body.strip()
+    if body.startswith("data:"):
+        body = body[5:].strip()
+    if not body:
+        return ""
+    try:
+        data = json.loads(body)
+    except json.JSONDecodeError:
+        return ""
+    text = ""
+    for cand in data.get("candidates", []) or []:
+        content = cand.get("content") or {}
+        for part in content.get("parts") or []:
+            if isinstance(part, dict):
+                t = part.get("text")
+                if isinstance(t, str):
+                    text += t
+    return text
+
+
+def _resolve_http_options(
+    overrides: Any,
+) -> HttpOptions | None:
+    """Reconstruct ``HttpOptions`` from serializable overrides, or None."""
+    if overrides is None:
+        return None
+    return HttpOptions.model_validate(overrides.model_dump(exclude_none=True))
+
+
+class GeminiApiCaller:
+    """Wraps a ``genai.Client`` and exposes Temporal activities for SDK calls.
+
+    The caller owns a reference to the user-provided ``genai.Client``.
+    All credential management, HTTP client configuration, etc. is the
+    responsibility of whoever constructs the client.
+    """
+
+    def __init__(
+        self,
+        client: GeminiClient,
+        credentials: google.auth.credentials.Credentials | None = None,
+    ) -> None:
+        """Initialize with a genai.Client and optional extra credentials."""
+        self._client = client
+        self._credentials = credentials
+
+    def activities(self) -> Sequence[Callable]:
+        """Return activities that route SDK calls through this client."""
+
+        @activity.defn
+        async def gemini_api_client_async_request(
+            req: _GeminiApiRequest,
+        ) -> _GeminiApiResponse:
+            """Execute a Gemini SDK API call with real credentials."""
+            response: SdkHttpResponse = (
+                await self._client.aio._api_client.async_request(
+                    http_method=req.http_method,
+                    path=req.path,
+                    request_dict=req.request_dict,
+                    http_options=_resolve_http_options(req.http_options_overrides),
+                )
+            )
+            return _GeminiApiResponse(
+                headers=response.headers or {},
+                body=response.body or "",
+            )
+
+        @activity.defn
+        async def gemini_api_client_async_request_streamed(
+            req: _GeminiApiRequest,
+        ) -> _GeminiApiStreamedResponse:
+            """Execute a streamed Gemini SDK API call, collecting all chunks.
+
+            When ``req.stream_turn_id`` is set, each chunk's text content is
+            republished as a ``reply_delta`` event on the parent workflow's
+            :class:`WorkflowStream` as soon as it arrives — fine-grained,
+            visible to the UI in real time. Function-call chunks (no text
+            parts) are silently skipped.
+            """
+            stream = await self._client.aio._api_client.async_request_streamed(
+                http_method=req.http_method,
+                path=req.path,
+                request_dict=req.request_dict,
+                http_options=_resolve_http_options(req.http_options_overrides),
+            )
+
+            chunks: list[_GeminiApiResponse] = []
+
+            # If a stream context rode in on the request, hand it to the
+            # harness's publisher helper. The activity never unpacks
+            # turn_id/turn_number itself — it just forwards the opaque
+            # carrier and calls ``publish`` per chunk.
+            async with AsyncExitStack() as stack:
+                publisher = None
+                if req.stream_context is not None:
+                    publisher = await stack.enter_async_context(
+                        AgentWorkflowRunner.publisher_from_activity(
+                            req.stream_context,
+                        )
+                    )
+
+                async for chunk in stream:
+                    body = chunk.body or ""
+                    chunks.append(
+                        _GeminiApiResponse(headers=chunk.headers or {}, body=body)
+                    )
+                    if publisher is not None:
+                        delta = _extract_text_delta(body)
+                        if delta:
+                            publisher.publish(ReplyDelta(text=delta))
+
+            return _GeminiApiStreamedResponse(chunks=chunks)
+
+        @activity.defn
+        async def gemini_files_upload(
+            req: _GeminiUploadFileRequest,
+        ) -> types.File:
+            """Upload a file using the real genai.Client on the worker."""
+            if req.file_bytes is not None:
+                import io
+
+                file_arg: Any = io.BytesIO(req.file_bytes)
+            else:
+                file_arg = req.file_path
+
+            return await self._client.aio.files.upload(file=file_arg, config=req.config)
+
+        @activity.defn
+        async def gemini_files_download(
+            req: _GeminiDownloadFileRequest,
+        ) -> bytes:
+            """Download a file using the real genai.Client on the worker."""
+            return await self._client.aio.files.download(
+                file=req.file, config=req.config
+            )
+
+        @activity.defn
+        async def gemini_files_register(
+            req: _GeminiRegisterFilesRequest,
+        ) -> types.RegisterFilesResponse:
+            """Register GCS files using the real genai.Client on the worker.
+
+            Uses ``credentials`` if provided at plugin init,
+            otherwise falls back to the client's own credentials.
+            Token refresh happens here on the worker side, so no auth
+            material enters the workflow event history.
+            """
+            auth = self._credentials or self._client._api_client._credentials
+            if auth is None:
+                raise ValueError(
+                    "No credentials available for register_files(). "
+                    "Pass extra_credentials to GoogleGenAIPlugin or initialize "
+                    "the genai.Client with credentials."
+                )
+            return await self._client.aio.files.register_files(
+                auth=auth,
+                uris=req.uris,
+                config=req.config,
+            )
+
+        @activity.defn
+        async def gemini_file_search_stores_upload(
+            req: _GeminiUploadToFileSearchStoreRequest,
+        ) -> types.UploadToFileSearchStoreOperation:
+            """Upload a file to a file search store on the worker."""
+            if req.file_bytes is not None:
+                import io
+
+                file_arg: Any = io.BytesIO(req.file_bytes)
+            else:
+                file_arg = req.file_path
+
+            return (
+                await self._client.aio.file_search_stores.upload_to_file_search_store(
+                    file_search_store_name=req.file_search_store_name,
+                    file=file_arg,
+                    config=req.config,
+                )
+            )
+
+        return [
+            gemini_api_client_async_request,
+            gemini_api_client_async_request_streamed,
+            gemini_files_upload,
+            gemini_files_download,
+            gemini_files_register,
+            gemini_file_search_stores_upload,
+            make_gemini_interactions_create_streamed(self._client),
+        ]

temporal_agent_harness/ai_sdks/google_genai_plugin/_google_genai_plugin.py

@@ -0,0 +1,98 @@
+"""Temporal plugin for Google Gemini SDK integration."""
+
+from __future__ import annotations
+
+import dataclasses
+
+import google.auth.credentials
+from google.genai import Client as GeminiClient
+
+from ._gemini_activity import GeminiApiCaller
+from temporalio.contrib.pydantic import PydanticPayloadConverter
+from temporalio.converter import DataConverter, DefaultPayloadConverter
+from temporalio.plugin import SimplePlugin
+from temporalio.worker import WorkflowRunner
+from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner
+
+
+def _data_converter(converter: DataConverter | None) -> DataConverter:
+    if converter is None:
+        return DataConverter(payload_converter_class=PydanticPayloadConverter)
+    elif converter.payload_converter_class is DefaultPayloadConverter:
+        return dataclasses.replace(
+            converter, payload_converter_class=PydanticPayloadConverter
+        )
+    return converter
+
+
+class GoogleGenAIPlugin(SimplePlugin):
+    """A Temporal Worker Plugin configured for the Google Gemini SDK.
+
+    .. warning::
+        This class is experimental and may change in future versions.
+        Use with caution in production environments.
+
+    This plugin registers the ``gemini_api_client_async_request`` activity
+    using the provided ``genai.Client`` with real credentials.  Workflows use
+    :func:`~temporalio.contrib.google_genai.workflow.google_genai_client` to
+    get an ``AsyncClient`` backed by a ``TemporalApiClient`` that routes all
+    API calls through this activity.
+
+    No credentials are passed to or from the workflow.  Auth material never
+    appears in Temporal's event history.
+
+    Example (Gemini Developer API)::
+
+        client = genai.Client(api_key=os.environ["GOOGLE_API_KEY"])
+        plugin = GoogleGenAIPlugin(client)
+
+    Example (Vertex AI)::
+
+        client = genai.Client(
+            vertexai=True, project="my-project", location="us-central1",
+        )
+        plugin = GoogleGenAIPlugin(client)
+
+    Example (with separate GCS credentials for file registration)::
+
+        client = genai.Client(api_key=os.environ["GOOGLE_API_KEY"])
+        gcs_creds, _ = google.auth.default()
+        plugin = GoogleGenAIPlugin(client, extra_credentials=gcs_creds)
+    """
+
+    def __init__(
+        self,
+        client: GeminiClient,
+        extra_credentials: google.auth.credentials.Credentials | None = None,
+    ) -> None:
+        """Initialize the Gemini plugin.
+
+        Args:
+            client: A fully configured ``genai.Client`` instance.
+                All credential management, HTTP client configuration, etc.
+                is the responsibility of the caller.
+            extra_credentials: Optional Google Cloud credentials used for
+                operations that require explicit auth (e.g.
+                ``files.register_files()``).  If not provided, the
+                client's own credentials are used.
+        """
+        self._api_caller = GeminiApiCaller(client, credentials=extra_credentials)
+
+        def workflow_runner(runner: WorkflowRunner | None) -> WorkflowRunner:
+            if not runner:
+                raise ValueError("No WorkflowRunner provided to GoogleGenAIPlugin.")
+            if isinstance(runner, SandboxedWorkflowRunner):
+                return dataclasses.replace(
+                    runner,
+                    restrictions=runner.restrictions.with_passthrough_modules(
+                        "google.genai"
+                    ),
+                )
+            return runner
+
+        super().__init__(
+            name="GoogleGenAIPlugin",
+            data_converter=_data_converter,
+            activities=self._api_caller.activities(),
+            workflow_runner=workflow_runner,
+        )