Skip to content

Commit 85d420f

Browse files
dmytrostrukdmytrostruk
authored
Chroma memory store - C# implementation (#1634)
### 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. --> This PR contains changes for Chroma C# memory store. For more information about Chroma, see: https://docs.trychroma.com/getting-started. To run integration tests, local Chroma server should be up and running. To start local Chroma server for testing (this information is also included in [README](https://github.com/microsoft/semantic-kernel/blob/0fb77e0c0f347e3f22d79b3c99e7f84d35340462/dotnet/src/Connectors/Connectors.Memory.Chroma/README.md)): 1. Clone Chroma: ```bash git clone https://github.com/chroma-core/chroma.git cd chroma ``` 2. Run local Chroma server with Docker within Chroma repository root: ```bash docker-compose up -d --build ``` 3. Use Semantic Kernel with Chroma, using server local endpoint `http://localhost:8000`: ```csharp const string endpoint = "http://localhost:8000"; ChromaMemoryStore memoryStore = new(endpoint); IKernel kernel = Kernel.Builder .WithLogger(logger) .WithOpenAITextEmbeddingGenerationService("text-embedding-ada-002", "OPENAI_API_KEY") .WithMemoryStorage(memoryStore) //.WithChromaMemoryStore(endpoint) // This method offers an alternative approach to registering Chroma memory store. .Build(); ``` ### Description <!-- Describe your changes, the overall approach, the underlying design. These notes will help understanding how your code works. Thanks! --> Main public classes: 1. `ChromaMemoryStore` - implementation of `IMemoryStore` interface. Uses `ChromaClient` to communicate with Chroma API. Throws `ChromaMemoryStoreException` in case of issues. 2. `ChromaClient` - responsible for communication with Chroma API. Throws `ChromaClientException` in case of issues. Chroma entity models: 1. `ChromaCollectionModel` - defines Chroma collection. 2. `ChromaEmbeddingsModel` - defines Chroma embeddings, constructed from `MemoryRecord`. 3. `ChromaQueryResultModel` - defines Chroma search result set. ### Contribution Checklist <!-- Before submitting this PR, please make sure: --> - [x] The code builds clean without any errors or warnings - [x] The PR follows SK Contribution Guidelines (https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md) - [x] The code follows the .NET coding conventions (https://learn.microsoft.com/dotnet/csharp/fundamentals/coding-style/coding-conventions) verified with `dotnet format` - [x] All unit tests pass, and I have added new tests where possible - [x] I didn't break anyone 😄
1 parent 527d378 commit 85d420f

28 files changed

+2047
-0
lines changed

‎dotnet/SK-dotnet.sln‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,8 @@ Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Connectors.Memory.Postgres"
7474
EndProject
7575
Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Connectors.Memory.Redis", "src\Connectors\Connectors.Memory.Redis\Connectors.Memory.Redis.csproj", "{3720F5ED-FB4D-485E-8A93-CDE60DEF0805}"
7676
EndProject
77+
Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "Connectors.Memory.Chroma", "src\Connectors\Connectors.Memory.Chroma\Connectors.Memory.Chroma.csproj", "{185E0CE8-C2DA-4E4C-A491-E8EB40316315}"
78+
EndProject
7779
Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Connectors.AI.OpenAI", "src\Connectors\Connectors.AI.OpenAI\Connectors.AI.OpenAI.csproj", "{AFA81EB7-F869-467D-8A90-744305D80AAC}"
7880
EndProject
7981
Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "SemanticKernel.Abstractions", "src\SemanticKernel.Abstractions\SemanticKernel.Abstractions.csproj", "{627742DB-1E52-468A-99BD-6FF1A542D25B}"
@@ -355,6 +357,12 @@ Global
355357
{50FAE231-6F24-4779-9D02-12ABBC9A49E2}.Publish|Any CPU.Build.0 = Publish|Any CPU
356358
{50FAE231-6F24-4779-9D02-12ABBC9A49E2}.Release|Any CPU.ActiveCfg = Release|Any CPU
357359
{50FAE231-6F24-4779-9D02-12ABBC9A49E2}.Release|Any CPU.Build.0 = Release|Any CPU
360+
{185E0CE8-C2DA-4E4C-A491-E8EB40316315}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
361+
{185E0CE8-C2DA-4E4C-A491-E8EB40316315}.Debug|Any CPU.Build.0 = Debug|Any CPU
362+
{185E0CE8-C2DA-4E4C-A491-E8EB40316315}.Publish|Any CPU.ActiveCfg = Debug|Any CPU
363+
{185E0CE8-C2DA-4E4C-A491-E8EB40316315}.Publish|Any CPU.Build.0 = Debug|Any CPU
364+
{185E0CE8-C2DA-4E4C-A491-E8EB40316315}.Release|Any CPU.ActiveCfg = Release|Any CPU
365+
{185E0CE8-C2DA-4E4C-A491-E8EB40316315}.Release|Any CPU.Build.0 = Release|Any CPU
358366
{0D0C4DAD-E6BC-4504-AE3A-EEA4E35920C1}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
359367
{0D0C4DAD-E6BC-4504-AE3A-EEA4E35920C1}.Debug|Any CPU.Build.0 = Debug|Any CPU
360368
{0D0C4DAD-E6BC-4504-AE3A-EEA4E35920C1}.Publish|Any CPU.ActiveCfg = Debug|Any CPU
@@ -418,6 +426,7 @@ Global
418426
{B00AD427-0047-4850-BEF9-BA8237EA9D8B} = {958AD708-F048-4FAF-94ED-D2F2B92748B9}
419427
{DB950192-30F1-48B1-88D7-F43FECCA1A1C} = {958AD708-F048-4FAF-94ED-D2F2B92748B9}
420428
{1C19D805-3573-4477-BF07-40180FCDE1BD} = {958AD708-F048-4FAF-94ED-D2F2B92748B9}
429+
{185E0CE8-C2DA-4E4C-A491-E8EB40316315} = {0247C2C9-86C3-45BA-8873-28B0948EDC0C}
421430
{0D0C4DAD-E6BC-4504-AE3A-EEA4E35920C1} = {9ECD1AA0-75B3-4E25-B0B5-9F0945B64974}
422431
{E6EDAB8F-3406-4DBF-9AAB-DF40DC2CA0FA} = {FA3720F1-C99A-49B2-9577-A940257098BF}
423432
EndGlobalSection

‎dotnet/samples/KernelSyntaxExamples/Example50_Chroma.cs‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
// Copyright (c) Microsoft. All rights reserved.
2+
3+
using System;
4+
using System.Threading.Tasks;
5+
using Microsoft.SemanticKernel;
6+
using Microsoft.SemanticKernel.Connectors.Memory.Chroma;
7+
using Microsoft.SemanticKernel.Memory;
8+
using RepoUtils;
9+
10+
// ReSharper disable once InconsistentNaming
11+
public static class Example50_Chroma
12+
{
13+
private const string MemoryCollectionName = "chroma-test";
14+
15+
public static async Task RunAsync()
16+
{
17+
string endpoint = Env.Var("CHROMA_ENDPOINT");
18+
19+
var memoryStore = new ChromaMemoryStore(endpoint);
20+
21+
IKernel kernel = Kernel.Builder
22+
.WithLogger(ConsoleLogger.Log)
23+
.WithOpenAITextCompletionService("text-davinci-003", Env.Var("OPENAI_API_KEY"))
24+
.WithOpenAITextEmbeddingGenerationService("text-embedding-ada-002", Env.Var("OPENAI_API_KEY"))
25+
.WithMemoryStorage(memoryStore)
26+
//.WithChromaMemoryStore(endpoint) // This method offers an alternative approach to registering Chroma memory store.
27+
.Build();
28+
29+
Console.WriteLine("== Printing Collections in DB ==");
30+
var collections = memoryStore.GetCollectionsAsync();
31+
await foreach (var collection in collections)
32+
{
33+
Console.WriteLine(collection);
34+
}
35+
36+
Console.WriteLine("== Adding Memories ==");
37+
38+
var key1 = await kernel.Memory.SaveInformationAsync(MemoryCollectionName, id: Guid.NewGuid().ToString(), text: "british short hair");
39+
var key2 = await kernel.Memory.SaveInformationAsync(MemoryCollectionName, id: Guid.NewGuid().ToString(), text: "orange tabby");
40+
var key3 = await kernel.Memory.SaveInformationAsync(MemoryCollectionName, id: Guid.NewGuid().ToString(), text: "norwegian forest cat");
41+
42+
Console.WriteLine("== Printing Collections in DB ==");
43+
collections = memoryStore.GetCollectionsAsync();
44+
await foreach (var collection in collections)
45+
{
46+
Console.WriteLine(collection);
47+
}
48+
49+
Console.WriteLine("== Retrieving Memories Through the Kernel ==");
50+
MemoryQueryResult? lookup = await kernel.Memory.GetAsync(MemoryCollectionName, key1);
51+
Console.WriteLine(lookup != null ? lookup.Metadata.Text : "ERROR: memory not found");
52+
53+
Console.WriteLine("== Similarity Searching Memories: My favorite color is orange ==");
54+
var searchResults = kernel.Memory.SearchAsync(MemoryCollectionName, "My favorite color is orange", limit: 3, minRelevanceScore: 0.6);
55+
56+
await foreach (var item in searchResults)
57+
{
58+
Console.WriteLine(item.Metadata.Text + " : " + item.Relevance);
59+
}
60+
61+
Console.WriteLine("== Removing Collection {0} ==", MemoryCollectionName);
62+
await memoryStore.DeleteCollectionAsync(MemoryCollectionName);
63+
64+
Console.WriteLine("== Printing Collections in DB ==");
65+
await foreach (var collection in collections)
66+
{
67+
Console.WriteLine(collection);
68+
}
69+
}
70+
}

‎dotnet/samples/KernelSyntaxExamples/KernelSyntaxExamples.csproj‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
<ProjectReference Include="..\..\src\Connectors\Connectors.AI.OpenAI\Connectors.AI.OpenAI.csproj" />
3030
<ProjectReference Include="..\..\src\Connectors\Connectors.AI.HuggingFace\Connectors.AI.HuggingFace.csproj" />
3131
<ProjectReference Include="..\..\src\Connectors\Connectors.Memory.AzureCognitiveSearch\Connectors.Memory.AzureCognitiveSearch.csproj" />
32+
<ProjectReference Include="..\..\src\Connectors\Connectors.Memory.Chroma\Connectors.Memory.Chroma.csproj" />
3233
<ProjectReference Include="..\..\src\Connectors\Connectors.Memory.Postgres\Connectors.Memory.Postgres.csproj" />
3334
<ProjectReference Include="..\..\src\Connectors\Connectors.Memory.Weaviate\Connectors.Memory.Weaviate.csproj" />
3435
<ProjectReference Include="..\..\src\Connectors\Connectors.Memory.Redis\Connectors.Memory.Redis.csproj" />

‎dotnet/samples/KernelSyntaxExamples/Program.cs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -154,5 +154,8 @@ public static async Task Main()
154154

155155
await Example49_LogitBias.RunAsync();
156156
Console.WriteLine("== DONE ==");
157+
158+
await Example50_Chroma.RunAsync();
159+
Console.WriteLine("== DONE ==");
157160
}
158161
}

‎dotnet/src/Connectors/Connectors.Memory.Chroma/ChromaClient.cs‎

Lines changed: 207 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,207 @@
1+
// Copyright (c) Microsoft. All rights reserved.
2+
3+
using System;
4+
using System.Collections.Generic;
5+
using System.Net.Http;
6+
using System.Runtime.CompilerServices;
7+
using System.Text;
8+
using System.Text.Json;
9+
using System.Threading;
10+
using System.Threading.Tasks;
11+
using Microsoft.Extensions.Logging;
12+
using Microsoft.Extensions.Logging.Abstractions;
13+
using Microsoft.SemanticKernel.Connectors.Memory.Chroma.Http.ApiSchema;
14+
using Microsoft.SemanticKernel.Connectors.Memory.Chroma.Http.ApiSchema.Internal;
15+
16+
namespace Microsoft.SemanticKernel.Connectors.Memory.Chroma;
17+
18+
/// <summary>
19+
/// An implementation of a client for the Chroma Vector DB. This class is used to
20+
/// create, delete, and get embeddings data from Chroma Vector DB instance.
21+
/// </summary>
22+
#pragma warning disable CA1001 // Types that own disposable fields should be disposable. Explanation - In this case, there is no need to dispose because either the NonDisposableHttpClientHandler or a custom HTTP client is being used.
23+
public class ChromaClient : IChromaClient
24+
#pragma warning restore CA1001 // Types that own disposable fields should be disposable. Explanation - In this case, there is no need to dispose because either the NonDisposableHttpClientHandler or a custom HTTP client is being used.
25+
{
26+
/// <summary>
27+
/// Initializes a new instance of the <see cref="ChromaClient"/> class.
28+
/// </summary>
29+
/// <param name="endpoint">Chroma server endpoint URL.</param>
30+
/// <param name="logger">Optional logger instance.</param>
31+
public ChromaClient(string endpoint, ILogger? logger = null)
32+
{
33+
this._httpClient = new HttpClient(NonDisposableHttpClientHandler.Instance, disposeHandler: false);
34+
this._endpoint = endpoint;
35+
this._logger = logger ?? NullLogger<ChromaClient>.Instance;
36+
}
37+
38+
/// <summary>
39+
/// Initializes a new instance of the <see cref="ChromaClient"/> class.
40+
/// </summary>
41+
/// <param name="httpClient">The <see cref="HttpClient"/> instance used for making HTTP requests.</param>
42+
/// <param name="endpoint">Chroma server endpoint URL.</param>
43+
/// <param name="logger">Optional logger instance.</param>
44+
/// <exception cref="ChromaClientException">Occurs when <see cref="HttpClient"/> doesn't have base address and endpoint parameter is not provided.</exception>
45+
public ChromaClient(HttpClient httpClient, string? endpoint = null, ILogger? logger = null)
46+
{
47+
if (string.IsNullOrEmpty(httpClient.BaseAddress?.AbsoluteUri) && string.IsNullOrEmpty(endpoint))
48+
{
49+
throw new ChromaClientException("The HttpClient BaseAddress and endpoint are both null or empty. Please ensure at least one is provided.");
50+
}
51+
52+
this._httpClient = httpClient;
53+
this._endpoint = endpoint;
54+
this._logger = logger ?? NullLogger<ChromaClient>.Instance;
55+
}
56+
57+
/// <inheritdoc />
58+
public async Task CreateCollectionAsync(string collectionName, CancellationToken cancellationToken = default)
59+
{
60+
this._logger.LogDebug("Creating collection {0}", collectionName);
61+
62+
using var request = CreateCollectionRequest.Create(collectionName).Build();
63+
64+
await this.ExecuteHttpRequestAsync(request, cancellationToken).ConfigureAwait(false);
65+
}
66+
67+
/// <inheritdoc />
68+
public async Task<ChromaCollectionModel?> GetCollectionAsync(string collectionName, CancellationToken cancellationToken = default)
69+
{
70+
this._logger.LogDebug("Getting collection {0}", collectionName);
71+
72+
using var request = GetCollectionRequest.Create(collectionName).Build();
73+
74+
(HttpResponseMessage response, string responseContent) = await this.ExecuteHttpRequestAsync(request, cancellationToken).ConfigureAwait(false);
75+
76+
var collection = JsonSerializer.Deserialize<ChromaCollectionModel>(responseContent);
77+
78+
return collection;
79+
}
80+
81+
/// <inheritdoc />
82+
public async Task DeleteCollectionAsync(string collectionName, CancellationToken cancellationToken = default)
83+
{
84+
this._logger.LogDebug("Deleting collection {0}", collectionName);
85+
86+
using var request = DeleteCollectionRequest.Create(collectionName).Build();
87+
88+
await this.ExecuteHttpRequestAsync(request, cancellationToken).ConfigureAwait(false);
89+
}
90+
91+
/// <inheritdoc />
92+
public async IAsyncEnumerable<string> ListCollectionsAsync([EnumeratorCancellation] CancellationToken cancellationToken = default)
93+
{
94+
this._logger.LogDebug("Listing collections");
95+
96+
using var request = ListCollectionsRequest.Create().Build();
97+
98+
(HttpResponseMessage response, string responseContent) = await this.ExecuteHttpRequestAsync(request, cancellationToken).ConfigureAwait(false);
99+
100+
var collections = JsonSerializer.Deserialize<List<ChromaCollectionModel>>(responseContent);
101+
102+
foreach (var collection in collections!)
103+
{
104+
yield return collection.Name;
105+
}
106+
}
107+
108+
/// <inheritdoc />
109+
public async Task AddEmbeddingsAsync(string collectionId, string[] ids, float[][] embeddings, object[]? metadatas = null, CancellationToken cancellationToken = default)
110+
{
111+
this._logger.LogDebug("Adding embeddings to collection with id: {0}", collectionId);
112+
113+
using var request = AddEmbeddingsRequest.Create(collectionId, ids, embeddings, metadatas).Build();
114+
115+
await this.ExecuteHttpRequestAsync(request, cancellationToken).ConfigureAwait(false);
116+
}
117+
118+
/// <inheritdoc />
119+
public async Task<ChromaEmbeddingsModel> GetEmbeddingsAsync(string collectionId, string[] ids, string[]? include = null, CancellationToken cancellationToken = default)
120+
{
121+
this._logger.LogDebug("Getting embeddings from collection with id: {0}", collectionId);
122+
123+
using var request = GetEmbeddingsRequest.Create(collectionId, ids, include).Build();
124+
125+
(HttpResponseMessage response, string responseContent) = await this.ExecuteHttpRequestAsync(request, cancellationToken).ConfigureAwait(false);
126+
127+
var embeddings = JsonSerializer.Deserialize<ChromaEmbeddingsModel>(responseContent);
128+
129+
return embeddings ?? new ChromaEmbeddingsModel();
130+
}
131+
132+
/// <inheritdoc />
133+
public async Task DeleteEmbeddingsAsync(string collectionId, string[] ids, CancellationToken cancellationToken = default)
134+
{
135+
this._logger.LogDebug("Deleting embeddings from collection with id: {0}", collectionId);
136+
137+
using var request = DeleteEmbeddingsRequest.Create(collectionId, ids).Build();
138+
139+
await this.ExecuteHttpRequestAsync(request, cancellationToken).ConfigureAwait(false);
140+
}
141+
142+
/// <inheritdoc />
143+
public async Task<ChromaQueryResultModel> QueryEmbeddingsAsync(string collectionId, float[][] queryEmbeddings, int nResults, string[]? include = null, CancellationToken cancellationToken = default)
144+
{
145+
this._logger.LogDebug("Query embeddings in collection with id: {0}", collectionId);
146+
147+
using var request = QueryEmbeddingsRequest.Create(collectionId, queryEmbeddings, nResults, include).Build();
148+
149+
(HttpResponseMessage response, string responseContent) = await this.ExecuteHttpRequestAsync(request, cancellationToken).ConfigureAwait(false);
150+
151+
var queryResult = JsonSerializer.Deserialize<ChromaQueryResultModel>(responseContent);
152+
153+
return queryResult ?? new ChromaQueryResultModel();
154+
}
155+
156+
#region private ================================================================================
157+
158+
private const string ApiRoute = "api/v1/";
159+
160+
private readonly ILogger _logger;
161+
private readonly HttpClient _httpClient;
162+
private readonly string? _endpoint = null;
163+
164+
private async Task<(HttpResponseMessage response, string responseContent)> ExecuteHttpRequestAsync(
165+
HttpRequestMessage request,
166+
CancellationToken cancellationToken = default)
167+
{
168+
string endpoint = this._endpoint ?? this._httpClient.BaseAddress.ToString();
169+
endpoint = this.SanitizeEndpoint(endpoint);
170+
171+
string operationName = request.RequestUri.ToString();
172+
173+
request.RequestUri = new Uri(new Uri(endpoint), operationName);
174+
175+
HttpResponseMessage response = await this._httpClient.SendAsync(request, cancellationToken).ConfigureAwait(false);
176+
177+
string responseContent = await response.Content.ReadAsStringAsync().ConfigureAwait(false);
178+
179+
try
180+
{
181+
response.EnsureSuccessStatusCode();
182+
}
183+
catch (HttpRequestException e)
184+
{
185+
this._logger.LogError(e, "{0} {1} operation failed: {2}, {3}", request.Method.Method, operationName, e.Message, responseContent);
186+
throw new ChromaClientException($"{request.Method.Method} {operationName} operation failed: {e.Message}, {responseContent}", e);
187+
}
188+
189+
return (response, responseContent);
190+
}
191+
192+
private string SanitizeEndpoint(string endpoint)
193+
{
194+
StringBuilder builder = new(endpoint);
195+
196+
if (!endpoint.EndsWith("/", StringComparison.Ordinal))
197+
{
198+
builder.Append('/');
199+
}
200+
201+
builder.Append(ApiRoute);
202+
203+
return builder.ToString();
204+
}
205+
206+
#endregion
207+
}

‎dotnet/src/Connectors/Connectors.Memory.Chroma/ChromaClientException.cs‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
// Copyright (c) Microsoft. All rights reserved.
2+
3+
using System;
4+
using System.Globalization;
5+
6+
namespace Microsoft.SemanticKernel.Connectors.Memory.Chroma;
7+
8+
/// <summary>
9+
/// Exception to identify issues in <see cref="ChromaClient"/> class.
10+
/// </summary>
11+
public class ChromaClientException : Exception
12+
{
13+
private const string CollectionDoesNotExistErrorFormat = "Collection {0} does not exist";
14+
private const string DeleteNonExistentCollectionErrorMessage = "list index out of range";
15+
16+
/// <summary>
17+
/// Initializes a new instance of the <see cref="ChromaClientException"/> class.
18+
/// </summary>
19+
public ChromaClientException() : base()
20+
{
21+
}
22+
23+
/// <summary>
24+
/// Initializes a new instance of the <see cref="ChromaClientException"/> class.
25+
/// </summary>
26+
/// <param name="message">The message that describes the error.</param>
27+
public ChromaClientException(string message) : base(message)
28+
{
29+
}
30+
31+
/// <summary>
32+
/// Initializes a new instance of the <see cref="ChromaClientException"/> class.
33+
/// </summary>
34+
/// <param name="message">The message that describes the error.</param>
35+
/// <param name="innerException">Instance of inner exception.</param>
36+
public ChromaClientException(string message, Exception innerException) : base(message, innerException)
37+
{
38+
}
39+
40+
/// <summary>
41+
/// Checks if Chroma API error means that collection does not exist.
42+
/// </summary>
43+
/// <param name="collectionName">Collection name.</param>
44+
public bool CollectionDoesNotExistException(string collectionName) =>
45+
this.Message.Contains(string.Format(CultureInfo.InvariantCulture, CollectionDoesNotExistErrorFormat, collectionName));
46+
47+
/// <summary>
48+
/// Checks if Chroma API error means that there was an attempt to delete non-existent collection.
49+
/// </summary>
50+
public bool DeleteNonExistentCollectionException() =>
51+
this.Message.Contains(DeleteNonExistentCollectionErrorMessage);
52+
}

0 commit comments

Comments
 (0)