GoogleGenAIChatGenerator

This component enables chat completion using Google Gemini models through Google Gen AI SDK.


Most common position in a pipeline	After a ChatPromptBuilder
Mandatory init variables	"api_key": A Google API key. Can be set with `GOOGLE_API_KEY` env var.
Mandatory run variables	“messages”: A list of `ChatMessage` objects representing the chat
Output variables	“replies”: A list of alternative replies of the model to the input chat
API reference	Google GenAI
GitHub link	https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/google_genai

Overview

GoogleGenAIChatGenerator supports gemini-2.0-flash (default), gemini-2.5-pro-exp-03-25, gemini-1.5-pro, and gemini-1.5-flash models.

Streaming

This Generator supports streaming the tokens from the LLM directly in output. To do so, pass a function to the streaming_callback init parameter.

Authentication

Google Gen AI is compatible with both the Gemini Developer API and the Vertex AI API.

To use this component with the Gemini Developer API and get an API key, visit Google AI Studio.
To use this component with the Vertex AI API, visit Google Cloud > Vertex AI.

The component uses a GOOGLE_API_KEY or GEMINI_API_KEY environment variable by default. Otherwise, you can pass an API key at initialization with a Secret and Secret.from_token static method:

embedder = GoogleGenAITextEmbedder(api_key=Secret.from_token("<your-api-key>"))

The following examples show how to use the component with the Gemini Developer API and the Vertex AI API.

Gemini Developer API (API Key Authentication)

from haystack_integrations.components.generators.google_genai import GoogleGenAIChatGenerator

# set the environment variable (GOOGLE_API_KEY or GEMINI_API_KEY)
chat_generator = GoogleGenAIChatGenerator()

Vertex AI (Application Default Credentials)

from haystack_integrations.components.generators.google_genai import GoogleGenAIChatGenerator

# Using Application Default Credentials (requires gcloud auth setup)
chat_generator = GoogleGenAIChatGenerator(
    api="vertex",
    vertex_ai_project="my-project",
    vertex_ai_location="us-central1",
)

Vertex AI (API Key Authentication)

from haystack_integrations.components.generators.google_genai import GoogleGenAIChatGenerator

# set the environment variable (GOOGLE_API_KEY or GEMINI_API_KEY)
chat_generator = GoogleGenAIChatGenerator(api="vertex")

Usage

To start using this integration, install the package with:

pip install google-genai-haystack

On its own

from haystack.dataclasses.chat_message import ChatMessage
from haystack_integrations.components.generators.google_genai import GoogleGenAIChatGenerator

# Initialize the chat generator
chat_generator = GoogleGenAIChatGenerator()

# Generate a response
messages = [ChatMessage.from_user("Tell me about movie Shawshank Redemption")]
response = chat_generator.run(messages=messages)
print(response["replies"][0].text)

You can also easily use function calls. First, define the function locally and convert into a Tool:

from typing import Annotated
from haystack.tools import create_tool_from_function

# example function to get the current weather
def get_current_weather(
    location: Annotated[str, "The city for which to get the weather, e.g. 'San Francisco'"] = "Munich",
    unit: Annotated[str, "The unit for the temperature, e.g. 'celsius'"] = "celsius",
) -> str:
    return f"The weather in {location} is sunny. The temperature is 20 {unit}."

tool = create_tool_from_function(get_current_weather)

Create a new instance of GoogleGenAIChatGenerator to set the tools and a ToolInvoker to invoke the tools.

import os
from haystack_integrations.components.generators.google_genai import GoogleGenAIChatGenerator
from haystack.components.tools import ToolInvoker

os.environ["GOOGLE_API_KEY"] = "<MY_API_KEY>"

genai_chat = GoogleGenAIChatGenerator(tools=[tool])

tool_invoker = ToolInvoker(tools=[tool])

And then ask a question:

from haystack.dataclasses import ChatMessage

messages = [ChatMessage.from_user("What is the temperature in celsius in Berlin?")]
res = genai_chat.run(messages=messages)

print(res["replies"][0].tool_calls)
>>> [ToolCall(tool_name='get_current_weather',
>>>           arguments={'unit': 'celsius', 'location': 'Berlin'}, id=None)]

tool_messages = tool_invoker.run(messages=replies)["tool_messages"]
messages = user_message + replies + tool_messages

messages += res["replies"][0] + [ChatMessage.from_function(content=weather, name="get_current_weather")]

final_replies = genai_chat.run(messages=messages)["replies"]
print(final_replies[0].text)
>>> The temperature in Berlin is 20 degrees Celsius.

With Streaming

from haystack.dataclasses.chat_message import ChatMessage
from haystack.dataclasses import StreamingChunk
from haystack_integrations.components.generators.google_genai import GoogleGenAIChatGenerator

def streaming_callback(chunk: StreamingChunk):
    print(chunk.content, end='', flush=True)

# Initialize with streaming callback
chat_generator = GoogleGenAIChatGenerator(
    streaming_callback=streaming_callback
)

# Generate a streaming response
messages = [ChatMessage.from_user("Write a short story")]
response = chat_generator.run(messages=messages)
# Text will stream in real-time through the callback

In a pipeline

import os
from haystack.components.builders import ChatPromptBuilder
from haystack.dataclasses import ChatMessage
from haystack import Pipeline
from haystack_integrations.components.generators.google_genai import GoogleGenAIChatGenerator

# no parameter init, we don't use any runtime template variables
prompt_builder = ChatPromptBuilder()

os.environ["GOOGLE_API_KEY"] = "<MY_API_KEY>"
genai_chat = GoogleGenAIChatGenerator()

pipe = Pipeline()
pipe.add_component("prompt_builder", prompt_builder)
pipe.add_component("genai", genai_chat)
pipe.connect("prompt_builder.prompt", "genai.messages")

location = "Rome"
messages = [ChatMessage.from_user("Tell me briefly about {{location}} history")]
res = pipe.run(data={"prompt_builder": {"template_variables":{"location": location}, "template": messages}})

print(res)

Updated 27 days ago