Add hosted search tool support

kzu · kzu · commit e41e0b85089d · 2025-12-05T12:43:16.000-03:00
diff --git a/src/Extensions.Grok/GrokChatClient.cs b/src/Extensions.Grok/GrokChatClient.cs
@@ -1,4 +1,5 @@
 using System.Text.Json;
+using System.Linq;
 using System.ClientModel.Primitives;
 
 using Microsoft.Extensions.AI;
@@ -30,15 +31,19 @@ public async Task<ChatResponse> GetResponseAsync(IEnumerable<ChatMessage> messag
         
         var protoResponse = await client.GetCompletionAsync(requestDto, cancellationToken: cancellationToken);
 
-        // Assuming single choice (n=1)
-        var output = protoResponse.Outputs[0];
-        
-        return new ChatResponse(new[] { MapToChatMessage(output.Message) })
+        var chatMessages = protoResponse.Outputs
+            .Select(x => MapToChatMessage(x.Message))
+            .Where(x => x.Contents.Count > 0)
+            .ToList();
+
+        var lastOutput = protoResponse.Outputs.LastOrDefault();
+
+        return new ChatResponse(chatMessages)
         {
             ResponseId = protoResponse.Id,
             ModelId = protoResponse.Model,
             CreatedAt = protoResponse.Created.ToDateTimeOffset(),
-            FinishReason = MapFinishReason(output.FinishReason),
+            FinishReason = lastOutput != null ? MapFinishReason(lastOutput.FinishReason) : null,
             Usage = MapToUsage(protoResponse.Usage),
         };
     }
@@ -110,6 +115,40 @@ GetCompletionsRequest MapToRequest(IEnumerable<ChatMessage> messages, ChatOption
                     };
                     request.Tools.Add(new Tool { Function = function });
                 }
+                else if (tool is HostedWebSearchTool webSearchTool)
+                {
+                    if (webSearchTool is GrokXSearchTool xSearch)
+                    {
+                        var toolProto = new XSearch
+                        {
+                            EnableImageUnderstanding = xSearch.EnableImageUnderstanding,
+                            EnableVideoUnderstanding = xSearch.EnableVideoUnderstanding,
+                        };
+
+                        if (xSearch.AllowedHandles is { } allowed) toolProto.AllowedXHandles.AddRange(allowed);
+                        if (xSearch.ExcludedHandles is { } excluded) toolProto.ExcludedXHandles.AddRange(excluded);
+                        if (xSearch.FromDate is { } from) toolProto.FromDate = Google.Protobuf.WellKnownTypes.Timestamp.FromDateTimeOffset(from.ToDateTime(TimeOnly.MinValue, DateTimeKind.Utc));
+                        if (xSearch.ToDate is { } to) toolProto.ToDate = Google.Protobuf.WellKnownTypes.Timestamp.FromDateTimeOffset(to.ToDateTime(TimeOnly.MinValue, DateTimeKind.Utc));
+
+                        request.Tools.Add(new Tool { XSearch = toolProto });
+                    }
+                    else if (webSearchTool is GrokSearchTool grokSearch)
+                    {
+                        var toolProto = new WebSearch
+                        {
+                            EnableImageUnderstanding = grokSearch.EnableImageUnderstanding,
+                        };
+
+                        if (grokSearch.AllowedDomains is { } allowed) toolProto.AllowedDomains.AddRange(allowed);
+                        if (grokSearch.ExcludedDomains is { } excluded) toolProto.ExcludedDomains.AddRange(excluded);
+
+                        request.Tools.Add(new Tool { WebSearch = toolProto });
+                    }
+                    else
+                    {
+                        request.Tools.Add(new Tool { WebSearch = new WebSearch() });
+                    }
+                }
             }
         }
 
@@ -120,10 +159,6 @@ GetCompletionsRequest MapToRequest(IEnumerable<ChatMessage> messages, ChatOption
                 FormatType = FormatType.JsonObject
             };
         }
-        else
-        {
-            request.ResponseFormat = new ResponseFormat { FormatType = FormatType.Text };
-        }
 
         return request;
     }
@@ -148,11 +183,18 @@ GetCompletionsRequest MapToRequest(IEnumerable<ChatMessage> messages, ChatOption
 
     static ChatMessage MapToChatMessage(CompletionMessage message)
     {
-        var chatMessage = new ChatMessage(ChatRole.Assistant, message.Content);
+        var chatMessage = new ChatMessage() { Role = MapRole(message.Role) };
+
+        if (!string.IsNullOrEmpty(message.Content))
+        {
+            chatMessage.Contents.Add(new TextContent(message.Content) { Annotations = [] });
+        }
 
         foreach (var toolCall in message.ToolCalls)
         {
-            if (toolCall.ToolCase == XaiApi.ToolCall.ToolOneofCase.Function)
+            // Only include client-side tools in the response messages
+            if (toolCall.ToolCase == XaiApi.ToolCall.ToolOneofCase.Function && 
+                toolCall.Type == ToolCallType.ClientSideTool)
             {
                 var arguments = !string.IsNullOrEmpty(toolCall.Function.Arguments)
                     ? JsonSerializer.Deserialize<IDictionary<string, object?>>(toolCall.Function.Arguments)
diff --git a/src/Extensions.Grok/GrokClient.cs b/src/Extensions.Grok/GrokClient.cs
@@ -1,33 +1,22 @@
-﻿using System.ClientModel;
-using System.ClientModel.Primitives;
-using System.Collections.Concurrent;
-using System.Net.Http;
+﻿using System.Collections.Concurrent;
 using System.Net.Http.Headers;
-using Grpc.Core;
 using Grpc.Net.Client;
 
 namespace Devlooped.Extensions.AI.Grok;
 
-public class GrokClient
+public class GrokClient(string apiKey, GrokClientOptions options)
 {
     static ConcurrentDictionary<(Uri, string), GrpcChannel> channels = [];
 
     public GrokClient(string apiKey) : this(apiKey, new GrokClientOptions()) { }
 
-    public GrokClient(string apiKey, GrokClientOptions options)
-    {
-        ApiKey = apiKey;
-        Options = options;
-        Endpoint = options.Endpoint;
-    }
-
-    public string ApiKey { get; }
+    public string ApiKey { get; } = apiKey;
 
     /// <summary>Gets or sets the endpoint for the service.</summary>
-    public Uri Endpoint { get; set; }
+    public Uri Endpoint { get; set; } = options.Endpoint;
 
     /// <summary>Gets the options used to configure the client.</summary>
-    public GrokClientOptions Options { get; }
+    public GrokClientOptions Options { get; } = options;
 
     internal GrpcChannel Channel => channels.GetOrAdd((Endpoint, ApiKey), key =>
     {
diff --git a/src/Extensions/Grok/GrokChatClient2.cs b/src/Extensions/Grok/GrokChatClient2.cs
diff --git a/src/Extensions/Grok/GrokClient2.cs b/src/Extensions/Grok/GrokClient2.cs
diff --git a/src/Extensions/Grok/GrokSearchTool.cs b/src/Extensions/Grok/GrokSearchTool.cs
@@ -1,47 +1,102 @@
-﻿using System.Text.Json.Serialization;
+﻿using System.ComponentModel;
+using System.Text.Json.Serialization;
 using Microsoft.Extensions.AI;
 
 namespace Devlooped.Extensions.AI.Grok;
 
 /// <summary>
-/// Enables or disables Grok's live search capabilities.
-/// See https://docs.x.ai/docs/guides/live-search#enabling-search
+/// Enables or disables Grok's search tool.
+/// See https://docs.x.ai/docs/guides/tools/search-tools
 /// </summary>
+[EditorBrowsable(EditorBrowsableState.Never)]
+[Obsolete("No longer supported in agent search too. Use either GrokSearchTool or GrokXSearchTool")]
 public enum GrokSearch
 {
     /// <summary>
-    /// (default): Live search is available to the model, but the model automatically decides whether to perform live search.
+    /// (default): Search tool is available to the model, but the model automatically decides whether to perform search.
     /// </summary>
     Auto,
     /// <summary>
-    /// Enables live search.
+    /// Enables search tool.
     /// </summary>
     On,
     /// <summary>
-    /// Disables search and uses the model without accessing additional information from data sources.
+    /// Disables search tool and uses the model without accessing additional information from data sources.
     /// </summary>
     Off
 }
 
-/// <summary>Configures Grok's live search capabilities. See https://docs.x.ai/docs/guides/live-search.</summary>
-public class GrokSearchTool(GrokSearch mode) : HostedWebSearchTool
+/// <summary>
+/// Configures Grok's agentic search tool. 
+/// See https://docs.x.ai/docs/guides/tools/search-tools.
+/// </summary>
+public class GrokSearchTool : HostedWebSearchTool
 {
-    /// <summary>Sets the search mode for Grok's live search capabilities.</summary>
-    public GrokSearch Mode { get; } = mode;
+    public GrokSearchTool() { }
+
     /// <inheritdoc/>
-    public override string Name => "Live Search";
+    public override string Name => "Search";
     /// <inheritdoc/>
-    public override string Description => "Performs live search using X.AI";
+    public override string Description => "Performs search using X.AI";
+
+    /// <summary>Use to make the web search only perform the search and web browsing on web pages that fall within the specified domains. Can include a maximum of five domains.</summary>
+    public IList<string>? AllowedDomains { get; set; }
+    /// <summary>Use to prevent the model from including the specified domains in any web search tool invocations and from browsing any pages on those domains. Can include a maximum of five domains.</summary>
+    public IList<string>? ExcludedDomains { get; set; }
+    /// <summary>See https://docs.x.ai/docs/guides/tools/search-tools#enable-image-understanding</summary>
+    public bool EnableImageUnderstanding { get; set; }
+
+    #region Legacy Live Search
+    [EditorBrowsable(EditorBrowsableState.Never)]
+    [Obsolete("Legacy live search mode is not available in new agentic search tool.")]
+    public GrokSearchTool(GrokSearch mode) => Mode = mode;
+    /// <summary>Sets the search mode for Grok's live search capabilities.</summary>
+    [EditorBrowsable(EditorBrowsableState.Never)]
+    [Obsolete("Legacy live search mode is not available in new agentic search tool.")]
+    public GrokSearch Mode { get; set; }
     /// <summary>See https://docs.x.ai/docs/guides/live-search#set-date-range-of-the-search-data</summary>
+    [EditorBrowsable(EditorBrowsableState.Never)]
+    [Obsolete("Date range can only be applied to X source.")]
     public DateOnly? FromDate { get; set; }
     /// <summary>See https://docs.x.ai/docs/guides/live-search#set-date-range-of-the-search-data</summary>
+    [EditorBrowsable(EditorBrowsableState.Never)]
+    [Obsolete("Date range can only be applied to X search tool.")]
     public DateOnly? ToDate { get; set; }
     /// <summary>See https://docs.x.ai/docs/guides/live-search#limit-the-maximum-amount-of-data-sources</summary>
+    [EditorBrowsable(EditorBrowsableState.Never)]
+    [Obsolete("No longer supported in search tool.")]
     public int? MaxSearchResults { get; set; }
     /// <summary>See https://docs.x.ai/docs/guides/live-search#data-sources-and-parameters</summary>
+    [EditorBrowsable(EditorBrowsableState.Never)]
+    [Obsolete("No longer supported in search tool.")]
     public IList<GrokSource>? Sources { get; set; }
     /// <summary>See https://docs.x.ai/docs/guides/live-search#returning-citations</summary>
+    [EditorBrowsable(EditorBrowsableState.Never)]
+    [Obsolete("No longer supported in search tool.")]
     public bool? ReturnCitations { get; set; }
+    #endregion
+}
+
+/// <summary>
+/// Configures Grok's agentic search tool for X. 
+/// See https://docs.x.ai/docs/guides/tools/search-tools.
+/// </summary>
+public class GrokXSearchTool : HostedWebSearchTool
+{
+    /// <summary>See https://docs.x.ai/docs/guides/tools/search-tools#only-consider-x-posts-from-specific-handles</summary>
+    [JsonPropertyName("allowed_x_handles")]
+    public IList<string>? AllowedHandles { get; set; }
+    /// <summary>See https://docs.x.ai/docs/guides/tools/search-tools#exclude-x-posts-from-specific-handles</summary>
+    [JsonPropertyName("excluded_x_handles")]
+    public IList<string>? ExcludedHandles { get; set; }
+    /// <summary>See https://docs.x.ai/docs/guides/tools/search-tools#date-range</summary>
+    public DateOnly? FromDate { get; set; }
+    /// <summary>See https://docs.x.ai/docs/guides/tools/search-tools#date-range</summary>
+    public DateOnly? ToDate { get; set; }
+    /// <summary>See https://docs.x.ai/docs/guides/tools/search-tools#enable-image-understanding-1</summary>
+    public bool EnableImageUnderstanding { get; set; }
+    /// <summary>See https://docs.x.ai/docs/guides/tools/search-tools#enable-video-understanding</summary>
+    public bool EnableVideoUnderstanding { get; set; }
 }
 
 /// <summary>Grok Live Search data source base type.</summary>
@@ -50,9 +105,13 @@ public class GrokSearchTool(GrokSearch mode) : HostedWebSearchTool
 [JsonDerivedType(typeof(GrokNewsSource), "news")]
 [JsonDerivedType(typeof(GrokRssSource), "rss")]
 [JsonDerivedType(typeof(GrokXSource), "x")]
+[EditorBrowsable(EditorBrowsableState.Never)]
+[Obsolete("No longer supported in agent search too. Use either GrokSearchTool or GrokXSearchTool")]
 public abstract class GrokSource { }
 
 /// <summary>Search-based data source base class providing common properties for `web` and `news` sources.</summary>
+[EditorBrowsable(EditorBrowsableState.Never)]
+[Obsolete("No longer supported in agent search too. Use either GrokSearchTool or GrokXSearchTool")]
 public abstract class GrokSearchSource : GrokSource
 {
     /// <summary>Include data from a specific country/region by specifying the ISO alpha-2 code of the country.</summary>
@@ -64,24 +123,32 @@ public abstract class GrokSearchSource : GrokSource
 }
 
 /// <summary>Web live search source.</summary>
+[EditorBrowsable(EditorBrowsableState.Never)]
+[Obsolete("No longer supported in agent search too. Use either GrokSearchTool or GrokXSearchTool")]
 public class GrokWebSource : GrokSearchSource
 {
     /// <summary>See https://docs.x.ai/docs/guides/live-search#parameter-allowed_websites-supported-by-web</summary>
     public IList<string>? AllowedWebsites { get; set; }
 }
 
 /// <summary>News live search source.</summary>
+[EditorBrowsable(EditorBrowsableState.Never)]
+[Obsolete("No longer supported in agent search too. Use either GrokSearchTool or GrokXSearchTool")]
 public class GrokNewsSource : GrokSearchSource { }
 
 /// <summary>RSS live search source.</summary>
 /// <param name="rss">The RSS feed to search.</param>
+[EditorBrowsable(EditorBrowsableState.Never)]
+[Obsolete("No longer supported in agent search too. Use either GrokSearchTool or GrokXSearchTool")]
 public class GrokRssSource(string rss) : GrokSource
 {
     /// <summary>See https://docs.x.ai/docs/guides/live-search#parameter-link-supported-by-rss</summary>
     public IList<string>? Links { get; set; } = [rss];
 }
 
 /// <summary>X live search source./summary>
+[EditorBrowsable(EditorBrowsableState.Never)]
+[Obsolete("No longer supported in agent search too. Use GrokXSearchTool")]
 public class GrokXSource : GrokSearchSource
 {
     /// <summary>See https://docs.x.ai/docs/guides/live-search#parameter-excluded_x_handles-supported-by-x</summary>
@@ -96,4 +163,9 @@ public class GrokXSource : GrokSearchSource
     /// <summary>See https://docs.x.ai/docs/guides/live-search#parameters-post_favorite_count-and-post_view_count-supported-by-x</summary>
     [JsonPropertyName("post_view_count")]
     public int? ViewCount { get; set; }
+    [JsonPropertyName("from_date")]
+    public DateOnly? FromDate { get; set; }
+    /// <summary>See https://docs.x.ai/docs/guides/live-search#set-date-range-of-the-search-data</summary>
+    [JsonPropertyName("to_date")]
+    public DateOnly? ToDate { get; set; }
 }
diff --git a/src/Tests/GrokTests.cs b/src/Tests/GrokTests.cs
@@ -302,4 +302,52 @@ public async Task CanAvoidCitations()
         Assert.NotNull(node);
         Assert.Null(node["citations"]);
     }
+
+    [SecretsFact("XAI_API_KEY")]
+    public async Task GrokGrpcInvokesHostedSearchTool()
+    {
+        var messages = new Chat()
+        {
+            { "system", "You are an AI assistant that knows how to search the web." },
+            { "user", "What's Tesla stock worth today? Search X and the news for latest info." },
+        };
+
+        var grok = new GrokClient(Configuration["XAI_API_KEY"]!).AsIChatClient("grok-4-fast");
+
+        var options = new ChatOptions
+        {
+            Tools = [new HostedWebSearchTool()]
+        };
+
+        var response = await grok.GetResponseAsync(messages, options);
+        var text = response.Text;
+
+        Assert.Contains("TSLA", text);
+        Assert.NotNull(response.ModelId);
+    }
+
+    [SecretsFact("XAI_API_KEY")]
+    public async Task GrokGrpcInvokesGrokSearchTool()
+    {
+        var messages = new Chat()
+        {
+            { "system", "You are an AI assistant that knows how to search the web." },
+            { "user", "What is the latest news about Microsoft?" },
+        };
+
+        var grok = new GrokClient(Configuration["XAI_API_KEY"]!).AsIChatClient("grok-4-fast");
+
+        var options = new ChatOptions
+        {
+            Tools = [new GrokSearchTool 
+            { 
+                AllowedDomains = ["microsoft.com", "news.microsoft.com"] 
+            }]
+        };
+
+        var response = await grok.GetResponseAsync(messages, options);
+        
+        Assert.NotNull(response.Text);
+        Assert.Contains("Microsoft", response.Text);
+    }
 }
