Skip to content

Commit 92eac7d

Browse files
kzukzu
committed
Add support for collection search results citations
1 parent 6722ced commit 92eac7d

File tree

6 files changed

+363
-90
lines changed

6 files changed

+363
-90
lines changed

‎readme.md‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -158,6 +158,50 @@ var options = new ChatOptions
158158
};
159159
```
160160

161+
To receive the actual search results and file references, include `CollectionsSearchCallOutput` in the options:
162+
163+
```csharp
164+
var options = new GrokChatOptions
165+
{
166+
Include = [IncludeOption.CollectionsSearchCallOutput],
167+
Tools = [new HostedFileSearchTool {
168+
Inputs = [new HostedVectorStoreContent("[collection_id]")]
169+
}]
170+
};
171+
172+
var response = await grok.GetResponseAsync(messages, options);
173+
174+
// Access the search results with file references
175+
var results = response.Messages
176+
.SelectMany(x => x.Contents)
177+
.OfType<CollectionSearchToolResultContent>();
178+
179+
foreach (var result in results)
180+
{
181+
// Each result contains files that were found and referenced
182+
var files = result.Outputs?.OfType<HostedFileContent>();
183+
foreach (var file in files ?? [])
184+
{
185+
Console.WriteLine($"File: {file.Name} (ID: {file.FileId})");
186+
187+
// Files include citation annotations with snippets
188+
foreach (var citation in file.Annotations?.OfType<CitationAnnotation>() ?? [])
189+
{
190+
Console.WriteLine($" Title: {citation.Title}");
191+
Console.WriteLine($" Snippet: {citation.Snippet}");
192+
Console.WriteLine($" URL: {citation.Url}"); // collections://[collection_id]/files/[file_id]
193+
}
194+
}
195+
}
196+
```
197+
198+
Citations from collection search include:
199+
- **Title**: Extracted from the first line of the chunk content (if available), typically the file name or heading
200+
- **Snippet**: The relevant text excerpt from the document
201+
- **FileId**: Identifier of the source file in the collection
202+
- **Url**: A `collections://` URI pointing to the specific file within the collection
203+
- **ToolName**: Always set to `"collections_search"`
204+
161205
Learn more about [collection search](https://docs.x.ai/docs/guides/tools/collections-search-tool).
162206

163207
## Remote MCP

‎src/xAI.Tests/ChatClientTests.cs‎

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,7 @@ public async Task GrokInvokesToolAndSearch()
9696
Assert.Equal(options.ModelId, response.ModelId);
9797

9898
var calls = response.Messages
99-
.SelectMany(x => x.Contents.OfType<HostedToolCallContent>())
99+
.SelectMany(x => x.Contents.OfType<CollectionSearchToolCallContent>())
100100
.Select(x => x.RawRepresentation as xAI.Protocol.ToolCall)
101101
.Where(x => x is not null)
102102
.ToList();
@@ -317,7 +317,6 @@ public async Task GrokInvokesHostedCollectionSearch()
317317

318318
var options = new GrokChatOptions
319319
{
320-
Include = { IncludeOption.CollectionsSearchCallOutput },
321320
Tools = [new HostedFileSearchTool {
322321
Inputs = [new HostedVectorStoreContent("collection_91559d9b-a55d-42fe-b2ad-ecf8904d9049")]
323322
}]
@@ -329,9 +328,21 @@ public async Task GrokInvokesHostedCollectionSearch()
329328
Assert.Contains("11,74", text);
330329
Assert.Contains(response.Messages
331330
.SelectMany(x => x.Contents)
332-
.OfType<HostedToolCallContent>()
331+
.OfType<CollectionSearchToolCallContent>()
333332
.Select(x => x.RawRepresentation as xAI.Protocol.ToolCall),
334333
x => x?.Type == xAI.Protocol.ToolCallType.CollectionsSearchTool);
334+
// No actual search results content since we didn't specify it in Include
335+
Assert.Empty(response.Messages.SelectMany(x => x.Contents).OfType<CollectionSearchToolResultContent>());
336+
337+
options.Include = [IncludeOption.CollectionsSearchCallOutput];
338+
response = await grok.GetResponseAsync(messages, options);
339+
340+
// Now it also contains the file reference as result content
341+
Assert.Contains(response.Messages
342+
.SelectMany(x => x.Contents)
343+
.OfType<CollectionSearchToolResultContent>()
344+
.SelectMany(x => (x.Outputs ?? []).OfType<HostedFileContent>()),
345+
x => x.Name == "LNS0004592.txt");
335346
}
336347

337348
[SecretsFact("XAI_API_KEY", "GITHUB_TOKEN")]
@@ -459,7 +470,7 @@ public async Task GrokStreamsUpdatesFromAllTools()
459470

460471
Assert.Contains(response.Messages
461472
.SelectMany(x => x.Contents)
462-
.OfType<HostedToolCallContent>()
473+
.OfType<CollectionSearchToolCallContent>()
463474
.Select(x => x.RawRepresentation as xAI.Protocol.ToolCall),
464475
x => x?.Type == xAI.Protocol.ToolCallType.WebSearchTool);
465476

src/xAI/HostedToolCallContent.cs renamed to src/xAI/CollectionSearchToolCallContent.cs

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,9 @@
1-
using System.Diagnostics.CodeAnalysis;
2-
using Microsoft.Extensions.AI;
1+
using Microsoft.Extensions.AI;
32

43
namespace xAI;
54

65
/// <summary>Represents a hosted tool agentic call.</summary>
7-
[Experimental("xAI001")]
8-
public class HostedToolCallContent : AIContent
6+
public class CollectionSearchToolCallContent : AIContent
97
{
108
/// <summary>Gets or sets the tool call ID.</summary>
119
public virtual string? CallId { get; set; }

src/xAI/HostedToolResultContent.cs renamed to src/xAI/CollectionSearchToolResultContent.cs

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,9 @@
1-
using System.Diagnostics;
2-
using System.Diagnostics.CodeAnalysis;
3-
using Microsoft.Extensions.AI;
1+
using Microsoft.Extensions.AI;
42

53
namespace xAI;
64

75
/// <summary>Represents a hosted tool agentic call.</summary>
8-
[DebuggerDisplay("{DebuggerDisplay,nq}")]
9-
[Experimental("xAI001")]
10-
public class HostedToolResultContent : AIContent
6+
public class CollectionSearchToolResultContent : AIContent
117
{
128
/// <summary>Gets or sets the tool call ID.</summary>
139
public virtual string? CallId { get; set; }

‎src/xAI/GrokChatClient.cs‎

Lines changed: 85 additions & 75 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
using System.Text.Json;
2+
using Google.Protobuf;
23
using Grpc.Core;
34
using Grpc.Net.Client;
45
using Microsoft.Extensions.AI;
@@ -39,79 +40,86 @@ public async Task<ChatResponse> GetResponseAsync(IEnumerable<ChatMessage> messag
3940
var response = await client.GetCompletionAsync(request, cancellationToken: cancellationToken);
4041
var lastOutput = response.Outputs.OrderByDescending(x => x.Index).FirstOrDefault();
4142

42-
if (lastOutput == null)
43-
{
44-
return new ChatResponse()
45-
{
46-
ResponseId = response.Id,
47-
ModelId = response.Model,
48-
CreatedAt = response.Created.ToDateTimeOffset(),
49-
Usage = MapToUsage(response.Usage),
50-
};
51-
}
52-
53-
var message = new ChatMessage(MapRole(lastOutput.Message.Role), default(string));
54-
var citations = response.Citations?.Distinct().Select(MapCitation).ToList<AIAnnotation>();
55-
56-
foreach (var output in response.Outputs.OrderBy(x => x.Index))
57-
{
58-
if (output.Message.Content is { Length: > 0 } text)
59-
{
60-
// Special-case output from tools
61-
if (output.Message.Role == MessageRole.RoleTool &&
62-
output.Message.ToolCalls.Count == 1 &&
63-
output.Message.ToolCalls[0] is { } toolCall)
64-
{
65-
if (toolCall.Type == ToolCallType.McpTool)
66-
{
67-
message.Contents.Add(new McpServerToolCallContent(toolCall.Id, toolCall.Function.Name, null)
68-
{
69-
RawRepresentation = toolCall
70-
});
71-
message.Contents.Add(new McpServerToolResultContent(toolCall.Id)
72-
{
73-
RawRepresentation = toolCall,
74-
Output = [new TextContent(text)]
75-
});
76-
continue;
77-
}
78-
else if (toolCall.Type == ToolCallType.CodeExecutionTool)
79-
{
80-
message.Contents.Add(new CodeInterpreterToolCallContent()
81-
{
82-
CallId = toolCall.Id,
83-
RawRepresentation = toolCall
84-
});
85-
message.Contents.Add(new CodeInterpreterToolResultContent()
86-
{
87-
CallId = toolCall.Id,
88-
RawRepresentation = toolCall,
89-
Outputs = [new TextContent(text)]
90-
});
91-
continue;
92-
}
93-
}
94-
95-
var content = new TextContent(text) { Annotations = citations };
96-
97-
foreach (var citation in output.Message.Citations)
98-
(content.Annotations ??= []).Add(MapInlineCitation(citation));
99-
100-
message.Contents.Add(content);
101-
}
102-
103-
foreach (var toolCall in output.Message.ToolCalls)
104-
message.Contents.Add(MapToolCall(toolCall));
105-
}
106-
107-
return new ChatResponse(message)
43+
var result = new ChatResponse()
10844
{
10945
ResponseId = response.Id,
11046
ModelId = response.Model,
11147
CreatedAt = response.Created?.ToDateTimeOffset(),
11248
FinishReason = lastOutput != null ? MapFinishReason(lastOutput.FinishReason) : null,
11349
Usage = MapToUsage(response.Usage),
11450
};
51+
52+
var citations = response.Citations?.Distinct().Select(MapCitation).ToList<AIAnnotation>();
53+
54+
((List<ChatMessage>)result.Messages).AddRange(response.Outputs.AsChatMessages(citations));
55+
56+
return result;
57+
58+
//foreach (var output in response.Outputs.OrderBy(x => x.Index))
59+
//{
60+
// if (output.Message.Content is { Length: > 0 } text)
61+
// {
62+
// // Special-case output from tools
63+
// if (output.Message.Role == MessageRole.RoleTool &&
64+
// output.Message.ToolCalls.Count == 1 &&
65+
// output.Message.ToolCalls[0] is { } toolCall)
66+
// {
67+
// if (toolCall.Type == ToolCallType.McpTool)
68+
// {
69+
// message.Contents.Add(new McpServerToolCallContent(toolCall.Id, toolCall.Function.Name, null)
70+
// {
71+
// RawRepresentation = toolCall
72+
// });
73+
// message.Contents.Add(new McpServerToolResultContent(toolCall.Id)
74+
// {
75+
// RawRepresentation = toolCall,
76+
// Output = [new TextContent(text)]
77+
// });
78+
// continue;
79+
// }
80+
// else if (toolCall.Type == ToolCallType.CodeExecutionTool)
81+
// {
82+
// message.Contents.Add(new CodeInterpreterToolCallContent()
83+
// {
84+
// CallId = toolCall.Id,
85+
// RawRepresentation = toolCall
86+
// });
87+
// message.Contents.Add(new CodeInterpreterToolResultContent()
88+
// {
89+
// CallId = toolCall.Id,
90+
// RawRepresentation = toolCall,
91+
// Outputs = [new TextContent(text)]
92+
// });
93+
// continue;
94+
// }
95+
// else if (toolCall.Type == ToolCallType.CollectionsSearchTool)
96+
// {
97+
// var settings = JsonParser.Settings.Default.WithIgnoreUnknownFields(true);
98+
// var parser = new JsonParser(settings);
99+
// var matches = parser.Parse<SearchResponse>(text);
100+
// }
101+
// }
102+
103+
// var content = new TextContent(text) { Annotations = citations };
104+
105+
// foreach (var citation in output.Message.Citations)
106+
// (content.Annotations ??= []).Add(MapInlineCitation(citation));
107+
108+
// message.Contents.Add(content);
109+
// }
110+
111+
// foreach (var toolCall in output.Message.ToolCalls)
112+
// message.Contents.Add(MapToolCall(toolCall));
113+
//}
114+
115+
//return new ChatResponse(message)
116+
//{
117+
// ResponseId = response.Id,
118+
// ModelId = response.Model,
119+
// CreatedAt = response.Created?.ToDateTimeOffset(),
120+
// FinishReason = lastOutput != null ? MapFinishReason(lastOutput.FinishReason) : null,
121+
// Usage = MapToUsage(response.Usage),
122+
//};
115123
}
116124

117125
AIContent MapToolCall(ToolCall toolCall) => toolCall.Type switch
@@ -134,7 +142,7 @@ public async Task<ChatResponse> GetResponseAsync(IEnumerable<ChatMessage> messag
134142
CallId = toolCall.Id,
135143
RawRepresentation = toolCall
136144
},
137-
_ => new HostedToolCallContent()
145+
_ => new CollectionSearchToolCallContent()
138146
{
139147
CallId = toolCall.Id,
140148
RawRepresentation = toolCall
@@ -188,15 +196,16 @@ async IAsyncEnumerable<ChatResponseUpdate> CompleteChatStreamingCore(IEnumerable
188196

189197
static CitationAnnotation MapInlineCitation(InlineCitation citation) => citation.CitationCase switch
190198
{
191-
InlineCitation.CitationOneofCase.WebCitation => new CitationAnnotation { Url = new(citation.WebCitation.Url) },
192-
InlineCitation.CitationOneofCase.XCitation => new CitationAnnotation { Url = new(citation.XCitation.Url) },
199+
InlineCitation.CitationOneofCase.WebCitation => new CitationAnnotation { Url = new(citation.WebCitation.Url), RawRepresentation = citation },
200+
InlineCitation.CitationOneofCase.XCitation => new CitationAnnotation { Url = new(citation.XCitation.Url), RawRepresentation = citation },
193201
InlineCitation.CitationOneofCase.CollectionsCitation => new CitationAnnotation
194202
{
195203
FileId = citation.CollectionsCitation.FileId,
196204
Snippet = citation.CollectionsCitation.ChunkContent,
197205
ToolName = "file_search",
206+
RawRepresentation = citation
198207
},
199-
_ => new CitationAnnotation()
208+
_ => new CitationAnnotation { RawRepresentation = citation }
200209
};
201210

202211
static CitationAnnotation MapCitation(string citation)
@@ -210,12 +219,13 @@ static CitationAnnotation MapCitation(string citation)
210219
var file = url.AbsolutePath[7..];
211220
return new CitationAnnotation
212221
{
213-
ToolName = "collections_search",
214-
FileId = file,
215222
AdditionalProperties = new AdditionalPropertiesDictionary
216-
{
217-
{ "collection_id", collection }
218-
}
223+
{
224+
{ "collection_id", collection }
225+
},
226+
FileId = file,
227+
ToolName = "collections_search",
228+
Url = new Uri($"collections://{collection}/files/{file}"),
219229
};
220230
}
221231

0 commit comments

Comments
 (0)