Skip to main content

VertexAICodeGenerator

Generate code using Google Vertex AI generative models.

Basic Information

  • Type: haystack_integrations.components.generators.google_vertex.code_generator.VertexAICodeGenerator
  • Components it can connect with:
    • PromptBuilder: Receives code prefix and suffix from PromptBuilder.
    • AnswerBuilder: Sends generated code snippets to AnswerBuilder.

Inputs

ParameterTypeDefaultDescription
prefixstrCode before the current point.
suffixOptional[str]NoneCode after the current point.

Outputs

ParameterTypeDefaultDescription
repliesList[str]A list of generated code snippets.

Overview

VertexAICodeGenerator generates code using Google Vertex AI generative models. It supports code-bison, code-bison-32k, and code-gecko models.

This component is useful for code completion and generation tasks where you provide a code prefix (and optionally a suffix) and the model generates the code in between.

Authorization

This component authenticates using Google Cloud Application Default Credentials (ADCs). For more information, see the official Google documentation.

Create secrets with the following keys: GCP_PROJECT_ID and optionally GCP_DEFAULT_REGION. For detailed instructions on creating secrets, see Create Secrets.

Usage Example

This pipeline uses VertexAICodeGenerator to generate code:

components:
VertexAICodeGenerator:
type: haystack_integrations.components.generators.google_vertex.code_generator.VertexAICodeGenerator
init_parameters:
model: code-bison
project_id:
location:
AnswerBuilder:
type: haystack.components.builders.answer_builder.AnswerBuilder
init_parameters:
pattern:
reference_pattern:
last_message_only: false
return_only_referenced_documents: true

connections:
- sender: VertexAICodeGenerator.replies
receiver: AnswerBuilder.replies

max_runs_per_component: 100

metadata: {}

inputs:
query:
- VertexAICodeGenerator.prefix
- AnswerBuilder.query

outputs:
answers: AnswerBuilder.answers

Parameters

Init Parameters

These are the parameters you can configure in Pipeline Builder:

ParameterTypeDefaultDescription
project_idOptional[str]NoneID of the GCP project to use. By default, it is set during Google Cloud authentication.
modelstrcode-bisonName of the model to use.
locationOptional[str]NoneThe default location to use when making API calls. If not set, uses us-central-1.
kwargsAnyAdditional keyword arguments to pass to the model. See the TextGenerationModel.predict() documentation.

Run Method Parameters

These are the parameters you can configure for the run() method. You can pass these parameters at query time through the API, in Playground, or when running a job.

ParameterTypeDefaultDescription
prefixstrCode before the current point.
suffixOptional[str]NoneCode after the current point.