Utiliser le moteur RAG Vertex AI dans l'API Gemini Live

La génération augmentée par récupération (RAG) est une technique utilisée pour récupérer et fournir des informations pertinentes aux LLM afin de générer des réponses vérifiables. Ces informations peuvent inclure des informations récentes, un thème et un contexte, ou une vérité terrain.

Cette page explique comment utiliser le moteur RAG Vertex AI avec l'API Gemini Live, qui vous permet de spécifier et de récupérer des informations à partir du corpus RAG.

Prérequis

Avant de pouvoir utiliser le moteur RAG de Vertex AI avec l'API Live multimodale, vous devez remplir les conditions préalables suivantes :

  1. Activez l'API RAG dans Vertex AI.

  2. Créez l'exemple de corpus RAG.

  3. Pour importer des fichiers dans le corpus RAG, consultez l'exemple d'API d'importation de fichiers RAG.

Configurer

Vous pouvez utiliser le moteur RAG Vertex AI avec l'API Live en spécifiant le moteur RAG Vertex AI comme outil. L'exemple de code suivant montre comment spécifier le moteur RAG de Vertex AI en tant qu'outil :

Remplacez les variables suivantes :

  • YOUR_PROJECT_ID : ID de votre projet Google Cloud .
  • YOUR_CORPUS_ID : ID de votre corpus.
  • YOUR_LOCATION : région dans laquelle traiter la requête.
PROJECT_ID = "YOUR_PROJECT_ID"
RAG_CORPUS_ID = "YOUR_CORPUS_ID"
LOCATION = "YOUR_LOCATION"

TOOLS = {
  "retrieval": {
    "vertex_rag_store": {
    "rag_resources": {
          "rag_corpus": "projects/${PROJECT_ID}/locations/${LOCATION}/ragCorpora/${RAG_CORPUS_ID}"
        }
     }
}

Utiliser Websocket pour communiquer en temps réel

Pour activer la communication en temps réel entre un client et un serveur, vous devez utiliser un Websocket. Ces exemples de code montrent comment utiliser un Websocket à l'aide de l'API Python et du SDK Python.

API Python 

CONFIG = {"response_modalities": ["TEXT"], "speech_config": { "language_code": "en-US" }}
headers = {
  "Content-Type": "application/json",
  "Authorization": f"Bearer {bearer_token[0]}",
}
HOST= "${LOCATION}-aiplatform.googleapis.com"
SERVICE_URL = f"wss://{HOST}/ws/google.cloud.aiplatform.v1beta1.LlmBidiService/BidiGenerateContent"
MODEL="gemini-2.0-flash-exp"

# Connect to the server
async with connect(SERVICE_URL, additional_headers=headers) as ws:
  # Setup the session
  await ws.send(
json.dumps(
          {
              "setup": {
                  "model": MODEL,
                  "generation_config": CONFIG,
                  # Setup RAG as a retrieval tool
                  "tools": TOOLS,
              }
          }
      )
  )

  # Receive setup response
  raw_response = await ws.recv(decode=False)
  setup_response = json.loads(raw_response.decode("ascii"))

  # Send text message
  text_input = "What are popular LLMs?"
  display(Markdown(f"**Input:** {text_input}"))

  msg = {
      "client_content": {
          "turns": [{"role": "user", "parts": [{"text": text_input}]}],
          "turn_complete": True,
      }
  }

  await ws.send(json.dumps(msg))

  responses = []

  # Receive chunks of server response
  async for raw_response in ws:
      response = json.loads(raw_response.decode())
      server_content = response.pop("serverContent", None)
      if server_content is None:
          break

      model_turn = server_content.pop("modelTurn", None)
      if model_turn is not None:
          parts = model_turn.pop("parts", None)
          if parts is not None:
              display(Markdown(f"**parts >** {parts}"))
              responses.append(parts[0]["text"])

      # End of turn
      turn_complete = server_content.pop("turnComplete", None)
      if turn_complete:
          grounding_metadata = server_content.pop("groundingMetadata", None)
          if grounding_metadata is not None:
            grounding_chunks = grounding_metadata.pop("groundingChunks", None)
            if grounding_chunks is not None:
              for chunk in grounding_chunks:
                display(Markdown(f"**grounding_chunk >** {chunk}"))
          break

  # Print the server response
  display(Markdown(f"**Response >** {''.join(responses)}"))

SDK Python

Pour savoir comment installer le SDK d'IA générative, consultez Installer une bibliothèque :

from google import genai
from google.genai import types
from google.genai.types import (Content, LiveConnectConfig, HttpOptions, Modality, Part,)
from IPython import display

MODEL="gemini-2.0-flash-exp"

client = genai.Client(
  vertexai=True,
  project=PROJECT_ID,
  location=LOCATION
)

async with client.aio.live.connect(
  model=MODEL,
  config=LiveConnectConfig(response_modalities=[Modality.TEXT],
                            tools=TOOLS),
) as session:
  text_input = "\'What are core LLM techniques?\'"
  print("> ", text_input, "\n")
  await session.send_client_content(
      turns=Content(role="user", parts=[Part(text=text_input)])
  )

  async for message in session.receive()
      if message.text:
          display.display(display.Markdown(message.text))
          continue

Utiliser le moteur RAG Vertex AI comme magasin de contexte

Vous pouvez utiliser le moteur RAG Vertex AI comme magasin de contexte pour l'API Gemini Live afin de stocker le contexte de session, de former et de récupérer les contextes passés liés à votre conversation, et d'enrichir le contexte actuel pour la génération de modèles. Vous pouvez également profiter de cette fonctionnalité pour partager des contextes entre vos différentes sessions Live API.

Le moteur RAG Vertex AI permet de stocker et d'indexer les formes de données suivantes à partir des contextes de session :

  • Texte
  • Discours audio

Créer un corpus de type MemoryCorpus

Pour stocker et indexer les textes de conversation à partir du contexte de session, vous devez créer un corpus RAG de type MemoryCorpus. Vous devez également spécifier un analyseur LLM dans la configuration de votre corpus de mémoire, qui est utilisé pour analyser les contextes de session stockés à partir de l'API Live afin de créer de la mémoire pour l'indexation.

Cet exemple de code montre comment créer un corpus. Toutefois, commencez par remplacer les variables par des valeurs.

# Currently supports Google first-party embedding models
EMBEDDING_MODEL = YOUR_EMBEDDING_MODEL  # Such as "publishers/google/models/text-embedding-005"
MEMORY_CORPUS_DISPLAY_NAME = YOUR_MEMORY_CORPUS_DISPLAY_NAME
LLM_PARSER_MODEL_NAME = YOUR_LLM_PARSER_MODEL_NAME  # Such as "projects/{project_id}/locations/{location}/publishers/google/models/gemini-2.5-pro-preview-05-06"

memory_corpus = rag.create_corpus(
   display_name=MEMORY_CORPUS_DISPLAY_NAME,
   corpus_type_config=rag.RagCorpusTypeConfig(
       corpus_type_config=rag.MemoryCorpus(
           llm_parser=rag.LlmParserConfig(
               model_name=LLM_PARSER_MODEL_NAME,
           )
       )
   ),
   backend_config=rag.RagVectorDbConfig(
       rag_embedding_model_config=rag.RagEmbeddingModelConfig(
           vertex_prediction_endpoint=rag.VertexPredictionEndpoint(
               publisher_model=EMBEDDING_MODEL
           )
       )
   ),
)

Spécifier votre corpus de mémoire pour stocker les contextes

Lorsque vous utilisez votre corpus de mémoire avec l'API Live, vous devez spécifier le corpus de mémoire comme outil de récupération, puis définir store_context sur true pour permettre à l'API Live de stocker les contextes de session.

Cet exemple de code montre comment spécifier votre corpus de mémoire pour stocker les contextes. Toutefois, commencez par remplacer les variables par des valeurs.

from google import genai
from google.genai import types
from google.genai.types import (Content, LiveConnectConfig, HttpOptions, Modality, Part)
from IPython import display

PROJECT_ID=YOUR_PROJECT_ID
LOCATION=YOUR_LOCATION
TEXT_INPUT=YOUR_TEXT_INPUT
MODEL_NAME=YOUR_MODEL_NAME  # Such as "gemini-2.0-flash-exp"

client = genai.Client(
   vertexai=True,
   project=PROJECT_ID,
   location=LOCATION,
)

memory_store=types.VertexRagStore(
   rag_resources=[
       types.VertexRagStoreRagResource(
           rag_corpus=memory_corpus.name
       )
   ],
   store_context=True
)

async with client.aio.live.connect(
   model=MODEL_NAME,
   config=LiveConnectConfig(response_modalities=[Modality.TEXT],
                            tools=[types.Tool(
                                retrieval=types.Retrieval(
                                    vertex_rag_store=memory_store))]),
) as session:
   text_input=TEXT_INPUT
   await session.send_client_content(
       turns=Content(role="user", parts=[Part(text=text_input)])
   )

   async for message in session.receive():
       if message.text:
           display.display(display.Markdown(message.text))
           continue

Étapes suivantes