Package org.springframework.ai.vertex.api

Class VertexAiApi

java.lang.Object

org.springframework.ai.vertex.api.VertexAiApi

public class VertexAiApi
extends Object

Vertex AI API client for the Generative Language model. https://developers.generativeai.google/api/rest/generativelanguage https://cloud.google.com/vertex-ai/docs/generative-ai/learn/streaming Provides methods to generate a response from the model given an input https://developers.generativeai.google/api/rest/generativelanguage/models/generateMessage as well as to generate embeddings for the input text: https://developers.generativeai.google/api/rest/generativelanguage/models/embedText Supported models:

                name=models/chat-bison-001,
                version=001,
                displayName=Chat Bison,
                description=Chat-optimized generative language model.,
                inputTokenLimit=4096,
                outputTokenLimit=1024,
                supportedGenerationMethods=[generateMessage, countMessageTokens],
                temperature=0.25,
                topP=0.95,
                topK=40

                name=models/text-bison-001,
                version=001,
                displayName=Text Bison,
                description=Model targeted for text generation.,
                inputTokenLimit=8196,
                outputTokenLimit=1024,
                supportedGenerationMethods=[generateText, countTextTokens, createTunedTextModel],
                temperature=0.7,
                topP=0.95,
                topK=40

                name=models/embedding-gecko-001,
                version=001,
                displayName=Embedding Gecko, description=Obtain a distributed representation of a text.,
                inputTokenLimit=1024,
                outputTokenLimit=1,
                supportedGenerationMethods=[embedText, countTextTokens],
                temperature=null,
                topP=null,
                topK=null
 

Author:

Christian Tzolov

Nested Class Summary

Nested Classes

Modifier and Type	Class	Description
static final record	VertexAiApi.Embedding	A list of floats representing the embedding.
static final record	VertexAiApi.GenerateMessageRequest	Message generation request body.
static final record	VertexAiApi.GenerateMessageResponse	The response from the model.
static final record	VertexAiApi.Message	The base unit of structured text.
static final record	VertexAiApi.MessagePrompt	All of the structured input text passed to the model as a prompt.
static final record	VertexAiApi.Model	Information about a Generative Language Model.
static final record	VertexAiApi.ResponseError	API error response.

Field Summary

Fields

Modifier and Type	Field	Description
static final String	DEFAULT_BASE_URL	The default base URL for accessing the Vertex AI API.
static final String	DEFAULT_EMBEDDING_MODEL	The default embedding model.
static final String	DEFAULT_GENERATE_MODEL	The default generation model.

Constructor Summary

Constructors

Constructor	Description
VertexAiApi(String apiKey)	Create a new chat completion api.
VertexAiApi(String baseUrl, String apiKey, String model, String embeddingModel, org.springframework.web.client.RestClient.Builder restClientBuilder)	Create a new chat completion api.

Method Summary

Modifier and Type	Method	Description
List<VertexAiApi.Embedding>	batchEmbedText(List<String> texts)	Generates a response from the model given an input.
Integer	countMessageTokens(VertexAiApi.MessagePrompt prompt)	Returns the number of tokens in the message prompt.
VertexAiApi.Embedding	embedText(String text)	Generates a response from the model given an input.
VertexAiApi.GenerateMessageResponse	generateMessage(VertexAiApi.GenerateMessageRequest request)	Generates a response from the model given an input.
VertexAiApi.Model	getModel(String modelName)	Returns the model details.
List<String>	listModels()	Returns the list of models available for use.

Methods inherited from class java.lang.Object

clone, equals, finalize, getClass, hashCode, notify, notifyAll, toString, wait, wait, wait

Field Details

DEFAULT_GENERATE_MODEL

public static final String DEFAULT_GENERATE_MODEL

The default generation model. This model is used to generate responses for the input text.

See Also:

• Constant Field Values

DEFAULT_EMBEDDING_MODEL

public static final String DEFAULT_EMBEDDING_MODEL

The default embedding model. This model is used to generate embeddings for the input text.

See Also:

• Constant Field Values

DEFAULT_BASE_URL

public static final String DEFAULT_BASE_URL

The default base URL for accessing the Vertex AI API.

See Also:

• Constant Field Values

Constructor Details

VertexAiApi

public VertexAiApi(String apiKey)

Create a new chat completion api.

Parameters:

apiKey - vertex apiKey.

VertexAiApi

public VertexAiApi(String baseUrl,
 String apiKey,
 String model,
 String embeddingModel,
 org.springframework.web.client.RestClient.Builder restClientBuilder)

Create a new chat completion api.

Parameters:

baseUrl - api base URL.

apiKey - vertex apiKey.

model - vertex model.

embeddingModel - vertex embedding model.

restClientBuilder - RestClient builder.

Method Details

generateMessage

public VertexAiApi.GenerateMessageResponse generateMessage(VertexAiApi.GenerateMessageRequest request)

Generates a response from the model given an input.

Parameters:

request - Request body.

Returns:

Response body.

embedText

public VertexAiApi.Embedding embedText(String text)

Generates a response from the model given an input.

Parameters:

text - Text to embed.

Returns:

Embedding response.

batchEmbedText

public List<VertexAiApi.Embedding> batchEmbedText(List<String> texts)

Generates a response from the model given an input.

Parameters:

texts - List of texts to embed.

Returns:

Embedding response containing a list of embeddings.

countMessageTokens

public Integer countMessageTokens(VertexAiApi.MessagePrompt prompt)

Returns the number of tokens in the message prompt.

Parameters:

prompt - Message prompt to count tokens for.

Returns:

Number of tokens in the message prompt.

listModels

public List<String> listModels()

Returns the list of models available for use.

Returns:

List of models available for use.

getModel

public VertexAiApi.Model getModel(String modelName)

Returns the model details.

Parameters:

modelName - Name of the model to get details for.

Returns:

Model details.