Add file search stores, register_files, chat example, and full integration test

- Add TemporalAsyncFileSearchStores: subclasses AsyncFileSearchStores,
  overriding upload_to_file_search_store() to dispatch through a
  Temporal activity. All other methods (create, get, delete, list,
  import_file, documents) work via async_request.

- Add TemporalAsyncFiles.register_files(): dispatches through
  gemini_files_register activity. The auth param from the caller is
  ignored; the activity uses extra_credentials from GeminiPlugin init
  or falls back to the client's own credentials.

- Add extra_credentials param to GeminiPlugin for operations that need
  explicit auth (e.g. GCS file registration).

- Wire TemporalAsyncFileSearchStores into TemporalAsyncClient.

- Add 12 new tests (31 total): file upload (str + bytes), download,
  file search store upload, multi-turn chat, TemporalAsyncClient
  wiring checks, low-level raise checks, and a full integration test
  that runs real activities with a mocked client covering
  generate_content, streaming, file upload/download, store upload,
  RAG query, and store deletion.

Author: JasonSteving99

temporalio/contrib/google_gemini_sdk/_gemini_activity.py

@@ -9,8 +9,9 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import Any, Callable
+from typing import Any, Callable, Optional
 
+import google.auth.credentials
 from google.genai import Client as GeminiClient, types
 from google.genai.types import HttpOptions
 from google.genai.types import HttpResponse as SdkHttpResponse
@@ -21,7 +22,9 @@
     _GeminiApiResponse,
     _GeminiApiStreamedResponse,
     _GeminiDownloadFileRequest,
+    _GeminiRegisterFilesRequest,
     _GeminiUploadFileRequest,
+    _GeminiUploadToFileSearchStoreRequest,
 )
 
 
@@ -35,15 +38,20 @@ def _resolve_http_options(
 
 
 class GeminiApiCaller:
-    """Wraps a ``genai.Client`` and exposes a Temporal activity for API calls.
+    """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) -> None:
+    def __init__(
+        self,
+        client: GeminiClient,
+        credentials: Optional[google.auth.credentials.Credentials] = None,
+    ) -> None:
         self._client = client
+        self._credentials = credentials
 
     def activities(self) -> Sequence[Callable]:
         """Return activities that route SDK calls through this client."""
@@ -112,9 +120,46 @@ async def gemini_files_download(
                 file=req.file, config=req.config
             )
 
+        @activity.defn(name="gemini_files_register")
+        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.
+            """
+            return await self._client.aio.files.register_files(
+                auth=self._credentials or self._client._api_client._credentials,
+                uris=req.uris,
+                config=req.config,
+            )
+
+        @activity.defn(name="gemini_file_search_stores_upload")
+        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,
         ]

temporalio/contrib/google_gemini_sdk/_gemini_plugin.py

@@ -3,7 +3,9 @@
 from __future__ import annotations
 
 import dataclasses
+from typing import Optional
 
+import google.auth.credentials
 from google.genai import Client as GeminiClient
 
 from temporalio.contrib.google_gemini_sdk._gemini_activity import GeminiApiCaller
@@ -51,17 +53,33 @@ class GeminiPlugin(SimplePlugin):
             vertexai=True, project="my-project", location="us-central1",
         )
         plugin = GeminiPlugin(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 = GeminiPlugin(client, extra_credentials=gcs_creds)
     """
 
-    def __init__(self, client: GeminiClient) -> None:
+    def __init__(
+        self,
+        client: GeminiClient,
+        extra_credentials: Optional[google.auth.credentials.Credentials] = 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)
+        self._api_caller = GeminiApiCaller(
+            client, credentials=extra_credentials
+        )
 
         def workflow_runner(runner: WorkflowRunner | None) -> WorkflowRunner:
             if not runner:

temporalio/contrib/google_gemini_sdk/_temporal_api_client.py

@@ -122,6 +122,22 @@ class _GeminiDownloadFileRequest(BaseModel):
     config: types.DownloadFileConfig | None = None
 
 
+class _GeminiRegisterFilesRequest(BaseModel):
+    """Serializable activity input for registering GCS files."""
+
+    uris: list[str]
+    config: types.RegisterFilesConfig | None = None
+
+
+class _GeminiUploadToFileSearchStoreRequest(BaseModel):
+    """Serializable activity input for uploading a file to a file search store."""
+
+    file_search_store_name: str
+    file_bytes: bytes | None = None
+    file_path: str | None = None
+    config: types.UploadToFileSearchStoreConfig | None = None
+
+
 class TemporalApiClient(BaseApiClient):
     """A ``BaseApiClient`` that routes all API calls through Temporal activities.
 

temporalio/contrib/google_gemini_sdk/_temporal_async_client.py

@@ -1,7 +1,8 @@
 """Temporal-aware AsyncClient shim.
 
 ``TemporalAsyncClient`` is an ``AsyncClient`` subclass that wires up
-``TemporalAsyncFiles`` in place of the default ``AsyncFiles``.
+Temporal-aware replacements for modules that need special handling
+(files, file search stores).
 """
 
 from __future__ import annotations
@@ -13,6 +14,9 @@
 from temporalio.contrib.google_gemini_sdk._temporal_api_client import (
     TemporalApiClient,
 )
+from temporalio.contrib.google_gemini_sdk._temporal_file_search_stores import (
+    TemporalAsyncFileSearchStores,
+)
 from temporalio.contrib.google_gemini_sdk._temporal_files import (
     TemporalAsyncFiles,
 )
@@ -21,10 +25,14 @@
 class TemporalAsyncClient(AsyncClient):
     """``AsyncClient`` subclass that uses Temporal-aware modules.
 
-    Replaces ``AsyncFiles`` with ``TemporalAsyncFiles`` so that file
-    upload/download operations run entirely inside Temporal activities.
-    Other modules (models, caches, etc.) are inherited unchanged and
-    work through ``TemporalApiClient``'s activity-backed HTTP methods.
+    Replaces ``AsyncFiles`` with ``TemporalAsyncFiles`` and
+    ``AsyncFileSearchStores`` with ``TemporalAsyncFileSearchStores``
+    so that file upload/download operations and file search store uploads
+    run entirely inside Temporal activities.
+
+    Other modules (models, tunings, caches, batches, live, tokens,
+    operations) are inherited unchanged and work through
+    ``TemporalApiClient``'s activity-backed HTTP methods.
     """
 
     def __init__(
@@ -34,3 +42,6 @@ def __init__(
     ) -> None:
         super().__init__(api_client)
         self._files = TemporalAsyncFiles(api_client, activity_config)
+        self._file_search_stores = TemporalAsyncFileSearchStores(
+            api_client, activity_config
+        )

temporalio/contrib/google_gemini_sdk/_temporal_file_search_stores.py

@@ -0,0 +1,100 @@
+"""Temporal-aware AsyncFileSearchStores shim.
+
+``TemporalAsyncFileSearchStores`` is an ``AsyncFileSearchStores`` subclass
+whose ``upload_to_file_search_store`` method dispatches through a Temporal
+activity so the entire upload (including filesystem access and resumable
+upload negotiation) runs on the activity worker.
+"""
+
+from __future__ import annotations
+
+import io
+from datetime import timedelta
+from typing import Optional, Union
+
+from google.genai import types
+from google.genai.file_search_stores import AsyncFileSearchStores
+from google.genai.types import HttpOptions
+
+from temporalio import workflow as temporal_workflow
+from temporalio.workflow import ActivityConfig
+
+from temporalio.contrib.google_gemini_sdk._temporal_api_client import (
+    TemporalApiClient,
+    _GeminiUploadToFileSearchStoreRequest,
+    _validate_http_options,
+)
+
+
+class TemporalAsyncFileSearchStores(AsyncFileSearchStores):
+    """``AsyncFileSearchStores`` subclass that routes ``upload_to_file_search_store`` through an activity.
+
+    The entire upload operation — including filesystem access, resumable
+    upload negotiation, and chunked transfer — runs inside a Temporal
+    activity on the worker.  All other methods (``create``, ``get``,
+    ``delete``, ``list``, ``import_file``, ``documents``) are inherited
+    and already work through the ``TemporalApiClient``'s ``async_request``
+    activity.
+    """
+
+    def __init__(
+        self,
+        api_client: TemporalApiClient,
+        activity_config: ActivityConfig | None = None,
+    ) -> None:
+        super().__init__(api_client)
+        self._activity_config = activity_config or ActivityConfig(
+            start_to_close_timeout=timedelta(seconds=60),
+        )
+
+    async def upload_to_file_search_store(
+        self,
+        *,
+        file_search_store_name: str,
+        file: Union[str, "os.PathLike[str]", io.IOBase],
+        config: Optional[types.UploadToFileSearchStoreConfigOrDict] = None,
+    ) -> types.UploadToFileSearchStoreOperation:
+        """Upload a file to a file search store via a Temporal activity.
+
+        Accepts a file path (resolved on the worker), ``os.PathLike``, or
+        an ``io.IOBase`` (bytes sent across the activity boundary).
+        """
+        act_config: ActivityConfig = {**self._activity_config}
+        if "summary" not in act_config:
+            act_config["summary"] = "file_search_stores.upload"
+
+        upload_config = None
+        if config is not None:
+            if isinstance(config, dict):
+                upload_config = types.UploadToFileSearchStoreConfig.model_validate(
+                    config
+                )
+            else:
+                upload_config = config
+            _validate_http_options(upload_config.http_options)
+
+        if isinstance(file, io.IOBase):
+            req = _GeminiUploadToFileSearchStoreRequest(
+                file_search_store_name=file_search_store_name,
+                file_bytes=file.read(),
+                config=upload_config,
+            )
+        elif isinstance(file, str):
+            req = _GeminiUploadToFileSearchStoreRequest(
+                file_search_store_name=file_search_store_name,
+                file_path=file,
+                config=upload_config,
+            )
+        else:
+            req = _GeminiUploadToFileSearchStoreRequest(
+                file_search_store_name=file_search_store_name,
+                file_path=file.__fspath__(),
+                config=upload_config,
+            )
+
+        return await temporal_workflow.execute_activity(
+            "gemini_file_search_stores_upload",
+            req,
+            result_type=types.UploadToFileSearchStoreOperation,
+            **act_config,
+        )

temporalio/contrib/google_gemini_sdk/_temporal_files.py

@@ -19,9 +19,12 @@
 from temporalio import workflow as temporal_workflow
 from temporalio.workflow import ActivityConfig
 
+import google.auth.credentials
+
 from temporalio.contrib.google_gemini_sdk._temporal_api_client import (
     TemporalApiClient,
     _GeminiDownloadFileRequest,
+    _GeminiRegisterFilesRequest,
     _GeminiUploadFileRequest,
     _validate_http_options,
 )
@@ -123,3 +126,38 @@ async def download(
             result_type=bytes,
             **act_config,
         )
+
+    async def register_files(
+        self,
+        *,
+        auth: google.auth.credentials.Credentials,
+        uris: list[str],
+        config: Optional[types.RegisterFilesConfigOrDict] = None,
+    ) -> types.RegisterFilesResponse:
+        """Register GCS files via a Temporal activity.
+
+        .. note::
+            The ``auth`` parameter is **ignored**.  The activity uses
+            ``credentials`` if provided to ``GeminiPlugin``,
+            otherwise falls back to the ``genai.Client``'s own credentials.
+            Either way, those credentials must have access to the GCS URIs
+            being registered.
+        """
+        act_config: ActivityConfig = {**self._activity_config}
+        if "summary" not in act_config:
+            act_config["summary"] = "files.register_files"
+
+        register_config = None
+        if config is not None:
+            if isinstance(config, dict):
+                register_config = types.RegisterFilesConfig.model_validate(config)
+            else:
+                register_config = config
+            _validate_http_options(register_config.http_options)
+
+        return await temporal_workflow.execute_activity(
+            "gemini_files_register",
+            _GeminiRegisterFilesRequest(uris=uris, config=register_config),
+            result_type=types.RegisterFilesResponse,
+            **act_config,
+        )