Skip to content

Commit 0ec3400

Browse files
lint
1 parent d62ea96 commit 0ec3400

6 files changed

Lines changed: 125 additions & 93 deletions

File tree

temporalio/contrib/google_gemini_sdk/_gemini_activity.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
from google.genai.types import HttpResponse as SdkHttpResponse
1919

2020
from temporalio import activity
21-
from temporalio.contrib.google_gemini_sdk._temporal_api_client import (
21+
from temporalio.contrib.google_gemini_sdk._models import (
2222
_GeminiApiRequest,
2323
_GeminiApiResponse,
2424
_GeminiApiStreamedResponse,
Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,107 @@
1+
"""Serializable Pydantic models for the Gemini SDK Temporal integration.
2+
3+
These models cross the activity boundary — they're constructed on the
4+
workflow side and deserialized on the activity side (or vice versa).
5+
"""
6+
7+
from __future__ import annotations
8+
9+
from typing import Any
10+
11+
from google.genai import types
12+
from pydantic import BaseModel
13+
14+
__all__ = [
15+
"_GeminiApiRequest",
16+
"_GeminiApiResponse",
17+
"_GeminiApiStreamedResponse",
18+
"_GeminiDownloadFileRequest",
19+
"_GeminiRegisterFilesRequest",
20+
"_GeminiUploadFileRequest",
21+
"_GeminiUploadToFileSearchStoreRequest",
22+
"_SerializableHttpOptions",
23+
]
24+
25+
26+
class _SerializableHttpOptions(BaseModel):
27+
"""Per-request HTTP options that can be serialized across the activity boundary.
28+
29+
Non-serializable fields (httpx_client, httpx_async_client, aiohttp_client,
30+
client_args, async_client_args) must be configured at GeminiPlugin init.
31+
32+
``timeout`` is excluded because Temporal owns timeouts/retries — configure
33+
via ``ActivityConfig`` instead.
34+
"""
35+
36+
base_url: str | None = None
37+
base_url_resource_scope: str | None = None
38+
api_version: str | None = None
39+
headers: dict[str, str] | None = None
40+
extra_body: dict[str, Any] | None = None
41+
42+
43+
# ── async_request models ──────────────────────────────────────────────────
44+
45+
46+
class _GeminiApiRequest(BaseModel):
47+
"""Serializable activity input for a Gemini SDK API call."""
48+
49+
http_method: str
50+
path: str
51+
request_dict: dict[str, object]
52+
http_options_overrides: _SerializableHttpOptions | None = None
53+
54+
55+
class _GeminiApiResponse(BaseModel):
56+
"""Serializable activity output for a Gemini SDK API call."""
57+
58+
headers: dict[str, str]
59+
body: str
60+
61+
62+
class _GeminiApiStreamedResponse(BaseModel):
63+
"""Serializable activity output for a batched streamed API call.
64+
65+
The activity collects all streamed chunks and returns them as a list.
66+
The ``TemporalApiClient`` then yields them one at a time to the SDK.
67+
"""
68+
69+
chunks: list[_GeminiApiResponse]
70+
71+
72+
# ── files upload/download models ──────────────────────────────────────────
73+
74+
75+
class _GeminiUploadFileRequest(BaseModel):
76+
"""Serializable activity input for a file upload.
77+
78+
For file path uploads the path is resolved on the worker. For
79+
in-memory uploads the raw bytes are sent across the activity boundary.
80+
"""
81+
82+
file_bytes: bytes | None = None
83+
file_path: str | None = None
84+
config: types.UploadFileConfig | None = None
85+
86+
87+
class _GeminiDownloadFileRequest(BaseModel):
88+
"""Serializable activity input for a file download."""
89+
90+
file: str
91+
config: types.DownloadFileConfig | None = None
92+
93+
94+
class _GeminiRegisterFilesRequest(BaseModel):
95+
"""Serializable activity input for registering GCS files."""
96+
97+
uris: list[str]
98+
config: types.RegisterFilesConfig | None = None
99+
100+
101+
class _GeminiUploadToFileSearchStoreRequest(BaseModel):
102+
"""Serializable activity input for uploading a file to a file search store."""
103+
104+
file_search_store_name: str
105+
file_bytes: bytes | None = None
106+
file_path: str | None = None
107+
config: types.UploadToFileSearchStoreConfig | None = None

temporalio/contrib/google_gemini_sdk/_temporal_api_client.py

Lines changed: 6 additions & 87 deletions
Original file line numberDiff line numberDiff line change
@@ -17,33 +17,19 @@
1717
from datetime import timedelta
1818
from typing import Any
1919

20-
from google.genai import types
2120
from google.genai._api_client import BaseApiClient
2221
from google.genai.types import HttpOptions, HttpOptionsOrDict
2322
from google.genai.types import HttpResponse as SdkHttpResponse
24-
from pydantic import BaseModel
2523

2624
from temporalio import workflow as temporal_workflow
25+
from temporalio.contrib.google_gemini_sdk._models import (
26+
_GeminiApiRequest,
27+
_GeminiApiResponse,
28+
_GeminiApiStreamedResponse,
29+
_SerializableHttpOptions,
30+
)
2731
from temporalio.workflow import ActivityConfig
2832

29-
30-
class _SerializableHttpOptions(BaseModel):
31-
"""Per-request HTTP options that can be serialized across the activity boundary.
32-
33-
Non-serializable fields (httpx_client, httpx_async_client, aiohttp_client,
34-
client_args, async_client_args) must be configured at GeminiPlugin init.
35-
36-
``timeout`` is excluded because Temporal owns timeouts/retries — configure
37-
via ``ActivityConfig`` instead.
38-
"""
39-
40-
base_url: str | None = None
41-
base_url_resource_scope: str | None = None
42-
api_version: str | None = None
43-
headers: dict[str, str] | None = None
44-
extra_body: dict[str, Any] | None = None
45-
46-
4733
# Fields on HttpOptions that cannot be serialized or should not be forwarded.
4834
_REJECTED_HTTP_OPTION_FIELDS = frozenset(
4935
{
@@ -72,73 +58,6 @@ def _validate_http_options(http_options: HttpOptions | None) -> None:
7258
)
7359

7460

75-
# ── async_request models ──────────────────────────────────────────────────
76-
77-
78-
class _GeminiApiRequest(BaseModel):
79-
"""Serializable activity input for a Gemini SDK API call."""
80-
81-
http_method: str
82-
path: str
83-
request_dict: dict[str, object]
84-
http_options_overrides: _SerializableHttpOptions | None = None
85-
86-
87-
class _GeminiApiResponse(BaseModel):
88-
"""Serializable activity output for a Gemini SDK API call."""
89-
90-
headers: dict[str, str]
91-
body: str
92-
93-
94-
class _GeminiApiStreamedResponse(BaseModel):
95-
"""Serializable activity output for a batched streamed API call.
96-
97-
The activity collects all streamed chunks and returns them as a list.
98-
The ``TemporalApiClient`` then yields them one at a time to the SDK.
99-
"""
100-
101-
chunks: list[_GeminiApiResponse]
102-
103-
104-
# ── files upload/download models ──────────────────────────────────────────
105-
106-
107-
class _GeminiUploadFileRequest(BaseModel):
108-
"""Serializable activity input for a file upload.
109-
110-
For file path uploads the path is resolved on the worker. For
111-
in-memory uploads the raw bytes are sent across the activity boundary.
112-
"""
113-
114-
file_bytes: bytes | None = None
115-
file_path: str | None = None
116-
config: types.UploadFileConfig | None = None
117-
118-
119-
class _GeminiDownloadFileRequest(BaseModel):
120-
"""Serializable activity input for a file download."""
121-
122-
file: str
123-
config: types.DownloadFileConfig | None = None
124-
125-
126-
class _GeminiRegisterFilesRequest(BaseModel):
127-
"""Serializable activity input for registering GCS files."""
128-
129-
uris: list[str]
130-
config: types.RegisterFilesConfig | None = None
131-
132-
133-
class _GeminiUploadToFileSearchStoreRequest(BaseModel):
134-
"""Serializable activity input for uploading a file to a file search store."""
135-
136-
file_search_store_name: str
137-
file_bytes: bytes | None = None
138-
file_path: str | None = None
139-
config: types.UploadToFileSearchStoreConfig | None = None
140-
141-
14261
class TemporalApiClient(BaseApiClient):
14362
"""A ``BaseApiClient`` that routes all API calls through Temporal activities.
14463

temporalio/contrib/google_gemini_sdk/_temporal_file_search_stores.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,11 @@
1616
from google.genai.file_search_stores import AsyncFileSearchStores
1717

1818
from temporalio import workflow as temporal_workflow
19+
from temporalio.contrib.google_gemini_sdk._models import (
20+
_GeminiUploadToFileSearchStoreRequest,
21+
)
1922
from temporalio.contrib.google_gemini_sdk._temporal_api_client import (
2023
TemporalApiClient,
21-
_GeminiUploadToFileSearchStoreRequest,
2224
_validate_http_options,
2325
)
2426
from temporalio.workflow import ActivityConfig

temporalio/contrib/google_gemini_sdk/_temporal_files.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,13 @@
1919
from google.genai.files import AsyncFiles
2020

2121
from temporalio import workflow as temporal_workflow
22-
from temporalio.contrib.google_gemini_sdk._temporal_api_client import (
23-
TemporalApiClient,
22+
from temporalio.contrib.google_gemini_sdk._models import (
2423
_GeminiDownloadFileRequest,
2524
_GeminiRegisterFilesRequest,
2625
_GeminiUploadFileRequest,
26+
)
27+
from temporalio.contrib.google_gemini_sdk._temporal_api_client import (
28+
TemporalApiClient,
2729
_validate_http_options,
2830
)
2931
from temporalio.workflow import ActivityConfig

tests/contrib/google_gemini_sdk/test_gemini.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,15 +38,17 @@
3838
activity_as_tool,
3939
gemini_client,
4040
)
41-
from temporalio.contrib.google_gemini_sdk._temporal_api_client import (
42-
TemporalApiClient,
41+
from temporalio.contrib.google_gemini_sdk._models import (
4342
_GeminiApiRequest,
4443
_GeminiApiResponse,
4544
_GeminiApiStreamedResponse,
4645
_GeminiDownloadFileRequest,
4746
_GeminiUploadFileRequest,
4847
_GeminiUploadToFileSearchStoreRequest,
4948
)
49+
from temporalio.contrib.google_gemini_sdk._temporal_api_client import (
50+
TemporalApiClient,
51+
)
5052
from temporalio.contrib.google_gemini_sdk._temporal_async_client import (
5153
TemporalAsyncClient,
5254
)

0 commit comments

Comments
 (0)