[GH-3723] Vertex AI Gemini logprobs support

Signed-off-by: Rodrigo Malara <[email protected]>

Author: rodrigomalara

models/spring-ai-vertex-ai-gemini/src/main/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatModel.java

@@ -49,6 +49,7 @@
 import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.springframework.ai.vertexai.gemini.api.VertexAiGeminiApi;
 import reactor.core.publisher.Flux;
 import reactor.core.scheduler.Schedulers;
 
@@ -580,8 +581,29 @@ protected List<Generation> responseCandidateToGeneration(Candidate candidate) {
 		int candidateIndex = candidate.getIndex();
 		FinishReason candidateFinishReason = candidate.getFinishReason();
 
+		// Convert from VertexAI protobuf to VertexAiGeminiApi DTOs
+		List<VertexAiGeminiApi.LogProbs.TopContent> topCandidates = candidate.getLogprobsResult()
+			.getTopCandidatesList()
+			.stream()
+			.filter(topCandidate -> !topCandidate.getCandidatesList().isEmpty())
+			.map(topCandidate -> new VertexAiGeminiApi.LogProbs.TopContent(topCandidate.getCandidatesList()
+				.stream()
+				.map(c -> new VertexAiGeminiApi.LogProbs.Content(c.getToken(), c.getLogProbability(), c.getTokenId()))
+				.toList()))
+			.toList();
+
+		// Convert from VertexAI protobuf to VertexAiGeminiApi DTOs
+		List<VertexAiGeminiApi.LogProbs.Content> chosenCandidates = candidate.getLogprobsResult()
+			.getChosenCandidatesList()
+			.stream()
+			.map(c -> new VertexAiGeminiApi.LogProbs.Content(c.getToken(), c.getLogProbability(), c.getTokenId()))
+			.toList();
+
+		VertexAiGeminiApi.LogProbs logprobs = new VertexAiGeminiApi.LogProbs(candidate.getAvgLogprobs(), topCandidates,
+				chosenCandidates);
+
 		Map<String, Object> messageMetadata = Map.of("candidateIndex", candidateIndex, "finishReason",
-				candidateFinishReason);
+				candidateFinishReason, "logprobs", logprobs);
 
 		ChatGenerationMetadata chatGenerationMetadata = ChatGenerationMetadata.builder()
 			.finishReason(candidateFinishReason.name())
@@ -737,6 +759,10 @@ private GenerationConfig toGenerationConfig(VertexAiGeminiChatOptions options) {
 		if (options.getPresencePenalty() != null) {
 			generationConfigBuilder.setPresencePenalty(options.getPresencePenalty().floatValue());
 		}
+		if (options.getLogprobs() != null) {
+			generationConfigBuilder.setLogprobs(options.getLogprobs());
+		}
+		generationConfigBuilder.setResponseLogprobs(options.getResponseLogprobs());
 
 		return generationConfigBuilder.build();
 	}

models/spring-ai-vertex-ai-gemini/src/main/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatOptions.java

@@ -64,6 +64,20 @@ public class VertexAiGeminiChatOptions implements ToolCallingChatOptions {
 	 */
 	private @JsonProperty("temperature") Double temperature;
 
+	/**
+	 * Optional. Enable returning the log probabilities of the top candidate tokens at each generation step.
+	 * The model's chosen token might not be the same as the top candidate token at each step.
+	 * Specify the number of candidates to return by using an integer value in the range of 1-20.
+	 * Should not be set unless responseLogprobs is set to true.
+	 */
+	private @JsonProperty("logprobs") Integer logprobs;
+
+	/**
+	 * Optional. If true, returns the log probabilities of the tokens that were chosen by the model at each step.
+	 * By default, this parameter is set to false.
+	 */
+	private @JsonProperty("responseLogprobs") boolean responseLogprobs;
+
 	/**
 	 * Optional. If specified, nucleus sampling will be used.
 	 */
@@ -162,6 +176,8 @@ public static VertexAiGeminiChatOptions fromOptions(VertexAiGeminiChatOptions fr
 		options.setSafetySettings(fromOptions.getSafetySettings());
 		options.setInternalToolExecutionEnabled(fromOptions.getInternalToolExecutionEnabled());
 		options.setToolContext(fromOptions.getToolContext());
+		options.setLogprobs(fromOptions.getLogprobs());
+		options.setResponseLogprobs(fromOptions.getResponseLogprobs());
 		return options;
 	}
 
@@ -183,6 +199,10 @@ public void setTemperature(Double temperature) {
 		this.temperature = temperature;
 	}
 
+	public void setResponseLogprobs(boolean responseLogprobs) {
+		this.responseLogprobs = responseLogprobs;
+	}
+
 	@Override
 	public Double getTopP() {
 		return this.topP;
@@ -326,6 +346,18 @@ public void setToolContext(Map<String, Object> toolContext) {
 		this.toolContext = toolContext;
 	}
 
+	public Integer getLogprobs() {
+		return logprobs;
+	}
+
+	public void setLogprobs(Integer logprobs) {
+		this.logprobs = logprobs;
+	}
+
+	public boolean getResponseLogprobs() {
+		return responseLogprobs;
+	}
+
 	@Override
 	public boolean equals(Object o) {
 		if (this == o) {
@@ -346,15 +378,17 @@ public boolean equals(Object o) {
 				&& Objects.equals(this.toolNames, that.toolNames)
 				&& Objects.equals(this.safetySettings, that.safetySettings)
 				&& Objects.equals(this.internalToolExecutionEnabled, that.internalToolExecutionEnabled)
-				&& Objects.equals(this.toolContext, that.toolContext);
+				&& Objects.equals(this.toolContext, that.toolContext) && Objects.equals(this.logprobs, that.logprobs)
+				&& Objects.equals(this.responseLogprobs, that.responseLogprobs);
 	}
 
 	@Override
 	public int hashCode() {
 		return Objects.hash(this.stopSequences, this.temperature, this.topP, this.topK, this.candidateCount,
 				this.frequencyPenalty, this.presencePenalty, this.maxOutputTokens, this.model, this.responseMimeType,
 				this.toolCallbacks, this.toolNames, this.googleSearchRetrieval, this.safetySettings,
-				this.internalToolExecutionEnabled, this.toolContext);
+				this.internalToolExecutionEnabled, this.toolContext, this.toolContext, this.logprobs,
+				this.responseLogprobs);
 	}
 
 	@Override
@@ -365,7 +399,8 @@ public String toString() {
 				+ this.candidateCount + ", maxOutputTokens=" + this.maxOutputTokens + ", model='" + this.model + '\''
 				+ ", responseMimeType='" + this.responseMimeType + '\'' + ", toolCallbacks=" + this.toolCallbacks
 				+ ", toolNames=" + this.toolNames + ", googleSearchRetrieval=" + this.googleSearchRetrieval
-				+ ", safetySettings=" + this.safetySettings + '}';
+				+ ", safetySettings=" + this.safetySettings + ", logProbs=" + this.logprobs + ", responseLogprobs="
+				+ this.responseLogprobs + '}';
 	}
 
 	@Override
@@ -488,6 +523,16 @@ public Builder toolContext(Map<String, Object> toolContext) {
 			return this;
 		}
 
+		public Builder logprobs(Integer logprobs) {
+			this.options.setLogprobs(logprobs);
+			return this;
+		}
+
+		public Builder responseLogprobs(Boolean responseLogprobs) {
+			this.options.setResponseLogprobs(responseLogprobs);
+			return this;
+		}
+
 		public VertexAiGeminiChatOptions build() {
 			return this.options;
 		}

models/spring-ai-vertex-ai-gemini/src/main/java/org/springframework/ai/vertexai/gemini/api/VertexAiGeminiApi.java

@@ -0,0 +1,16 @@
+package org.springframework.ai.vertexai.gemini.api;
+
+import java.util.List;
+
+public class VertexAiGeminiApi {
+
+	public record LogProbs(Double avgLogprobs, List<TopContent> topCandidates,
+			List<LogProbs.Content> chosenCandidates) {
+		public record Content(String token, Float logprob, Integer id) {
+		}
+
+		public record TopContent(List<Content> candidates) {
+		}
+	}
+
+}

models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/CreateGeminiRequestTests.java

@@ -262,6 +262,8 @@ public void createRequestWithGenerationConfigOptions() {
 				.stopSequences(List.of("stop1", "stop2"))
 				.candidateCount(1)
 				.responseMimeType("application/json")
+				.responseLogprobs(true)
+				.logprobs(2)
 				.build())
 			.build();
 
@@ -280,6 +282,8 @@ public void createRequestWithGenerationConfigOptions() {
 		assertThat(request.model().getGenerationConfig().getStopSequences(0)).isEqualTo("stop1");
 		assertThat(request.model().getGenerationConfig().getStopSequences(1)).isEqualTo("stop2");
 		assertThat(request.model().getGenerationConfig().getResponseMimeType()).isEqualTo("application/json");
+		assertThat(request.model().getGenerationConfig().getLogprobs()).isEqualTo(2);
+		assertThat(request.model().getGenerationConfig().getResponseLogprobs()).isEqualTo(true);
 	}
 
 }

models/spring-ai-vertex-ai-gemini/src/test/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatModelIT.java

@@ -47,6 +47,7 @@
 import org.springframework.ai.model.tool.ToolCallingManager;
 import org.springframework.ai.tool.annotation.Tool;
 import org.springframework.ai.vertexai.gemini.VertexAiGeminiChatModel.ChatModel;
+import org.springframework.ai.vertexai.gemini.api.VertexAiGeminiApi;
 import org.springframework.ai.vertexai.gemini.common.VertexAiGeminiSafetySetting;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Value;
@@ -226,6 +227,26 @@ void textStream() {
 		assertThat(generationTextFromStream).isNotEmpty();
 	}
 
+	@Test
+	void logprobs() {
+		VertexAiGeminiChatOptions chatOptions = VertexAiGeminiChatOptions.builder()
+			.logprobs(1)
+			.responseLogprobs(true)
+			.build();
+
+		var logprobs = (VertexAiGeminiApi.LogProbs) this.chatModel
+			.call(new Prompt("Explain Bulgaria? Answer in 10 paragraphs.", chatOptions))
+			.getResult()
+			.getOutput()
+			.getMetadata()
+			.get("logprobs");
+
+		assertThat(logprobs).isNotNull();
+		assertThat(logprobs.avgLogprobs()).isNotZero();
+		assertThat(logprobs.topCandidates()).isNotEmpty();
+		assertThat(logprobs.chosenCandidates()).isNotEmpty();
+	}
+
 	@Test
 	void beanStreamOutputConverterRecords() {