mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
.NET: Add rag samples with sample TextSearchStore (#1664)
* Port store for adding text to a vector store to AF * Fix typo. * Change TextSearchStore to sample, and add sample to use it and do rag with a custom schema * Add more tests and fix broken ones * Fix merge issue * Fix sample after merge. * Convert TextSearchStore to use Dynamic mode to be AOT compatible. * Add some more clarification on when to use assistant messages in rag searches.
This commit is contained in:
westey
+22
@@ -0,0 +1,22 @@
|
||||
<Project Sdk="Microsoft.NET.Sdk">
|
||||
|
||||
<PropertyGroup>
|
||||
<OutputType>Exe</OutputType>
|
||||
<TargetFramework>net9.0</TargetFramework>
|
||||
|
||||
<Nullable>enable</Nullable>
|
||||
<ImplicitUsings>enable</ImplicitUsings>
|
||||
</PropertyGroup>
|
||||
|
||||
<ItemGroup>
|
||||
<PackageReference Include="Azure.AI.OpenAI" />
|
||||
<PackageReference Include="Azure.Identity" />
|
||||
<PackageReference Include="Microsoft.Extensions.AI.OpenAI" />
|
||||
<PackageReference Include="Microsoft.SemanticKernel.Connectors.InMemory" />
|
||||
</ItemGroup>
|
||||
|
||||
<ItemGroup>
|
||||
<ProjectReference Include="..\..\..\..\src\Microsoft.Agents.AI.OpenAI\Microsoft.Agents.AI.OpenAI.csproj" />
|
||||
</ItemGroup>
|
||||
|
||||
</Project>
|
||||
+107
@@ -0,0 +1,107 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
// This sample shows how to use TextSearchProvider to add retrieval augmented generation (RAG) capabilities to an AI agent.
|
||||
// The sample uses an In-Memory vector store, which can easily be replaced with any other vector store that implements the Microsoft.Extensions.VectorData abstractions.
|
||||
// The TextSearchProvider runs a search against the vector store via the TextSearchStore before each model invocation and injects the results into the model context.
|
||||
// The TextSearchStore is a sample store implementation that hardcodes a storage schema and uses the vector store to store and retrieve documents.
|
||||
|
||||
using Azure.AI.OpenAI;
|
||||
using Azure.Identity;
|
||||
using Microsoft.Agents.AI;
|
||||
using Microsoft.Agents.AI.Data;
|
||||
using Microsoft.Agents.AI.Samples;
|
||||
using Microsoft.Extensions.AI;
|
||||
using Microsoft.Extensions.VectorData;
|
||||
using Microsoft.SemanticKernel.Connectors.InMemory;
|
||||
using OpenAI;
|
||||
|
||||
var endpoint = Environment.GetEnvironmentVariable("AZURE_OPENAI_ENDPOINT") ?? throw new InvalidOperationException("AZURE_OPENAI_ENDPOINT is not set.");
|
||||
var deploymentName = Environment.GetEnvironmentVariable("AZURE_OPENAI_DEPLOYMENT_NAME") ?? "gpt-4o-mini";
|
||||
var embeddingDeploymentName = Environment.GetEnvironmentVariable("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME") ?? "text-embedding-3-large";
|
||||
|
||||
AzureOpenAIClient azureOpenAIClient = new(
|
||||
new Uri(endpoint),
|
||||
new AzureCliCredential());
|
||||
|
||||
// Create an In-Memory vector store that uses the Azure OpenAI embedding model to generate embeddings.
|
||||
VectorStore vectorStore = new InMemoryVectorStore(new()
|
||||
{
|
||||
EmbeddingGenerator = azureOpenAIClient.GetEmbeddingClient(embeddingDeploymentName).AsIEmbeddingGenerator()
|
||||
});
|
||||
|
||||
// Create a store that defines a storage schema, and uses the vector store to store and retrieve documents.
|
||||
TextSearchStore textSearchStore = new(vectorStore, "product-and-policy-info", 3072);
|
||||
|
||||
// Upload sample documents into the store.
|
||||
await textSearchStore.UpsertDocumentsAsync(GetSampleDocuments());
|
||||
|
||||
// Create an adapter function that the TextSearchProvider can use to run searches against the TextSearchStore.
|
||||
Func<string, CancellationToken, Task<IEnumerable<TextSearchProvider.TextSearchResult>>> SearchAdapter = async (text, ct) =>
|
||||
{
|
||||
// Here we are limiting the search results to the single top result to demonstrate that we are accurately matching
|
||||
// specific search results for each question, but in a real world case, more results should be used.
|
||||
var searchResults = await textSearchStore.SearchAsync(text, 1, ct);
|
||||
return searchResults.Select(r => new TextSearchProvider.TextSearchResult
|
||||
{
|
||||
SourceName = r.SourceName,
|
||||
SourceLink = r.SourceLink,
|
||||
Text = r.Text ?? string.Empty,
|
||||
RawRepresentation = r
|
||||
});
|
||||
};
|
||||
|
||||
// Configure the options for the TextSearchProvider.
|
||||
TextSearchProviderOptions textSearchOptions = new()
|
||||
{
|
||||
// Run the search prior to every model invocation.
|
||||
SearchTime = TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke,
|
||||
};
|
||||
|
||||
// Create the AI agent with the TextSearchProvider as the AI context provider.
|
||||
AIAgent agent = azureOpenAIClient
|
||||
.GetChatClient(deploymentName)
|
||||
.CreateAIAgent(new ChatClientAgentOptions
|
||||
{
|
||||
Instructions = "You are a helpful support specialist for Contoso Outdoors. Answer questions using the provided context and cite the source document when available.",
|
||||
AIContextProviderFactory = ctx => ctx.SerializedState.ValueKind is not System.Text.Json.JsonValueKind.Null and not System.Text.Json.JsonValueKind.Undefined
|
||||
? new TextSearchProvider(SearchAdapter, ctx.SerializedState, ctx.JsonSerializerOptions, textSearchOptions)
|
||||
: new TextSearchProvider(SearchAdapter, textSearchOptions)
|
||||
});
|
||||
|
||||
AgentThread thread = agent.GetNewThread();
|
||||
|
||||
Console.WriteLine(">> Asking about returns\n");
|
||||
Console.WriteLine(await agent.RunAsync("Hi! I need help understanding the return policy.", thread));
|
||||
|
||||
Console.WriteLine("\n>> Asking about shipping\n");
|
||||
Console.WriteLine(await agent.RunAsync("How long does standard shipping usually take?", thread));
|
||||
|
||||
Console.WriteLine("\n>> Asking about product care\n");
|
||||
Console.WriteLine(await agent.RunAsync("What is the best way to maintain the TrailRunner tent fabric?", thread));
|
||||
|
||||
// Produces some sample search documents.
|
||||
// Each one contains a source name and link, which the agent can use to cite sources in its responses.
|
||||
static IEnumerable<TextSearchDocument> GetSampleDocuments()
|
||||
{
|
||||
yield return new TextSearchDocument
|
||||
{
|
||||
SourceId = "return-policy-001",
|
||||
SourceName = "Contoso Outdoors Return Policy",
|
||||
SourceLink = "https://contoso.com/policies/returns",
|
||||
Text = "Customers may return any item within 30 days of delivery. Items should be unused and include original packaging. Refunds are issued to the original payment method within 5 business days of inspection."
|
||||
};
|
||||
yield return new TextSearchDocument
|
||||
{
|
||||
SourceId = "shipping-guide-001",
|
||||
SourceName = "Contoso Outdoors Shipping Guide",
|
||||
SourceLink = "https://contoso.com/help/shipping",
|
||||
Text = "Standard shipping is free on orders over $50 and typically arrives in 3-5 business days within the continental United States. Expedited options are available at checkout."
|
||||
};
|
||||
yield return new TextSearchDocument
|
||||
{
|
||||
SourceId = "tent-care-001",
|
||||
SourceName = "TrailRunner Tent Care Instructions",
|
||||
SourceLink = "https://contoso.com/manuals/trailrunner-tent",
|
||||
Text = "Clean the tent fabric with lukewarm water and a non-detergent soap. Allow it to air dry completely before storage and avoid prolonged UV exposure to extend the lifespan of the waterproof coating."
|
||||
};
|
||||
}
|
||||
+51
@@ -0,0 +1,51 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
namespace Microsoft.Agents.AI.Samples;
|
||||
|
||||
/// <summary>
|
||||
/// Represents a document that can be used for Retrieval Augmented Generation (RAG) that stores textual data.
|
||||
/// </summary>
|
||||
public sealed class TextSearchDocument
|
||||
{
|
||||
/// <summary>
|
||||
/// Gets or sets an optional list of namespaces that the document should belong to.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// A namespace is a logical grouping of documents, e.g. may include a group id to scope the document to a specific group of users.
|
||||
/// </remarks>
|
||||
public IList<string> Namespaces { get; set; } = [];
|
||||
|
||||
/// <summary>
|
||||
/// Gets or sets the content as text.
|
||||
/// </summary>
|
||||
public string? Text { get; set; }
|
||||
|
||||
/// <summary>
|
||||
/// Gets or sets an optional source ID for the document.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// This ID should be unique within the collection that the document is stored in, and can
|
||||
/// be used to map back to the source artifact for this document.
|
||||
/// If updates need to be made later or the source document was deleted and this document
|
||||
/// also needs to be deleted, this id can be used to find the document again.
|
||||
/// </remarks>
|
||||
public string? SourceId { get; set; }
|
||||
|
||||
/// <summary>
|
||||
/// Gets or sets an optional name for the source document.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// This can be used to provide display names for citation links when the document is referenced as
|
||||
/// part of a response to a query.
|
||||
/// </remarks>
|
||||
public string? SourceName { get; set; }
|
||||
|
||||
/// <summary>
|
||||
/// Gets or sets an optional link back to the source of the document.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// This can be used to provide citation links when the document is referenced as
|
||||
/// part of a response to a query.
|
||||
/// </remarks>
|
||||
public string? SourceLink { get; set; }
|
||||
}
|
||||
+392
@@ -0,0 +1,392 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
using System.Linq.Expressions;
|
||||
using System.Text.RegularExpressions;
|
||||
using Microsoft.Extensions.VectorData;
|
||||
|
||||
namespace Microsoft.Agents.AI.Samples;
|
||||
|
||||
/// <summary>
|
||||
/// A class that allows for easy storage and retrieval of documents in a Vector Store for Retrieval Augmented Generation (RAG).
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// <para>
|
||||
/// This class provides an opinionated schema for storing documents in a vector store. It is valuable for simple scenarios
|
||||
/// where you want to store text + embedding, or a reference to an external document + embedding without needing to customize the schema.
|
||||
/// If you want to control the schema yourself, use an implementation of <see cref="VectorStoreCollection{TKey, TRecord}"/> directly instead.
|
||||
/// </para>
|
||||
/// <para>
|
||||
/// This class and its related types are currently provided as a sample implementation, but may be promoted to a first-class supported API in future releases.
|
||||
/// </para>
|
||||
/// </remarks>
|
||||
public sealed partial class TextSearchStore : IDisposable
|
||||
{
|
||||
#if NET
|
||||
[GeneratedRegex(@"\p{L}+", RegexOptions.IgnoreCase, "en-US")]
|
||||
private static partial Regex AnyLanguageWordRegex();
|
||||
|
||||
private static readonly Func<string, ICollection<string>> s_defaultWordSegmenter = text => AnyLanguageWordRegex().Matches(text).Select(x => x.Value).ToList();
|
||||
#else
|
||||
private static readonly Regex s_anyLanguageWordRegex = new(@"\p{L}+", RegexOptions.Compiled);
|
||||
private static Regex AnyLanguageWordRegex() => s_anyLanguageWordRegex;
|
||||
|
||||
private static readonly Func<string, ICollection<string>> s_defaultWordSegmenter = text =>
|
||||
{
|
||||
List<string> words = new();
|
||||
foreach (Match word in AnyLanguageWordRegex().Matches(text))
|
||||
{
|
||||
words.Add(word.Value);
|
||||
}
|
||||
return words;
|
||||
};
|
||||
#endif
|
||||
|
||||
private readonly VectorStore _vectorStore;
|
||||
private readonly TextSearchStoreOptions _options;
|
||||
private readonly Func<string, ICollection<string>> _wordSegmenter;
|
||||
|
||||
private readonly VectorStoreCollection<object, Dictionary<string, object?>> _vectorStoreRecordCollection;
|
||||
private readonly SemaphoreSlim _collectionInitializationLock = new(1, 1);
|
||||
private bool _collectionInitialized;
|
||||
private bool _disposedValue;
|
||||
|
||||
/// <summary>
|
||||
/// Initializes a new instance of the <see cref="TextSearchStore"/> class.
|
||||
/// </summary>
|
||||
/// <param name="vectorStore">The vector store to store and read the memories from.</param>
|
||||
/// <param name="collectionName">The name of the collection in the vector store to store and read the memories from.</param>
|
||||
/// <param name="vectorDimensions">The number of dimensions to use for the memory embeddings.</param>
|
||||
/// <param name="options">Options to configure the behavior of this class.</param>
|
||||
/// <exception cref="NotSupportedException">Thrown if the key type provided is not supported.</exception>
|
||||
public TextSearchStore(
|
||||
VectorStore vectorStore,
|
||||
string collectionName,
|
||||
int vectorDimensions,
|
||||
TextSearchStoreOptions? options = default)
|
||||
{
|
||||
// Verify
|
||||
if (vectorStore is null)
|
||||
{
|
||||
throw new ArgumentNullException(nameof(vectorStore));
|
||||
}
|
||||
|
||||
if (string.IsNullOrWhiteSpace(collectionName))
|
||||
{
|
||||
throw new ArgumentException("Collection name cannot be null or whitespace.", nameof(collectionName));
|
||||
}
|
||||
|
||||
if (vectorDimensions < 1)
|
||||
{
|
||||
throw new ArgumentOutOfRangeException(nameof(vectorDimensions), "Vector dimensions must be greater than zero.");
|
||||
}
|
||||
|
||||
if (options?.KeyType is not null && options.KeyType != typeof(string) && options.KeyType != typeof(Guid))
|
||||
{
|
||||
throw new NotSupportedException($"Unsupported key of type '{options.KeyType.Name}'");
|
||||
}
|
||||
|
||||
if (options?.KeyType is not null && options.KeyType != typeof(string) && options?.UseSourceIdAsPrimaryKey is true)
|
||||
{
|
||||
throw new NotSupportedException($"The {nameof(TextSearchStoreOptions.UseSourceIdAsPrimaryKey)} option can only be used when the key type is 'string'.");
|
||||
}
|
||||
|
||||
// Assign
|
||||
this._vectorStore = vectorStore;
|
||||
this._options = options ?? new TextSearchStoreOptions();
|
||||
this._wordSegmenter = this._options.WordSegmenter ?? s_defaultWordSegmenter;
|
||||
|
||||
// Create a definition so that we can use the dimensions provided at runtime.
|
||||
VectorStoreCollectionDefinition ragDocumentDefinition = new()
|
||||
{
|
||||
Properties = new List<VectorStoreProperty>()
|
||||
{
|
||||
new VectorStoreKeyProperty("Key", this._options.KeyType ?? typeof(string)),
|
||||
new VectorStoreDataProperty("Namespaces", typeof(List<string>)) { IsIndexed = true },
|
||||
new VectorStoreDataProperty("SourceId", typeof(string)) { IsIndexed = true },
|
||||
new VectorStoreDataProperty("Text", typeof(string)) { IsFullTextIndexed = true },
|
||||
new VectorStoreDataProperty("SourceName", typeof(string)),
|
||||
new VectorStoreDataProperty("SourceLink", typeof(string)),
|
||||
new VectorStoreVectorProperty("TextEmbedding", typeof(string), vectorDimensions),
|
||||
}
|
||||
};
|
||||
|
||||
this._vectorStoreRecordCollection = this._vectorStore.GetDynamicCollection(collectionName, ragDocumentDefinition);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Upserts a batch of text chunks into the vector store.
|
||||
/// </summary>
|
||||
/// <param name="textChunks">The text chunks to upload.</param>
|
||||
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
|
||||
/// <returns>A task that completes when the documents have been upserted.</returns>
|
||||
public async Task UpsertTextAsync(IEnumerable<string> textChunks, CancellationToken cancellationToken = default)
|
||||
{
|
||||
if (textChunks == null)
|
||||
{
|
||||
throw new ArgumentNullException(nameof(textChunks));
|
||||
}
|
||||
|
||||
var vectorStoreRecordCollection = await this.EnsureCollectionExistsAsync(cancellationToken).ConfigureAwait(false);
|
||||
|
||||
var storageDocuments = textChunks.Select(textChunk =>
|
||||
{
|
||||
// Without text we cannot generate a vector.
|
||||
if (string.IsNullOrWhiteSpace(textChunk))
|
||||
{
|
||||
throw new ArgumentException("One of the provided text chunks is null.", nameof(textChunks));
|
||||
}
|
||||
|
||||
return new Dictionary<string, object?>
|
||||
{
|
||||
{ "Key", this.GenerateUniqueKey(null) },
|
||||
{ "Namespaces", new List<string>() },
|
||||
{ "Text", textChunk },
|
||||
{ "TextEmbedding", textChunk },
|
||||
};
|
||||
});
|
||||
|
||||
await vectorStoreRecordCollection.UpsertAsync(storageDocuments, cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Upserts a batch of documents into the vector store.
|
||||
/// </summary>
|
||||
/// <param name="documents">The documents to upload.</param>
|
||||
/// <param name="options">Optional options to control the upsert behavior.</param>
|
||||
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
|
||||
/// <returns>A task that completes when the documents have been upserted.</returns>
|
||||
public async Task UpsertDocumentsAsync(IEnumerable<TextSearchDocument> documents, TextSearchStoreUpsertOptions? options = null, CancellationToken cancellationToken = default)
|
||||
{
|
||||
if (documents is null)
|
||||
{
|
||||
throw new ArgumentNullException(nameof(documents));
|
||||
}
|
||||
|
||||
var vectorStoreRecordCollection = await this.EnsureCollectionExistsAsync(cancellationToken).ConfigureAwait(false);
|
||||
|
||||
var storageDocuments = documents.Select(document =>
|
||||
{
|
||||
if (document is null)
|
||||
{
|
||||
throw new ArgumentNullException(nameof(documents), "One of the provided documents is null.");
|
||||
}
|
||||
|
||||
// Without text we cannot generate a vector.
|
||||
if (string.IsNullOrWhiteSpace(document.Text))
|
||||
{
|
||||
throw new ArgumentException($"The {nameof(TextSearchDocument.Text)} property must be set.", nameof(document));
|
||||
}
|
||||
|
||||
// If we aren't persisting the text, we need a source id or link to refer back to the original document.
|
||||
if (options?.DoNotPersistSourceText is true && string.IsNullOrWhiteSpace(document.SourceId) && string.IsNullOrWhiteSpace(document.SourceLink))
|
||||
{
|
||||
throw new ArgumentException($"Either the {nameof(TextSearchDocument.SourceId)} or {nameof(TextSearchDocument.SourceLink)} properties must be set when the {nameof(TextSearchStoreUpsertOptions.DoNotPersistSourceText)} setting is true.", nameof(document));
|
||||
}
|
||||
|
||||
var key = this.GenerateUniqueKey(this._options.UseSourceIdAsPrimaryKey ?? false ? document.SourceId : null);
|
||||
|
||||
return new Dictionary<string, object?>()
|
||||
{
|
||||
{ "Key", key },
|
||||
{ "Namespaces", document.Namespaces.ToList() },
|
||||
{ "SourceId", document.SourceId },
|
||||
{ "Text", options?.DoNotPersistSourceText is true ? null : document.Text },
|
||||
{ "SourceName", document.SourceName },
|
||||
{ "SourceLink", document.SourceLink },
|
||||
{ "TextEmbedding", document.Text },
|
||||
};
|
||||
});
|
||||
|
||||
await vectorStoreRecordCollection.UpsertAsync(storageDocuments, cancellationToken).ConfigureAwait(false);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Search the database for documents similar to the provided query.
|
||||
/// </summary>
|
||||
/// <param name="query">The text query to find similar documents to.</param>
|
||||
/// <param name="top">The maximum number of results to return.</param>
|
||||
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
|
||||
/// <returns>The search results.</returns>
|
||||
public async Task<IEnumerable<TextSearchDocument>> SearchAsync(string query, int top, CancellationToken cancellationToken = default)
|
||||
{
|
||||
var searchResult = await this.SearchCoreAsync(query, top, cancellationToken).ConfigureAwait(false);
|
||||
|
||||
return searchResult.Select(x => new TextSearchDocument()
|
||||
{
|
||||
Namespaces = (List<string>)x["Namespaces"]!,
|
||||
Text = (string?)x["Text"],
|
||||
SourceId = (string?)x["SourceId"],
|
||||
SourceName = (string?)x["SourceName"],
|
||||
SourceLink = (string?)x["SourceLink"],
|
||||
});
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Internal search implementation with hydration of id / link only storage.
|
||||
/// </summary>
|
||||
/// <param name="query">The text query to find similar documents to.</param>
|
||||
/// <param name="top">The maximum number of results to return.</param>
|
||||
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
|
||||
/// <returns>The search results.</returns>
|
||||
private async Task<IEnumerable<Dictionary<string, object?>>> SearchCoreAsync(string query, int top, CancellationToken cancellationToken = default)
|
||||
{
|
||||
// Short circuit if the query is empty.
|
||||
if (string.IsNullOrWhiteSpace(query))
|
||||
{
|
||||
return [];
|
||||
}
|
||||
|
||||
var vectorStoreRecordCollection = await this.EnsureCollectionExistsAsync(cancellationToken).ConfigureAwait(false);
|
||||
|
||||
// If the user has not opted out of hybrid search, check if the vector store supports it.
|
||||
var hybridSearchCollection = this._options.UseHybridSearch ?? true ?
|
||||
vectorStoreRecordCollection.GetService(typeof(IKeywordHybridSearchable<Dictionary<string, object?>>)) as IKeywordHybridSearchable<Dictionary<string, object?>> :
|
||||
null;
|
||||
|
||||
// Optional filter to limit the search to a specific namespace.
|
||||
Expression<Func<Dictionary<string, object?>, bool>>? filter = string.IsNullOrWhiteSpace(this._options.SearchNamespace) ? null : x => ((List<string>)x["Namespaces"]!).Contains(this._options.SearchNamespace);
|
||||
|
||||
// Execute a hybrid search if possible, otherwise perform a regular vector search.
|
||||
var searchResult = hybridSearchCollection is null
|
||||
? vectorStoreRecordCollection.SearchAsync(
|
||||
query,
|
||||
top,
|
||||
options: new()
|
||||
{
|
||||
Filter = filter,
|
||||
},
|
||||
cancellationToken: cancellationToken)
|
||||
: hybridSearchCollection.HybridSearchAsync(
|
||||
query,
|
||||
this._wordSegmenter(query),
|
||||
top,
|
||||
options: new()
|
||||
{
|
||||
Filter = filter,
|
||||
},
|
||||
cancellationToken: cancellationToken);
|
||||
|
||||
// Retrieve the documents from the search results.
|
||||
List<Dictionary<string, object?>> searchResponseDocs = new();
|
||||
await foreach (var searchResponseDoc in searchResult.WithCancellation(cancellationToken).ConfigureAwait(false))
|
||||
{
|
||||
searchResponseDocs.Add(searchResponseDoc.Record);
|
||||
}
|
||||
|
||||
// Find any source ids and links for which the text needs to be retrieved.
|
||||
var sourceIdsToRetrieve = searchResponseDocs
|
||||
.Where(x => string.IsNullOrWhiteSpace((string?)x["Text"]))
|
||||
.Select(x => new TextSearchStoreOptions.SourceRetrievalRequest((string?)x["SourceId"], (string?)x["SourceLink"]))
|
||||
.ToList();
|
||||
|
||||
// If we have none, we can return early.
|
||||
if (sourceIdsToRetrieve.Count == 0)
|
||||
{
|
||||
return searchResponseDocs;
|
||||
}
|
||||
|
||||
if (this._options.SourceRetrievalCallback is null)
|
||||
{
|
||||
throw new InvalidOperationException($"The {nameof(TextSearchStoreOptions.SourceRetrievalCallback)} option must be set if retrieving documents without stored text.");
|
||||
}
|
||||
|
||||
// Retrieve the source text for the documents that need it.
|
||||
var retrievalResponses = await this._options.SourceRetrievalCallback(sourceIdsToRetrieve).ConfigureAwait(false);
|
||||
|
||||
if (retrievalResponses is null)
|
||||
{
|
||||
throw new InvalidOperationException($"The {nameof(TextSearchStoreOptions.SourceRetrievalCallback)} must return a non-null value.");
|
||||
}
|
||||
|
||||
// Update the retrieved documents with the retrieved text.
|
||||
return searchResponseDocs.GroupJoin(
|
||||
retrievalResponses,
|
||||
searchResponseDoc => (searchResponseDoc["SourceId"], searchResponseDoc["SourceLink"]),
|
||||
retrievalResponse => (retrievalResponse.SourceId, retrievalResponse.SourceLink),
|
||||
(searchResponseDoc, textRetrievalResponse) => (searchResponseDoc, textRetrievalResponse))
|
||||
.SelectMany(
|
||||
joinedSet => joinedSet.textRetrievalResponse.DefaultIfEmpty(),
|
||||
(combined, textRetrievalResponse) =>
|
||||
{
|
||||
combined.searchResponseDoc["Text"] = textRetrievalResponse?.Text ?? combined.searchResponseDoc["Text"];
|
||||
return combined.searchResponseDoc;
|
||||
});
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Thread safe method to get the collection and ensure that it is created at least once.
|
||||
/// </summary>
|
||||
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
|
||||
/// <returns>The created collection.</returns>
|
||||
private async Task<VectorStoreCollection<object, Dictionary<string, object?>>> EnsureCollectionExistsAsync(CancellationToken cancellationToken)
|
||||
{
|
||||
// Return immediately if the collection is already created, no need to do any locking in this case.
|
||||
if (this._collectionInitialized)
|
||||
{
|
||||
return this._vectorStoreRecordCollection;
|
||||
}
|
||||
|
||||
// Wait on a lock to ensure that only one thread can create the collection.
|
||||
await this._collectionInitializationLock.WaitAsync(cancellationToken).ConfigureAwait(false);
|
||||
|
||||
// If multiple threads waited on the lock, and the first already created the collection,
|
||||
// we can return immediately without doing any work in subsequent threads.
|
||||
if (this._collectionInitialized)
|
||||
{
|
||||
this._collectionInitializationLock.Release();
|
||||
return this._vectorStoreRecordCollection;
|
||||
}
|
||||
|
||||
// Only the winning thread should reach this point and create the collection.
|
||||
try
|
||||
{
|
||||
await this._vectorStoreRecordCollection.EnsureCollectionExistsAsync(cancellationToken).ConfigureAwait(false);
|
||||
this._collectionInitialized = true;
|
||||
}
|
||||
finally
|
||||
{
|
||||
this._collectionInitializationLock.Release();
|
||||
}
|
||||
|
||||
return this._vectorStoreRecordCollection;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Generates a unique key for the RAG document.
|
||||
/// </summary>
|
||||
/// <param name="sourceId">Source id of the source document for this RAG document.</param>
|
||||
/// <returns>A new unique key.</returns>
|
||||
/// <exception cref="NotSupportedException">Thrown if the requested key type is not supported.</exception>
|
||||
private object GenerateUniqueKey(string? sourceId)
|
||||
=> this._options.KeyType switch
|
||||
{
|
||||
_ when (this._options.KeyType == null || this._options.KeyType == typeof(string)) && !string.IsNullOrWhiteSpace(sourceId) => sourceId!,
|
||||
_ when this._options.KeyType == null || this._options.KeyType == typeof(string) => Guid.NewGuid().ToString(),
|
||||
_ when this._options.KeyType == typeof(Guid) => Guid.NewGuid(),
|
||||
|
||||
_ => throw new NotSupportedException($"Unsupported key of type '{this._options.KeyType.Name}'")
|
||||
};
|
||||
|
||||
/// <inheritdoc/>
|
||||
private void Dispose(bool disposing)
|
||||
{
|
||||
if (!this._disposedValue)
|
||||
{
|
||||
if (disposing)
|
||||
{
|
||||
this._vectorStoreRecordCollection.Dispose();
|
||||
this._collectionInitializationLock.Dispose();
|
||||
}
|
||||
|
||||
this._disposedValue = true;
|
||||
}
|
||||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
public void Dispose()
|
||||
{
|
||||
// Do not change this code. Put cleanup code in 'Dispose(bool disposing)' method
|
||||
this.Dispose(disposing: true);
|
||||
GC.SuppressFinalize(this);
|
||||
}
|
||||
}
|
||||
+140
@@ -0,0 +1,140 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
namespace Microsoft.Agents.AI.Samples;
|
||||
|
||||
/// <summary>
|
||||
/// Contains options for the <see cref="TextSearchStore"/>.
|
||||
/// </summary>
|
||||
public sealed class TextSearchStoreOptions
|
||||
{
|
||||
/// <summary>
|
||||
/// Gets or sets an optional namespace to pre-filter the possible
|
||||
/// records with when doing a vector search.
|
||||
/// </summary>
|
||||
public string? SearchNamespace { get; init; }
|
||||
|
||||
/// <summary>
|
||||
/// Gets or sets a value indicating whether to use the source ID as the primary key for records.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// <para>
|
||||
/// Using the source ID as the primary key allows for easy updates from the source for any changed
|
||||
/// records, since those records can just be upserted again, and will overwrite the previous version
|
||||
/// of the same record.
|
||||
/// </para>
|
||||
/// <para>
|
||||
/// This setting can only be used when the chosen key type is a string.
|
||||
/// </para>
|
||||
/// </remarks>
|
||||
/// <value>
|
||||
/// Defaults to <c>false</c> if not set.
|
||||
/// </value>
|
||||
public bool? UseSourceIdAsPrimaryKey { get; init; }
|
||||
|
||||
/// <summary>
|
||||
/// Gets or sets a value indicating whether to use hybrid search if it is available for the provided vector store.
|
||||
/// </summary>
|
||||
/// <value>
|
||||
/// Defaults to <c>true</c> if not set.
|
||||
/// </value>
|
||||
public bool? UseHybridSearch { get; init; }
|
||||
|
||||
/// <summary>
|
||||
/// Gets or sets a word segmenter function to split search text into separate words for the purposes of hybrid search.
|
||||
/// This will not be used if <see cref="UseHybridSearch"/> is set to <c>false</c>.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Defaults to a simple text-character-based segmenter that splits the text by any character that is not a text character.
|
||||
/// </remarks>
|
||||
public Func<string, ICollection<string>>? WordSegmenter { get; init; }
|
||||
|
||||
/// <summary>
|
||||
/// Gets or sets the type of key to use for records in the text search store.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// Make sure to pick a key type that is supported by the underlying vector store.
|
||||
/// Note that you have to choose <see cref="string"/> when using <see cref="UseSourceIdAsPrimaryKey"/>.
|
||||
/// </remarks>
|
||||
/// <value>Defaults to <see cref="string"/> if not set. Only <see cref="string"/> and <see cref="Guid"/> is currently supported.</value>
|
||||
public Type? KeyType { get; init; }
|
||||
|
||||
/// <summary>
|
||||
/// Gets or sets an optional callback to load the source text using the source id or source link
|
||||
/// if the source text is not persisted in the database.
|
||||
/// </summary>
|
||||
/// <remarks>
|
||||
/// The response should include the source id or source link, as provided in the request,
|
||||
/// plus the source text loaded from the source.
|
||||
/// </remarks>
|
||||
public Func<List<SourceRetrievalRequest>, Task<IEnumerable<SourceRetrievalResponse>>>? SourceRetrievalCallback { get; init; }
|
||||
|
||||
/// <summary>
|
||||
/// Represents a request to the <see cref="SourceRetrievalCallback"/>.
|
||||
/// </summary>
|
||||
public sealed class SourceRetrievalRequest
|
||||
{
|
||||
/// <summary>
|
||||
/// Initializes a new instance of the <see cref="SourceRetrievalRequest"/> class.
|
||||
/// </summary>
|
||||
/// <param name="sourceId">The source ID of the document to retrieve.</param>
|
||||
/// <param name="sourceLink">The source link of the document to retrieve.</param>
|
||||
public SourceRetrievalRequest(string? sourceId, string? sourceLink)
|
||||
{
|
||||
this.SourceId = sourceId;
|
||||
this.SourceLink = sourceLink;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Gets or sets the source ID of the document to retrieve.
|
||||
/// </summary>
|
||||
public string? SourceId { get; set; }
|
||||
|
||||
/// <summary>
|
||||
/// Gets or sets the source link of the document to retrieve.
|
||||
/// </summary>
|
||||
public string? SourceLink { get; set; }
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Represents a response from the <see cref="SourceRetrievalCallback"/>.
|
||||
/// </summary>
|
||||
public sealed class SourceRetrievalResponse
|
||||
{
|
||||
/// <summary>
|
||||
/// Initializes a new instance of the <see cref="SourceRetrievalResponse"/> class.
|
||||
/// </summary>
|
||||
/// <param name="request">The request matching this response.</param>
|
||||
/// <param name="text">The source text that was retrieved.</param>
|
||||
public SourceRetrievalResponse(SourceRetrievalRequest request, string text)
|
||||
{
|
||||
if (request == null)
|
||||
{
|
||||
throw new ArgumentNullException(nameof(request));
|
||||
}
|
||||
|
||||
if (text == null)
|
||||
{
|
||||
throw new ArgumentNullException(nameof(text));
|
||||
}
|
||||
|
||||
this.SourceId = request.SourceId;
|
||||
this.SourceLink = request.SourceLink;
|
||||
this.Text = text;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
/// Gets or sets the source ID of the document that was retrieved.
|
||||
/// </summary>
|
||||
public string? SourceId { get; set; }
|
||||
|
||||
/// <summary>
|
||||
/// Gets or sets the source link of the document that was retrieved.
|
||||
/// </summary>
|
||||
public string? SourceLink { get; set; }
|
||||
|
||||
/// <summary>
|
||||
/// Gets or sets the source text of the document that was retrieved.
|
||||
/// </summary>
|
||||
public string Text { get; set; }
|
||||
}
|
||||
}
|
||||
+17
@@ -0,0 +1,17 @@
|
||||
// Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
namespace Microsoft.Agents.AI.Samples;
|
||||
|
||||
/// <summary>
|
||||
/// Contains options for <see cref="TextSearchStore.UpsertDocumentsAsync(IEnumerable{TextSearchDocument}, TextSearchStoreUpsertOptions?, CancellationToken)"/>.
|
||||
/// </summary>
|
||||
public sealed class TextSearchStoreUpsertOptions
|
||||
{
|
||||
/// <summary>
|
||||
/// Gets or sets a value indicating whether the source text should be persisted in the database.
|
||||
/// </summary>
|
||||
/// <value>
|
||||
/// Defaults to <see langword="false"/> if not set.
|
||||
/// </value>
|
||||
public bool DoNotPersistSourceText { get; init; }
|
||||
}
|
||||
Reference in New Issue
Block a user