Skip to content

Commit

Permalink
Improve ACS memory connector performance/cost (#776)
Browse files Browse the repository at this point in the history 
### Motivation and Context

* Checking if an index exists adds extra latency the first time an index
is being used.
* Some indexing can be removed to reduce the quota used, and CPU usage
on the service side.


### Description

* Create indexes/collections only on write only when an index doesn't
exist.
* Remove some indexing from the record schema
  • Loading branch information
@dluc
dluc committed May 3, 2023
1 parent 34a4e6b commit c6dec8e
Show file tree
Hide file tree
Showing 2 changed files with 94 additions and 44 deletions.
102 changes: 62 additions & 40 deletions dotnet/src/Connectors/Connectors.Memory.AzureCognitiveSearch/AzureCognitiveSearchMemory.cs
View file
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
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
}

0 comments on commit c6dec8e

Please sign in to comment.