VertexAICodeGenerator

Generate code using Google Vertex AI generative models.

Key Features

• Code generation and completion using Google Vertex AI models (code-bison, code-bison-32k, code-gecko)
• Code completion task: provide a code prefix (and optionally a suffix) and the model fills in the code in between
• Authenticates using Google Cloud Application Default Credentials (ADCs)

Configuration

1. Drag the VertexAICodeGenerator component onto the canvas from the Component Library.
2. Click on the component to open the configuration panel.
3. On the General tab:
   1. Enter your GCP project ID. Create a secret with the key GCP_PROJECT_ID. For detailed instructions, see Create Secrets.
   2. Select a model. Supported models: code-bison, code-bison-32k, code-gecko.
   3. Optionally, enter the location. If not set, uses us-central1.
4. Go to the Advanced tab to configure additional model keyword arguments.

Connections

VertexAICodeGenerator accepts a prefix (str) input and an optional suffix (Optional[str]) input. It outputs generated code snippets as replies (a list of strings).

Connect the prefix input to the query output from the Input component or from PromptBuilder. Connect the replies output to AnswerBuilder.

Source Code

To check this component's source code, open code_generator.py in the Haystack Core Integrations repository.

Usage Examples

Basic Configuration

  VertexAICodeGenerator:
    type: haystack_integrations.components.generators.google_vertex.code_generator.VertexAICodeGenerator
    init_parameters:
      model: code-bison

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

Inputs

Parameter	Type	Default	Description
prefix	str		Code before the current point.
suffix	Optional[str]	None	Code after the current point.

Outputs

Parameter	Type	Description
replies	List[str]	A list of generated code snippets.

Init Parameters

These are the parameters you can configure in Pipeline Builder:

Parameter	Type	Default	Description
project_id	Optional[str]	None	ID of the GCP project to use. By default, it is set during Google Cloud authentication.
model	str	code-bison	Name of the model to use.
location	Optional[str]	None	The default location to use when making API calls. If not set, uses us-central-1.
kwargs	Any		Additional keyword arguments to pass to the model. See the TextGenerationModel.predict() documentation.

Run Method Parameters

These are the parameters you can configure for the component's run() method. This means you can pass these parameters at query time through the API, in Playground, or when running a job. For details, see Modify Pipeline Parameters at Query Time.

Parameter	Type	Default	Description
prefix	str		Code before the current point.
suffix	Optional[str]	None	Code after the current point.

Related Information

• Create Secrets
• Google Cloud Application Default Credentials