Add AITool.GetService (#6830)

Following the same pattern as elsewhere in M.E.AI, this enables a consumer to reach through layers of delegating tools to grab information from inner ones, such as whether they were marked as requiring approval.

Author: stephentoub

src/Libraries/Microsoft.Extensions.AI.Abstractions/CHANGELOG.md

@@ -3,6 +3,7 @@
 ## NOT YET RELEASED
 
 - Added new `ChatResponseFormat.ForJsonSchema` overloads that export a JSON schema from a .NET type.
+- Added new `AITool.GetService` virtual method.
 - Updated `TextReasoningContent` to include `ProtectedData` for representing encrypted/redacted content.
 - Fixed `MinLength`/`MaxLength`/`Length` attribute mapping in nullable string properties during schema export.
 

src/Libraries/Microsoft.Extensions.AI.Abstractions/Functions/DelegatingAIFunction.cs

@@ -58,4 +58,14 @@ protected DelegatingAIFunction(AIFunction innerFunction)
     /// <inheritdoc />
     protected override ValueTask<object?> InvokeCoreAsync(AIFunctionArguments arguments, CancellationToken cancellationToken) =>
         InnerFunction.InvokeAsync(arguments, cancellationToken);
+
+    /// <inheritdoc />
+    public override object? GetService(Type serviceType, object? serviceKey = null)
+    {
+        _ = Throw.IfNull(serviceType);
+
+        return
+            serviceKey is null && serviceType.IsInstanceOfType(this) ? this :
+            InnerFunction.GetService(serviceType, serviceKey);
+    }
 }

src/Libraries/Microsoft.Extensions.AI.Abstractions/Functions/DelegatingAIFunctionDeclaration.cs

@@ -45,4 +45,14 @@ protected DelegatingAIFunctionDeclaration(AIFunctionDeclaration innerFunction)
 
     /// <inheritdoc />
     public override string ToString() => InnerFunction.ToString();
+
+    /// <inheritdoc />
+    public override object? GetService(Type serviceType, object? serviceKey = null)
+    {
+        _ = Throw.IfNull(serviceType);
+
+        return
+            serviceKey is null && serviceType.IsInstanceOfType(this) ? this :
+            InnerFunction.GetService(serviceType, serviceKey);
+    }
 }

src/Libraries/Microsoft.Extensions.AI.Abstractions/Microsoft.Extensions.AI.Abstractions.json

@@ -685,6 +685,14 @@
           "Member": "Microsoft.Extensions.AI.AITool.AITool();",
           "Stage": "Stable"
         },
+        {
+          "Member": "virtual object? Microsoft.Extensions.AI.AITool.GetService(System.Type serviceType, object? serviceKey = null);",
+          "Stage": "Stable"
+        },
+        {
+          "Member": "TService? Microsoft.Extensions.AI.AITool.GetService<TService>(object? serviceKey = null);",
+          "Stage": "Stable"
+        },
         {
           "Member": "override string Microsoft.Extensions.AI.AITool.ToString();",
           "Stage": "Stable"
@@ -1477,6 +1485,10 @@
           "Member": "Microsoft.Extensions.AI.DelegatingAIFunction.DelegatingAIFunction(Microsoft.Extensions.AI.AIFunction innerFunction);",
           "Stage": "Stable"
         },
+        {
+          "Member": "override object? Microsoft.Extensions.AI.DelegatingAIFunction.GetService(System.Type serviceType, object? serviceKey = null);",
+          "Stage": "Stable"
+        },
         {
           "Member": "override System.Threading.Tasks.ValueTask<object?> Microsoft.Extensions.AI.DelegatingAIFunction.InvokeCoreAsync(Microsoft.Extensions.AI.AIFunctionArguments arguments, System.Threading.CancellationToken cancellationToken);",
           "Stage": "Stable"

src/Libraries/Microsoft.Extensions.AI.Abstractions/Tools/AITool.cs

@@ -1,10 +1,12 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System;
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Text;
 using Microsoft.Shared.Collections;
+using Microsoft.Shared.Diagnostics;
 
 namespace Microsoft.Extensions.AI;
 
@@ -31,6 +33,35 @@ protected AITool()
     /// <inheritdoc/>
     public override string ToString() => Name;
 
+    /// <summary>Asks the <see cref="AITool"/> for an object of the specified type <paramref name="serviceType"/>.</summary>
+    /// <param name="serviceType">The type of object being requested.</param>
+    /// <param name="serviceKey">An optional key that can be used to help identify the target service.</param>
+    /// <returns>The found object, otherwise <see langword="null"/>.</returns>
+    /// <exception cref="ArgumentNullException"><paramref name="serviceType"/> is <see langword="null"/>.</exception>
+    /// <remarks>
+    /// The purpose of this method is to allow for the retrieval of strongly-typed services that might be provided by the <see cref="AITool"/>,
+    /// including itself or any services it might be wrapping.
+    /// </remarks>
+    public virtual object? GetService(Type serviceType, object? serviceKey = null)
+    {
+        _ = Throw.IfNull(serviceType);
+
+        return
+            serviceKey is null && serviceType.IsInstanceOfType(this) ? this :
+            null;
+    }
+
+    /// <summary>Asks the <see cref="AITool"/> for an object of type <typeparamref name="TService"/>.</summary>
+    /// <typeparam name="TService">The type of the object to be retrieved.</typeparam>
+    /// <param name="serviceKey">An optional key that can be used to help identify the target service.</param>
+    /// <returns>The found object, otherwise <see langword="null"/>.</returns>
+    /// <remarks>
+    /// The purpose of this method is to allow for the retrieval of strongly typed services that may be provided by the <see cref="AITool"/>,
+    /// including itself or any services it might be wrapping.
+    /// </remarks>
+    public TService? GetService<TService>(object? serviceKey = null) =>
+        GetService(typeof(TService), serviceKey) is TService service ? service : default;
+
     /// <summary>Gets the string to display in the debugger for this instance.</summary>
     [DebuggerBrowsable(DebuggerBrowsableState.Never)]
     private string DebuggerDisplay

src/Libraries/Microsoft.Extensions.AI.AzureAIInference/CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## NOT YET RELEASED
 
+- Updated to accommodate the additions in `Microsoft.Extensions.AI.Abstractions`.
+
 ## 9.9.0-preview.1.25458.4
 
 - Updated tool mapping to recognize any `AIFunctionDeclaration`.

src/Libraries/Microsoft.Extensions.AI.OpenAI/OpenAIResponsesChatClient.cs

@@ -897,5 +897,15 @@ internal sealed class ResponseToolAITool(ResponseTool tool) : AITool
     {
         public ResponseTool Tool => tool;
         public override string Name => Tool.GetType().Name;
+
+        /// <inheritdoc />
+        public override object? GetService(Type serviceType, object? serviceKey = null)
+        {
+            _ = Throw.IfNull(serviceType);
+
+            return
+                serviceKey is null && serviceType.IsInstanceOfType(Tool) ? Tool :
+                base.GetService(serviceType, serviceKey);
+        }
     }
 }

src/Libraries/Microsoft.Extensions.AI/CHANGELOG.md

@@ -2,8 +2,9 @@
 
 ## NOT YET RELEASED
 
-- Updated the EnableSensitiveData properties on OpenTelemetryChatClient/EmbeddingGenerator to respect a OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT environment variable.
-- Added OpenTelemetryImageGenerator to provide OpenTelemetry instrumentation for IImageGenerator implementations.
+- Updated the `EnableSensitiveData` properties on `OpenTelemetryChatClient/EmbeddingGenerator` to respect a `OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT` environment variable.
+- Updated `OpenTelemetryChatClient/EmbeddingGenerator` to emit recent additions to the OpenTelemetry Semantic Conventions for Generative AI systems.
+- Added `OpenTelemetryImageGenerator` to provide OpenTelemetry instrumentation for `IImageGenerator` implementations.
 
 ## 9.9.0
 

src/Libraries/Microsoft.Extensions.AI/ChatCompletion/FunctionInvokingChatClient.cs

@@ -430,7 +430,7 @@ public override async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseA
         List<ChatMessage> originalMessages = [.. messages];
         messages = originalMessages;
 
-        ApprovalRequiredAIFunction[]? approvalRequiredFunctions = null; // available tools that require approval
+        AITool[]? approvalRequiredFunctions = null; // available tools that require approval
         List<ChatMessage>? augmentedHistory = null; // the actual history of messages sent on turns other than the first
         List<FunctionCallContent>? functionCallContents = null; // function call contents that need responding to in the current turn
         List<ChatMessage>? responseMessages = null; // tracked list of messages, across multiple turns, to be used in fallback cases to reconstitute history
@@ -539,7 +539,7 @@ public override async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseA
                     approvalRequiredFunctions =
                         (options?.Tools ?? Enumerable.Empty<AITool>())
                         .Concat(AdditionalTools ?? Enumerable.Empty<AITool>())
-                        .OfType<ApprovalRequiredAIFunction>()
+                        .Where(t => t.GetService<ApprovalRequiredAIFunction>() is not null)
                         .ToArray();
                 }
 
@@ -741,7 +741,7 @@ private static (Dictionary<string, AITool>? ToolMap, bool AnyRequireApproval) Cr
                 for (int i = 0; i < count; i++)
                 {
                     AITool tool = toolList[i];
-                    anyRequireApproval |= tool is ApprovalRequiredAIFunction;
+                    anyRequireApproval |= tool.GetService<ApprovalRequiredAIFunction>() is not null;
                     map[tool.Name] = tool;
                 }
             }
@@ -1475,7 +1475,7 @@ private static ChatMessage ConvertToFunctionCallContentMessage(ApprovalResultWit
     /// </summary>
     private static (bool hasApprovalRequiringFcc, int lastApprovalCheckedFCCIndex) CheckForApprovalRequiringFCC(
         List<FunctionCallContent>? functionCallContents,
-        ApprovalRequiredAIFunction[] approvalRequiredFunctions,
+        AITool[] approvalRequiredFunctions,
         bool hasApprovalRequiringFcc,
         int lastApprovalCheckedFCCIndex)
     {
@@ -1556,7 +1556,7 @@ private static IList<ChatMessage> ReplaceFunctionCallsWithApprovalRequests(
                     {
                         foreach (var t in toolMap)
                         {
-                            if (t.Value is ApprovalRequiredAIFunction araf && araf.Name == functionCall.Name)
+                            if (t.Value.GetService<ApprovalRequiredAIFunction>() is { } araf && araf.Name == functionCall.Name)
                             {
                                 anyApprovalRequired = true;
                                 break;

test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/Functions/DelegatingAIFunctionTests.cs

@@ -20,7 +20,7 @@ public void Constructor_NullInnerFunction_ThrowsArgumentNullException()
     [Fact]
     public void DefaultOverrides_DelegateToInnerFunction()
     {
-        AIFunction expected = AIFunctionFactory.Create(() => 42);
+        AIFunction expected = new ApprovalRequiredAIFunction(AIFunctionFactory.Create(() => 42));
         DerivedFunction actual = new(expected);
 
         Assert.Same(expected, actual.InnerFunction);
@@ -32,6 +32,7 @@ public void DefaultOverrides_DelegateToInnerFunction()
         Assert.Same(expected.UnderlyingMethod, actual.UnderlyingMethod);
         Assert.Same(expected.AdditionalProperties, actual.AdditionalProperties);
         Assert.Equal(expected.ToString(), actual.ToString());
+        Assert.Same(expected, actual.GetService<ApprovalRequiredAIFunction>());
     }
 
     private sealed class DerivedFunction(AIFunction innerFunction) : DelegatingAIFunction(innerFunction)