Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 2 additions & 8 deletions integrations/google_genai/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ name = "google-genai-haystack"
dynamic = ["version"]
description = 'Use models like Gemini via Google Gen AI SDK'
readme = "README.md"
requires-python = ">=3.9"
requires-python = ">=3.10"
license = "Apache-2.0"
keywords = []
authors = [
Expand All @@ -18,15 +18,14 @@ classifiers = [
"License :: OSI Approved :: Apache Software License",
"Development Status :: 4 - Beta",
"Programming Language :: Python",
"Programming Language :: Python :: 3.9",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
"Programming Language :: Python :: 3.13",
"Programming Language :: Python :: Implementation :: CPython",
"Programming Language :: Python :: Implementation :: PyPy",
]
dependencies = ["haystack-ai>=2.19.0", "google-genai[aiohttp]>=1.18.0", "jsonref>=1.0.0"]
dependencies = ["haystack-ai>=2.22.0", "google-genai[aiohttp]>=1.18.0", "jsonref>=1.0.0"]

[project.urls]
Documentation = "https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/google_genai#readme"
Expand Down Expand Up @@ -87,7 +86,6 @@ module = [
ignore_missing_imports = true

[tool.ruff]
target-version = "py39"
line-length = 120

[tool.ruff.lint]
Expand Down Expand Up @@ -140,10 +138,6 @@ ignore = [
# Allow function call argument defaults e.g. `Secret.from_env_var`
"B008",
]
unfixable = [
# Don't touch unused imports
"F401",
]

[tool.ruff.lint.isort]
known-first-party = ["haystack_integrations"]
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
#
# SPDX-License-Identifier: Apache-2.0

from typing import Literal, Optional
from typing import Literal

from google.genai import Client
from haystack import logging
Expand All @@ -14,8 +14,8 @@
def _get_client(
api_key: Secret,
api: Literal["gemini", "vertex"],
vertex_ai_project: Optional[str],
vertex_ai_location: Optional[str],
vertex_ai_project: str | None,
vertex_ai_location: str | None,
) -> Client:
"""
Internal utility function to get a Google GenAI client.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
#
# SPDX-License-Identifier: Apache-2.0

from typing import Any, Literal, Optional, Union
from typing import Any, Literal

from google.genai import types
from haystack import Document, component, default_from_dict, default_to_dict, logging
Expand Down Expand Up @@ -75,16 +75,16 @@ def __init__(
*,
api_key: Secret = Secret.from_env_var(["GOOGLE_API_KEY", "GEMINI_API_KEY"], strict=False),
api: Literal["gemini", "vertex"] = "gemini",
vertex_ai_project: Optional[str] = None,
vertex_ai_location: Optional[str] = None,
vertex_ai_project: str | None = None,
vertex_ai_location: str | None = None,
model: str = "text-embedding-004",
prefix: str = "",
suffix: str = "",
batch_size: int = 32,
progress_bar: bool = True,
meta_fields_to_embed: Optional[list[str]] = None,
meta_fields_to_embed: list[str] | None = None,
embedding_separator: str = "\n",
config: Optional[dict[str, Any]] = None,
config: dict[str, Any] | None = None,
) -> None:
"""
Creates an GoogleGenAIDocumentEmbedder component.
Expand Down Expand Up @@ -195,7 +195,7 @@ def _prepare_texts_to_embed(self, documents: list[Document]) -> list[str]:

def _embed_batch(
self, texts_to_embed: list[str], batch_size: int
) -> tuple[list[Optional[list[float]]], dict[str, Any]]:
) -> tuple[list[list[float] | None], dict[str, Any]]:
"""
Embed a list of texts in batches.
"""
Expand Down Expand Up @@ -227,7 +227,7 @@ def _embed_batch(

async def _embed_batch_async(
self, texts_to_embed: list[str], batch_size: int
) -> tuple[list[Optional[list[float]]], dict[str, Any]]:
) -> tuple[list[list[float] | None], dict[str, Any]]:
"""
Embed a list of texts in batches asynchronously.
"""
Expand Down Expand Up @@ -257,7 +257,7 @@ async def _embed_batch_async(
return all_embeddings, meta

@component.output_types(documents=list[Document], meta=dict[str, Any])
def run(self, documents: list[Document]) -> Union[dict[str, list[Document]], dict[str, Any]]:
def run(self, documents: list[Document]) -> dict[str, list[Document]] | dict[str, Any]:
"""
Embeds a list of documents.

Expand All @@ -281,13 +281,13 @@ def run(self, documents: list[Document]) -> Union[dict[str, list[Document]], dic
meta: dict[str, Any]
embeddings, meta = self._embed_batch(texts_to_embed=texts_to_embed, batch_size=self._batch_size)

for doc, emb in zip(documents, embeddings):
for doc, emb in zip(documents, embeddings, strict=True):
doc.embedding = emb

return {"documents": documents, "meta": meta}

@component.output_types(documents=list[Document], meta=dict[str, Any])
async def run_async(self, documents: list[Document]) -> Union[dict[str, list[Document]], dict[str, Any]]:
async def run_async(self, documents: list[Document]) -> dict[str, list[Document]] | dict[str, Any]:
"""
Embeds a list of documents asynchronously.

Expand All @@ -310,7 +310,7 @@ async def run_async(self, documents: list[Document]) -> Union[dict[str, list[Doc

embeddings, meta = await self._embed_batch_async(texts_to_embed=texts_to_embed, batch_size=self._batch_size)

for doc, emb in zip(documents, embeddings):
for doc, emb in zip(documents, embeddings, strict=True):
doc.embedding = emb

return {"documents": documents, "meta": meta}
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
#
# SPDX-License-Identifier: Apache-2.0

from typing import Any, Literal, Optional, Union
from typing import Any, Literal

from google.genai import types
from haystack import component, default_from_dict, default_to_dict, logging
Expand Down Expand Up @@ -76,12 +76,12 @@ def __init__(
*,
api_key: Secret = Secret.from_env_var(["GOOGLE_API_KEY", "GEMINI_API_KEY"], strict=False),
api: Literal["gemini", "vertex"] = "gemini",
vertex_ai_project: Optional[str] = None,
vertex_ai_location: Optional[str] = None,
vertex_ai_project: str | None = None,
vertex_ai_location: str | None = None,
model: str = "text-embedding-004",
prefix: str = "",
suffix: str = "",
config: Optional[dict[str, Any]] = None,
config: dict[str, Any] | None = None,
) -> None:
"""
Creates an GoogleGenAITextEmbedder component.
Expand Down Expand Up @@ -177,7 +177,7 @@ def _prepare_output(self, result: types.EmbedContentResponse) -> dict[str, Any]:
return {"embedding": embedding, "meta": {"model": self._model_name}}

@component.output_types(embedding=list[float], meta=dict[str, Any])
def run(self, text: str) -> Union[dict[str, list[float]], dict[str, Any]]:
def run(self, text: str) -> dict[str, list[float]] | dict[str, Any]:
"""
Embeds a single string.

Expand All @@ -194,7 +194,7 @@ def run(self, text: str) -> Union[dict[str, list[float]], dict[str, Any]]:
return self._prepare_output(result=response)

@component.output_types(embedding=list[float], meta=dict[str, Any])
async def run_async(self, text: str) -> Union[dict[str, list[float]], dict[str, Any]]:
async def run_async(self, text: str) -> dict[str, list[float]] | dict[str, Any]:
"""
Asynchronously embed a single string.

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
import json
from collections.abc import AsyncIterator, Iterator
from datetime import datetime, timezone
from typing import Any, Literal, Optional
from typing import Any, Literal

from google.genai import types
from haystack import logging
Expand Down Expand Up @@ -456,13 +456,13 @@ def __init__(
*,
api_key: Secret = Secret.from_env_var(["GOOGLE_API_KEY", "GEMINI_API_KEY"], strict=False),
api: Literal["gemini", "vertex"] = "gemini",
vertex_ai_project: Optional[str] = None,
vertex_ai_location: Optional[str] = None,
vertex_ai_project: str | None = None,
vertex_ai_location: str | None = None,
model: str = "gemini-2.5-flash",
generation_kwargs: Optional[dict[str, Any]] = None,
safety_settings: Optional[list[dict[str, Any]]] = None,
streaming_callback: Optional[StreamingCallbackT] = None,
tools: Optional[ToolsType] = None,
generation_kwargs: dict[str, Any] | None = None,
safety_settings: list[dict[str, Any]] | None = None,
streaming_callback: StreamingCallbackT | None = None,
tools: ToolsType | None = None,
):
"""
Initialize a GoogleGenAIChatGenerator instance.
Expand Down Expand Up @@ -801,10 +801,10 @@ def _process_thinking_config(self, generation_kwargs: dict[str, Any]) -> dict[st
def run(
self,
messages: list[ChatMessage],
generation_kwargs: Optional[dict[str, Any]] = None,
safety_settings: Optional[list[dict[str, Any]]] = None,
streaming_callback: Optional[StreamingCallbackT] = None,
tools: Optional[ToolsType] = None,
generation_kwargs: dict[str, Any] | None = None,
safety_settings: list[dict[str, Any]] | None = None,
streaming_callback: StreamingCallbackT | None = None,
tools: ToolsType | None = None,
) -> dict[str, Any]:
"""
Run the Google Gen AI chat generator on the given input data.
Expand Down Expand Up @@ -909,10 +909,10 @@ def run(
async def run_async(
self,
messages: list[ChatMessage],
generation_kwargs: Optional[dict[str, Any]] = None,
safety_settings: Optional[list[dict[str, Any]]] = None,
streaming_callback: Optional[StreamingCallbackT] = None,
tools: Optional[ToolsType] = None,
generation_kwargs: dict[str, Any] | None = None,
safety_settings: list[dict[str, Any]] | None = None,
streaming_callback: StreamingCallbackT | None = None,
tools: ToolsType | None = None,
) -> dict[str, Any]:
"""
Async version of the run method. Run the Google Gen AI chat generator on the given input data.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,12 @@
#
# SPDX-License-Identifier: Apache-2.0

from typing import Any, Union
from typing import Any


def remove_key_from_schema(
schema: Union[dict[str, Any], list[Any], Any], target_key: str
) -> Union[dict[str, Any], list[Any], Any]:
schema: dict[str, Any] | list[Any] | Any, target_key: str
) -> dict[str, Any] | list[Any] | Any:
"""
Recursively traverse a schema and remove all occurrences of the target key.

Expand Down
Loading