Skip to content

Commit 0d8dbf1

Browse files
GovindhKishoreGovindhKishore
authored
feat(google_genai): add timeout and max_retries to embedder components (#3412)
1 parent 718517d commit 0d8dbf1

6 files changed

Lines changed: 178 additions & 2 deletions

File tree

‎integrations/google_genai/src/haystack_integrations/components/embedders/google_genai/document_embedder.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,8 @@ def __init__(
8585
meta_fields_to_embed: list[str] | None = None,
8686
embedding_separator: str = "\n",
8787
config: dict[str, Any] | None = None,
88+
timeout: float | None = None,
89+
max_retries: int | None = None,
8890
) -> None:
8991
"""
9092
Creates an GoogleGenAIDocumentEmbedder component.
@@ -122,6 +124,11 @@ def __init__(
122124
Specifying task types in `config` does not take effect for `gemini-embedding-2`.
123125
See [Gemini documentation](https://ai.google.dev/gemini-api/docs/embeddings#task-types) for more
124126
information.
127+
:param timeout:
128+
The timeout in seconds for the underlying Google GenAI client network requests.
129+
:param max_retries:
130+
The maximum number of retries for the underlying Google GenAI client network requests.
131+
125132
"""
126133
self._api_key = api_key
127134
self._api = api
@@ -135,12 +142,16 @@ def __init__(
135142
self._meta_fields_to_embed = meta_fields_to_embed or []
136143
self._embedding_separator = embedding_separator
137144
self._config = config
145+
self._timeout = timeout
146+
self._max_retries = max_retries
138147

139148
self._client = _get_client(
140149
api_key=api_key,
141150
api=api,
142151
vertex_ai_project=vertex_ai_project,
143152
vertex_ai_location=vertex_ai_location,
153+
timeout=timeout,
154+
max_retries=max_retries,
144155
)
145156

146157
def to_dict(self) -> dict[str, Any]:
@@ -164,6 +175,8 @@ def to_dict(self) -> dict[str, Any]:
164175
vertex_ai_project=self._vertex_ai_project,
165176
vertex_ai_location=self._vertex_ai_location,
166177
config=self._config,
178+
timeout=self._timeout,
179+
max_retries=self._max_retries,
167180
)
168181

169182
@classmethod

‎integrations/google_genai/src/haystack_integrations/components/embedders/google_genai/multimodal_document_embedder.py‎

Lines changed: 49 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,14 +9,14 @@
99
from typing import Any, Literal
1010

1111
from google.genai.types import Content, EmbedContentConfig, Part
12-
from haystack import Document, component, logging
12+
from haystack import Document, component, default_from_dict, default_to_dict, logging
1313
from haystack.components.converters.image.image_utils import (
1414
_batch_convert_pdf_pages_to_images,
1515
_encode_image_to_base64,
1616
_PDFPageInfo,
1717
)
1818
from haystack.dataclasses import ByteStream
19-
from haystack.utils.auth import Secret
19+
from haystack.utils import Secret, deserialize_secrets_inplace
2020
from tqdm import tqdm
2121
from tqdm.asyncio import tqdm as async_tqdm
2222
from typing_extensions import NotRequired, TypedDict
@@ -181,6 +181,8 @@ def __init__(
181181
batch_size: int = 6,
182182
progress_bar: bool = True,
183183
config: dict[str, Any] | None = None,
184+
timeout: float | None = None,
185+
max_retries: int | None = None,
184186
) -> None:
185187
"""
186188
Creates an GoogleGenAIMultimodalDocumentEmbedder component.
@@ -217,6 +219,10 @@ def __init__(
217219
You can for example set the output dimensionality of the embedding: `{"output_dimensionality": 768}`.
218220
See [Google API documentation](https://googleapis.github.io/python-genai/genai.html#genai.types.EmbedContentConfig)
219221
for the available options.
222+
:param timeout:
223+
The timeout in seconds for the underlying Google GenAI client network requests.
224+
:param max_retries:
225+
The maximum number of retries for the underlying Google GenAI client network requests.
220226
"""
221227
self._api_key = api_key
222228
self._api = api
@@ -229,14 +235,55 @@ def __init__(
229235
self._batch_size = batch_size
230236
self._progress_bar = progress_bar
231237
self._config = config
238+
self._timeout = timeout
239+
self._max_retries = max_retries
232240

233241
self._client = _get_client(
234242
api_key=api_key,
235243
api=api,
236244
vertex_ai_project=vertex_ai_project,
237245
vertex_ai_location=vertex_ai_location,
246+
timeout=timeout,
247+
max_retries=max_retries,
238248
)
239249

250+
def to_dict(self) -> dict[str, Any]:
251+
"""
252+
Serializes the component to a dictionary.
253+
254+
:returns:
255+
Dictionary with serialized data.
256+
"""
257+
return default_to_dict(
258+
self,
259+
model=self._model,
260+
file_path_meta_field=self._file_path_meta_field,
261+
root_path=self._root_path if self._root_path != "" else None,
262+
image_size=self._image_size,
263+
batch_size=self._batch_size,
264+
progress_bar=self._progress_bar,
265+
api_key=self._api_key.to_dict(),
266+
api=self._api,
267+
vertex_ai_project=self._vertex_ai_project,
268+
vertex_ai_location=self._vertex_ai_location,
269+
config=self._config,
270+
timeout=self._timeout,
271+
max_retries=self._max_retries,
272+
)
273+
274+
@classmethod
275+
def from_dict(cls, data: dict[str, Any]) -> "GoogleGenAIMultimodalDocumentEmbedder":
276+
"""
277+
Deserializes the component from a dictionary.
278+
279+
:param data:
280+
Dictionary to deserialize from.
281+
:returns:
282+
Deserialized component.
283+
"""
284+
deserialize_secrets_inplace(data["init_parameters"], keys=["api_key"])
285+
return default_from_dict(cls, data)
286+
240287
def _extract_parts_to_embed(self, documents: list[Document]) -> list[Part]:
241288
"""
242289
Validates the input documents and extracts the files to embed in the format expected by the Google AI API.

‎integrations/google_genai/src/haystack_integrations/components/embedders/google_genai/text_embedder.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,8 @@ def __init__(
8282
prefix: str = "",
8383
suffix: str = "",
8484
config: dict[str, Any] | None = None,
85+
timeout: float | None = None,
86+
max_retries: int | None = None,
8587
) -> None:
8688
"""
8789
Creates an GoogleGenAITextEmbedder component.
@@ -111,6 +113,10 @@ def __init__(
111113
Specifying task types in `config` does not take effect for `gemini-embedding-2`.
112114
See [Gemini documentation](https://ai.google.dev/gemini-api/docs/embeddings#task-types) for more
113115
information.
116+
:param timeout:
117+
The timeout in seconds for the underlying Google GenAI client network requests.
118+
:param max_retries:
119+
The maximum number of retries for the underlying Google GenAI client network requests.
114120
"""
115121

116122
self._api_key = api_key
@@ -121,11 +127,15 @@ def __init__(
121127
self._prefix = prefix
122128
self._suffix = suffix
123129
self._config = config
130+
self._timeout = timeout
131+
self._max_retries = max_retries
124132
self._client = _get_client(
125133
api_key=api_key,
126134
api=api,
127135
vertex_ai_project=vertex_ai_project,
128136
vertex_ai_location=vertex_ai_location,
137+
timeout=timeout,
138+
max_retries=max_retries,
129139
)
130140

131141
def to_dict(self) -> dict[str, Any]:
@@ -145,6 +155,8 @@ def to_dict(self) -> dict[str, Any]:
145155
prefix=self._prefix,
146156
suffix=self._suffix,
147157
config=self._config,
158+
timeout=self._timeout,
159+
max_retries=self._max_retries,
148160
)
149161

150162
@classmethod

‎integrations/google_genai/tests/test_document_embedder.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,8 @@ def test_init_default(self, monkeypatch):
4040
assert embedder._meta_fields_to_embed == []
4141
assert embedder._embedding_separator == "\n"
4242
assert embedder._config is None
43+
assert embedder._timeout is None
44+
assert embedder._max_retries is None
4345

4446
def test_init_with_parameters(self, monkeypatch):
4547
embedder = GoogleGenAIDocumentEmbedder(
@@ -52,6 +54,8 @@ def test_init_with_parameters(self, monkeypatch):
5254
meta_fields_to_embed=["test_field"],
5355
embedding_separator=" | ",
5456
config={"task_type": "CLASSIFICATION"},
57+
timeout=30.0,
58+
max_retries=3,
5559
)
5660
assert embedder._api_key.resolve_value() == "fake-api-key-2"
5761
assert embedder._model == "model"
@@ -62,6 +66,8 @@ def test_init_with_parameters(self, monkeypatch):
6266
assert embedder._meta_fields_to_embed == ["test_field"]
6367
assert embedder._embedding_separator == " | "
6468
assert embedder._config == {"task_type": "CLASSIFICATION"}
69+
assert embedder._timeout == 30.0
70+
assert embedder._max_retries == 3
6571

6672
def test_init_fail_wo_api_key(self, monkeypatch):
6773
monkeypatch.delenv("GOOGLE_API_KEY", raising=False)
@@ -90,6 +96,8 @@ def test_to_dict(self, monkeypatch):
9096
"api": "gemini",
9197
"vertex_ai_project": None,
9298
"vertex_ai_location": None,
99+
"timeout": None,
100+
"max_retries": None,
93101
},
94102
}
95103

@@ -105,6 +113,8 @@ def test_to_dict_with_custom_init_parameters(self, monkeypatch):
105113
meta_fields_to_embed=["test_field"],
106114
embedding_separator=" | ",
107115
config={"task_type": "CLASSIFICATION"},
116+
timeout=30.0,
117+
max_retries=3,
108118
)
109119
data = component.to_dict()
110120
assert data == {
@@ -124,6 +134,8 @@ def test_to_dict_with_custom_init_parameters(self, monkeypatch):
124134
"api": "gemini",
125135
"vertex_ai_project": None,
126136
"vertex_ai_location": None,
137+
"timeout": 30.0,
138+
"max_retries": 3,
127139
},
128140
}
129141

@@ -145,6 +157,8 @@ def test_from_dict(self, monkeypatch):
145157
"api": "gemini",
146158
"vertex_ai_project": None,
147159
"vertex_ai_location": None,
160+
"timeout": 30.0,
161+
"max_retries": 3,
148162
},
149163
}
150164
monkeypatch.setenv("GOOGLE_API_KEY", "fake-api-key")
@@ -161,6 +175,8 @@ def test_from_dict(self, monkeypatch):
161175
assert embedder._api == "gemini"
162176
assert embedder._vertex_ai_project is None
163177
assert embedder._vertex_ai_location is None
178+
assert embedder._timeout == 30.0
179+
assert embedder._max_retries == 3
164180

165181
def test_prepare_texts_to_embed_w_metadata(self):
166182
documents = [

‎integrations/google_genai/tests/test_multimodal_document_embedder.py‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,8 @@ def test_init_with_parameters(self):
7676
root_path="root_path",
7777
image_size=(1024, 1024),
7878
config={"task_type": "CLASSIFICATION"},
79+
timeout=30.0,
80+
max_retries=3,
7981
)
8082
assert embedder._api_key.resolve_value() == "fake-api-key-2"
8183
assert embedder._model == "model"
@@ -85,6 +87,8 @@ def test_init_with_parameters(self):
8587
assert embedder._batch_size == 64
8688
assert embedder._progress_bar is False
8789
assert embedder._config == {"task_type": "CLASSIFICATION"}
90+
assert embedder._timeout == 30.0
91+
assert embedder._max_retries == 3
8892

8993
def test_init_fail_wo_api_key(self, monkeypatch):
9094
monkeypatch.delenv("GOOGLE_API_KEY", raising=False)
@@ -114,6 +118,8 @@ def test_to_dict(self, monkeypatch):
114118
"vertex_ai_project": None,
115119
"vertex_ai_location": None,
116120
"config": None,
121+
"timeout": None,
122+
"max_retries": None,
117123
},
118124
}
119125

@@ -134,6 +140,8 @@ def test_from_dict(self, monkeypatch):
134140
"vertex_ai_project": None,
135141
"vertex_ai_location": None,
136142
"config": {"task_type": "CLASSIFICATION"},
143+
"timeout": 30.0,
144+
"max_retries": 3,
137145
},
138146
}
139147
monkeypatch.setenv("GOOGLE_API_KEY", "fake-api-key")
@@ -150,6 +158,70 @@ def test_from_dict(self, monkeypatch):
150158
assert embedder._api == "gemini"
151159
assert embedder._vertex_ai_project is None
152160
assert embedder._vertex_ai_location is None
161+
assert embedder._timeout == 30.0
162+
assert embedder._max_retries == 3
163+
164+
def test_to_dict_with_custom_init_parameters(self, monkeypatch):
165+
monkeypatch.setenv("GOOGLE_API_KEY", "fake-api-key")
166+
component = GoogleGenAIMultimodalDocumentEmbedder(
167+
model="model",
168+
batch_size=12,
169+
timeout=30.0,
170+
max_retries=3,
171+
config={"output_dimensionality": 768},
172+
root_path="some_root_path",
173+
image_size=(512, 512),
174+
progress_bar=False,
175+
)
176+
data = component.to_dict()
177+
assert data == {
178+
"type": (
179+
"haystack_integrations.components.embedders.google_genai.multimodal_document_embedder.GoogleGenAIMultimodalDocumentEmbedder"
180+
),
181+
"init_parameters": {
182+
"model": "model",
183+
"file_path_meta_field": "file_path",
184+
"root_path": "some_root_path",
185+
"image_size": (512, 512),
186+
"batch_size": 12,
187+
"progress_bar": False,
188+
"api_key": {"type": "env_var", "env_vars": ["GOOGLE_API_KEY", "GEMINI_API_KEY"], "strict": False},
189+
"api": "gemini",
190+
"vertex_ai_project": None,
191+
"vertex_ai_location": None,
192+
"config": {"output_dimensionality": 768},
193+
"timeout": 30.0,
194+
"max_retries": 3,
195+
},
196+
}
197+
198+
def test_to_dict_from_dict_roundtrip(self, monkeypatch):
199+
monkeypatch.setenv("GOOGLE_API_KEY", "fake-api-key")
200+
original = GoogleGenAIMultimodalDocumentEmbedder(
201+
model="model",
202+
batch_size=12,
203+
timeout=30.0,
204+
max_retries=3,
205+
image_size=(512, 512),
206+
progress_bar=False,
207+
config={"output_dimensionality": 768},
208+
root_path="some_root_path",
209+
)
210+
data = original.to_dict()
211+
restored = component_from_dict(GoogleGenAIMultimodalDocumentEmbedder, data, "embedder")
212+
assert restored._api_key.resolve_value() == "fake-api-key"
213+
assert restored._model == "model"
214+
assert restored._batch_size == 12
215+
assert restored._timeout == 30.0
216+
assert restored._max_retries == 3
217+
assert restored._config == {"output_dimensionality": 768}
218+
assert restored._root_path == "some_root_path"
219+
assert restored._image_size == (512, 512)
220+
assert restored._progress_bar is False
221+
assert restored._file_path_meta_field == "file_path"
222+
assert restored._api == "gemini"
223+
assert restored._vertex_ai_project is None
224+
assert restored._vertex_ai_location is None
153225

154226
def test_extract_parts_to_embed_images(self, test_files_path):
155227
embedder = GoogleGenAIMultimodalDocumentEmbedder(api_key=Secret.from_token("fake-api-key"))

0 commit comments

Comments
 (0)