-
Notifications
You must be signed in to change notification settings - Fork 3.4k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
.Net: Added vector search implementation for Azure CosmosDB for Mongo…
…DB (#8887) ### 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. --> Related: #6522 - Implemented `VectorizedSearchAsync` method in Azure CosmosDB for MongoDB connector. - Added unit and integration tests. ### 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
f35d051
commit 8dc7dc3
Showing
5 changed files
with
414 additions
and
5 deletions.
There are no files selected for viewing
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
78 changes: 78 additions & 0 deletions
78
...ors.Memory.AzureCosmosDBMongoDB/AzureCosmosDBMongoDBVectorStoreCollectionSearchMapping.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,78 @@ | ||
// Copyright (c) Microsoft. All rights reserved. | ||
|
||
using Microsoft.SemanticKernel.Data; | ||
using MongoDB.Bson; | ||
|
||
namespace Microsoft.SemanticKernel.Connectors.AzureCosmosDBMongoDB; | ||
|
||
/// <summary> | ||
/// Contains mapping helpers to use when searching for documents using Azure CosmosDB MongoDB. | ||
/// </summary> | ||
internal sealed class AzureCosmosDBMongoDBVectorStoreCollectionSearchMapping | ||
{ | ||
/// <summary>Returns search part of the search query for <see cref="IndexKind.Hnsw"/> index kind.</summary> | ||
public static BsonDocument GetSearchQueryForHnswIndex<TVector>( | ||
TVector vector, | ||
string vectorPropertyName, | ||
int limit, | ||
int efSearch) | ||
{ | ||
return new BsonDocument | ||
{ | ||
{ "$search", | ||
new BsonDocument | ||
{ | ||
{ "cosmosSearch", | ||
new BsonDocument | ||
{ | ||
{ "vector", BsonArray.Create(vector) }, | ||
{ "path", vectorPropertyName }, | ||
{ "k", limit }, | ||
{ "efSearch", efSearch } | ||
} | ||
} | ||
} | ||
} | ||
}; | ||
} | ||
|
||
/// <summary>Returns search part of the search query for <see cref="IndexKind.IvfFlat"/> index kind.</summary> | ||
public static BsonDocument GetSearchQueryForIvfIndex<TVector>( | ||
TVector vector, | ||
string vectorPropertyName, | ||
int limit) | ||
{ | ||
return new BsonDocument | ||
{ | ||
{ "$search", | ||
new BsonDocument | ||
{ | ||
{ "cosmosSearch", | ||
new BsonDocument | ||
{ | ||
{ "vector", BsonArray.Create(vector) }, | ||
{ "path", vectorPropertyName }, | ||
{ "k", limit }, | ||
} | ||
}, | ||
{ "returnStoredSource", true } | ||
} | ||
} | ||
}; | ||
} | ||
|
||
/// <summary>Returns projection part of the search query to return similarity score together with document.</summary> | ||
public static BsonDocument GetProjectionQuery(string scorePropertyName, string documentPropertyName) | ||
{ | ||
return new BsonDocument | ||
{ | ||
{ "$project", | ||
new BsonDocument | ||
{ | ||
{ scorePropertyName, new BsonDocument { { "$meta", "searchScore" } } }, | ||
{ documentPropertyName, "$$ROOT" } | ||
} | ||
} | ||
}; | ||
} | ||
} |
Oops, something went wrong.