in src/WebJobs.Extensions.OpenAI.CosmosDBSearch/CosmosDBSearchProvider.cs [108:168]
public async Task<SearchResponse> SearchAsync(SearchRequest request)
{
if (request.Embeddings.IsEmpty)
{
throw new ArgumentException("Embeddings must be provided.");
}
try
{
MongoClient cosmosClient = this.cosmosDBClients.GetOrAdd(
request.ConnectionInfo.ConnectionName,
_ => this.CreateMongoClient(request.ConnectionInfo.ConnectionName)
);
IMongoCollection<BsonDocument> collection = cosmosClient
.GetDatabase(this.databaseName)
.GetCollection<BsonDocument>(this.collectionName);
// Search Azure Cosmos DB for MongoDB vCore collection for similar embeddings and project fields.
BsonDocument[]? pipeline = null;
switch (this.cosmosDBSearchConfigOptions.Value.Kind)
{
case "vector-ivf":
pipeline = this.GetVectorIVFSearchPipeline(request);
break;
case "vector-hnsw":
pipeline = this.GetVectorHNSWSearchPipeline(request);
break;
case "vector-diskann":
pipeline = this.GetVectorDiskANNSearchPipeline(request);
break;
}
IAsyncCursor<BsonDocument> cursor = await collection.AggregateAsync<BsonDocument>(
pipeline
);
List<BsonDocument> documents = await cursor.ToListAsync();
List<OutputDocument> resultFromDocuments = documents
.ToList()
.ConvertAll(bsonDocument =>
BsonSerializer.Deserialize<OutputDocument>(bsonDocument)
);
List<SearchResult> searchResults = new();
foreach (OutputDocument doc in resultFromDocuments)
{
searchResults.Add(new SearchResult(doc.Title, doc.Text));
}
SearchResponse response = new(searchResults);
return response;
}
catch (MongoException ex)
{
this.logger.LogError(ex, "CosmosDBSearchProvider:SearchAsync error");
throw;
}
}