Add allow for citations

Author: kzu

src/Extensions.Grok/GrokChatClient.cs

@@ -32,7 +32,7 @@ public async Task<ChatResponse> GetResponseAsync(IEnumerable<ChatMessage> messag
         var protoResponse = await client.GetCompletionAsync(requestDto, cancellationToken: cancellationToken);
 
         var chatMessages = protoResponse.Outputs
-            .Select(x => MapToChatMessage(x.Message))
+            .Select(x => MapToChatMessage(x.Message, protoResponse.Citations))
             .Where(x => x.Contents.Count > 0)
             .ToList();
 
@@ -63,7 +63,7 @@ async IAsyncEnumerable<ChatResponseUpdate> CompleteChatStreamingCore(IEnumerable
                 var outputChunk = chunk.Outputs[0];
 
                 // Use positional arguments for ChatResponseUpdate
-                yield return new ChatResponseUpdate(
+                var update = new ChatResponseUpdate(
                     outputChunk.Delta.Role != MessageRole.InvalidRole ? MapRole(outputChunk.Delta.Role) : null,
                     outputChunk.Delta.Content
                 )
@@ -73,6 +73,23 @@ async IAsyncEnumerable<ChatResponseUpdate> CompleteChatStreamingCore(IEnumerable
                     CreatedAt = chunk.Created.ToDateTimeOffset(),
                     FinishReason = outputChunk.FinishReason != FinishReason.ReasonInvalid ? MapFinishReason(outputChunk.FinishReason) : null,
                 };
+
+                if (chunk.Citations is { Count: > 0 } citations)
+                {
+                    var textContent = update.Contents.OfType<TextContent>().FirstOrDefault();
+                    if (textContent == null)
+                    {
+                        textContent = new TextContent(string.Empty);
+                        update.Contents.Add(textContent);
+                    }
+                    
+                    foreach (var citation in citations.Distinct())
+                    {
+                        (textContent.Annotations ??= []).Add(new CitationAnnotation { Url = new(citation) });
+                    }
+                }
+
+                yield return update;
             }
         }
     }
@@ -181,13 +198,21 @@ GetCompletionsRequest MapToRequest(IEnumerable<ChatMessage> messages, ChatOption
         _ => ChatRole.Assistant
     };
 
-    static ChatMessage MapToChatMessage(CompletionMessage message)
+    static ChatMessage MapToChatMessage(CompletionMessage message, IList<string>? citations = null)
     {
         var chatMessage = new ChatMessage() { Role = MapRole(message.Role) };
 
-        if (!string.IsNullOrEmpty(message.Content))
+        if (!string.IsNullOrEmpty(message.Content) || (citations is { Count: > 0 }))
         {
-            chatMessage.Contents.Add(new TextContent(message.Content) { Annotations = [] });
+            var textContent = new TextContent(message.Content ?? string.Empty);
+            if (citations is { Count: > 0 })
+            {
+                foreach (var citation in citations.Distinct())
+                {
+                    (textContent.Annotations ??= []).Add(new CitationAnnotation { Url = new(citation) });
+                }
+            }
+            chatMessage.Contents.Add(textContent);
         }
 
         foreach (var toolCall in message.ToolCalls)

src/Tests/GrokTests.cs

@@ -324,10 +324,14 @@ public async Task GrokGrpcInvokesHostedSearchTool()
 
         Assert.Contains("TSLA", text);
         Assert.NotNull(response.ModelId);
+        Assert.Contains(new Uri("https://finance.yahoo.com/quote/TSLA/news/"), response.Messages
+            .SelectMany(x => x.Contents)
+            .SelectMany(x => x.Annotations?.OfType<CitationAnnotation>() ?? [])
+            .Select(x => x.Url));
     }
 
     [SecretsFact("XAI_API_KEY")]
-    public async Task GrokGrpcInvokesGrokSearchTool()
+    public async Task GrokGrpcInvokesGrokSearchToolIncludesDomain()
     {
         var messages = new Chat()
         {
@@ -341,13 +345,66 @@ public async Task GrokGrpcInvokesGrokSearchTool()
         {
             Tools = [new GrokSearchTool 
             { 
-                AllowedDomains = ["microsoft.com", "news.microsoft.com"] 
+                AllowedDomains = ["microsoft.com", "news.microsoft.com"],
             }]
         };
 
         var response = await grok.GetResponseAsync(messages, options);
 
         Assert.NotNull(response.Text);
         Assert.Contains("Microsoft", response.Text);
+
+        var urls = response.Messages
+            .SelectMany(x => x.Contents)
+            .SelectMany(x => x.Annotations?.OfType<CitationAnnotation>() ?? [])
+            .Where(x => x.Url is not null)
+            .Select(x => x.Url!)
+            .ToList();
+
+        foreach (var url in urls)
+        {
+            output.WriteLine(url.ToString());
+        }
+
+        Assert.All(urls, x => x.Host.EndsWith(".microsoft.com"));
+    }
+
+    [SecretsFact("XAI_API_KEY")]
+    public async Task GrokGrpcInvokesGrokSearchToolExcludesDomain()
+    {
+        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
+            {
+                ExcludedDomains = ["blogs.microsoft.com"]
+            }]
+        };
+
+        var response = await grok.GetResponseAsync(messages, options);
+
+        Assert.NotNull(response.Text);
+        Assert.Contains("Microsoft", response.Text);
+
+        var urls = response.Messages
+            .SelectMany(x => x.Contents)
+            .SelectMany(x => x.Annotations?.OfType<CitationAnnotation>() ?? [])
+            .Where(x => x.Url is not null)
+            .Select(x => x.Url!)
+            .ToList();
+
+        foreach (var url in urls)
+        {
+            output.WriteLine(url.ToString());
+        }
+
+        Assert.DoesNotContain(urls, x => x.Host == "blogs.microsoft.com");
     }
 }