Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Improve ACS memory connector performance/cost #776

Merged
merged 1 commit into from
May 3, 2023
Merged

Improve ACS memory connector performance/cost #776

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
102 changes: 62 additions & 40 deletions dotnet/src/Connectors/Connectors.Memory.AzureCognitiveSearch/AzureCognitiveSearchMemory.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -106,12 +106,22 @@ public AzureCognitiveSearchMemory(string endpoint, TokenCredential credentials)
{
collection = NormalizeIndexName(collection);

var client = await this.GetSearchClientAsync(collection, cancellationToken).ConfigureAwait(false);
var client = this.GetSearchClient(collection);

Response<AzureCognitiveSearchRecord>? result = await client.GetDocumentAsync<AzureCognitiveSearchRecord>(
EncodeId(key), cancellationToken: cancellationToken).ConfigureAwait(false);
Response<AzureCognitiveSearchRecord>? result;
try
{
result = await client
.GetDocumentAsync<AzureCognitiveSearchRecord>(EncodeId(key), cancellationToken: cancellationToken)
.ConfigureAwait(false);
}
catch (RequestFailedException e) when (e.Status == 404)
{
// Index not found, no data to return
return null;
}

if (result == null || result.Value == null)
if (result?.Value == null)
{
throw new AzureCognitiveSearchMemoryException("Memory read returned null");
}
Expand All @@ -130,25 +140,37 @@ public AzureCognitiveSearchMemory(string endpoint, TokenCredential credentials)
{
collection = NormalizeIndexName(collection);

var client = await this.GetSearchClientAsync(collection, cancellationToken).ConfigureAwait(false);
var client = this.GetSearchClient(collection);

// TODO: use vectors
var options = new SearchOptions
{
QueryType = SearchQueryType.Semantic,
SemanticConfigurationName = "default",
QueryLanguage = "en-us", // TODO: this shouldn't be required
QueryLanguage = "en-us",
Size = limit,
};

Response<SearchResults<AzureCognitiveSearchRecord>>? searchResult = await client
.SearchAsync<AzureCognitiveSearchRecord>(query, options, cancellationToken: cancellationToken)
.ConfigureAwait(false);
Response<SearchResults<AzureCognitiveSearchRecord>>? searchResult = null;
try
{
searchResult = await client
.SearchAsync<AzureCognitiveSearchRecord>(query, options, cancellationToken: cancellationToken)
.ConfigureAwait(false);
}
catch (RequestFailedException e) when (e.Status == 404)
{
// Index not found, no data to return
}

await foreach (SearchResult<AzureCognitiveSearchRecord>? doc in searchResult.Value.GetResultsAsync())
if (searchResult != null)
{
if (doc.RerankerScore < minRelevanceScore) { break; }
await foreach (SearchResult<AzureCognitiveSearchRecord>? doc in searchResult.Value.GetResultsAsync())
{
if (doc.RerankerScore < minRelevanceScore) { break; }

yield return new MemoryQueryResult(ToMemoryRecordMetadata(doc.Document), doc.RerankerScore ?? 1, null);
yield return new MemoryQueryResult(ToMemoryRecordMetadata(doc.Document), doc.RerankerScore ?? 1, null);
}
}
}

Expand All @@ -159,9 +181,15 @@ public async Task RemoveAsync(string collection, string key, CancellationToken c

var records = new List<AzureCognitiveSearchRecord> { new() { Id = EncodeId(key) } };

var client = await this.GetSearchClientAsync(collection, cancellationToken).ConfigureAwait(false);

await client.DeleteDocumentsAsync(records, cancellationToken: cancellationToken).ConfigureAwait(false);
var client = this.GetSearchClient(collection);
try
{
await client.DeleteDocumentsAsync(records, cancellationToken: cancellationToken).ConfigureAwait(false);
}
catch (RequestFailedException e) when (e.Status == 404)
{
// Index not found, no data to delete
}
}

/// <inheritdoc />
Expand Down Expand Up @@ -192,30 +220,12 @@ await foreach (var index in indexes)

/// <summary>
/// Get a search client for the index specified.
/// Note: the index might not exist, but we avoid checking everytime and the extra latency.
/// </summary>
/// <param name="indexName">Index name</param>
/// <param name="cancellationToken">Task cancellation token</param>
/// <returns>Search client ready to read/write</returns>
private async Task<SearchClient> GetSearchClientAsync(
string indexName,
CancellationToken cancellationToken = default)
private SearchClient GetSearchClient(string indexName)
{
Response<SearchIndex>? existingIndex = null;
try
{
// Search the index
existingIndex = await this._adminClient.GetIndexAsync(indexName, cancellationToken).ConfigureAwait(false);
}
catch (RequestFailedException e) when (e.Status == 404)
{
}

// Create the index if it doesn't exist
if (existingIndex == null || existingIndex.Value == null)
{
await this.CreateIndexAsync(indexName, cancellationToken).ConfigureAwait(false);
}

// Search an available client from the local cache
if (!this._clientsByIndex.TryGetValue(indexName, out SearchClient client))
{
Expand Down Expand Up @@ -243,6 +253,7 @@ await foreach (var index in indexes)
{
Configurations =
{
// TODO: replace with vector search
new SemanticConfiguration("default", new PrioritizedFields
{
TitleField = new SemanticField { FieldName = "Description" },
Expand All @@ -264,12 +275,23 @@ await foreach (var index in indexes)
AzureCognitiveSearchRecord record,
CancellationToken cancellationToken = default)
{
var client = await this.GetSearchClientAsync(indexName, cancellationToken).ConfigureAwait(false);
var client = this.GetSearchClient(indexName);

Task<Response<IndexDocumentsResult>> UpsertCode() => client
.MergeOrUploadDocumentsAsync(new List<AzureCognitiveSearchRecord> { record },
new IndexDocumentsOptions { ThrowOnAnyError = true },
cancellationToken);

Response<IndexDocumentsResult>? result = await client.MergeOrUploadDocumentsAsync(
new List<AzureCognitiveSearchRecord> { record },
new IndexDocumentsOptions { ThrowOnAnyError = true },
cancellationToken).ConfigureAwait(false);
Response<IndexDocumentsResult>? result;
try
{
result = await UpsertCode().ConfigureAwait(false);
}
catch (RequestFailedException e) when (e.Status == 404)
{
await this.CreateIndexAsync(indexName, cancellationToken).ConfigureAwait(false);
result = await UpsertCode().ConfigureAwait(false);
}

if (result == null || result.Value.Results.Count == 0)
{
Expand Down
36 changes: 32 additions & 4 deletions dotnet/src/Connectors/Connectors.Memory.AzureCognitiveSearch/AzureCognitiveSearchRecord.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,23 +5,51 @@

namespace Microsoft.SemanticKernel.Connectors.Memory.AzureCognitiveSearch;

/// <summary>
/// Azure Cognitive Search record and index definition.
/// Note: once defined, index cannot be modified.
/// </summary>
public class AzureCognitiveSearchRecord
{
[SimpleField(IsKey = true, IsFilterable = true)]
/// <summary>
/// Record Id.
/// The record is not filterable to save quota, also SK uses only semantic search.
/// </summary>
[SimpleField(IsKey = true, IsFilterable = false)]
public string Id { get; set; } = string.Empty;

[SearchableField(AnalyzerName = LexicalAnalyzerName.Values.EnMicrosoft)]
/// <summary>
/// Content is stored here.
/// </summary>
[SearchableField(AnalyzerName = LexicalAnalyzerName.Values.StandardLucene)]
public string? Text { get; set; } = string.Empty;

[SearchableField(AnalyzerName = LexicalAnalyzerName.Values.EnMicrosoft)]
/// <summary>
/// Optional description of the content, e.g. a title. This can be useful when
/// indexing external data without pulling in the entire content.
/// </summary>
[SearchableField(AnalyzerName = LexicalAnalyzerName.Values.StandardLucene)]
public string? Description { get; set; } = string.Empty;

[SearchableField(AnalyzerName = LexicalAnalyzerName.Values.EnMicrosoft)]
/// <summary>
/// Additional metadata. Currently this is a string, where you could store serialized data as JSON.
/// In future the design might change to allow storing named values and leverage filters.
/// </summary>
[SearchableField(AnalyzerName = LexicalAnalyzerName.Values.StandardLucene)]
public string? AdditionalMetadata { get; set; } = string.Empty;

/// <summary>
/// Name of the external source, in cases where the content and the Id are
/// referenced to external information.
/// </summary>
[SimpleField(IsFilterable = false)]
public string ExternalSourceName { get; set; } = string.Empty;

/// <summary>
/// Whether the record references external information.
/// </summary>
[SimpleField(IsFilterable = false)]
public bool IsReference { get; set; } = false;

// TODO: add one more field with the vector, float array, mark it as searchable
}