datapizza-ai-clients-google 0.0.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,207 @@
1
+ # Byte-compiled / optimized / DLL files
2
+ __pycache__/
3
+ *.py[codz]
4
+ *$py.class
5
+
6
+ # C extensions
7
+ *.so
8
+
9
+ # Distribution / packaging
10
+ .Python
11
+ build/
12
+ develop-eggs/
13
+ dist/
14
+ downloads/
15
+ eggs/
16
+ .eggs/
17
+ lib/
18
+ lib64/
19
+ parts/
20
+ sdist/
21
+ var/
22
+ wheels/
23
+ share/python-wheels/
24
+ *.egg-info/
25
+ .installed.cfg
26
+ *.egg
27
+ MANIFEST
28
+
29
+ # PyInstaller
30
+ # Usually these files are written by a python script from a template
31
+ # before PyInstaller builds the exe, so as to inject date/other infos into it.
32
+ *.manifest
33
+ *.spec
34
+
35
+ # Installer logs
36
+ pip-log.txt
37
+ pip-delete-this-directory.txt
38
+
39
+ # Unit test / coverage reports
40
+ htmlcov/
41
+ .tox/
42
+ .nox/
43
+ .coverage
44
+ .coverage.*
45
+ .cache
46
+ nosetests.xml
47
+ coverage.xml
48
+ *.cover
49
+ *.py.cover
50
+ .hypothesis/
51
+ .pytest_cache/
52
+ cover/
53
+
54
+ # Translations
55
+ *.mo
56
+ *.pot
57
+
58
+ # Django stuff:
59
+ *.log
60
+ local_settings.py
61
+ db.sqlite3
62
+ db.sqlite3-journal
63
+
64
+ # Flask stuff:
65
+ instance/
66
+ .webassets-cache
67
+
68
+ # Scrapy stuff:
69
+ .scrapy
70
+
71
+ # Sphinx documentation
72
+ docs/_build/
73
+
74
+ # PyBuilder
75
+ .pybuilder/
76
+ target/
77
+
78
+ # Jupyter Notebook
79
+ .ipynb_checkpoints
80
+
81
+ # IPython
82
+ profile_default/
83
+ ipython_config.py
84
+
85
+ # pyenv
86
+ # For a library or package, you might want to ignore these files since the code is
87
+ # intended to run in multiple environments; otherwise, check them in:
88
+ # .python-version
89
+
90
+ # pipenv
91
+ # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
92
+ # However, in case of collaboration, if having platform-specific dependencies or dependencies
93
+ # having no cross-platform support, pipenv may install dependencies that don't work, or not
94
+ # install all needed dependencies.
95
+ #Pipfile.lock
96
+
97
+ # UV
98
+ # Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
99
+ # This is especially recommended for binary packages to ensure reproducibility, and is more
100
+ # commonly ignored for libraries.
101
+ uv.lock
102
+
103
+ # poetry
104
+ # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
105
+ # This is especially recommended for binary packages to ensure reproducibility, and is more
106
+ # commonly ignored for libraries.
107
+ # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
108
+ #poetry.lock
109
+ #poetry.toml
110
+
111
+ # pdm
112
+ # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
113
+ # pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
114
+ # https://pdm-project.org/en/latest/usage/project/#working-with-version-control
115
+ #pdm.lock
116
+ #pdm.toml
117
+ .pdm-python
118
+ .pdm-build/
119
+
120
+ # pixi
121
+ # Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
122
+ #pixi.lock
123
+ # Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
124
+ # in the .venv directory. It is recommended not to include this directory in version control.
125
+ .pixi
126
+
127
+ # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
128
+ __pypackages__/
129
+
130
+ # Celery stuff
131
+ celerybeat-schedule
132
+ celerybeat.pid
133
+
134
+ # SageMath parsed files
135
+ *.sage.py
136
+
137
+ # Environments
138
+ .env
139
+ .envrc
140
+ .venv
141
+ env/
142
+ venv/
143
+ ENV/
144
+ env.bak/
145
+ venv.bak/
146
+
147
+ # Spyder project settings
148
+ .spyderproject
149
+ .spyproject
150
+
151
+ # Rope project settings
152
+ .ropeproject
153
+
154
+ # mkdocs documentation
155
+ /site
156
+
157
+ # mypy
158
+ .mypy_cache/
159
+ .dmypy.json
160
+ dmypy.json
161
+
162
+ # Pyre type checker
163
+ .pyre/
164
+
165
+ # pytype static type analyzer
166
+ .pytype/
167
+
168
+ # Cython debug symbols
169
+ cython_debug/
170
+
171
+ # PyCharm
172
+ # JetBrains specific template is maintained in a separate JetBrains.gitignore that can
173
+ # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
174
+ # and can be added to the global gitignore or merged into this file. For a more nuclear
175
+ # option (not recommended) you can uncomment the following to ignore the entire idea folder.
176
+ #.idea/
177
+
178
+ # Abstra
179
+ # Abstra is an AI-powered process automation framework.
180
+ # Ignore directories containing user credentials, local state, and settings.
181
+ # Learn more at https://abstra.io/docs
182
+ .abstra/
183
+
184
+ # Visual Studio Code
185
+ # Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
186
+ # that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
187
+ # and can be added to the global gitignore or merged into this file. However, if you prefer,
188
+ # you could uncomment the following to ignore the entire vscode folder
189
+ # .vscode/
190
+
191
+ # Ruff stuff:
192
+ .ruff_cache/
193
+
194
+ # PyPI configuration file
195
+ .pypirc
196
+
197
+ # Cursor
198
+ # Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to
199
+ # exclude from AI features like autocomplete and code analysis. Recommended for sensitive data
200
+ # refer to https://docs.cursor.com/context/ignore-files
201
+ .cursorignore
202
+ .cursorindexingignore
203
+
204
+ # Marimo
205
+ marimo/_static/
206
+ marimo/_lsp/
207
+ __marimo__/
@@ -0,0 +1,12 @@
1
+ Metadata-Version: 2.4
2
+ Name: datapizza-ai-clients-google
3
+ Version: 0.0.1
4
+ Summary: Google (Gemini) client for the datapizza-ai framework
5
+ License: MIT
6
+ Classifier: License :: OSI Approved :: MIT License
7
+ Classifier: Programming Language :: Python :: 3
8
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
9
+ Classifier: Topic :: Software Development :: Libraries :: Application Frameworks
10
+ Requires-Python: <3.13.0,>=3.10.0
11
+ Requires-Dist: datapizza-ai-core>=0.0.1
12
+ Requires-Dist: google-genai<2.0.0,>=1.3.0
File without changes
@@ -0,0 +1,3 @@
1
+ from .google_client import GoogleClient
2
+
3
+ __all__ = ["GoogleClient"]
@@ -0,0 +1,659 @@
1
+ from collections.abc import AsyncIterator, Iterator
2
+ from typing import Literal
3
+
4
+ from datapizza.core.cache import Cache
5
+ from datapizza.core.clients import Client, ClientResponse
6
+ from datapizza.memory import Memory
7
+ from datapizza.tools import Tool
8
+ from datapizza.tools.tool_converter import ToolConverter
9
+ from datapizza.type import (
10
+ FunctionCallBlock,
11
+ Model,
12
+ StructuredBlock,
13
+ TextBlock,
14
+ ThoughtBlock,
15
+ )
16
+
17
+ from google import genai
18
+ from google.genai import types
19
+ from google.oauth2 import service_account
20
+
21
+ from .memory_adapter import GoogleMemoryAdapter
22
+
23
+
24
+ class GoogleClient(Client):
25
+ """A client for interacting with Google's Generative AI APIs.
26
+
27
+ This class provides methods for invoking the Google GenAI API to generate responses
28
+ based on given input data. It extends the Client class.
29
+ """
30
+
31
+ def __init__(
32
+ self,
33
+ api_key: str | None = None,
34
+ model: str = "gemini-2.0-flash",
35
+ system_prompt: str = "",
36
+ temperature: float | None = None,
37
+ cache: Cache | None = None,
38
+ project_id: str | None = None,
39
+ location: str | None = None,
40
+ credentials_path: str | None = None,
41
+ use_vertexai: bool = False,
42
+ ):
43
+ """
44
+ Args:
45
+ api_key: The API key for the Google API.
46
+ model: The model to use for the Google API.
47
+ system_prompt: The system prompt to use for the Google API.
48
+ temperature: The temperature to use for the Google API.
49
+ cache: The cache to use for the Google API.
50
+ project_id: The project ID for the Google API.
51
+ location: The location for the Google API.
52
+ credentials_path: The path to the credentials for the Google API.
53
+ use_vertexai: Whether to use Vertex AI for the Google API.
54
+ """
55
+ if temperature and not 0 <= temperature <= 2:
56
+ raise ValueError("Temperature must be between 0 and 2")
57
+
58
+ super().__init__(
59
+ model_name=model,
60
+ system_prompt=system_prompt,
61
+ temperature=temperature,
62
+ cache=cache,
63
+ )
64
+ self.memory_adapter = GoogleMemoryAdapter()
65
+
66
+ try:
67
+ if use_vertexai:
68
+ if not credentials_path:
69
+ raise ValueError("credentials_path must be provided")
70
+ if not project_id:
71
+ raise ValueError("project_id must be provided")
72
+ if not location:
73
+ raise ValueError("location must be provided")
74
+
75
+ credentials = service_account.Credentials.from_service_account_file(
76
+ credentials_path,
77
+ scopes=["https://www.googleapis.com/auth/cloud-platform"],
78
+ )
79
+ self.client = genai.Client(
80
+ vertexai=True,
81
+ project=project_id,
82
+ location=location,
83
+ credentials=credentials,
84
+ )
85
+ else:
86
+ if not api_key:
87
+ raise ValueError("api_key must be provided")
88
+
89
+ self.client = genai.Client(api_key=api_key)
90
+
91
+ except Exception as e:
92
+ raise RuntimeError(
93
+ f"Failed to initialize Google GenAI client: {e!s}"
94
+ ) from None
95
+
96
+ def _convert_tool(self, tool: Tool) -> dict:
97
+ """Convert tools to Google function format"""
98
+ return ToolConverter.to_google_format(tool)
99
+
100
+ def _prepare_tools(self, tools: list[Tool] | None) -> list[types.Tool] | None:
101
+ if not tools:
102
+ return None
103
+
104
+ google_tools = []
105
+ function_declarations = []
106
+ has_google_search = False
107
+
108
+ for tool in tools:
109
+ # Check if tool has google search capability
110
+ if hasattr(tool, "name") and "google_search" in tool.name.lower():
111
+ has_google_search = True
112
+ elif isinstance(tool, Tool):
113
+ function_declarations.append(self._convert_tool(tool))
114
+ elif isinstance(tool, dict):
115
+ google_tools.append(tool)
116
+ else:
117
+ raise ValueError(f"Unknown tool type: {type(tool)}")
118
+
119
+ if function_declarations:
120
+ google_tools.append(types.Tool(function_declarations=function_declarations))
121
+
122
+ if has_google_search:
123
+ google_tools.append(types.Tool(google_search=types.GoogleSearch()))
124
+
125
+ return google_tools if google_tools else None
126
+
127
+ def _convert_tool_choice(
128
+ self, tool_choice: Literal["auto", "required", "none"] | list[str]
129
+ ) -> types.ToolConfig:
130
+ adjusted_tool_choice: types.ToolConfig
131
+ if isinstance(tool_choice, list):
132
+ adjusted_tool_choice = types.ToolConfig(
133
+ function_calling_config=types.FunctionCallingConfig(
134
+ mode="ANY", # type: ignore
135
+ allowed_function_names=tool_choice,
136
+ )
137
+ )
138
+ elif tool_choice == "required":
139
+ adjusted_tool_choice = types.ToolConfig(
140
+ function_calling_config=types.FunctionCallingConfig(mode="ANY") # type: ignore
141
+ )
142
+ elif tool_choice == "none":
143
+ adjusted_tool_choice = types.ToolConfig(
144
+ function_calling_config=types.FunctionCallingConfig(mode="NONE") # type: ignore
145
+ )
146
+ elif tool_choice == "auto":
147
+ adjusted_tool_choice = types.ToolConfig(
148
+ function_calling_config=types.FunctionCallingConfig(mode="AUTO") # type: ignore
149
+ )
150
+ return adjusted_tool_choice
151
+
152
+ def _invoke(
153
+ self,
154
+ *,
155
+ input: str,
156
+ tools: list[Tool] | None,
157
+ memory: Memory | None,
158
+ tool_choice: Literal["auto", "required", "none"] | list[str],
159
+ temperature: float | None,
160
+ max_tokens: int,
161
+ system_prompt: str | None,
162
+ **kwargs,
163
+ ) -> ClientResponse:
164
+ """Implementation of the abstract _invoke method"""
165
+ if tools is None:
166
+ tools = []
167
+ contents = self._memory_to_contents(None, input, memory)
168
+
169
+ tool_map = {tool.name: tool for tool in tools if isinstance(tool, Tool)}
170
+
171
+ prepared_tools = self._prepare_tools(tools)
172
+ config = types.GenerateContentConfig(
173
+ temperature=temperature or self.temperature,
174
+ system_instruction=system_prompt or self.system_prompt,
175
+ max_output_tokens=max_tokens or None,
176
+ tools=prepared_tools, # type: ignore
177
+ tool_config=self._convert_tool_choice(tool_choice)
178
+ if tools and any(isinstance(tool, Tool) for tool in tools)
179
+ else None,
180
+ **kwargs,
181
+ )
182
+
183
+ response = self.client.models.generate_content(
184
+ model=self.model_name,
185
+ contents=contents, # type: ignore
186
+ config=config, # type: ignore
187
+ )
188
+ return self._response_to_client_response(response, tool_map)
189
+
190
+ async def _a_invoke(
191
+ self,
192
+ *,
193
+ input: str,
194
+ tools: list[Tool] | None,
195
+ memory: Memory | None,
196
+ tool_choice: Literal["auto", "required", "none"] | list[str],
197
+ temperature: float | None,
198
+ max_tokens: int,
199
+ system_prompt: str | None,
200
+ **kwargs,
201
+ ) -> ClientResponse:
202
+ """Implementation of the abstract _invoke method"""
203
+ if tools is None:
204
+ tools = []
205
+ contents = self._memory_to_contents(None, input, memory)
206
+
207
+ tool_map = {tool.name: tool for tool in tools if isinstance(tool, Tool)}
208
+
209
+ prepared_tools = self._prepare_tools(tools)
210
+ config = types.GenerateContentConfig(
211
+ temperature=temperature or self.temperature,
212
+ system_instruction=system_prompt or self.system_prompt,
213
+ max_output_tokens=max_tokens or None,
214
+ tools=prepared_tools, # type: ignore
215
+ tool_config=self._convert_tool_choice(tool_choice)
216
+ if tools and any(isinstance(tool, Tool) for tool in tools)
217
+ else None,
218
+ **kwargs,
219
+ )
220
+
221
+ response = await self.client.aio.models.generate_content(
222
+ model=self.model_name,
223
+ contents=contents, # type: ignore
224
+ config=config, # type: ignore
225
+ )
226
+ return self._response_to_client_response(response, tool_map)
227
+
228
+ def _stream_invoke(
229
+ self,
230
+ input: str,
231
+ tools: list[Tool] | None,
232
+ memory: Memory | None,
233
+ tool_choice: Literal["auto", "required", "none"] | list[str],
234
+ temperature: float | None,
235
+ max_tokens: int,
236
+ system_prompt: str | None,
237
+ **kwargs,
238
+ ) -> Iterator[ClientResponse]:
239
+ """Implementation of the abstract _stream_invoke method"""
240
+ if tools is None:
241
+ tools = []
242
+ contents = self._memory_to_contents(None, input, memory)
243
+
244
+ prepared_tools = self._prepare_tools(tools)
245
+ config = types.GenerateContentConfig(
246
+ temperature=temperature or self.temperature,
247
+ system_instruction=system_prompt or self.system_prompt,
248
+ max_output_tokens=max_tokens or None,
249
+ tools=prepared_tools, # type: ignore
250
+ tool_config=self._convert_tool_choice(tool_choice)
251
+ if tools and any(isinstance(tool, Tool) for tool in tools)
252
+ else None,
253
+ **kwargs,
254
+ )
255
+
256
+ message_text = ""
257
+ thought_block = ThoughtBlock(content="")
258
+
259
+ for chunk in self.client.models.generate_content_stream(
260
+ model=self.model_name,
261
+ contents=contents, # type: ignore
262
+ config=config,
263
+ ):
264
+ if not chunk.candidates:
265
+ raise ValueError("No candidates in response")
266
+
267
+ finish_reason = chunk.candidates[0].finish_reason
268
+ stop_reason = (
269
+ finish_reason.value.lower()
270
+ if finish_reason is not None
271
+ else finish_reason
272
+ )
273
+
274
+ if not chunk.candidates[0].content:
275
+ raise ValueError("No content in response")
276
+
277
+ if not chunk.candidates[0].content.parts:
278
+ yield ClientResponse(
279
+ content=[],
280
+ delta=chunk.text or "",
281
+ stop_reason=stop_reason,
282
+ prompt_tokens_used=(
283
+ chunk.usage_metadata.prompt_token_count
284
+ if chunk.usage_metadata
285
+ and chunk.usage_metadata.prompt_token_count
286
+ else 0
287
+ ),
288
+ completion_tokens_used=(
289
+ chunk.usage_metadata.candidates_token_count
290
+ if chunk.usage_metadata
291
+ and chunk.usage_metadata.candidates_token_count
292
+ else 0
293
+ ),
294
+ cached_tokens_used=(
295
+ chunk.usage_metadata.cached_content_token_count
296
+ if chunk.usage_metadata
297
+ and chunk.usage_metadata.cached_content_token_count
298
+ else 0
299
+ ),
300
+ )
301
+ continue
302
+
303
+ for part in chunk.candidates[0].content.parts:
304
+ if not part.text:
305
+ continue
306
+ elif hasattr(part, "thought") and part.thought:
307
+ thought_block.content += part.text
308
+ else: # If it's not a thought, it's a message
309
+ if part.text:
310
+ message_text += str(chunk.text or "")
311
+
312
+ yield ClientResponse(
313
+ content=[],
314
+ delta=chunk.text or "",
315
+ stop_reason=stop_reason,
316
+ prompt_tokens_used=(
317
+ chunk.usage_metadata.prompt_token_count
318
+ if chunk.usage_metadata
319
+ and chunk.usage_metadata.prompt_token_count
320
+ else 0
321
+ ),
322
+ completion_tokens_used=(
323
+ chunk.usage_metadata.candidates_token_count
324
+ if chunk.usage_metadata
325
+ and chunk.usage_metadata.candidates_token_count
326
+ else 0
327
+ ),
328
+ cached_tokens_used=(
329
+ chunk.usage_metadata.cached_content_token_count
330
+ if chunk.usage_metadata
331
+ and chunk.usage_metadata.cached_content_token_count
332
+ else 0
333
+ ),
334
+ )
335
+
336
+ async def _a_stream_invoke(
337
+ self,
338
+ input: str,
339
+ tools: list[Tool] | None = None,
340
+ memory: Memory | None = None,
341
+ tool_choice: Literal["auto", "required", "none"] | list[str] = "auto",
342
+ temperature: float | None = None,
343
+ max_tokens: int | None = None,
344
+ system_prompt: str | None = None,
345
+ **kwargs,
346
+ ) -> AsyncIterator[ClientResponse]:
347
+ """Implementation of the abstract _a_stream_invoke method for Google"""
348
+ if tools is None:
349
+ tools = []
350
+ contents = self._memory_to_contents(None, input, memory)
351
+
352
+ prepared_tools = self._prepare_tools(tools)
353
+ config = types.GenerateContentConfig(
354
+ temperature=temperature or self.temperature,
355
+ system_instruction=system_prompt or self.system_prompt,
356
+ max_output_tokens=max_tokens or None,
357
+ tools=prepared_tools, # type: ignore
358
+ tool_config=self._convert_tool_choice(tool_choice)
359
+ if tools and any(isinstance(tool, Tool) for tool in tools)
360
+ else None,
361
+ **kwargs,
362
+ )
363
+
364
+ message_text = ""
365
+ thought_block = ThoughtBlock(content="")
366
+ async for chunk in await self.client.aio.models.generate_content_stream(
367
+ model=self.model_name,
368
+ contents=contents, # type: ignore
369
+ config=config,
370
+ ): # type: ignore
371
+ finish_reason = chunk.candidates[0].finish_reason
372
+ stop_reason = (
373
+ finish_reason.value.lower()
374
+ if finish_reason is not None
375
+ else finish_reason
376
+ )
377
+
378
+ # Handle the case where the response has no parts
379
+ if not chunk.candidates[0].content.parts:
380
+ yield ClientResponse(
381
+ content=[],
382
+ delta=chunk.text or "",
383
+ stop_reason=stop_reason,
384
+ prompt_tokens_used=chunk.usage_metadata.prompt_token_count
385
+ if chunk.usage_metadata
386
+ else 0,
387
+ completion_tokens_used=chunk.usage_metadata.candidates_token_count
388
+ if chunk.usage_metadata
389
+ else 0,
390
+ cached_tokens_used=chunk.usage_metadata.cached_content_token_count
391
+ if chunk.usage_metadata
392
+ else 0,
393
+ )
394
+ continue
395
+
396
+ for part in chunk.candidates[0].content.parts:
397
+ if not part.text:
398
+ continue
399
+ elif hasattr(part, "thought") and part.thought:
400
+ thought_block.content += part.text
401
+ else: # If it's not a thought, it's a message
402
+ if part.text:
403
+ message_text += chunk.text or ""
404
+ yield ClientResponse(
405
+ content=[],
406
+ delta=chunk.text or "",
407
+ stop_reason=stop_reason,
408
+ prompt_tokens_used=chunk.usage_metadata.prompt_token_count
409
+ if chunk.usage_metadata
410
+ and chunk.usage_metadata.prompt_token_count
411
+ else 0,
412
+ completion_tokens_used=chunk.usage_metadata.candidates_token_count
413
+ if chunk.usage_metadata
414
+ and chunk.usage_metadata.candidates_token_count
415
+ else 0,
416
+ cached_tokens_used=chunk.usage_metadata.cached_content_token_count
417
+ if chunk.usage_metadata
418
+ and chunk.usage_metadata.cached_content_token_count
419
+ else 0,
420
+ )
421
+
422
+ def _structured_response(
423
+ self,
424
+ input: str,
425
+ output_cls: type[Model],
426
+ memory: Memory | None,
427
+ temperature: float | None,
428
+ max_tokens: int,
429
+ system_prompt: str | None,
430
+ tools: list[Tool] | None,
431
+ tool_choice: Literal["auto", "required", "none"] | list[str] = "auto",
432
+ **kwargs,
433
+ ) -> ClientResponse:
434
+ """Implementation of the abstract _structured_response method"""
435
+ contents = self._memory_to_contents(self.system_prompt, input, memory)
436
+
437
+ prepared_tools = self._prepare_tools(tools)
438
+ response = self.client.models.generate_content(
439
+ model=self.model_name,
440
+ contents=contents, # type: ignore
441
+ config=types.GenerateContentConfig(
442
+ system_instruction=system_prompt,
443
+ temperature=temperature,
444
+ max_output_tokens=max_tokens,
445
+ response_mime_type="application/json",
446
+ tools=prepared_tools, # type: ignore
447
+ tool_config=self._convert_tool_choice(tool_choice)
448
+ if tools and any(isinstance(tool, Tool) for tool in tools)
449
+ else None,
450
+ response_schema=(
451
+ output_cls.model_json_schema()
452
+ if hasattr(output_cls, "model_json_schema")
453
+ else output_cls
454
+ ),
455
+ ),
456
+ )
457
+ if not response or not response.candidates:
458
+ raise ValueError("No response from Google GenAI")
459
+
460
+ structured_data = output_cls.model_validate_json(str(response.text))
461
+ return ClientResponse(
462
+ content=[StructuredBlock(content=structured_data)],
463
+ stop_reason=response.candidates[0].finish_reason.value.lower()
464
+ if response.candidates[0].finish_reason
465
+ else None,
466
+ prompt_tokens_used=(
467
+ response.usage_metadata.prompt_token_count
468
+ if response.usage_metadata
469
+ and response.usage_metadata.prompt_token_count
470
+ else 0
471
+ ),
472
+ completion_tokens_used=(
473
+ response.usage_metadata.candidates_token_count
474
+ if response.usage_metadata
475
+ and response.usage_metadata.candidates_token_count
476
+ else 0
477
+ ),
478
+ cached_tokens_used=(
479
+ response.usage_metadata.cached_content_token_count
480
+ if response.usage_metadata
481
+ and response.usage_metadata.cached_content_token_count
482
+ else 0
483
+ ),
484
+ )
485
+
486
+ async def _a_structured_response(
487
+ self,
488
+ input: str,
489
+ output_cls: type[Model],
490
+ memory: Memory | None,
491
+ temperature: float | None,
492
+ max_tokens: int,
493
+ system_prompt: str | None,
494
+ tools: list[Tool] | None,
495
+ tool_choice: Literal["auto", "required", "none"] | list[str] = "auto",
496
+ **kwargs,
497
+ ) -> ClientResponse:
498
+ """Implementation of the abstract _structured_response method"""
499
+ contents = self._memory_to_contents(self.system_prompt, input, memory)
500
+ prepared_tools = self._prepare_tools(tools)
501
+ response = await self.client.aio.models.generate_content(
502
+ model=self.model_name,
503
+ contents=contents, # type: ignore
504
+ config=types.GenerateContentConfig(
505
+ system_instruction=system_prompt,
506
+ temperature=temperature,
507
+ max_output_tokens=max_tokens,
508
+ response_mime_type="application/json",
509
+ tools=prepared_tools, # type: ignore
510
+ tool_config=self._convert_tool_choice(tool_choice)
511
+ if tools and any(isinstance(tool, Tool) for tool in tools)
512
+ else None,
513
+ response_schema=(
514
+ output_cls.model_json_schema()
515
+ if hasattr(output_cls, "model_json_schema")
516
+ else output_cls
517
+ ),
518
+ ),
519
+ )
520
+
521
+ if not response or not response.candidates:
522
+ raise ValueError("No response from Google GenAI")
523
+
524
+ structured_data = output_cls.model_validate_json(str(response.text))
525
+ return ClientResponse(
526
+ content=[StructuredBlock(content=structured_data)],
527
+ stop_reason=response.candidates[0].finish_reason.value.lower()
528
+ if response.candidates[0].finish_reason
529
+ else None,
530
+ prompt_tokens_used=(
531
+ response.usage_metadata.prompt_token_count
532
+ if response.usage_metadata
533
+ and response.usage_metadata.prompt_token_count
534
+ else 0
535
+ ),
536
+ completion_tokens_used=(
537
+ response.usage_metadata.candidates_token_count
538
+ if response.usage_metadata
539
+ and response.usage_metadata.candidates_token_count
540
+ else 0
541
+ ),
542
+ cached_tokens_used=(
543
+ response.usage_metadata.cached_content_token_count
544
+ if response.usage_metadata
545
+ and response.usage_metadata.cached_content_token_count
546
+ else 0
547
+ ),
548
+ )
549
+
550
+ def _embed(
551
+ self,
552
+ text: str | list[str],
553
+ model_name: str | None,
554
+ task_type: str = "RETRIEVAL_DOCUMENT",
555
+ output_dimensionality: int = 768,
556
+ title: str | None = None,
557
+ **kwargs,
558
+ ) -> list[float] | list[list[float] | None]:
559
+ """Embed a text using the model"""
560
+ response = self.client.models.embed_content(
561
+ model=model_name or self.model_name,
562
+ contents=text, # type: ignore
563
+ config=types.EmbedContentConfig(
564
+ task_type=task_type,
565
+ output_dimensionality=output_dimensionality,
566
+ title=title,
567
+ **kwargs,
568
+ ),
569
+ )
570
+ # Extract the embedding values from the response
571
+ if not response.embeddings:
572
+ return []
573
+
574
+ embeddings = [embedding.values for embedding in response.embeddings]
575
+
576
+ if isinstance(text, str) and embeddings[0]:
577
+ return embeddings[0]
578
+
579
+ return embeddings
580
+
581
+ async def _a_embed(
582
+ self,
583
+ text: str | list[str],
584
+ model_name: str | None,
585
+ task_type: str = "RETRIEVAL_DOCUMENT",
586
+ output_dimensionality: int = 768,
587
+ title: str | None = None,
588
+ **kwargs,
589
+ ) -> list[float] | list[list[float] | None]:
590
+ """Embed a text using the model"""
591
+ response = await self.client.aio.models.embed_content(
592
+ model=model_name or self.model_name,
593
+ contents=text, # type: ignore
594
+ config=types.EmbedContentConfig(
595
+ task_type=task_type,
596
+ output_dimensionality=output_dimensionality,
597
+ title=title,
598
+ **kwargs,
599
+ ),
600
+ )
601
+ # Extract the embedding values from the response
602
+ if not response.embeddings:
603
+ return []
604
+ embeddings = [embedding.values for embedding in response.embeddings]
605
+
606
+ if isinstance(text, str) and embeddings[0]:
607
+ return embeddings[0]
608
+
609
+ return embeddings
610
+
611
+ def _response_to_client_response(
612
+ self, response, tool_map: dict[str, Tool] | None = None
613
+ ) -> ClientResponse:
614
+ blocks = []
615
+ # Handle function calls if present
616
+ if hasattr(response, "function_calls") and response.function_calls:
617
+ for fc in response.function_calls:
618
+ if not tool_map:
619
+ raise ValueError("Tool map is required")
620
+
621
+ tool = tool_map.get(fc.name, None)
622
+ if not tool:
623
+ raise ValueError(f"Tool {fc.name} not found in tool map")
624
+
625
+ blocks.append(
626
+ FunctionCallBlock(
627
+ name=fc.name,
628
+ arguments=fc.args,
629
+ id=f"fc_{id(fc)}",
630
+ tool=tool,
631
+ )
632
+ )
633
+ else:
634
+ if hasattr(response, "text") and response.text:
635
+ blocks.append(TextBlock(content=response.text))
636
+
637
+ if hasattr(response, "candidates") and response.candidates:
638
+ for part in response.candidates[0].content.parts:
639
+ if not part.text:
640
+ continue
641
+ if hasattr(part, "thought") and part.thought:
642
+ blocks.append(ThoughtBlock(content=part.text))
643
+
644
+ usage_metadata = getattr(response, "usage_metadata", None)
645
+ return ClientResponse(
646
+ content=blocks,
647
+ stop_reason=(response.candidates[0].finish_reason.value.lower())
648
+ if hasattr(response, "candidates") and response.candidates
649
+ else None,
650
+ prompt_tokens_used=usage_metadata.prompt_token_count
651
+ if usage_metadata
652
+ else 0,
653
+ completion_tokens_used=usage_metadata.candidates_token_count
654
+ if usage_metadata
655
+ else 0,
656
+ cached_tokens_used=usage_metadata.cached_content_token_count
657
+ if usage_metadata
658
+ else 0,
659
+ )
@@ -0,0 +1,140 @@
1
+ import base64
2
+
3
+ from datapizza.memory.memory import Turn
4
+ from datapizza.memory.memory_adapter import MemoryAdapter
5
+ from datapizza.type import (
6
+ ROLE,
7
+ FunctionCallBlock,
8
+ FunctionCallResultBlock,
9
+ MediaBlock,
10
+ StructuredBlock,
11
+ TextBlock,
12
+ )
13
+
14
+ from google.genai import types
15
+
16
+
17
+ class GoogleMemoryAdapter(MemoryAdapter):
18
+ def _turn_to_message(self, turn: Turn) -> dict:
19
+ content = []
20
+ for block in turn:
21
+ block_dict = {}
22
+
23
+ match block:
24
+ case TextBlock():
25
+ block_dict = {"text": block.content}
26
+ case FunctionCallBlock():
27
+ block_dict = {
28
+ "function_call": {"name": block.name, "args": block.arguments}
29
+ }
30
+ case FunctionCallResultBlock():
31
+ block_dict = types.Part.from_function_response(
32
+ name=block.tool.name,
33
+ response={"result": block.result},
34
+ )
35
+ case StructuredBlock():
36
+ block_dict = {"text": str(block.content)}
37
+ case MediaBlock():
38
+ match block.media.media_type:
39
+ case "image":
40
+ block_dict = self._process_image_block(block)
41
+ case "pdf":
42
+ block_dict = self._process_pdf_block(block)
43
+
44
+ case "audio":
45
+ block_dict = self._process_audio_block(block)
46
+
47
+ case _:
48
+ raise NotImplementedError(
49
+ f"Unsupported media type: {block.media.media_type}"
50
+ )
51
+
52
+ content.append(block_dict)
53
+
54
+ return {
55
+ "role": turn.role.google_role,
56
+ "parts": (content),
57
+ }
58
+
59
+ def _process_audio_block(self, block: MediaBlock) -> types.Part:
60
+ match block.media.source_type:
61
+ case "raw":
62
+ return types.Part.from_bytes(
63
+ data=block.media.source,
64
+ mime_type="audio/mp3",
65
+ )
66
+
67
+ case "path":
68
+ with open(block.media.source, "rb") as f:
69
+ audio_bytes = f.read()
70
+
71
+ return types.Part.from_bytes(
72
+ data=audio_bytes,
73
+ mime_type="audio/mp3",
74
+ )
75
+
76
+ case _:
77
+ raise NotImplementedError(
78
+ f"Unsupported media source type: {block.media.source_type} for audio, source type supported: raw, path"
79
+ )
80
+
81
+ def _process_pdf_block(self, block: MediaBlock) -> types.Part | dict:
82
+ match block.media.source_type:
83
+ case "raw":
84
+ return types.Part.from_bytes(
85
+ data=block.media.source,
86
+ mime_type="application/pdf",
87
+ )
88
+ case "base64":
89
+ return {
90
+ "inline_data": {
91
+ "mime_type": "application/pdf",
92
+ "data": block.media.source,
93
+ }
94
+ }
95
+ case "path":
96
+ with open(block.media.source, "rb") as f:
97
+ pdf_bytes = f.read()
98
+
99
+ return {
100
+ "inline_data": {
101
+ "mime_type": "application/pdf",
102
+ "data": pdf_bytes,
103
+ }
104
+ }
105
+
106
+ case _:
107
+ raise NotImplementedError(
108
+ f"Unsupported media source type: {block.media.source_type} only supported: raw, base64, path"
109
+ )
110
+
111
+ def _process_image_block(self, block: MediaBlock) -> dict:
112
+ match block.media.source_type:
113
+ case "url":
114
+ return types.Part.from_uri(
115
+ file_uri=block.media.source,
116
+ mime_type=f"image/{block.media.extension}",
117
+ ) # type: ignore
118
+ case "base64":
119
+ return {
120
+ "inline_data": {
121
+ "mime_type": f"image/{block.media.extension}",
122
+ "data": block.media.source,
123
+ }
124
+ }
125
+ case "path":
126
+ with open(block.media.source, "rb") as image_file:
127
+ base64_image = base64.b64encode(image_file.read()).decode("utf-8")
128
+ return {
129
+ "inline_data": {
130
+ "mime_type": f"image/{block.media.extension}",
131
+ "data": base64_image,
132
+ }
133
+ }
134
+ case _:
135
+ raise NotImplementedError(
136
+ f"Unsupported media source type: {block.media.source_type} for image, only url, base64, path are supported"
137
+ )
138
+
139
+ def _text_to_message(self, text: str, role: ROLE) -> dict:
140
+ return {"role": role.google_role, "parts": [{"text": text}]}
@@ -0,0 +1,57 @@
1
+ # Build system configuration
2
+ [build-system]
3
+ requires = ["hatchling"]
4
+ build-backend = "hatchling.build"
5
+
6
+ # Project metadata
7
+ [project]
8
+ name = "datapizza-ai-clients-google"
9
+ version = "0.0.1"
10
+ description = "Google (Gemini) client for the datapizza-ai framework"
11
+ readme = "README.md"
12
+ license = {text = "MIT"}
13
+
14
+ requires-python = ">=3.10.0,<3.13.0"
15
+ classifiers = [
16
+ "Programming Language :: Python :: 3",
17
+ "License :: OSI Approved :: MIT License",
18
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
19
+ "Topic :: Software Development :: Libraries :: Application Frameworks",
20
+ ]
21
+ dependencies = [
22
+ "datapizza-ai-core>=0.0.1",
23
+ "google-genai>=1.3.0,<2.0.0",
24
+ ]
25
+
26
+ # Development dependencies
27
+ [dependency-groups]
28
+ dev = [
29
+ "deptry>=0.23.0",
30
+ "pytest",
31
+ "ruff>=0.11.5",
32
+ ]
33
+
34
+ # Hatch build configuration
35
+ [tool.hatch.build.targets.sdist]
36
+ include = ["datapizza"]
37
+ exclude = ["**/BUILD"]
38
+
39
+ [tool.hatch.build.targets.wheel]
40
+ include = ["datapizza"]
41
+ exclude = ["**/BUILD"]
42
+
43
+ # Ruff configuration
44
+ [tool.ruff]
45
+ line-length = 88
46
+
47
+ [tool.ruff.lint]
48
+ select = [
49
+ "W", # pycodestyle warnings
50
+ "F", # pyflakes
51
+ "B", # flake8-bugbear
52
+ "I", # isort
53
+ "UP", # pyupgrade
54
+ "SIM", # flake8-simplify
55
+ "RUF", # Ruff-specific rules
56
+ "C4", # flake8-comprehensions
57
+ ]