Refactored to constructorArg for mandatory args in GoogleVertexAiUnifiedStreamingProcessor

Author: lhoet-google

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/GoogleVertexAiUnifiedStreamingProcessor.java

@@ -192,17 +192,17 @@ private static StreamingUnifiedChatCompletionResults.ChatCompletionChunk.Choice
 
         static {
             PARSER.declareObjectArray(
-                ConstructingObjectParser.optionalConstructorArg(),
+                ConstructingObjectParser.constructorArg(),
                 (p, c) -> CandidateParser.parse(p),
                 new ParseField(CANDIDATES_FIELD)
             );
             PARSER.declareObject(
-                ConstructingObjectParser.optionalConstructorArg(),
+                ConstructingObjectParser.constructorArg(),
                 (p, c) -> UsageMetadataParser.parse(p),
                 new ParseField(USAGE_METADATA_FIELD)
             );
-            PARSER.declareString(ConstructingObjectParser.optionalConstructorArg(), new ParseField(MODEL_VERSION_FIELD));
-            PARSER.declareString(ConstructingObjectParser.optionalConstructorArg(), new ParseField(RESPONSE_ID_FIELD));
+            PARSER.declareString(ConstructingObjectParser.constructorArg(), new ParseField(MODEL_VERSION_FIELD));
+            PARSER.declareString(ConstructingObjectParser.constructorArg(), new ParseField(RESPONSE_ID_FIELD));
         }
 
         public static StreamingUnifiedChatCompletionResults.ChatCompletionChunk parse(XContentParser parser) throws IOException {
@@ -224,7 +224,7 @@ private static class CandidateParser {
 
         static {
             PARSER.declareObject(
-                ConstructingObjectParser.optionalConstructorArg(),
+                ConstructingObjectParser.constructorArg(),
                 (p, c) -> ContentParser.parse(p),
                 new ParseField(CONTENT_FIELD)
             );
@@ -248,9 +248,9 @@ private static class ContentParser {
         );
 
         static {
-            PARSER.declareString(ConstructingObjectParser.optionalConstructorArg(), new ParseField(ROLE_FIELD));
+            PARSER.declareString(ConstructingObjectParser.constructorArg(), new ParseField(ROLE_FIELD));
             PARSER.declareObjectArray(
-                ConstructingObjectParser.optionalConstructorArg(),
+                ConstructingObjectParser.constructorArg(),
                 (p, c) -> PartParser.parse(p),
                 new ParseField(PARTS_FIELD)
             );

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/googlevertexai/GoogleVertexAiUnifiedStreamingProcessorTests.java

@@ -81,7 +81,7 @@ public void testJsonLiteral() {
         }
     }
 
-    public void testJsonLiteral_optionalTopLevelFieldsMissing() {
+    public void testJsonLiteral_usageMetadataTokenCountMissing() {
         String json = """
                 {
                   "candidates" : [ {
@@ -90,7 +90,12 @@ public void testJsonLiteral_optionalTopLevelFieldsMissing() {
                       "parts" : [ { "text" : "Hello" } ]
                     },
                     "finishReason": "STOP"
-                  } ]
+                  } ],
+                  "usageMetadata" : {
+                    "trafficType" : "ON_DEMAND"
+                  },
+                  "modelVersion": "gemini-2.0-flash-001",
+                  "responseId": "responseId"
                 }
             """;
 
@@ -101,12 +106,11 @@ public void testJsonLiteral_optionalTopLevelFieldsMissing() {
         try (XContentParser parser = XContentFactory.xContent(XContentType.JSON).createParser(parserConfig, json)) {
             var chunk = GoogleVertexAiUnifiedStreamingProcessor.GoogleVertexAiChatCompletionChunkParser.parse(parser);
 
-            assertNull(chunk.id());
+            assertEquals("responseId", chunk.id());
             assertEquals(1, chunk.choices().size());
             var choice = chunk.choices().getFirst();
             assertEquals("Hello", choice.delta().content());
             assertEquals("model", choice.delta().role());
-            assertNull(chunk.model());
             assertEquals("STOP", choice.finishReason());
             assertEquals(0, choice.index());
             assertNull(choice.delta().toolCalls());
@@ -131,7 +135,14 @@ public void testJsonLiteral_functionCallArgsMissing() {
                       ]
                     }
                   } ],
-                  "responseId" : "resId789"
+                  "responseId" : "resId789",
+                  "modelVersion": "gemini-2.0-flash-00",
+                  "usageMetadata" : {
+                    "promptTokenCount": 10,
+                    "candidatesTokenCount": 20,
+                    "totalTokenCount": 30,
+                    "trafficType" : "ON_DEMAND"
+                  }
                 }
             """;
         XContentParserConfiguration parserConfig = XContentParserConfiguration.EMPTY.withDeprecationHandler(
@@ -171,7 +182,14 @@ public void testJsonLiteral_multipleTextParts() {
                     },
                     "finishReason": "STOP"
                   } ],
-                  "responseId" : "multiTextId"
+                  "responseId" : "multiTextId",
+                  "usageMetadata" : {
+                    "promptTokenCount": 10,
+                    "candidatesTokenCount": 20,
+                    "totalTokenCount": 30,
+                    "trafficType" : "ON_DEMAND"
+                  },
+                  "modelVersion": "gemini-2.0-flash-001"
                 }
             """;
 
@@ -192,8 +210,7 @@ public void testJsonLiteral_multipleTextParts() {
             assertEquals("STOP", choice.finishReason());
             assertEquals(0, choice.index());
             assertNull(choice.delta().toolCalls());
-            assertNull(chunk.model());
-            assertNull(chunk.usage());
+            assertEquals("gemini-2.0-flash-001", chunk.model());
         } catch (IOException e) {
             fail("IOException during test: " + e.getMessage());
         }