Google GenAI integration for Haystack
Module haystack_integrations.components.generators.google_genai.chat.chat_generator
GoogleGenAIChatGenerator
A component for generating chat completions using Google's Gemini models via the Google Gen AI SDK.
This component provides an interface to Google's Gemini models through the new google-genai SDK, supporting models like gemini-2.0-flash and other Gemini variants.
Usage example
from haystack.dataclasses.chat_message import ChatMessage
from haystack.tools import Tool, Toolset
from haystack_integrations.components.generators.google_genai import GoogleGenAIChatGenerator
# Initialize the chat generator
chat_generator = GoogleGenAIChatGenerator(model="gemini-2.0-flash")
# Generate a response
messages = [ChatMessage.from_user("Tell me about the future of AI")]
response = chat_generator.run(messages=messages)
print(response["replies"][0].text)
# Tool usage example
def weather_function(city: str):
return f"The weather in {city} is sunny and 25°C"
weather_tool = Tool(
name="weather",
description="Get weather information for a city",
parameters={"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
function=weather_function
)
# Can use either List[Tool] or Toolset
chat_generator_with_tools = GoogleGenAIChatGenerator(
model="gemini-2.0-flash",
tools=[weather_tool] # or tools=Toolset([weather_tool])
)
messages = [ChatMessage.from_user("What's the weather in Paris?")]
response = chat_generator_with_tools.run(messages=messages)
GoogleGenAIChatGenerator.__init__
def __init__(*,
api_key: Secret = Secret.from_env_var(
["GOOGLE_API_KEY", "GEMINI_API_KEY"]),
model: str = "gemini-2.0-flash",
generation_kwargs: Optional[Dict[str, Any]] = None,
safety_settings: Optional[List[Dict[str, Any]]] = None,
streaming_callback: Optional[StreamingCallbackT] = None,
tools: Optional[Union[List[Tool], Toolset]] = None)
Initialize a GoogleGenAIChatGenerator instance.
Arguments:
api_key
: Google API key, defaults to theGOOGLE_API_KEY
andGEMINI_API_KEY
environment variables, see https://ai.google.dev/gemini-api/docs/api-key for more information.model
: Name of the model to use (e.g., "gemini-2.0-flash")generation_kwargs
: Configuration for generation (temperature, max_tokens, etc.)safety_settings
: Safety settings for content filteringstreaming_callback
: A callback function that is called when a new token is received from the stream.tools
: A list of Tool objects or a Toolset that the model can use. Each tool should have a unique name.
GoogleGenAIChatGenerator.to_dict
def to_dict() -> Dict[str, Any]
Serializes the component to a dictionary.
Returns:
Dictionary with serialized data.
GoogleGenAIChatGenerator.from_dict
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "GoogleGenAIChatGenerator"
Deserializes the component from a dictionary.
Arguments:
data
: Dictionary to deserialize from.
Returns:
Deserialized component.
GoogleGenAIChatGenerator.run
@component.output_types(replies=List[ChatMessage])
def run(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[Union[List[Tool], Toolset]] = None) -> Dict[str, Any]
Run the Google Gen AI chat generator on the given input data.
Arguments:
messages
: A list of ChatMessage instances representing the input messages.generation_kwargs
: Configuration for generation. If provided, it will override the default config.safety_settings
: Safety settings for content filtering. If provided, it will override the default settings.streaming_callback
: A callback function that is called when a new token is received from the stream.tools
: A list of Tool objects or a Toolset that the model can use. If provided, it will override the tools set during initialization.
Raises:
RuntimeError
: If there is an error in the Google Gen AI chat generation.ValueError
: If a ChatMessage does not contain at least one of TextContent, ToolCall, or ToolCallResult or if the role in ChatMessage is different from User, System, Assistant.
Returns:
A dictionary with the following keys:
replies
: A list containing the generated ChatMessage responses.
GoogleGenAIChatGenerator.run_async
@component.output_types(replies=List[ChatMessage])
async def run_async(
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[Union[List[Tool], Toolset]] = None) -> Dict[str, Any]
Async version of the run method. Run the Google Gen AI chat generator on the given input data.
Arguments:
messages
: A list of ChatMessage instances representing the input messages.generation_kwargs
: Configuration for generation. If provided, it will override the default config.safety_settings
: Safety settings for content filtering. If provided, it will override the default settings.streaming_callback
: A callback function that is called when a new token is received from the stream.tools
: A list of Tool objects or a Toolset that the model can use. If provided, it will override the tools set during initialization.
Raises:
RuntimeError
: If there is an error in the Google Gen AI chat generation.ValueError
: If a ChatMessage does not contain at least one of TextContent, ToolCall, or ToolCallResult or if the role in ChatMessage is different from User, System, Assistant.
Returns:
A dictionary with the following keys:
replies
: A list containing the generated ChatMessage responses.
Module haystack_integrations.components.embedders.google_genai.document_embedder
GoogleGenAIDocumentEmbedder
Computes document embeddings using Google AI models.
Usage example
from haystack import Document
from haystack_integrations.components.embedders.google_genai import GoogleGenAIDocumentEmbedder
doc = Document(content="I love pizza!")
document_embedder = GoogleGenAIDocumentEmbedder()
result = document_embedder.run([doc])
print(result['documents'][0].embedding)
# [0.017020374536514282, -0.023255806416273117, ...]
GoogleGenAIDocumentEmbedder.__init__
def __init__(*,
api_key: Secret = Secret.from_env_var(
["GOOGLE_API_KEY", "GEMINI_API_KEY"]),
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,
embedding_separator: str = "\n",
config: Optional[Dict[str, Any]] = None) -> None
Creates an GoogleGenAIDocumentEmbedder component.
Arguments:
api_key
: The Google API key. You can set it with the environment variableGOOGLE_API_KEY
orGEMINI_API_KEY
, or pass it via this parameter during initialization.model
: The name of the model to use for calculating embeddings. The default model istext-embedding-ada-002
.prefix
: A string to add at the beginning of each text.suffix
: A string to add at the end of each text.batch_size
: Number of documents to embed at once.progress_bar
: IfTrue
, shows a progress bar when running.meta_fields_to_embed
: List of metadata fields to embed along with the document text.embedding_separator
: Separator used to concatenate the metadata fields to the document text.config
: A dictionary of keyword arguments to configure embedding content configurationtypes.EmbedContentConfig
. If not specified, it defaults to {"task_type": "SEMANTIC_SIMILARITY"}. For more information, see the Google AI Task types.
GoogleGenAIDocumentEmbedder.to_dict
def to_dict() -> Dict[str, Any]
Serializes the component to a dictionary.
Returns:
Dictionary with serialized data.
GoogleGenAIDocumentEmbedder.from_dict
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "GoogleGenAIDocumentEmbedder"
Deserializes the component from a dictionary.
Arguments:
data
: Dictionary to deserialize from.
Returns:
Deserialized component.
GoogleGenAIDocumentEmbedder.run
@component.output_types(documents=List[Document], meta=Dict[str, Any])
def run(
documents: List[Document]
) -> Union[Dict[str, List[Document]], Dict[str, Any]]
Embeds a list of documents.
Arguments:
documents
: A list of documents to embed.
Returns:
A dictionary with the following keys:
documents
: A list of documents with embeddings.meta
: Information about the usage of the model.
GoogleGenAIDocumentEmbedder.run_async
@component.output_types(documents=List[Document], meta=Dict[str, Any])
async def run_async(
documents: List[Document]
) -> Union[Dict[str, List[Document]], Dict[str, Any]]
Embeds a list of documents asynchronously.
Arguments:
documents
: A list of documents to embed.
Returns:
A dictionary with the following keys:
documents
: A list of documents with embeddings.meta
: Information about the usage of the model.
Module haystack_integrations.components.embedders.google_genai.text_embedder
GoogleGenAITextEmbedder
Embeds strings using Google AI models.
You can use it to embed user query and send it to an embedding Retriever.
Usage example
from haystack_integrations.components.embedders.google_genai import GoogleGenAITextEmbedder
text_to_embed = "I love pizza!"
text_embedder = GoogleGenAITextEmbedder()
print(text_embedder.run(text_to_embed))
# {'embedding': [0.017020374536514282, -0.023255806416273117, ...],
# 'meta': {'model': 'text-embedding-004-v2',
# 'usage': {'prompt_tokens': 4, 'total_tokens': 4}}}
GoogleGenAITextEmbedder.__init__
def __init__(*,
api_key: Secret = Secret.from_env_var(
["GOOGLE_API_KEY", "GEMINI_API_KEY"]),
model: str = "text-embedding-004",
prefix: str = "",
suffix: str = "",
config: Optional[Dict[str, Any]] = None) -> None
Creates an GoogleGenAITextEmbedder component.
Arguments:
api_key
: The Google API key. You can set it with the environment variableGOOGLE_API_KEY
orGEMINI_API_KEY
, or pass it via this parameter during initialization.model
: The name of the model to use for calculating embeddings. The default model istext-embedding-004
.prefix
: A string to add at the beginning of each text to embed.suffix
: A string to add at the end of each text to embed.config
: A dictionary of keyword arguments to configure embedding content configurationtypes.EmbedContentConfig
. If not specified, it defaults to {"task_type": "SEMANTIC_SIMILARITY"}. For more information, see the Google AI Task types.
GoogleGenAITextEmbedder.to_dict
def to_dict() -> Dict[str, Any]
Serializes the component to a dictionary.
Returns:
Dictionary with serialized data.
GoogleGenAITextEmbedder.from_dict
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "GoogleGenAITextEmbedder"
Deserializes the component from a dictionary.
Arguments:
data
: Dictionary to deserialize from.
Returns:
Deserialized component.
GoogleGenAITextEmbedder.run
@component.output_types(embedding=List[float], meta=Dict[str, Any])
def run(text: str) -> Union[Dict[str, List[float]], Dict[str, Any]]
Embeds a single string.
Arguments:
text
: Text to embed.
Returns:
A dictionary with the following keys:
embedding
: The embedding of the input text.meta
: Information about the usage of the model.
GoogleGenAITextEmbedder.run_async
@component.output_types(embedding=List[float], meta=Dict[str, Any])
async def run_async(
text: str) -> Union[Dict[str, List[float]], Dict[str, Any]]
Asynchronously embed a single string.
This is the asynchronous version of the run
method. It has the same parameters and return values
but can be used with await
in async code.
Arguments:
text
: Text to embed.
Returns:
A dictionary with the following keys:
embedding
: The embedding of the input text.meta
: Information about the usage of the model.