Define API to capture metadata from AI responses

pom.xml

@@ -79,15 +79,15 @@
 		<maven.compiler.source>17</maven.compiler.source>
         <maven.compiler.target>17</maven.compiler.target>
 
-		<!-- prodution dependencies -->
+		<!-- production dependencies -->
 		<spring-boot.version>3.1.3</spring-boot.version>
 		<stringtemplate.version>4.0.2</stringtemplate.version>
 		<open-ai-client.version>0.16.0</open-ai-client.version>
 		<azure-open-ai-client.version>1.0.0-beta.3</azure-open-ai-client.version>
 		<jtokkit.version>0.6.1</jtokkit.version>
 		<victools.version>4.31.1</victools.version>
 
-		<!-- readers/writer/stores dependecies-->
+		<!-- readers/writer/stores dependencies-->
 		<pdfbox.version>3.0.0</pdfbox.version>
 		<pgvector.version>0.1.3</pgvector.version>
 		<postgresql.version>42.6.0</postgresql.version>
@@ -97,7 +97,7 @@
 		<fastjson.version>2.0.42</fastjson.version>
 		<azure-search.version>11.6.0</azure-search.version>
 
-		<!-- testing dependecies -->
+		<!-- testing dependencies -->
 		<testcontainers.version>1.19.0</testcontainers.version>
 
 		<!-- documentation dependencies -->
@@ -409,4 +409,4 @@
 		</plugins>
 	</reporting>
 
-</project>
+</project>

spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/client/AzureOpenAiClient.java

@@ -24,9 +24,12 @@
 import com.azure.ai.openai.models.ChatRole;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.springframework.ai.azure.openai.metadata.AzureOpenAiGenerationMetadata;
 import org.springframework.ai.client.AiClient;
 import org.springframework.ai.client.AiResponse;
 import org.springframework.ai.client.Generation;
+import org.springframework.ai.metadata.PromptMetadata;
+import org.springframework.ai.metadata.PromptMetadata.PromptFilterMetadata;
 import org.springframework.ai.prompt.Prompt;
 import org.springframework.ai.prompt.messages.Message;
 import org.springframework.util.Assert;
@@ -124,13 +127,22 @@ public AiResponse generate(Prompt prompt) {
 
 		for (ChatChoice choice : chatCompletions.getChoices()) {
 			ChatMessage choiceMessage = choice.getMessage();
-			// TODO investigate mapping of additional metadata/runtime info to the general
-			// model.
 			Generation generation = new Generation(choiceMessage.getContent());
 			generations.add(generation);
 		}
 
-		return new AiResponse(generations);
+		return new AiResponse(generations, AzureOpenAiGenerationMetadata.from(chatCompletions))
+			.withPromptMetadata(generatePromptMetadata(chatCompletions));
+	}
+
+	private PromptMetadata generatePromptMetadata(ChatCompletions chatCompletions) {
+
+		return PromptMetadata.of(chatCompletions.getPromptFilterResults()
+			.stream()
+			.map(promptFilterResult -> PromptFilterMetadata.from(promptFilterResult.getPromptIndex(),
+					promptFilterResult.getContentFilterResults()))
+			.toList());
+
 	}
 
 }

spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/metadata/AzureOpenAiGenerationMetadata.java

@@ -0,0 +1,69 @@
+/*
+ * Copyright 2023 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.azure.openai.metadata;
+
+import com.azure.ai.openai.models.ChatCompletions;
+
+import org.springframework.ai.metadata.GenerationMetadata;
+import org.springframework.ai.metadata.Usage;
+import org.springframework.util.Assert;
+
+/**
+ * {@link GenerationMetadata} implementation for
+ * {@literal Microsoft Azure OpenAI Service}.
+ *
+ * @author John Blum
+ * @see org.springframework.ai.metadata.GenerationMetadata
+ * @since 0.7.1
+ */
+public class AzureOpenAiGenerationMetadata implements GenerationMetadata {
+
+	protected static final String AI_METADATA_STRING = "{ @type: %1$s, id: %2$s, usage: %3$s, rateLimit: %4$s }";
+
+	@SuppressWarnings("all")
+	public static AzureOpenAiGenerationMetadata from(ChatCompletions chatCompletions) {
+		Assert.notNull(chatCompletions, "Azure OpenAI ChatCompletions must not be null");
+		String id = chatCompletions.getId();
+		AzureOpenAiUsage usage = AzureOpenAiUsage.from(chatCompletions);
+		AzureOpenAiGenerationMetadata generationMetadata = new AzureOpenAiGenerationMetadata(id, usage);
+		return generationMetadata;
+	}
+
+	private final String id;
+
+	private final Usage usage;
+
+	protected AzureOpenAiGenerationMetadata(String id, AzureOpenAiUsage usage) {
+		this.id = id;
+		this.usage = usage;
+	}
+
+	public String getId() {
+		return this.id;
+	}
+
+	@Override
+	public Usage getUsage() {
+		return this.usage;
+	}
+
+	@Override
+	public String toString() {
+		return AI_METADATA_STRING.formatted(getClass().getTypeName(), getId(), getUsage(), getRateLimit());
+	}
+
+}

spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/metadata/AzureOpenAiUsage.java

@@ -0,0 +1,74 @@
+/*
+ * Copyright 2023 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.azure.openai.metadata;
+
+import com.azure.ai.openai.models.ChatCompletions;
+import com.azure.ai.openai.models.CompletionsUsage;
+
+import org.springframework.ai.metadata.Usage;
+import org.springframework.util.Assert;
+
+/**
+ * {@link Usage} implementation for {@literal Microsoft Azure OpenAI Service}.
+ *
+ * @author John Blum
+ * @see com.azure.ai.openai.models.CompletionsUsage
+ * @since 0.7.0
+ */
+public class AzureOpenAiUsage implements Usage {
+
+	public static AzureOpenAiUsage from(ChatCompletions chatCompletions) {
+		Assert.notNull(chatCompletions, "ChatCompletions must not be null");
+		return from(chatCompletions.getUsage());
+	}
+
+	public static AzureOpenAiUsage from(CompletionsUsage usage) {
+		return new AzureOpenAiUsage(usage);
+	}
+
+	private final CompletionsUsage usage;
+
+	public AzureOpenAiUsage(CompletionsUsage usage) {
+		Assert.notNull(usage, "CompletionUsage must not be null");
+		this.usage = usage;
+	}
+
+	protected CompletionsUsage getUsage() {
+		return this.usage;
+	}
+
+	@Override
+	public Long getPromptTokens() {
+		return Integer.valueOf(getUsage().getPromptTokens()).longValue();
+	}
+
+	@Override
+	public Long getGenerationTokens() {
+		return Integer.valueOf(getUsage().getCompletionTokens()).longValue();
+	}
+
+	@Override
+	public Long getTotalTokens() {
+		return Integer.valueOf(getUsage().getTotalTokens()).longValue();
+	}
+
+	@Override
+	public String toString() {
+		return getUsage().toString();
+	}
+
+}

spring-ai-core/src/main/java/org/springframework/ai/client/AiResponse.java

@@ -19,53 +19,108 @@
 import java.util.List;
 import java.util.Map;
 
+import org.springframework.ai.metadata.GenerationMetadata;
+import org.springframework.ai.metadata.PromptMetadata;
+import org.springframework.lang.Nullable;
+
 public class AiResponse {
 
+	private final GenerationMetadata metadata;
+
 	private final List<Generation> generations;
 
-	private Map<String, Object> providerOutput;
+	private final Map<String, Object> providerOutput;
+
+	private final Map<String, Object> runInfo;
 
-	private Map<String, Object> runInfo;
+	private PromptMetadata promptMetadata;
 
 	public AiResponse(List<Generation> generations) {
-		this(generations, Collections.emptyMap(), Collections.emptyMap());
+		this(generations, Collections.emptyMap(), Collections.emptyMap(), GenerationMetadata.NULL);
+	}
+
+	public AiResponse(List<Generation> generations, GenerationMetadata metadata) {
+		this(generations, Collections.emptyMap(), Collections.emptyMap(), metadata);
 	}
 
 	public AiResponse(List<Generation> generations, Map<String, Object> providerOutput) {
-		this(generations, providerOutput, Collections.emptyMap());
+		this(generations, providerOutput, Collections.emptyMap(), GenerationMetadata.NULL);
 	}
 
 	public AiResponse(List<Generation> generations, Map<String, Object> providerOutput, Map<String, Object> runInfo) {
+		this(generations, providerOutput, runInfo, GenerationMetadata.NULL);
+	}
+
+	public AiResponse(List<Generation> generations, Map<String, Object> providerOutput, Map<String, Object> runInfo,
+			GenerationMetadata metadata) {
+
+		this.metadata = metadata;
 		this.generations = List.copyOf(generations);
 		this.providerOutput = Map.copyOf(providerOutput);
 		this.runInfo = Map.copyOf(runInfo);
 	}
 
 	/**
-	 * The list of generated outputs. It is a list of lists because the Prompt could
-	 * request multiple output generations.
-	 * @return
+	 * The {@link List} of {@link Generation generated outputs}.
+	 * <p>
+	 * It is a {@link List} of {@link List lists} because the Prompt could request
+	 * multiple output {@link Generation generations}.
+	 * @return the {@link List} of {@link Generation generated outputs}.
 	 */
 	public List<Generation> getGenerations() {
-		return Collections.unmodifiableList(generations);
+		return this.generations;
 	}
 
 	public Generation getGeneration() {
 		return this.generations.get(0);
 	}
 
+	/**
+	 * Returns {@link GenerationMetadata} containing information about the use of the AI
+	 * provider's API.
+	 * @return {@link GenerationMetadata} containing information about the use of the AI
+	 * provider's API.
+	 */
+	public GenerationMetadata getGenerationMetadata() {
+		return this.metadata;
+	}
+
+	/**
+	 * Returns {@link PromptMetadata} containing information on prompt processing by the
+	 * AI.
+	 * @return {@link PromptMetadata} containing information on prompt processing by the
+	 * AI.
+	 */
+	public PromptMetadata getPromptMetadata() {
+		PromptMetadata promptMetadata = this.promptMetadata;
+		return promptMetadata != null ? promptMetadata : PromptMetadata.empty();
+	}
+
 	/**
 	 * Arbitrary model provider specific output
 	 */
 	public Map<String, Object> getProviderOutput() {
-		return Collections.unmodifiableMap(providerOutput);
+		return this.providerOutput;
 	}
 
 	/**
 	 * The run metadata information
 	 */
 	public Map<String, Object> getRunInfo() {
-		return Collections.unmodifiableMap(runInfo);
+		return this.runInfo;
+	}
+
+	/**
+	 * Builder method used to include {@link PromptMetadata} returned in the AI response
+	 * when processing the prompt.
+	 * @param promptMetadata {@link PromptMetadata} returned by the AI in the response
+	 * when processing the prompt.
+	 * @return this {@link AiResponse}.
+	 * @see #getPromptMetadata()
+	 */
+	public AiResponse withPromptMetadata(@Nullable PromptMetadata promptMetadata) {
+		this.promptMetadata = promptMetadata;
+		return this;
 	}
 
 }

spring-ai-core/src/main/java/org/springframework/ai/client/Generation.java

@@ -19,28 +19,43 @@
 import java.util.Collections;
 import java.util.Map;
 
+import org.springframework.ai.metadata.ChoiceMetadata;
+import org.springframework.lang.Nullable;
+
 public class Generation {
 
 	// Just text for now
 	private final String text;
 
 	private Map<String, Object> info;
 
+	private ChoiceMetadata choiceMetadata;
+
 	public Generation(String text) {
 		this(text, Collections.emptyMap());
 	}
 
 	public Generation(String text, Map<String, Object> info) {
 		this.text = text;
-		this.info = info;
+		this.info = Map.copyOf(info);
 	}
 
 	public String getText() {
 		return this.text;
 	}
 
 	public Map<String, Object> getInfo() {
-		return Collections.unmodifiableMap(this.info);
+		return this.info;
+	}
+
+	public ChoiceMetadata getChoiceMetadata() {
+		ChoiceMetadata choiceMetadata = this.choiceMetadata;
+		return choiceMetadata != null ? choiceMetadata : ChoiceMetadata.NULL;
+	}
+
+	public Generation withChoiceMetadata(@Nullable ChoiceMetadata choiceMetadata) {
+		this.choiceMetadata = choiceMetadata;
+		return this;
 	}
 
 	@Override