|
17 | 17 | from datetime import timedelta |
18 | 18 | from typing import Any |
19 | 19 |
|
20 | | -from google.genai import types |
21 | 20 | from google.genai._api_client import BaseApiClient |
22 | 21 | from google.genai.types import HttpOptions, HttpOptionsOrDict |
23 | 22 | from google.genai.types import HttpResponse as SdkHttpResponse |
24 | | -from pydantic import BaseModel |
25 | 23 |
|
26 | 24 | 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 | +) |
27 | 31 | from temporalio.workflow import ActivityConfig |
28 | 32 |
|
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 | | - |
47 | 33 | # Fields on HttpOptions that cannot be serialized or should not be forwarded. |
48 | 34 | _REJECTED_HTTP_OPTION_FIELDS = frozenset( |
49 | 35 | { |
@@ -72,73 +58,6 @@ def _validate_http_options(http_options: HttpOptions | None) -> None: |
72 | 58 | ) |
73 | 59 |
|
74 | 60 |
|
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 | | - |
142 | 61 | class TemporalApiClient(BaseApiClient): |
143 | 62 | """A ``BaseApiClient`` that routes all API calls through Temporal activities. |
144 | 63 |
|
|
0 commit comments