Cómo usar una base de datos de Weaviate con el motor de RAG de Vertex AI

Organiza tus páginas con colecciones Guarda y categoriza el contenido según tus preferencias.

En esta página, se muestra cómo conectar tu corpus de RAG Engine a tu base de datos de Weaviate.

También puedes seguir este notebook RAG Engine con Weaviate.

Puedes usar tu instancia de la base de datos de Weaviate, que es una base de datos de código abierto, con RAG Engine para indexar y realizar una búsqueda de similitud basada en vectores. Una búsqueda de similitud es una forma de encontrar fragmentos de texto que sean similares al texto que buscas, lo que requiere el uso de un modelo de embedding. El modelo de embedding produce datos de vectores para cada fragmento de texto que se compara. La búsqueda de similitud se usa para recuperar contextos semánticos para la fundamentación y mostrar el contenido más preciso de tu LLM.

Con RAG Engine, puedes seguir usando tu instancia de base de datos de vectores completamente administrada, de la que eres responsable de aprovisionar. RAG Engine usa la base de datos de vectores para el almacenamiento, la administración de índices y la búsqueda.

Consideraciones

Ten en cuenta los siguientes pasos antes de usar la base de datos de Weaviate:

1. Debes crear, configurar y, luego, implementar la instancia y la colección de la base de datos de Weaviate. Sigue las instrucciones en Crea tu colección de Weaviate para configurar una colección según tu esquema.
2. Debes proporcionar una clave de API de Weaviate, que permite que RAG Engine interactúe con la base de datos de Weaviate. RAG Engine admite AuthN y AuthZ basados en claves de API, que se conectan a tu base de datos de Weaviate y admiten una conexión HTTPS.
3. RAG Engine no almacena ni administra tu clave de API de Weaviate. En su lugar, debes hacer lo siguiente:
   1. Almacena tu clave en Google Cloud Secret Manager.
   2. Otorga permisos a la cuenta de servicio de tu proyecto para acceder a tu secreto.
   3. Proporciona acceso a RAG Engine al nombre del recurso de tu secreto.
   4. Cuando interactúas con tu base de datos de Weaviate, RAG Engine accede a tu recurso de secreto con tu cuenta de servicio.
4. El corpus de RAG Engine y la colección de Weaviate tienen una asignación de uno a uno. Los archivos RAG se almacenan en una colección de bases de datos de Weaviate. Cuando se realiza una llamada a la API de CreateRagCorpus o a la API de UpdateRagCorpus, el corpus de RAG se asocia a la colección de la base de datos.
5. Además de las búsquedas semánticas densas basadas en embeddings, la búsqueda híbrida también es compatible con el motor de RAG a través de una base de datos de Weaviate. También puedes ajustar el peso entre la similitud de vectores densa y dispersa en una búsqueda híbrida.

Aprovisiona la base de datos de Weaviate

Antes de usar la base de datos de Weaviate con RAG Engine, debes hacer lo siguiente:

1. Configura y, luego, implementa tu instancia de la base de datos de Weaviate.
2. Prepara el extremo HTTPS.
3. Crea tu colección de Weaviate.
4. Usa tu clave de API para aprovisionar Weaviate con AuthN y AuthZ.
5. Aprovisiona tu cuenta de servicio de RAG Engine.

Configura y, luego, implementa tu instancia de la base de datos de Weaviate

Debes seguir la guía de inicio rápido de la guía oficial de Weaviate. Sin embargo, puedes usar la guía deGoogle Cloud Marketplace, que es opcional.

Puedes configurar tu instancia de Weaviate en cualquier lugar, siempre que se pueda acceder al extremo de Weaviate para configurarlo y, luego, implementarlo en tu proyecto. Luego, puedes administrar por completo tu instancia de la base de datos de Weaviate.

Dado que RAG Engine no participa en ninguna etapa del ciclo de vida de la instancia de la base de datos de Weaviate, es tu responsabilidad otorgar permisos a RAG Engine para que pueda almacenar y buscar datos en tu base de datos de Weaviate. También es tu responsabilidad asegurarte de que el motor de RAG pueda usar los datos de tu base de datos. Por ejemplo, si cambias tus datos, RAG Engine no es responsable de ningún comportamiento inesperado debido a esos cambios.

Prepara el extremo HTTPS

Durante el aprovisionamiento de Weaviate, asegúrate de crear un extremo HTTPS. Aunque se admiten las conexiones HTTP, preferimos que el tráfico de la base de datos de RAG Engine y Weaviate use una conexión HTTPS.

Crea tu colección de Weaviate

Debido a que el corpus de RAG Engine y la colección de Weaviate tienen una asignación de uno a uno, debes crear una colección en tu base de datos de Weaviate antes de asociarla con el corpus de RAG Engine. Esta asociación única se realiza cuando llamas a la API de CreateRagCorpus o a la API de UpdateRagCorpus.

Cuando crees una colección en Weaviate, debes usar el siguiente esquema:

Usa tu clave de API para aprovisionar Weaviate con AuthN y AuthZ

Aprovisionar la clave de API de Weaviate implica los siguientes pasos:

1. Crea la clave de API de Weaviate.
2. Configura Weaviate con tu clave de API de Weaviate.
3. Almacena tu clave de API de Weaviate en Secret Manager.

Crea la clave de API

El motor de RAG solo puede conectarse a tus instancias de la base de datos de Weaviate con tu clave de API para la autenticación y autorización. Debes seguir la guía oficial de autenticación de Weaviate para configurar la autenticación basada en claves de API en tu instancia de la base de datos de Weaviate.

Si la creación de la clave de API de Weaviate requiere información de identidad para asociarse con la que proviene de RAG Engine, debes crear tu primer corpus y usar tu cuenta de servicio de RAG Engine como identidad.

Almacena tu clave de API en Secret Manager

Una clave de API contiene información de identificación personal sensible (IIPS), que está sujeta a requisitos legales. Si los datos de la SPII se ven comprometidos o se usan de forma inadecuada, una persona podría sufrir un riesgo o daño importante. Para minimizar los riesgos para una persona mientras usa RAG Engine, no almacene ni administre su clave de API, y evite compartir la clave de API sin encriptar.

Para proteger los SPII, haz lo siguiente:

1. Almacena tu clave de API en Secret Manager.
2. Otorga a tu cuenta de servicio de RAG Engine los permisos a tus secretos y administra el control de acceso a nivel del recurso secreto.
   1. Navega a los permisos de tu proyecto.
   2. Habilita la opción Incluir asignaciones de roles proporcionadas por Google.
   3. Busca la cuenta de servicio, que tiene el formato

      service-{project number}@gcp-sa-vertex-rag.iam.gserviceaccount.com

   4. Edita los principales de la cuenta de servicio.
   5. Agrega el rol de Descriptor de acceso a secretos de Secret Manager a la cuenta de servicio.
3. Durante la creación o actualización del corpus de RAG, pasa el nombre del recurso secreto a RAG Engine y almacénalo.

Cuando realizas solicitudes a la API a tus instancias de la base de datos de Weaviate, RAG Engine usa cada cuenta de servicio para leer la clave de API que corresponde a tus recursos secretos en Secret Manager desde tus proyectos.

Aprovisiona tu cuenta de servicio de RAG Engine

Cuando creas el primer recurso de tu proyecto, RAG Engine crea una cuenta de servicio dedicada. Puedes encontrar tu cuenta de servicio en la página de IAM de tu proyecto. La cuenta de servicio sigue este formato:

service-{project number}@gcp-sa-vertex-rag.iam.gserviceaccount.com

Por ejemplo, service-123456789@gcp-sa-vertex-rag.iam.gserviceaccount.com.

Cuando se realiza la integración con la base de datos de Weaviate, se usa tu cuenta de servicio en las siguientes situaciones:

• Puedes usar tu cuenta de servicio para generar tu clave de API de Weaviate para la autenticación. En algunos casos, generar la clave de API no requiere información del usuario, lo que significa que no se requiere una cuenta de servicio cuando se genera la clave de API.
• Puedes vincular tu cuenta de servicio con la clave de API en tu base de datos de Weaviate para configurar la autenticación (AuthN) y la autorización (AuthZ). Sin embargo, no es obligatoria.
• Puedes almacenar el Secret Manager de la clave de API en tu proyecto y otorgarle permisos a tu cuenta de servicio para estos recursos secretos.
• RAG Engine usa cuentas de servicio para acceder a la clave de API desde el Secret Manager en tus proyectos.

Configura tu entorno de la consola de Google Cloud

Prepara tu corpus de RAG

Para acceder a los datos de tu base de datos de Weaviate, RAG Engine debe tener acceso a un corpus de RAG. En esta sección, se proporcionan los pasos para crear un solo corpus RAG y corpus RAG adicionales.

Usa las APIs de CreateRagCorpus y UpdateRagCorpus

Debes especificar los siguientes campos cuando llames a las APIs de CreateRagCorpus y UpdateRagCorpus:

• rag_vector_db_config.weaviate: Después de llamar a la API de CreateRagCorpus, se elige la configuración de la base de datos de vectores. La configuración de la base de datos de vectores contiene todos los campos de configuración. Si no se establece el campo rag_vector_db_config.weaviate, rag_vector_db_config.rag_managed_db se establece de forma predeterminada.
• weaviate.http_endpoint: El extremo HTTPS o HTTP de Weaviate se crea durante el aprovisionamiento de la instancia de la base de datos de Weaviate.
• weaviate.collection_name: Es el nombre de la colección que se crea durante el aprovisionamiento de la instancia de Weaviate. El nombre debe comenzar con una letra mayúscula.
• api_auth.api_key_config: La configuración especifica que se debe usar una clave de API para autorizar tu acceso a la base de datos de vectores.
• api_key_config.api_key_secret_version: Es el nombre del recurso del secreto que se almacena en Secret Manager, que contiene tu clave de API de Weaviate.

Puedes crear y asociar tu corpus de RAG a la colección de Weaviate en tu instancia de base de datos. Sin embargo, es posible que necesites la cuenta de servicio para generar tu clave de API y configurar tu instancia de la base de datos de Weaviate. Cuando creas tu primer corpus de RAG, se genera la cuenta de servicio. Después de crear tu primer corpus de RAG, es posible que la asociación entre la base de datos de Weaviate y la clave de API no esté lista para usarse en la creación de otro corpus de RAG.

En caso de que tu base de datos y clave no estén listas para asociarse a tu corpus de RAG, haz lo siguiente:

1. Configura el campo weaviate en rag_vector_db_config.

   • No puedes cambiar la base de datos de vectores asociada.
   • Deja los campos http_endpoint y collection_name vacíos. Ambos campos se pueden actualizar más adelante.

2. Si no tienes tu clave de API almacenada en Secret Manager, puedes dejar el campo api_auth en blanco. Cuando llames a la API de UpdateRagCorpus, podrás actualizar el campo api_auth. Weaviate requiere que se haga lo siguiente:

   1. Configura api_key_config en el campo api_auth.

   2. Establece el api_key_secret_version de tu clave de API de Weaviate en Secret Manager. El campo api_key_secret_version usa el siguiente formato:

      projects/{project}/secrets/{secret}/versions/{version}

3. Si especificas campos que solo se pueden establecer una vez, como http_endpoint o collection_name, no podrás cambiarlos, a menos que borres el corpus de RAG y lo vuelvas a crear. Se pueden actualizar otros campos, como el campo de clave de API, api_key_secret_version.

4. Cuando llames a UpdateRagCorpus, puedes configurar el campo vector_db. Tu llamada a la API de CreateRagCorpus debe establecer vector_db en weaviate. De lo contrario, el sistema elige la opción Base de datos administrada por RAG, que es la predeterminada. No se puede cambiar esta opción cuando llamas a la API de UpdateRagCorpus. Cuando llamas a UpdateRagCorpus y el campo vector_db está configurado de forma parcial, puedes actualizar los campos que están marcados como Cambiables (también conocidos como mutables).

En esta tabla, se enumeran los campos WeaviateConfig mutables e inmutables que se usan en tu código.

Crea el primer corpus de RAG

Cuando no exista la cuenta de servicio de RAG Engine, haz lo siguiente:

1. Crea un corpus de RAG en RAG Engine con una configuración vacía de Weaviate, que inicia el aprovisionamiento de RAG Engine para crear una cuenta de servicio.
2. Elige un nombre para tu cuenta de servicio de RAG Engine que siga este formato:
   

   service-{project number}@gcp-sa-vertex-rag.iam.gserviceaccount.com

   Por ejemplo, service-123456789@gcp-sa-vertex-rag.iam.gserviceaccount.com.

3. Con tu cuenta de servicio, accede al secreto que se almacena en Secret Manager de tu proyecto, que contiene tu clave de API de Weaviate.
4. Obtén la siguiente información después de que se complete el aprovisionamiento de Weaviate:
   • Tu extremo HTTPS o HTTP de Weaviate.
   • Es el nombre de tu colección de Weaviate.
5. Llama a la API de CreateRagCorpus para crear un corpus de RAG con una configuración de Weaviate vacía y llama a la API de UpdateRagCorpus para actualizar el corpus de RAG con la siguiente información:
   • Tu extremo HTTPS o HTTP de Weaviate.
   • Es el nombre de tu colección de Weaviate.
   • Es el nombre del recurso de la clave de API.

Crea otro corpus de RAG

Cuando exista la cuenta de servicio de RAG Engine, haz lo siguiente:

1. Obtén tu cuenta de servicio de RAG Engine de los permisos de tu proyecto.
2. Habilita la opción "Incluir asignaciones de roles proporcionadas por Google".
3. Elige un nombre para tu cuenta de servicio de RAG Engine que siga este formato:
   

   service-{project number}@gcp-sa-vertex-rag.iam.gserviceaccount.com

4. Con tu cuenta de servicio, accede al secreto que se almacena en Secret Manager de tu proyecto, que contiene tu clave de API de Weaviate.
5. Durante el aprovisionamiento de Weaviate, obtén la siguiente información:
   • El extremo HTTPS o HTTP de Weaviate
   • Es el nombre de tu colección de Weaviate.
6. Crea un corpus de RAG en RAG Engine y conéctate a tu colección de Weaviate haciendo una de las siguientes acciones:
   1. Realiza una llamada a la API de CreateRagCorpus para crear un corpus RAG con una configuración de Weaviate propagada, que es la opción preferida.
   2. Realiza una llamada a la API de CreateRagCorpus para crear un corpus de RAG con una configuración de Weaviate vacía y realiza una llamada a la API de UpdateRagCorpus para actualizar el corpus de RAG con la siguiente información:
      • Extremo HTTP de la base de datos de Weaviate
      • Nombre de la colección de Weaviate
      • Clave de API

Ejemplos

En esta sección, se presenta un código de muestra que muestra cómo configurar tu base de datos de Weaviate, Secret Manager, el corpus de RAG y el archivo RAG. También se proporciona código de muestra para demostrar cómo importar archivos, recuperar contexto, generar contenido y borrar el corpus y los archivos de RAG.

Para usar el notebook de la API de RAG de Model Garden, consulta Cómo usar Weaviate con Llama 3.

Configura tu base de datos de Weaviate

En esta muestra de código, se muestra cómo configurar tus datos de Weaviate y Secret Manager.

# TODO(developer): Update the variables.
# The HTTPS/HTTP Weaviate endpoint you created during provisioning.
HTTP_ENDPOINT_NAME="https://your.weaviate.endpoint.com"

# Your Weaviate API Key.
WEAVIATE_API_KEY="example-api-key"

# Select your Weaviate collection name, which roughly corresponds to a Vertex AI Knowledge Engine Corpus.
# For example, "MyCollectionName"
# Note that the first letter needs to be capitalized.
# Otherwise, Weavaite will capitalize it for you.
WEAVIATE_COLLECTION_NAME="MyCollectionName"

# Create a collection in Weaviate which includes the required schema fields shown below.
echo '{
  "class": "'${WEAVIATE_COLLECTION_NAME}'",
  "properties": [
    { "name": "fileId", "dataType": [ "string" ] },
    { "name": "corpusId", "dataType": [ "string" ] },
    { "name": "chunkId", "dataType": [ "string" ] },
    { "name": "chunkDataType", "dataType": [ "string" ] },
    { "name": "chunkData", "dataType": [ "string" ] },
    { "name": "fileOriginalUri", "dataType": [ "string" ] }
  ]
}' | curl \
    -X POST \
    -H 'Content-Type: application/json' \
    -H "Authorization: Bearer "${WEAVIATE_API_KEY} \
    -d @- \
    ${HTTP_ENDPOINT_NAME}/v1/schema

Configura Secret Manager

Para configurar Secret Manager, debes habilitarlo y establecer permisos.

Crear secreto

Para habilitar tu Secret Manager, haz lo siguiente:

# Create a secret in SecretManager.
curl "https://secretmanager.googleapis.com/v1/projects/${PROJECT_ID}/secrets?secretId=${SECRET_NAME}" \
    --request "POST" \
    --header "authorization: Bearer $(gcloud auth print-access-token)" \
    --header "content-type: application/json" \
    --data "{\"replication\": {\"automatic\": {}}}"

def create_secret(
    project_id: str, secret_id: str, ttl: Optional[str] = None
) -> secretmanager.Secret:
    """
    Create a new secret with the given name. A secret is a logical wrapper
    around a collection of secret versions. Secret versions hold the actual
    secret material.

     Args:
        project_id (str): The project ID where the secret is to be created.
        secret_id (str): The ID to assign to the new secret. This ID must be unique within the project.
        ttl (Optional[str]): An optional string that specifies the secret's time-to-live in seconds with
                             format (e.g., "900s" for 15 minutes). If specified, the secret
                             versions will be automatically deleted upon reaching the end of the TTL period.

    Returns:
        secretmanager.Secret: An object representing the newly created secret, containing details like the
                              secret's name, replication settings, and optionally its TTL.

    Example:
        # Create a secret with automatic replication and no TTL
        new_secret = create_secret("my-project", "my-new-secret")

        # Create a secret with a TTL of 30 days
        new_secret_with_ttl = create_secret("my-project", "my-timed-secret", "7776000s")
    """

    # Import the Secret Manager client library.
    from google.cloud import secretmanager

    # Create the Secret Manager client.
    client = secretmanager.SecretManagerServiceClient()

    # Build the resource name of the parent project.
    parent = f"projects/{project_id}"

    # Create the secret.
    response = client.create_secret(
        request={
            "parent": parent,
            "secret_id": secret_id,
            "secret": {"replication": {"automatic": {}}, "ttl": ttl},
        }
    )

    # Print the new secret name.
    print(f"Created secret: {response.name}")

Configurar permisos

Debes otorgar permisos de Secret Manager a tu cuenta de servicio.

def iam_grant_access(
    project_id: str, secret_id: str, member: str
) -> iam_policy_pb2.SetIamPolicyRequest:
    """
    Grant the given member access to a secret.
    """

    # Import the Secret Manager client library.
    from google.cloud import secretmanager

    # Create the Secret Manager client.
    client = secretmanager.SecretManagerServiceClient()

    # Build the resource name of the secret.
    name = client.secret_path(project_id, secret_id)

    # Get the current IAM policy.
    policy = client.get_iam_policy(request={"resource": name})

    # Add the given member with access permissions.
    policy.bindings.add(role="roles/secretmanager.secretAccessor", members=[member])

    # Update the IAM Policy.
    new_policy = client.set_iam_policy(request={"resource": name, "policy": policy})

    # Print data about the secret.
    print(f"Updated IAM policy on {secret_id}")

Agrega una versión del Secret

# TODO(developer): Update the variables.
# Select a resource name for your Secret, which contains your API Key.
SECRET_NAME="MyWeaviateApiKeySecret"

# Your Weaviate API Key.
WEAVIATE_API_KEY="example-api-key"
# Encode your WEAVIATE_API_KEY using base 64.
SECRET_DATA=$(echo ${WEAVIATE_API_KEY} | base64)

# Create a new version of your secret which uses SECRET_DATA as payload
curl "https://secretmanager.googleapis.com/v1/projects/${PROJECT_ID}/secrets/${SECRET_NAME}:addVersion" \
    --request "POST" \
    --header "authorization: Bearer $(gcloud auth print-access-token)" \
    --header "content-type: application/json" \
    --data "{\"payload\": {\"data\": \"${SECRET_DATA}\"}}"

from google.cloud import secretmanager
import google_crc32c  # type: ignore


def add_secret_version(
    project_id: str, secret_id: str, payload: str
) -> secretmanager.SecretVersion:
    """
    Add a new secret version to the given secret with the provided payload.
    """

    # Create the Secret Manager client.
    client = secretmanager.SecretManagerServiceClient()

    # Build the resource name of the parent secret.
    parent = client.secret_path(project_id, secret_id)

    # Convert the string payload into a bytes. This step can be omitted if you
    # pass in bytes instead of a str for the payload argument.
    payload_bytes = payload.encode("UTF-8")

    # Calculate payload checksum. Passing a checksum in add-version request
    # is optional.
    crc32c = google_crc32c.Checksum()
    crc32c.update(payload_bytes)

    # Add the secret version.
    response = client.add_secret_version(
        request={
            "parent": parent,
            "payload": {
                "data": payload_bytes,
                "data_crc32c": int(crc32c.hexdigest(), 16),
            },
        }
    )

    # Print the new secret version name.
    print(f"Added secret version: {response.name}")

Usa Weaviate con Llama 3

En el notebook de la API de RAG de Model Garden, se muestra cómo usar el SDK de Vertex AI para Python con un corpus de Weaviate y un modelo de Llama 3. Para usar el notebook, debes hacer lo siguiente:

1. Configura tu base de datos de Weaviate.

2. Configura tu Secret Manager.

3. Usa el notebook de la API de RAG de Model Garden.

Para obtener más ejemplos, consulta Ejemplos.

Crea un corpus RAG

En esta muestra de código, se muestra cómo crear un corpus de RAG y se establece la instancia de Weaviate como su base de datos de vectores.

  # TODO(developer): Update the variables.
  PROJECT_ID = "YOUR_PROJECT_ID"
  # The HTTPS/HTTP Weaviate endpoint you created during provisioning.
  HTTP_ENDPOINT_NAME="https://your.weaviate.endpoint.com"

  # Your Weaviate collection name, which roughly corresponds to a Vertex AI Knowledge Engine Corpus.
  # For example, "MyCollectionName"
  # Note that the first letter needs to be capitalized.
  # Otherwise, Weaviate will capitalize it for you.
  WEAVIATE_COLLECTION_NAME="MyCollectionName"

  # The resource name of your Weaviate API Key your Secret.
  SECRET_NAME="MyWeaviateApiKeySecret"
  # The Secret Manager resource name containing the API Key for your Weaviate endpoint.
  # For example, projects/{project}/secrets/{secret}/versions/latest
  APIKEY_SECRET_VERSION="projects/${PROJECT_ID}/secrets/${SECRET_NAME}/versions/latest"

  # Select a Corpus display name.
  CORPUS_DISPLAY_NAME="SpecialCorpus"

  # Call CreateRagCorpus API and set all Vector DB Config parameters for Weaviate to create a new corpus associated to your selected Weaviate collection.
  curl -X POST \
  -H "Authorization: Bearer $(gcloud auth print-access-token)" \
  -H "Content-Type: application/json" \
  https://us-central1-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/us-central1/ragCorpora \
  -d '{
        "display_name" : '\""${CORPUS_DISPLAY_NAME}"\"',
        "rag_vector_db_config" : {
                "weaviate": {
                      "http_endpoint": '\""${HTTP_ENDPOINT_NAME}"\"',
                      "collection_name": '\""${WEAVIATE_COLLECTION_NAME}"\"'
                },
          "api_auth" : {
                  "api_key_config": {
                        "api_key_secret_version": '\""${APIKEY_SECRET_VERSION}"\"'
                  }
          }
        }
    }'

  # TODO(developer): Update the variables.
  # Get operation_id returned in CreateRagCorpus.
  OPERATION_ID="your-operation-id"

  # Poll Operation status until done = true in the response.
  curl -X GET \
  -H "Authorization: Bearer $(gcloud auth print-access-token)" \
    -H "Content-Type: application/json" \
  https://us-central1-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/us-central1/operations/${OPERATION_ID}

  # Call ListRagCorpora API to verify the RAG corpus is created successfully.
  curl -sS -X GET \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer $(gcloud auth print-access-token)" \
    "https://us-central1-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/us-central1/ragCorpora"

from vertexai.preview import rag
import vertexai

# TODO(developer): Update and un-comment below lines
# PROJECT_ID = "your-project-id"
# weaviate_http_endpoint = "weaviate-http-endpoint"
# weaviate_collection_name = "weaviate-collection-name"
# weaviate_api_key_secret_manager_version = "projects/{PROJECT_ID}/secrets/{SECRET_NAME}/versions/latest"
# display_name = "test_corpus"
# description = "Corpus Description"

# Initialize Vertex AI API once per session
vertexai.init(project=PROJECT_ID, location="us-central1")

# Configure embedding model (Optional)
embedding_model_config = rag.EmbeddingModelConfig(
    publisher_model="publishers/google/models/text-embedding-004"
)

# Configure Vector DB
vector_db = rag.Weaviate(
    weaviate_http_endpoint=weaviate_http_endpoint,
    collection_name=weaviate_collection_name,
    api_key=weaviate_api_key_secret_manager_version,
)

corpus = rag.create_corpus(
    display_name=display_name,
    description=description,
    embedding_model_config=embedding_model_config,
    vector_db=vector_db,
)
print(corpus)
# Example response:
# RagCorpus(name='projects/1234567890/locations/us-central1/ragCorpora/1234567890',
# display_name='test_corpus', description='Corpus Description', embedding_model_config=...
# ...

Usa el archivo RAG

La API de RAG controla la carga, la importación, la publicación y la eliminación de archivos.

POST https://LOCATION-aiplatform.googleapis.com/upload/v1beta1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles:upload

{
 "rag_file": {
  "display_name": "FILE_DISPLAY_NAME",
  "description": "RAG_FILE_DESCRIPTION"
 }
}

curl -X POST \
     -H "Content-Type: application/json; charset=utf-8" \
     -d @INPUT_FILE \
     "https://LOCATION-aiplatform.googleapis.com/upload/v1beta1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles:upload"

$headers = @{  }

Invoke-WebRequest `
    -Method POST `
    -Headers $headers `
    -ContentType: "application/json; charset=utf-8" `
    -InFile INPUT_FILE `
    -Uri "https://LOCATION-aiplatform.googleapis.com/upload/v1beta1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles:upload" | Select-Object -Expand Content

from vertexai import rag
import vertexai

# TODO(developer): Update and un-comment below lines
# PROJECT_ID = "your-project-id"
# corpus_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}"
# path = "path/to/local/file.txt"
# display_name = "file_display_name"
# description = "file description"

# Initialize Vertex AI API once per session
vertexai.init(project=PROJECT_ID, location="us-central1")

rag_file = rag.upload_file(
    corpus_name=corpus_name,
    path=path,
    display_name=display_name,
    description=description,
)
print(rag_file)
# RagFile(name='projects/[PROJECT_ID]/locations/us-central1/ragCorpora/1234567890/ragFiles/09876543',
#  display_name='file_display_name', description='file description')

Importar archivos RAG

Los archivos y las carpetas se pueden importar desde Drive o Cloud Storage.

POST https://LOCATION-aiplatform.googleapis.com/upload/v1beta1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles:import

{
  "import_rag_files_config": {
    "gcs_source": {
      "uris": GCS_URIS
    },
    "google_drive_source": {
      "resource_ids": {
        "resource_id": DRIVE_RESOURCE_ID,
        "resource_type": DRIVE_RESOURCE_TYPE
      },
    }
  }
}

curl -X POST \
     -H "Content-Type: application/json; charset=utf-8" \
     -d @request.json \
     "https://LOCATION-aiplatform.googleapis.com/upload/v1beta1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles:import"

$headers = @{  }

Invoke-WebRequest `
    -Method POST `
    -Headers $headers `
    -ContentType: "application/json; charset=utf-8" `
    -InFile request.json `
    -Uri "https://LOCATION-aiplatform.googleapis.com/upload/v1beta1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles:import" | Select-Object -Expand Content

// Cloud Storage bucket/file location.
// Such as "gs://rag-e2e-test/"
GCS_URIS=YOUR_GCS_LOCATION

// Enter the QPM rate to limit RAG's access to your embedding model
// Example: 1000
EMBEDDING_MODEL_QPM_RATE=MAX_EMBEDDING_REQUESTS_PER_MIN_LIMIT

// ImportRagFiles
// Import a single Cloud Storage file or all files in a Cloud Storage bucket.
// Input: ENDPOINT, PROJECT_ID, RAG_CORPUS_ID, GCS_URIS
// Output: ImportRagFilesOperationMetadataNumber
// Use ListRagFiles to find the server-generated rag_file_id.
curl -X POST \
-H "Authorization: Bearer $(gcloud auth print-access-token)" \
-H "Content-Type: application/json" \
https://${ENDPOINT}/v1beta1/projects/${PROJECT_ID}/locations/${LOCATION}/ragCorpora/${RAG_CORPUS_ID}/ragFiles:import \
-d '{
  "import_rag_files_config": {
    "gcs_source": {
      "uris": '\""${GCS_URIS}"\"'
    },
    "rag_file_chunking_config": {
      "chunk_size": 512
    },
    "max_embedding_requests_per_min": '"${EMBEDDING_MODEL_QPM_RATE}"'
  }
}'

// Poll the operation status.
// The response contains the number of files imported.
OPERATION_ID=OPERATION_ID
poll_op_wait ${OPERATION_ID}

// Google Drive folder location.
FOLDER_RESOURCE_ID=YOUR_GOOGLE_DRIVE_FOLDER_RESOURCE_ID

// Enter the QPM rate to limit RAG's access to your embedding model
// Example: 1000
EMBEDDING_MODEL_QPM_RATE=MAX_EMBEDDING_REQUESTS_PER_MIN_LIMIT

// ImportRagFiles
// Import all files in a Google Drive folder.
// Input: ENDPOINT, PROJECT_ID, RAG_CORPUS_ID, FOLDER_RESOURCE_ID
// Output: ImportRagFilesOperationMetadataNumber
// Use ListRagFiles to find the server-generated rag_file_id.
curl -X POST \
-H "Authorization: Bearer $(gcloud auth print-access-token)" \
-H "Content-Type: application/json" \
https://${ENDPOINT}/v1beta1/projects/${PROJECT_ID}/locations/${LOCATION}/ragCorpora/${RAG_CORPUS_ID}/ragFiles:import \
-d '{
  "import_rag_files_config": {
    "google_drive_source": {
      "resource_ids": {
        "resource_id": '\""${FOLDER_RESOURCE_ID}"\"',
        "resource_type": "RESOURCE_TYPE_FOLDER"
      }
    },
    "max_embedding_requests_per_min": '"${EMBEDDING_MODEL_QPM_RATE}"'
  }
}'

// Poll the operation status.
// The response contains the number of files imported.
OPERATION_ID=OPERATION_ID
poll_op_wait ${OPERATION_ID}

from vertexai import rag
import vertexai

# TODO(developer): Update and un-comment below lines
# PROJECT_ID = "your-project-id"
# corpus_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}"
# paths = ["https://drive.google.com/file/123", "gs://my_bucket/my_files_dir"]  # Supports Google Cloud Storage and Google Drive Links

# Initialize Vertex AI API once per session
vertexai.init(project=PROJECT_ID, location="us-central1")

response = rag.import_files(
    corpus_name=corpus_name,
    paths=paths,
    transformation_config=rag.TransformationConfig(
        rag.ChunkingConfig(chunk_size=512, chunk_overlap=100)
    ),
    max_embedding_requests_per_min=900,  # Optional
)
print(f"Imported {response.imported_rag_files_count} files.")
# Example response:
# Imported 2 files.

Obtén un archivo RAG

GET https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles/RAG_FILE_ID

curl -X GET \
     "https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles/RAG_FILE_ID"

$headers = @{  }

Invoke-WebRequest `
    -Method GET `
    -Headers $headers `
    -Uri "https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles/RAG_FILE_ID" | Select-Object -Expand Content

from vertexai import rag
import vertexai

# TODO(developer): Update and un-comment below lines
# PROJECT_ID = "your-project-id"
# file_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}/ragFiles/{rag_file_id}"

# Initialize Vertex AI API once per session
vertexai.init(project=PROJECT_ID, location="us-central1")

rag_file = rag.get_file(name=file_name)
print(rag_file)
# Example response:
# RagFile(name='projects/1234567890/locations/us-central1/ragCorpora/11111111111/ragFiles/22222222222',
# display_name='file_display_name', description='file description')

Enumera archivos Rag

GET https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles?page_size=PAGE_SIZE&page_token=PAGE_TOKEN

curl -X GET \
     "https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles?page_size=PAGE_SIZE&page_token=PAGE_TOKEN"

$headers = @{  }

Invoke-WebRequest `
    -Method GET `
    -Headers $headers `
    -Uri "https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles?page_size=PAGE_SIZE&page_token=PAGE_TOKEN" | Select-Object -Expand Content

from vertexai import rag
import vertexai

# TODO(developer): Update and un-comment below lines
# PROJECT_ID = "your-project-id"
# corpus_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}"

# Initialize Vertex AI API once per session
vertexai.init(project=PROJECT_ID, location="us-central1")

files = rag.list_files(corpus_name=corpus_name)
for file in files:
    print(file.display_name)
    print(file.name)
# Example response:
# g-drive_file.txt
# projects/1234567890/locations/us-central1/ragCorpora/111111111111/ragFiles/222222222222
# g_cloud_file.txt
# projects/1234567890/locations/us-central1/ragCorpora/111111111111/ragFiles/333333333333

Borra un archivo RAG

DELETE https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles/RAG_FILE_ID

curl -X DELETE \
     "https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles/RAG_FILE_ID"

$headers = @{  }

Invoke-WebRequest `
    -Method DELETE `
    -Headers $headers `
    -Uri "https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION/ragCorpora/RAG_CORPUS_ID/ragFiles/RAG_FILE_ID" | Select-Object -Expand Content

from vertexai import rag
import vertexai

# TODO(developer): Update and un-comment below lines
# PROJECT_ID = "your-project-id"
# file_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}/ragFiles/{rag_file_id}"

# Initialize Vertex AI API once per session
vertexai.init(project=PROJECT_ID, location="us-central1")

rag.delete_file(name=file_name)
print(f"File {file_name} deleted.")
# Example response:
# Successfully deleted the RagFile.
# File projects/1234567890/locations/us-central1/ragCorpora/1111111111/ragFiles/2222222222 deleted.

Recupera el contexto

Cuando un usuario hace una pregunta o proporciona una instrucción, el componente de recuperación en RAG busca en su base de conocimiento la información relevante para la consulta.

POST https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION:retrieveContexts

{
 "vertex_rag_store": {
    "rag_resources": {
      "rag_corpus": "RAG_CORPUS_RESOURCE",
    },
    "vector_distance_threshold": 0.8
  },
  "query": {
   "text": "TEXT",
   "similarity_top_k": SIMILARITY_TOP_K
  }
 }

curl -X POST \
     -H "Content-Type: application/json; charset=utf-8" \
     -d @request.json \
     "https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION:retrieveContexts"

$headers = @{  }

Invoke-WebRequest `
    -Method POST `
    -Headers $headers `
    -ContentType: "application/json; charset=utf-8" `
    -InFile request.json `
    -Uri "https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION:retrieveContexts" | Select-Object -Expand Content

from vertexai import rag
import vertexai

# TODO(developer): Update and un-comment below lines
# PROJECT_ID = "your-project-id"
# corpus_name = "projects/[PROJECT_ID]/locations/us-central1/ragCorpora/[rag_corpus_id]"

# Initialize Vertex AI API once per session
vertexai.init(project=PROJECT_ID, location="us-central1")

response = rag.retrieval_query(
    rag_resources=[
        rag.RagResource(
            rag_corpus=corpus_name,
            # Optional: supply IDs from `rag.list_files()`.
            # rag_file_ids=["rag-file-1", "rag-file-2", ...],
        )
    ],
    text="Hello World!",
    rag_retrieval_config=rag.RagRetrievalConfig(
        top_k=10,
        filter=rag.utils.resources.Filter(vector_distance_threshold=0.5),
    ),
)
print(response)
# Example response:
# contexts {
#   contexts {
#     source_uri: "gs://your-bucket-name/file.txt"
#     text: "....
#   ....

Genera contenido

Una predicción controla el método LLM que genera contenido.

POST https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION/publishers/google/models/MODEL_ID:GENERATION_METHOD

{
 "contents": {
  "role": "user",
  "parts": {
    "text": "INPUT_PROMPT"
  }
 },
 "tools": {
  "retrieval": {
   "disable_attribution": false,
   "vertex_rag_store": {
    "rag_resources": {
      "rag_corpus": "RAG_CORPUS_RESOURCE",
    },
    "similarity_top_k": SIMILARITY_TOP_K,
    "vector_distance_threshold": VECTOR_DISTANCE_THRESHOLD
   }
  }
 }
}

curl -X POST \
     -H "Content-Type: application/json; charset=utf-8" \
     -d @request.json \
     "https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION/publishers/google/models/MODEL_ID:GENERATION_METHOD"

$headers = @{  }

Invoke-WebRequest `
    -Method POST `
    -Headers $headers `
    -ContentType: "application/json; charset=utf-8" `
    -InFile request.json `
    -Uri "https://LOCATION-aiplatform.googleapis.com/v1beta1/projects/PROJECT_ID/locations/LOCATION/publishers/google/models/MODEL_ID:GENERATION_METHOD" | Select-Object -Expand Content

from vertexai import rag
from vertexai.generative_models import GenerativeModel, Tool
import vertexai

# TODO(developer): Update and un-comment below lines
# PROJECT_ID = "your-project-id"
# corpus_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}"

# Initialize Vertex AI API once per session
vertexai.init(project=PROJECT_ID, location="us-central1")

rag_retrieval_tool = Tool.from_retrieval(
    retrieval=rag.Retrieval(
        source=rag.VertexRagStore(
            rag_resources=[
                rag.RagResource(
                    rag_corpus=corpus_name,
                    # Optional: supply IDs from `rag.list_files()`.
                    # rag_file_ids=["rag-file-1", "rag-file-2", ...],
                )
            ],
            rag_retrieval_config=rag.RagRetrievalConfig(
                top_k=10,
                filter=rag.utils.resources.Filter(vector_distance_threshold=0.5),
            ),
        ),
    )
)

rag_model = GenerativeModel(
    model_name="gemini-2.0-flash-001", tools=[rag_retrieval_tool]
)
response = rag_model.generate_content("Why is the sky blue?")
print(response.text)
# Example response:
#   The sky appears blue due to a phenomenon called Rayleigh scattering.
#   Sunlight, which contains all colors of the rainbow, is scattered
#   by the tiny particles in the Earth's atmosphere....
#   ...

Búsqueda híbrida

La búsqueda híbrida es compatible con la base de datos de Weaviate, que combina las búsquedas semánticas y de palabras clave para mejorar la relevancia de los resultados de la búsqueda. Durante la recuperación de los resultados de la búsqueda, una combinación de puntuaciones de similitud de la semántica (un vector denso) y la concordancia de palabras clave (un vector disperso) produce los resultados finales clasificados.

Búsqueda híbrida con la API de recuperación de RAG Engine

Este es un ejemplo de cómo habilitar una búsqueda híbrida con la API de recuperación de RAG Engine.

  # TODO(developer): Update the variables.
  PROJECT_ID = "YOUR_PROJECT_ID"
  # The HTTPS/HTTP Weaviate endpoint you created during provisioning.
  HTTP_ENDPOINT_NAME="https://your.weaviate.endpoint.com"

  # Your Weaviate collection name, which roughly corresponds to a Vertex AI Knowledge Engine Corpus.
  # For example, "MyCollectionName"
  # Note that the first letter needs to be capitalized.
  # Otherwise, Weaviate will capitalize it for you.
  WEAVIATE_COLLECTION_NAME="MyCollectionName"

  # The resource name of your Weaviate API Key your Secret.
  SECRET_NAME="MyWeaviateApiKeySecret"
  # The Secret Manager resource name containing the API Key for your Weaviate endpoint.
  # For example, projects/{project}/secrets/{secret}/versions/latest
  APIKEY_SECRET_VERSION="projects/${PROJECT_ID}/secrets/${SECRET_NAME}/versions/latest"

  # Select a Corpus display name.
  CORPUS_DISPLAY_NAME="SpecialCorpus"

  # Call CreateRagCorpus API and set all Vector DB Config parameters for Weaviate to create a new corpus associated to your selected Weaviate collection.
  curl -X POST \
  -H "Authorization: Bearer $(gcloud auth print-access-token)" \
  -H "Content-Type: application/json" \
  https://us-central1-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/us-central1/ragCorpora \
  -d '{
        "display_name" : '\""${CORPUS_DISPLAY_NAME}"\"',
        "rag_vector_db_config" : {
                "weaviate": {
                      "http_endpoint": '\""${HTTP_ENDPOINT_NAME}"\"',
                      "collection_name": '\""${WEAVIATE_COLLECTION_NAME}"\"'
                },
          "api_auth" : {
                  "api_key_config": {
                        "api_key_secret_version": '\""${APIKEY_SECRET_VERSION}"\"'
                  }
          }
        }
    }'

  # TODO(developer): Update the variables.
  # Get operation_id returned in CreateRagCorpus.
  OPERATION_ID="your-operation-id"

  # Poll Operation status until done = true in the response.
  curl -X GET \
  -H "Authorization: Bearer $(gcloud auth print-access-token)" \
    -H "Content-Type: application/json" \
  https://us-central1-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/us-central1/operations/${OPERATION_ID}

  # Call ListRagCorpora API to verify the RAG corpus is created successfully.
  curl -sS -X GET \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer $(gcloud auth print-access-token)" \
    "https://us-central1-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/us-central1/ragCorpora"

from vertexai import rag
import vertexai

# TODO(developer): Update and un-comment below lines
# PROJECT_ID = "your-project-id"
# corpus_name = "projects/[PROJECT_ID]/locations/us-central1/ragCorpora/[rag_corpus_id]"

# Initialize Vertex AI API once per session
vertexai.init(project=PROJECT_ID, location="us-central1")

response = rag.retrieval_query(
    rag_resources=[
        rag.RagResource(
            rag_corpus=corpus_name,
            # Optional: supply IDs from `rag.list_files()`.
            # rag_file_ids=["rag-file-1", "rag-file-2", ...],
        )
    ],
    text="Hello World!",
    rag_retrieval_config=rag.RagRetrievalConfig(
        top_k=10,
        filter=rag.utils.resources.Filter(vector_distance_threshold=0.5),
    ),
)
print(response)
# Example response:
# contexts {
#   contexts {
#     source_uri: "gs://your-bucket-name/file.txt"
#     text: "....
#   ....

Usa la búsqueda híbrida y RAG Engine para la generación fundamentada

Este es un ejemplo de cómo usar la búsqueda híbrida y el motor de RAG para la generación fundamentada.

  # TODO(developer): Update the variables.
  PROJECT_ID = "YOUR_PROJECT_ID"
  # The HTTPS/HTTP Weaviate endpoint you created during provisioning.
  HTTP_ENDPOINT_NAME="https://your.weaviate.endpoint.com"

  # Your Weaviate collection name, which roughly corresponds to a Vertex AI Knowledge Engine Corpus.
  # For example, "MyCollectionName"
  # Note that the first letter needs to be capitalized.
  # Otherwise, Weaviate will capitalize it for you.
  WEAVIATE_COLLECTION_NAME="MyCollectionName"

  # The resource name of your Weaviate API Key your Secret.
  SECRET_NAME="MyWeaviateApiKeySecret"
  # The Secret Manager resource name containing the API Key for your Weaviate endpoint.
  # For example, projects/{project}/secrets/{secret}/versions/latest
  APIKEY_SECRET_VERSION="projects/${PROJECT_ID}/secrets/${SECRET_NAME}/versions/latest"

  # Select a Corpus display name.
  CORPUS_DISPLAY_NAME="SpecialCorpus"

  # Call CreateRagCorpus API and set all Vector DB Config parameters for Weaviate to create a new corpus associated to your selected Weaviate collection.
  curl -X POST \
  -H "Authorization: Bearer $(gcloud auth print-access-token)" \
  -H "Content-Type: application/json" \
  https://us-central1-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/us-central1/ragCorpora \
  -d '{
        "display_name" : '\""${CORPUS_DISPLAY_NAME}"\"',
        "rag_vector_db_config" : {
                "weaviate": {
                      "http_endpoint": '\""${HTTP_ENDPOINT_NAME}"\"',
                      "collection_name": '\""${WEAVIATE_COLLECTION_NAME}"\"'
                },
          "api_auth" : {
                  "api_key_config": {
                        "api_key_secret_version": '\""${APIKEY_SECRET_VERSION}"\"'
                  }
          }
        }
    }'

  # TODO(developer): Update the variables.
  # Get operation_id returned in CreateRagCorpus.
  OPERATION_ID="your-operation-id"

  # Poll Operation status until done = true in the response.
  curl -X GET \
  -H "Authorization: Bearer $(gcloud auth print-access-token)" \
    -H "Content-Type: application/json" \
  https://us-central1-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/us-central1/operations/${OPERATION_ID}

  # Call ListRagCorpora API to verify the RAG corpus is created successfully.
  curl -sS -X GET \
  -H "Content-Type: application/json" \
  -H "Authorization: Bearer $(gcloud auth print-access-token)" \
    "https://us-central1-aiplatform.googleapis.com/v1beta1/projects/${PROJECT_ID}/locations/us-central1/ragCorpora"

from vertexai import rag
from vertexai.generative_models import GenerativeModel, Tool
import vertexai

# TODO(developer): Update and un-comment below lines
# PROJECT_ID = "your-project-id"
# corpus_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}"

# Initialize Vertex AI API once per session
vertexai.init(project=PROJECT_ID, location="us-central1")

rag_retrieval_tool = Tool.from_retrieval(
    retrieval=rag.Retrieval(
        source=rag.VertexRagStore(
            rag_resources=[
                rag.RagResource(
                    rag_corpus=corpus_name,
                    # Optional: supply IDs from `rag.list_files()`.
                    # rag_file_ids=["rag-file-1", "rag-file-2", ...],
                )
            ],
            rag_retrieval_config=rag.RagRetrievalConfig(
                top_k=10,
                filter=rag.utils.resources.Filter(vector_distance_threshold=0.5),
            ),
        ),
    )
)

rag_model = GenerativeModel(
    model_name="gemini-2.0-flash-001", tools=[rag_retrieval_tool]
)
response = rag_model.generate_content("Why is the sky blue?")
print(response.text)
# Example response:
#   The sky appears blue due to a phenomenon called Rayleigh scattering.
#   Sunlight, which contains all colors of the rainbow, is scattered
#   by the tiny particles in the Earth's atmosphere....
#   ...

¿Qué sigue?

• Cómo usar Pinecone con el motor de RAG de Vertex AI