Add Tool Argument Augmentation for dynamic tool schema enhancement

Introduce utilities to augment tool input schemas with additional
arguments (e.g., reasoning, metadata) without modifying tool
implementations. Includes AugmentedToolCallbackProvider,
AugmentedToolCallback, and ToolInputSchemaAugmenter components.

Signed-off-by: Christian Tzolov <[email protected]>

Author: tzolov

spring-ai-docs/src/main/antora/modules/ROOT/pages/api/tools.adoc

@@ -1328,6 +1328,74 @@ ToolCallbackResolver toolCallbackResolver(List<FunctionCallback> toolCallbacks)
 
 The `ToolCallbackResolver` is used internally by the `ToolCallingManager` to resolve tools dynamically at runtime, supporting both xref:_framework_controlled_tool_execution[] and xref:_user_controlled_tool_execution[].
 
+[[tool-argument-augmentation]]
+== Tool Argument Augmentation
+
+Spring AI provides a utility for **dynamic augmentation of tool input schemas** with additional arguments. This allows capturing extra information from the model—such as reasoning or metadata—without modifying the underlying tool implementation.
+
+Common use cases include:
+
+* **Inner Thinking/Reasoning**: Capture the model's step-by-step reasoning before executing a tool
+* **Memory Enhancement**: Extract insights to store in long-term memory
+* **Analytics & Tracking**: Collect metadata, user intent, or usage patterns
+* **Multi-Agent Coordination**: Pass agent identifiers or coordination signals
+
+=== Quick Start
+
+1. **Define augmented arguments** as a Java Record:
+
+[source,java]
+----
+public record AgentThinking(
+    @ToolParam(description = "Your reasoning for calling this tool", required = true)
+    String innerThought,
+
+    @ToolParam(description = "Confidence level (low, medium, high)", required = false)
+    String confidence
+) {}
+----
+
+2. **Wrap your tool** with `AugmentedToolCallbackProvider`:
+
+[source,java]
+----
+AugmentedToolCallbackProvider<AgentThinking> provider = AugmentedToolCallbackProvider
+    .<AgentThinking>builder()
+    .toolObject(new MyTools())  // Your @Tool annotated class
+    .argumentType(AgentThinking.class)
+    .argumentConsumer(event -> {
+        AgentThinking thinking = event.arguments();
+        log.info("Tool: {} | Reasoning: {}", event.toolDefinition().name(), thinking.innerThought());
+    })
+    .removeExtraArgumentsAfterProcessing(true)
+    .build();
+----
+
+3. **Use with ChatClient**:
+
+[source,java]
+----
+ChatClient chatClient = ChatClient.builder(chatModel)
+    .defaultToolCallbacks(provider)
+    .build();
+----
+
+The LLM sees the augmented schema with your additional fields. Your consumer receives the `AgentThinking` record, while the original tool receives only its expected arguments.
+
+=== Core Components
+
+* `AugmentedToolCallbackProvider<T>` - Wraps tool objects or providers, augmenting all tools with the specified Record type
+* `AugmentedToolCallback<T>` - Wraps individual `ToolCallback` instances
+* `AugmentedArgumentEvent<T>` - Contains `toolDefinition()`, `rawInput()`, and `arguments()` for consumers
+* `ToolInputSchemaAugmenter` - Low-level utility for schema manipulation
+
+=== Configuration
+
+The `removeExtraArgumentsAfterProcessing` option controls whether augmented arguments are passed to the original tool:
+
+* `true` (default) - Remove augmented arguments before calling the tool
+* `false` - Preserve augmented arguments in the input (if the tool can ignore extra fields)
+
 == Observability
 
 Tool calling includes observability support with spring.ai.tool observations that measure completion time and propagate tracing information. See xref:observability/index.adoc#_tool_calling[Tool Calling Observability].

spring-ai-model/src/main/java/org/springframework/ai/model/ModelOptionsUtils.java

@@ -378,6 +378,18 @@ private static String toGetName(String name) {
 	 */
 	public static String getJsonSchema(Type inputType, boolean toUpperCaseTypeValues) {
 
+		ObjectNode node = getJsonSchema(inputType);
+
+		if (toUpperCaseTypeValues) { // Required for OpenAPI 3.0 (at least Vertex AI
+			// version of it).
+			toUpperCaseTypeValues(node);
+		}
+
+		return node.toPrettyString();
+	}
+
+	public static ObjectNode getJsonSchema(Type inputType) {
+
 		if (SCHEMA_GENERATOR_CACHE.get() == null) {
 
 			JacksonModule jacksonModule = new JacksonModule(JacksonOption.RESPECT_JSONPROPERTY_REQUIRED);
@@ -405,12 +417,7 @@ public static String getJsonSchema(Type inputType, boolean toUpperCaseTypeValues
 			node.putObject("properties");
 		}
 
-		if (toUpperCaseTypeValues) { // Required for OpenAPI 3.0 (at least Vertex AI
-			// version of it).
-			toUpperCaseTypeValues(node);
-		}
-
-		return node.toPrettyString();
+		return node;
 	}
 
 	public static void toUpperCaseTypeValues(ObjectNode node) {

spring-ai-model/src/main/java/org/springframework/ai/tool/augment/AugmentedArgumentEvent.java

@@ -0,0 +1,32 @@
+/*
+ * 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.tool.augment;
+
+import org.springframework.ai.tool.definition.ToolDefinition;
+
+/**
+ * An event that encapsulates the augmented arguments extracted from a tool input, along
+ * with the associated tool definition and raw input data.
+ *
+ * @param <T> The type of the augmented arguments record.
+ * @param toolDefinition The tool definition associated with the event.
+ * @param rawInput The raw input data as a string.
+ * @param arguments The augmented arguments extracted from the input.
+ * @author Christian Tzolov
+ */
+public record AugmentedArgumentEvent<T>(ToolDefinition toolDefinition, String rawInput, T arguments) {
+}

spring-ai-model/src/main/java/org/springframework/ai/tool/augment/AugmentedToolCallback.java

@@ -0,0 +1,146 @@
+/*
+ * 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.tool.augment;
+
+import java.util.List;
+import java.util.Map;
+import java.util.function.Consumer;
+
+import com.fasterxml.jackson.core.type.TypeReference;
+
+import org.springframework.ai.chat.model.ToolContext;
+import org.springframework.ai.tool.ToolCallback;
+import org.springframework.ai.tool.augment.ToolInputSchemaAugmenter.AugmentedArgumentType;
+import org.springframework.ai.tool.definition.ToolDefinition;
+import org.springframework.ai.util.json.JsonParser;
+import org.springframework.lang.Nullable;
+import org.springframework.util.Assert;
+
+/**
+ * This class wraps an existing {@link ToolCallback} and modifies its input schema to
+ * include additional fields defined in the provided Record type. It also provides a
+ * mechanism to handle these extended arguments, either by consuming them via a provided
+ * {@link Consumer} or by removing them from the input after processing.
+ *
+ * @author Christian Tzolov
+ */
+public class AugmentedToolCallback<T extends Record> implements ToolCallback {
+
+	/**
+	 * The delegate ToolCallback that this class extends.
+	 */
+	private final ToolCallback delegate;
+
+	/**
+	 * The augmented ToolDefinition that includes the augmented input schema.
+	 */
+	private ToolDefinition augmentedToolDefinition;
+
+	/**
+	 * The record class type that defines the structure of the augmented arguments.
+	 */
+	private Class<T> augmentedArgumentsClass;
+
+	/**
+	 * A consumer that processes the augmented arguments extracted from the tool input.
+	 */
+	private Consumer<AugmentedArgumentEvent<T>> augmentedArgumentsConsumer;
+
+	/**
+	 * The list of tool argument types that have been added to the tool input schema.
+	 */
+	private List<AugmentedArgumentType> augmentedArgumentTypes;
+
+	/**
+	 * A flag indicating whether to remove the augmented arguments from the tool input
+	 * after they have been processed. If the arguments are not removed, they will remain
+	 * in the tool input for the delegate to process. In many cases this could be useful.
+	 */
+	private boolean removeAugmentedArgumentsAfterProcessing = false;
+
+	public AugmentedToolCallback(ToolCallback delegate, Class<T> augmentedArgumentsClass,
+			Consumer<AugmentedArgumentEvent<T>> augmentedArgumentsConsumer,
+			boolean removeExtraArgumentsAfterProcessing) {
+		Assert.notNull(delegate, "Delegate ToolCallback must not be null");
+		Assert.notNull(augmentedArgumentsClass, "Argument types must not be null");
+		Assert.isTrue(augmentedArgumentsClass.isRecord(), "Argument types must be a Record type");
+		Assert.isTrue(augmentedArgumentsClass.getRecordComponents().length > 0,
+				"Argument types must have at least one field");
+
+		this.delegate = delegate;
+		this.augmentedArgumentTypes = ToolInputSchemaAugmenter.toAugmentedArgumentTypes(augmentedArgumentsClass);
+		String originalSchema = this.delegate.getToolDefinition().inputSchema();
+		String augmentedSchema = ToolInputSchemaAugmenter.augmentToolInputSchema(originalSchema,
+				this.augmentedArgumentTypes);
+		this.augmentedToolDefinition = ToolDefinition.builder()
+			.name(this.delegate.getToolDefinition().name())
+			.description(this.delegate.getToolDefinition().description())
+			.inputSchema(augmentedSchema)
+			.build();
+
+		this.augmentedArgumentsClass = augmentedArgumentsClass;
+		this.augmentedArgumentsConsumer = augmentedArgumentsConsumer;
+		this.removeAugmentedArgumentsAfterProcessing = removeExtraArgumentsAfterProcessing;
+	}
+
+	@Override
+	public ToolDefinition getToolDefinition() {
+		return this.augmentedToolDefinition;
+	}
+
+	@Override
+	public String call(String toolInput) {
+		return this.delegate.call(this.handleAugmentedArguments(toolInput));
+	}
+
+	@Override
+	public String call(String toolInput, @Nullable ToolContext tooContext) {
+		return this.delegate.call(this.handleAugmentedArguments(toolInput), tooContext);
+	}
+
+	/**
+	 * Handles the augmented arguments in the tool input. It extracts the augmented
+	 * arguments from the tool input, processes them using the provided consumer, and
+	 * optionally removes them from the tool input.
+	 * @param toolInput the input as received from the LLM.
+	 * @return the input to send to the delegate ToolCallback
+	 */
+	private String handleAugmentedArguments(String toolInput) {
+
+		// Extract the augmented arguments from the toolInput and send them to the
+		// consumer if provided.
+		if (this.augmentedArgumentsConsumer != null) {
+			T augmentedArguments = JsonParser.fromJson(toolInput, this.augmentedArgumentsClass);
+			this.augmentedArgumentsConsumer
+				.accept(new AugmentedArgumentEvent<>(this.augmentedToolDefinition, toolInput, augmentedArguments));
+		}
+
+		// Optionally remove the extra arguments from the toolInput
+		if (this.removeAugmentedArgumentsAfterProcessing) {
+			var args = JsonParser.fromJson(toolInput, new TypeReference<Map<String, Object>>() {
+			});
+
+			for (AugmentedArgumentType newFieldType : this.augmentedArgumentTypes) {
+				args.remove(newFieldType.name());
+			}
+			toolInput = JsonParser.toJson(args);
+		}
+
+		return toolInput;
+	}
+
+}