Fixed remaining tests

Signed-off-by: ddobrin <[email protected]>

Author: ddobrin

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

@@ -16,8 +16,10 @@
 
 package org.springframework.ai.vertexai.gemini;
 
+import java.net.URI;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
@@ -311,7 +313,20 @@ private static List<Part> mediaToParts(Collection<Media> media) {
 		List<Part> parts = new ArrayList<>();
 
 		List<Part> mediaParts = media.stream()
-			.map(mediaData -> Part.fromBytes((byte[]) mediaData.getData(), mediaData.getMimeType().toString()))
+			.map(mediaData -> {
+				Object data = mediaData.getData();
+				String mimeType = mediaData.getMimeType().toString();
+				
+				if (data instanceof byte[]) {
+					return Part.fromBytes((byte[]) data, mimeType);
+				} else if (data instanceof URI || data instanceof String) {
+					// Handle URI or String URLs
+					String uri = data.toString();
+					return Part.fromUri(uri, mimeType);
+				} else {
+					throw new IllegalArgumentException("Unsupported media data type: " + data.getClass());
+				}
+			})
 			.toList();
 
 		if (!CollectionUtils.isEmpty(mediaParts)) {
@@ -324,7 +339,22 @@ private static List<Part> mediaToParts(Collection<Media> media) {
 	// Helper methods for JSON/Map conversion
 	private static Map<String, Object> parseJsonToMap(String json) {
 		try {
-			return ModelOptionsUtils.OBJECT_MAPPER.readValue(json, Map.class);
+			// First, try to parse as an array
+			Object parsed = ModelOptionsUtils.OBJECT_MAPPER.readValue(json, Object.class);
+			if (parsed instanceof List) {
+				// It's an array, wrap it in a map with "result" key
+				Map<String, Object> wrapper = new HashMap<>();
+				wrapper.put("result", parsed);
+				return wrapper;
+			} else if (parsed instanceof Map) {
+				// It's already a map, return it
+				return (Map<String, Object>) parsed;
+			} else {
+				// It's a primitive or other type, wrap it
+				Map<String, Object> wrapper = new HashMap<>();
+				wrapper.put("result", parsed);
+				return wrapper;
+			}
 		}
 		catch (Exception e) {
 			throw new RuntimeException("Failed to parse JSON: " + json, e);
@@ -949,7 +979,7 @@ public enum ChatModel implements ChatModelDescription {
 		 * See: <a href=
 		 * "https://cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/2-5-pro">gemini-2.5-pro</a>
 		 */
-		GEMINI_2_5_PRO("gemini-2.5-pro-preview-05-06"),
+		GEMINI_2_5_PRO("gemini-2.5-pro"),
 
 		/**
 		 * <b>gemini-2.5-flash</b> is a thinking model that offers great, well-rounded
@@ -965,7 +995,7 @@ public enum ChatModel implements ChatModelDescription {
 		 * See: <a href=
 		 * "https://cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/2-5-flash">gemini-2.5-flash</a>
 		 */
-		GEMINI_2_5_FLASH("gemini-2.5-flash-preview-04-17");
+		GEMINI_2_5_FLASH("gemini-2.5-flash");
 
 		public final String value;
 

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

@@ -233,7 +233,7 @@ public void defaultOptionsTools() {
 		assertThat(tool.functionDeclarations()).isPresent();
 		assertThat(tool.functionDeclarations().get()).hasSize(1);
 		assertThat(tool.functionDeclarations().get().get(0).name()).as("Explicitly enabled function")
-			.isEqualTo(TOOL_FUNCTION_NAME);
+			.isEqualTo(String.format(TOOL_FUNCTION_NAME));
 
 		// Override the default options function with one from the prompt
 		requestPrompt = client.buildRequestPrompt(new Prompt("Test message content",

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

@@ -91,8 +91,6 @@ void testMessageHistory() {
 		assertThat(response.getResult().getOutput().getText()).containsAnyOf("Blackbeard", "Bartholomew");
 	}
 
-	// Disabled until Gemini 2.5 PRO has an official release
-	@Disabled
 	@Test
 	void googleSearchToolPro() {
 		Prompt prompt = createPrompt(VertexAiGeminiChatOptions.builder()
@@ -131,7 +129,7 @@ void testSafetySettings() {
 
 	@NonNull
 	private Prompt createPrompt(VertexAiGeminiChatOptions chatOptions) {
-		String request = "Tell me about 3 famous pirates from the Golden Age of Piracy and why they did.";
+		String request = "Name 3 famous pirates from the Golden Age of Piracy and what they did.";
 		String name = "Bob";
 		String voice = "pirate";
 		UserMessage userMessage = new UserMessage(request);