
Most common position in a pipeline	After a ChatPromptBuilder
Mandatory init variables	"api_key": A Google AI Studio 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 AI
GitHub link	https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/google_ai

GoogleAIGeminiChatGenerator supports gemini-pro, gemini-1.5-pro, and gemini-1.5-flash models.

For available models, see https://ai.google.dev/gemini-api/docs/models/gemini.

Parameters Overview

GoogleAIGeminiChatGenerator uses a Google Makersuite API key for authentication. You can write this key in an api_key parameter or as a GOOGLE_API_KEY environment variable (recommended).

To get an API key, visit the Google Makersuite website.

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.

Usage

To begin working with GoogleAIGeminiChatGenerator, install the google-ai-haystack package:

pip install google-ai-haystack

On its own

Basic usage:

import os
from haystack.dataclasses import ChatMessage
from haystack_integrations.components.generators.google_ai import GoogleAIGeminiChatGenerator

os.environ["GOOGLE_API_KEY"] = "<MY_API_KEY>"
gemini_chat = GoogleAIGeminiChatGenerator()

messages = [ChatMessage.from_user("Tell me the name of a movie")]
res = gemini_chat.run(messages)

print(res["replies"][0].content)
>>> The Shawshank Redemption

messages += [res["replies"], ChatMessage.from_user("Who's the main actor?")]
res = gemini_chat.run(messages)

print(res["replies"][0].content)
>>> Tim Robbins

When chatting with Gemini Pro, you can also easily use function calls. First, define the function locally:

from google.ai.generativelanguage import FunctionDeclaration, Tool

get_current_weather_func = FunctionDeclaration(
  name="get_current_weather",
  description="Get the current weather in a given location",
  parameters={
    "type_": "OBJECT",
    "properties": {
      "location": {"type_": "STRING", "description": "The city and state, e.g. San Francisco, CA"},
      "unit": {
        "type_": "STRING",
        "enum": [
          "celsius",
          "fahrenheit",
        ],
      },
    },
    "required": ["location"],
  },
)

tool = Tool(function_declarations=[get_current_weather_func])

For demonstration purposes, we’re also creating a function that will always tell us the same thing:

def get_current_weather(location: str, unit: str = "celsius"):
  return {"weather": "sunny", "temperature": 21.8, "unit": unit}

Create a new instance of GoogleAIGeminiChatGenerator to set the tools:

import os
from haystack_integrations.components.generators.google_ai import GoogleAIGeminiChatGenerator

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

gemini_chat = GoogleAIGeminiChatGenerator(model="gemini-pro", 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 = gemini_chat.run(messages=messages)

print(res["replies"][0].content)
>>> {'location': 'Berlin, Germany', 'unit': 'metric'}

weather = get_current_weather(**res["replies"][0].content)

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

res = gemini_chat.run(messages=messages)
print(res["replies"][0].content)
>>> 21.8 Celsius

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_ai import GoogleAIGeminiChatGenerator


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

os.environ["GOOGLE_API_KEY"] = "<MY_API_KEY>"
gemini_chat = GoogleAIGeminiChatGenerator()

pipe = Pipeline()
pipe.add_component("prompt_builder", prompt_builder)
pipe.add_component("gemini", gemini_chat)
pipe.connect("prompt_builder.prompt", "gemini.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)

>>> - **753 B.C.:** Traditional date of the founding of Rome by Romulus and Remus.
>>> - **509 B.C.:** Establishment of the Roman Republic, replacing the Etruscan monarchy.
>>> - **492-264 B.C.:** Series of wars against neighboring tribes, resulting in the expansion of the Roman Republic's territory.
>>> - **264-146 B.C.:** Three Punic Wars against Carthage, resulting in the destruction of Carthage and the Roman Republic becoming the dominant power in the Mediterranean.
>>> - **133-73 B.C.:** Series of civil wars and slave revolts, leading to the rise of Julius Caesar.
>>> - **49 B.C.:** Julius Caesar crosses the Rubicon River, starting the Roman Civil War.
>>> - **44 B.C.:** Julius Caesar is assassinated, leading to the Second Triumvirate of Octavian, Mark Antony, and Lepidus.
>>> - **31 B.C.:** Battle of Actium, where Octavian defeats Mark Antony and Cleopatra, becoming the sole ruler of Rome.
>>> - **27 B.C.:** The Roman Republic is transformed into the Roman Empire, with Octavian becoming the first Roman emperor, known as Augustus.
>>> - **1st century A.D.:** The Roman Empire reaches its greatest extent, stretching from Britain to Egypt.
>>> - **3rd century A.D.:** The Roman Empire begins to decline, facing internal instability, invasions by Germanic tribes, and the rise of Christianity.
>>> - **476 A.D.:** The last Western Roman emperor, Romulus Augustulus, is overthrown by the Germanic leader Odoacer, marking the end of the Roman Empire in the West.