-
Notifications
You must be signed in to change notification settings - Fork 2.8k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
.Net: Added example of Azure AI Search as Plugin with custom schema a…
…nd configurable search fields (#5093) ### Motivation and Context <!-- Thank you for your contribution to the semantic-kernel repo! Please help reviewers and future users, providing the following information: 1. Why is this change required? 2. What problem does it solve? 3. What scenario does it contribute to? 4. If it fixes an open issue, please link to the issue here. --> Today, Azure AI Search connector has two limitations: 1. It's working only with predefined schema, which is not preferred approach, because index in Azure AI Search is customizable and may use different schema from the one that exists in SK today. 2. Search operation uses only `Embedding` index field to perform search, while it's possible to specify multiple vector fields with different names. In order to resolve these limitations, major refactoring on abstraction level is required, which will impact not only Azure AI Search connector, but all other memory connectors as well, and it will take some time. This PR contains example that shows how to register Azure AI Search functionality in Kernel as Plugin and use it to perform search operations and communicate with AI based on your data. In this example, index schema is custom and can be changed when needed, as well as `searchFields` parameter is configurable via `KernelArguments`. The usage of Azure AI Search plugin is very similar to already existing approach with `AzureAISearchMemoryStore` connector. ### Contribution Checklist <!-- Before submitting this PR, please make sure: --> - [x] The code builds clean without any errors or warnings - [x] The PR follows the [SK Contribution Guidelines](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md) and the [pre-submission formatting script](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md#development-scripts) raises no violations - [x] All unit tests pass, and I have added new tests where possible - [x] I didn't break anyone 😄
- Loading branch information
1 parent
3b14d7c
commit 6954e97
Showing
1 changed file
with
207 additions
and
0 deletions.
There are no files selected for viewing
207 changes: 207 additions & 0 deletions
207
dotnet/samples/KernelSyntaxExamples/Example84_AzureAISearchPlugin.cs
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,207 @@ | ||
// Copyright (c) Microsoft. All rights reserved. | ||
|
||
using System; | ||
using System.Collections.Generic; | ||
using System.Linq; | ||
using System.Text.Json; | ||
using System.Text.Json.Serialization; | ||
using System.Threading; | ||
using System.Threading.Tasks; | ||
using Azure; | ||
using Azure.Search.Documents; | ||
using Azure.Search.Documents.Indexes; | ||
using Azure.Search.Documents.Models; | ||
using Microsoft.Extensions.DependencyInjection; | ||
using Microsoft.SemanticKernel; | ||
using Microsoft.SemanticKernel.Embeddings; | ||
using Xunit; | ||
using Xunit.Abstractions; | ||
|
||
namespace Examples; | ||
|
||
public class Example84_AzureAISearchPlugin : BaseTest | ||
{ | ||
/// <summary> | ||
/// Shows how to register Azure AI Search service as a plugin and work with custom index schema. | ||
/// </summary> | ||
[Fact] | ||
public async Task AzureAISearchPluginAsync() | ||
{ | ||
// Azure AI Search configuration | ||
Uri endpoint = new(TestConfiguration.AzureAISearch.Endpoint); | ||
AzureKeyCredential keyCredential = new(TestConfiguration.AzureAISearch.ApiKey); | ||
|
||
// Create kernel builder | ||
IKernelBuilder kernelBuilder = Kernel.CreateBuilder(); | ||
|
||
// SearchIndexClient from Azure .NET SDK to perform search operations. | ||
kernelBuilder.Services.AddSingleton<SearchIndexClient>((_) => new SearchIndexClient(endpoint, keyCredential)); | ||
|
||
// Custom AzureAISearchService to configure request parameters and make a request. | ||
kernelBuilder.Services.AddSingleton<IAzureAISearchService, AzureAISearchService>(); | ||
|
||
// Embedding generation service to convert string query to vector | ||
kernelBuilder.AddOpenAITextEmbeddingGeneration("text-embedding-ada-002", TestConfiguration.OpenAI.ApiKey); | ||
|
||
// Chat completion service to ask questions based on data from Azure AI Search index. | ||
kernelBuilder.AddOpenAIChatCompletion("gpt-4", TestConfiguration.OpenAI.ApiKey); | ||
|
||
// Register Azure AI Search Plugin | ||
kernelBuilder.Plugins.AddFromType<AzureAISearchPlugin>(); | ||
|
||
// Create kernel | ||
var kernel = kernelBuilder.Build(); | ||
|
||
// Query with index name | ||
// The final prompt will look like this "Emily and David are...(more text based on data). Who is David?". | ||
var result1 = await kernel.InvokePromptAsync( | ||
"{{search 'David' collection='index-1'}} Who is David?"); | ||
|
||
WriteLine(result1); | ||
|
||
// Query with index name and search fields. | ||
// Search fields are optional. Since one index may contain multiple searchable fields, | ||
// it's possible to specify which fields should be used during search for each request. | ||
var arguments = new KernelArguments { ["searchFields"] = JsonSerializer.Serialize(new List<string> { "vector" }) }; | ||
|
||
// The final prompt will look like this "Elara is...(more text based on data). Who is Elara?". | ||
var result2 = await kernel.InvokePromptAsync( | ||
"{{search 'Story' collection='index-2' searchFields=$searchFields}} Who is Elara?", | ||
arguments); | ||
|
||
WriteLine(result2); | ||
} | ||
|
||
public Example84_AzureAISearchPlugin(ITestOutputHelper output) : base(output) | ||
{ | ||
} | ||
|
||
#region Index Schema | ||
|
||
/// <summary> | ||
/// Custom index schema. It may contain any fields that exist in search index. | ||
/// </summary> | ||
private sealed class IndexSchema | ||
{ | ||
[JsonPropertyName("chunk_id")] | ||
public string ChunkId { get; set; } | ||
|
||
[JsonPropertyName("parent_id")] | ||
public string ParentId { get; set; } | ||
|
||
[JsonPropertyName("chunk")] | ||
public string Chunk { get; set; } | ||
|
||
[JsonPropertyName("title")] | ||
public string Title { get; set; } | ||
|
||
[JsonPropertyName("vector")] | ||
public ReadOnlyMemory<float> Vector { get; set; } | ||
} | ||
|
||
#endregion | ||
|
||
#region Azure AI Search Service | ||
|
||
/// <summary> | ||
/// Abstraction for Azure AI Search service. | ||
/// </summary> | ||
private interface IAzureAISearchService | ||
{ | ||
Task<string?> SearchAsync( | ||
string collectionName, | ||
ReadOnlyMemory<float> vector, | ||
List<string>? searchFields = null, | ||
CancellationToken cancellationToken = default); | ||
} | ||
|
||
/// <summary> | ||
/// Implementation of Azure AI Search service. | ||
/// </summary> | ||
private sealed class AzureAISearchService : IAzureAISearchService | ||
{ | ||
private readonly List<string> _defaultVectorFields = new() { "vector" }; | ||
|
||
private readonly SearchIndexClient _indexClient; | ||
|
||
public AzureAISearchService(SearchIndexClient indexClient) | ||
{ | ||
this._indexClient = indexClient; | ||
} | ||
|
||
public async Task<string?> SearchAsync( | ||
string collectionName, | ||
ReadOnlyMemory<float> vector, | ||
List<string>? searchFields = null, | ||
CancellationToken cancellationToken = default) | ||
{ | ||
// Get client for search operations | ||
SearchClient searchClient = this._indexClient.GetSearchClient(collectionName); | ||
|
||
// Use search fields passed from Plugin or default fields configured in this class. | ||
List<string> fields = searchFields is { Count: > 0 } ? searchFields : this._defaultVectorFields; | ||
|
||
// Configure request parameters | ||
VectorizedQuery vectorQuery = new(vector); | ||
fields.ForEach(field => vectorQuery.Fields.Add(field)); | ||
|
||
SearchOptions searchOptions = new() { VectorSearch = new() { Queries = { vectorQuery } } }; | ||
|
||
// Perform search request | ||
Response<SearchResults<IndexSchema>> response = await searchClient.SearchAsync<IndexSchema>(searchOptions, cancellationToken); | ||
|
||
List<IndexSchema> results = new(); | ||
|
||
// Collect search results | ||
await foreach (SearchResult<IndexSchema> result in response.Value.GetResultsAsync()) | ||
{ | ||
results.Add(result.Document); | ||
} | ||
|
||
// Return text from first result. | ||
// In real applications, the logic can check document score, sort and return top N results | ||
// or aggregate all results in one text. | ||
// The logic and decision which text data to return should be based on business scenario. | ||
return results.FirstOrDefault()?.Chunk; | ||
} | ||
} | ||
|
||
#endregion | ||
|
||
#region Azure AI Search SK Plugin | ||
|
||
/// <summary> | ||
/// Azure AI Search SK Plugin. | ||
/// It uses <see cref="ITextEmbeddingGenerationService"/> to convert string query to vector. | ||
/// It uses <see cref="IAzureAISearchService"/> to perform a request to Azure AI Search. | ||
/// </summary> | ||
private sealed class AzureAISearchPlugin | ||
{ | ||
private readonly ITextEmbeddingGenerationService _textEmbeddingGenerationService; | ||
private readonly IAzureAISearchService _searchService; | ||
|
||
public AzureAISearchPlugin( | ||
ITextEmbeddingGenerationService textEmbeddingGenerationService, | ||
IAzureAISearchService searchService) | ||
{ | ||
this._textEmbeddingGenerationService = textEmbeddingGenerationService; | ||
this._searchService = searchService; | ||
} | ||
|
||
[KernelFunction("Search")] | ||
public async Task<string> SearchAsync( | ||
string query, | ||
string collection, | ||
List<string>? searchFields = null, | ||
CancellationToken cancellationToken = default) | ||
{ | ||
// Convert string query to vector | ||
ReadOnlyMemory<float> embedding = await this._textEmbeddingGenerationService.GenerateEmbeddingAsync(query, cancellationToken: cancellationToken); | ||
|
||
// Perform search | ||
return await this._searchService.SearchAsync(collection, embedding, searchFields, cancellationToken) ?? string.Empty; | ||
} | ||
} | ||
|
||
#endregion | ||
} |
6954e97
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Hello @dmytrostruk, maybe this is a super basic question (I'm new to the topic) but how is this a plugin? Isn't a plugin a config json and a prompt txt?