DEV Community

Johnny Z
Johnny Z

Posted on

Kernel Memory with Cosmos DB for NoSQL vector search.

Officially announced in Microsoft Build 2024, Cosmos DB for NoSQL now support vector search. It also means Kernel Memory can be integrated with Cosmos DB for NoSQL.

Enable Cosmos DB for NoSQL to support vector search.

Azure CosmosDB configurations

Implement IMemoryDb for kernel memory with cosmos client

The key is VectorDistance function to match against embeddings.

var sql =
    $"""
    SELECT Top @topN
      x.id, x.tags, x.payload, x.embedding, x.similarityScore
    FROM (
      SELECT
        c.id, c.tags, c.payload, c.embedding,VectorDistance(c.embedding, @embedding) AS similarityScore 
      FROM
        c
    ) AS x
    WHERE x.similarityScore > @similarityScore
    ORDER BY x.similarityScore desc
    """;

var queryDefinition = new QueryDefinition(sql)
    .WithParameter("@topN", limit)
    .WithParameter("@embedding", textEmbedding.Data)
    .WithParameter("@similarityScore", minRelevance);

// Index name as cosmos container name
var feedIterator = _cosmosClient
    .GetDatabase(DatabaseName)
    .GetContainer(index)
    .GetItemQueryIterator<MemoryRecordResult>(queryDefinition)
Enter fullscreen mode Exit fullscreen mode

Sample code here

Please feel free to reach out on twitter @roamingcode

Top comments (0)