Skip to content

Commit

Permalink
.Net: Added example of Azure AI Search as Plugin with custom schema a…
Browse files Browse the repository at this point in the history
…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
dmytrostruk committed Feb 21, 2024
1 parent 3b14d7c commit 6954e97
Showing 1 changed file with 207 additions and 0 deletions.
207 changes: 207 additions & 0 deletions dotnet/samples/KernelSyntaxExamples/Example84_AzureAISearchPlugin.cs
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
}

1 comment on commit 6954e97

@DavidCamposDSB
Copy link

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?

Please sign in to comment.