From 0c95cc396f0e739c52ca07cb0106c90f630c2a99 Mon Sep 17 00:00:00 2001 From: jonghoon park Date: Tue, 13 May 2025 02:28:00 +0900 Subject: [PATCH] GH-3118: Fix NPE in OpenAIChatModel compatible streaming api with vertex ai gemini Fixes: #3118 Signed-off-by: jonghoon park --- .../ai/openai/OpenAiChatModel.java | 7 +- .../VertexAIGeminiWithOpenAiChatModelIT.java | 121 ++++++++++++++++++ 2 files changed, 125 insertions(+), 3 deletions(-) create mode 100644 models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/proxy/VertexAIGeminiWithOpenAiChatModelIT.java diff --git a/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatModel.java b/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatModel.java index d7719b141..ae840a073 100644 --- a/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatModel.java +++ b/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatModel.java @@ -104,6 +104,7 @@ import org.springframework.util.StringUtils; * @author Ilayaperumal Gopinathan * @author Alexandros Pappas * @author Soby Chacko + * @author Jonghoon Park * @see ChatModel * @see StreamingChatModel * @see OpenAiApi @@ -304,15 +305,15 @@ public class OpenAiChatModel implements ChatModel { Flux chatResponse = completionChunks.map(this::chunkToChatCompletion) .switchMap(chatCompletion -> Mono.just(chatCompletion).map(chatCompletion2 -> { try { - @SuppressWarnings("null") - String id = chatCompletion2.id(); + // If an id is not provided, set to "NO_ID" (for compatible APIs). + String id = chatCompletion2.id() == null ? "NO_ID" : chatCompletion2.id(); List generations = chatCompletion2.choices().stream().map(choice -> { // @formatter:off if (choice.message().role() != null) { roleMap.putIfAbsent(id, choice.message().role().name()); } Map metadata = Map.of( - "id", chatCompletion2.id(), + "id", id, "role", roleMap.getOrDefault(id, ""), "index", choice.index(), "finishReason", choice.finishReason() != null ? choice.finishReason().name() : "", diff --git a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/proxy/VertexAIGeminiWithOpenAiChatModelIT.java b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/proxy/VertexAIGeminiWithOpenAiChatModelIT.java new file mode 100644 index 000000000..7b3ed862f --- /dev/null +++ b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/proxy/VertexAIGeminiWithOpenAiChatModelIT.java @@ -0,0 +1,121 @@ +/* + * Copyright 2025-2025 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.ai.openai.chat.proxy; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; +import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.messages.Message; +import org.springframework.ai.chat.messages.UserMessage; +import org.springframework.ai.chat.model.ChatResponse; +import org.springframework.ai.chat.model.Generation; +import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.ai.chat.prompt.SystemPromptTemplate; +import org.springframework.ai.model.SimpleApiKey; +import org.springframework.ai.model.tool.ToolCallingManager; +import org.springframework.ai.openai.OpenAiChatModel; +import org.springframework.ai.openai.OpenAiChatOptions; +import org.springframework.ai.openai.api.OpenAiApi; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.beans.factory.annotation.Value; +import org.springframework.boot.SpringBootConfiguration; +import org.springframework.boot.test.context.SpringBootTest; +import org.springframework.context.annotation.Bean; +import org.springframework.core.io.Resource; +import reactor.core.publisher.Flux; + +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * @author Jonghoon Park + */ +@SpringBootTest(classes = VertexAIGeminiWithOpenAiChatModelIT.Config.class) +@EnabledIfEnvironmentVariable(named = "GEMINI_API_KEY", matches = ".+") +class VertexAIGeminiWithOpenAiChatModelIT { + + private static final String VERTEX_AI_GEMINI_BASE_URL = "https://generativelanguage.googleapis.com"; + + private static final String VERTEX_AI_GEMINI_DEFAULT_MODEL = "gemini-2.0-flash"; + + @Value("classpath:/prompts/system-message.st") + private Resource systemResource; + + @Autowired + private OpenAiChatModel chatModel; + + @Test + void roleTest() { + UserMessage userMessage = new UserMessage( + "Tell me about 3 famous pirates from the Golden Age of Piracy and what they did."); + SystemPromptTemplate systemPromptTemplate = new SystemPromptTemplate(this.systemResource); + Message systemMessage = systemPromptTemplate.createMessage(Map.of("name", "Bob", "voice", "pirate")); + Prompt prompt = new Prompt(List.of(userMessage, systemMessage)); + ChatResponse response = this.chatModel.call(prompt); + assertThat(response.getResults()).hasSize(1); + assertThat(response.getResults().get(0).getOutput().getText()).contains("Blackbeard"); + } + + @Test + void streamRoleTest() { + UserMessage userMessage = new UserMessage( + "Tell me about 3 famous pirates from the Golden Age of Piracy and what they did."); + SystemPromptTemplate systemPromptTemplate = new SystemPromptTemplate(this.systemResource); + Message systemMessage = systemPromptTemplate.createMessage(Map.of("name", "Bob", "voice", "pirate")); + Prompt prompt = new Prompt(List.of(userMessage, systemMessage)); + Flux flux = this.chatModel.stream(prompt); + + List responses = flux.collectList().block(); + assertThat(responses.size()).isGreaterThan(1); + + String stitchedResponseContent = responses.stream() + .map(ChatResponse::getResults) + .flatMap(List::stream) + .map(Generation::getOutput) + .map(AssistantMessage::getText) + .collect(Collectors.joining()); + + assertThat(stitchedResponseContent).contains("Blackbeard"); + } + + @SpringBootConfiguration + static class Config { + + @Bean + public OpenAiApi chatCompletionApi() { + return OpenAiApi.builder() + .baseUrl(VERTEX_AI_GEMINI_BASE_URL) + .completionsPath("/v1beta/openai/chat/completions") + .apiKey(new SimpleApiKey(System.getenv("GEMINI_API_KEY"))) + .build(); + } + + @Bean + public OpenAiChatModel openAiClient(OpenAiApi openAiApi) { + return OpenAiChatModel.builder() + .openAiApi(openAiApi) + .toolCallingManager(ToolCallingManager.builder().build()) + .defaultOptions(OpenAiChatOptions.builder().model(VERTEX_AI_GEMINI_DEFAULT_MODEL).build()) + .build(); + } + + } + +}