|
17 | 17 | package com.google.cloud.vertexai.generativeai.preview;
|
18 | 18 |
|
19 | 19 | import static com.google.common.truth.Truth.assertThat;
|
| 20 | +import static org.junit.Assert.assertThrows; |
20 | 21 |
|
21 | 22 | import com.google.cloud.vertexai.api.Part;
|
22 | 23 | import com.google.protobuf.ByteString;
|
| 24 | +import com.google.protobuf.Struct; |
| 25 | +import com.google.protobuf.Value; |
23 | 26 | import java.net.URI;
|
24 | 27 | import java.net.URISyntaxException;
|
| 28 | +import java.util.HashMap; |
| 29 | +import java.util.Map; |
25 | 30 | import org.junit.Test;
|
26 | 31 | import org.junit.runner.RunWith;
|
27 | 32 | import org.junit.runners.JUnit4;
|
@@ -72,4 +77,56 @@ public void fromMimeTypeAndData_dataInURI() throws URISyntaxException {
|
72 | 77 | assertThat(part.getFileData().getMimeType()).isEqualTo("image/png");
|
73 | 78 | assertThat(part.getFileData().getFileUri()).isEqualTo(fileUri.toString());
|
74 | 79 | }
|
| 80 | + |
| 81 | + @Test |
| 82 | + public void testFromFunctionResponseWithStruct() { |
| 83 | + String functionName = "getCurrentWeather"; |
| 84 | + Struct functionResponse = |
| 85 | + Struct.newBuilder() |
| 86 | + .putFields("currentWeather", Value.newBuilder().setStringValue("Super nice!").build()) |
| 87 | + .putFields("currentTemperature", Value.newBuilder().setNumberValue(85.0).build()) |
| 88 | + .putFields("isRaining", Value.newBuilder().setBoolValue(false).build()) |
| 89 | + .build(); |
| 90 | + |
| 91 | + Part part = PartMaker.fromFunctionResponse(functionName, functionResponse); |
| 92 | + |
| 93 | + assertThat(part.getFunctionResponse().getName()).isEqualTo("getCurrentWeather"); |
| 94 | + assertThat(part.getFunctionResponse().getResponse()).isEqualTo(functionResponse); |
| 95 | + } |
| 96 | + |
| 97 | + @Test |
| 98 | + public void testFromFunctionResponseWithMap() { |
| 99 | + String functionName = "getCurrentWeather"; |
| 100 | + Map<String, Object> functionResponse = new HashMap<>(); |
| 101 | + functionResponse.put("currentWeather", "Super nice!"); |
| 102 | + functionResponse.put("currentTemperature", 85.0); |
| 103 | + functionResponse.put("isRaining", false); |
| 104 | + functionResponse.put("other", null); |
| 105 | + |
| 106 | + Part part = PartMaker.fromFunctionResponse(functionName, functionResponse); |
| 107 | + |
| 108 | + assertThat(part.getFunctionResponse().getName()).isEqualTo("getCurrentWeather"); |
| 109 | + |
| 110 | + Map<String, Value> fieldsMap = part.getFunctionResponse().getResponse().getFieldsMap(); |
| 111 | + assertThat(fieldsMap.get("currentWeather").getStringValue()).isEqualTo("Super nice!"); |
| 112 | + assertThat(fieldsMap.get("currentTemperature").getNumberValue()).isEqualTo(85.0); |
| 113 | + assertThat(fieldsMap.get("isRaining").getBoolValue()).isEqualTo(false); |
| 114 | + assertThat(fieldsMap.get("other").hasNullValue()).isEqualTo(true); |
| 115 | + } |
| 116 | + |
| 117 | + @Test |
| 118 | + public void testFromFunctionResponseWithInvalidMap() { |
| 119 | + String functionName = "getCurrentWeather"; |
| 120 | + Map<String, Object> invalidResponse = new HashMap<>(); |
| 121 | + invalidResponse.put("currentWeather", new byte[] {1, 2, 3}); |
| 122 | + IllegalArgumentException thrown = |
| 123 | + assertThrows( |
| 124 | + IllegalArgumentException.class, |
| 125 | + () -> PartMaker.fromFunctionResponse(functionName, invalidResponse)); |
| 126 | + assertThat(thrown) |
| 127 | + .hasMessageThat() |
| 128 | + .isEqualTo( |
| 129 | + "The value in the map can only be one of the following format: " |
| 130 | + + "String, Double, Boolean, null."); |
| 131 | + } |
75 | 132 | }
|
0 commit comments