apache_beam.ml.rag.embeddings.vertex_ai module

RAG-specific embedding implementations using Vertex AI models.

class apache_beam.ml.rag.embeddings.vertex_ai.VertexAITextEmbeddings(model_name: str, *, title: str | None = None, task_type: str = 'RETRIEVAL_DOCUMENT', project: str | None = None, location: str | None = None, credentials: Credentials | None = None, **kwargs)

Bases: EmbeddingsManager

Utilizes Vertex AI text embeddings for semantic search and RAG pipelines.

Parameters:

• model_name – Name of the Vertex AI text embedding model

• title – Optional title for the text content

• task_type – Task type for embeddings (default: RETRIEVAL_DOCUMENT)

• project – GCP project ID

• location – GCP location

• credentials – Optional GCP credentials

• **kwargs – Additional arguments passed to EmbeddingsManager including

• inference_args. (ModelHandler)

get_model_handler()

Returns model handler configured with RAG adapter.

get_ptransform_for_processing(**kwargs) → PTransform[PCollection[Chunk], PCollection[Chunk]]

Returns PTransform that uses the RAG adapter.