Add AITool.GetService

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/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.OpenAI/OpenAIResponsesChatClient.cs

@@ -910,5 +910,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/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;
                 }
             }
@@ -1455,7 +1455,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)
     {
@@ -1536,7 +1536,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)

test/Libraries/Microsoft.Extensions.AI.Abstractions.Tests/Tools/AIToolTests.cs

@@ -1,6 +1,7 @@
 // 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 Xunit;
 
 namespace Microsoft.Extensions.AI;
@@ -17,5 +18,25 @@ public void Constructor_Roundtrips()
         Assert.Empty(tool.AdditionalProperties);
     }
 
+    [Fact]
+    public void GetService_ReturnsExpectedObject()
+    {
+        DerivedAITool tool = new();
+
+        Assert.Throws<ArgumentNullException>("serviceType", () => tool.GetService(null!));
+
+        Assert.Same(tool, tool.GetService(typeof(object)));
+        Assert.Same(tool, tool.GetService(typeof(AITool)));
+        Assert.Same(tool, tool.GetService(typeof(DerivedAITool)));
+
+        Assert.Same(tool, tool.GetService<object>());
+        Assert.Same(tool, tool.GetService<AITool>());
+        Assert.Same(tool, tool.GetService<DerivedAITool>());
+
+        Assert.Null(tool.GetService<object>("key"));
+        Assert.Null(tool.GetService<AITool>("key"));
+        Assert.Null(tool.GetService<DerivedAITool>("key"));
+    }
+
     private sealed class DerivedAITool : AITool;
 }

test/Libraries/Microsoft.Extensions.AI.OpenAI.Tests/OpenAIConversionTests.cs

@@ -1193,12 +1193,17 @@ public void ListAddResponseTool_AddsToolCorrectly()
         Assert.Single(options.Tools);
         Assert.NotNull(options.Tools[0]);
 
+        var rawSearchTool = ResponseTool.CreateWebSearchTool();
         options = new()
         {
-            Tools = [ResponseTool.CreateWebSearchTool().AsAITool()],
+            Tools = [rawSearchTool.AsAITool()],
         };
         Assert.Single(options.Tools);
         Assert.NotNull(options.Tools[0]);
+
+        Assert.Same(rawSearchTool, options.Tools[0].GetService<ResponseTool>());
+        Assert.Same(rawSearchTool, options.Tools[0].GetService<WebSearchTool>());
+        Assert.Null(options.Tools[0].GetService<ResponseTool>("key"));
     }
 
     private static async IAsyncEnumerable<T> CreateAsyncEnumerable<T>(IEnumerable<T> source)