microsoft/semantic-kernel
1
0
Fork 0
mirror of https://github.com/microsoft/semantic-kernel.git synced 2026-10-12 02:19:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Chroma memory store - C# implementation (#1634)

Browse Source 
### 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 😄
This commit is contained in:
Dmytro Struk
2023-06-26 19:24:30 +00:00
committed by GitHub
parent 527d3782f8
commit 85d420f77c
28 changed files with 2047 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dotnet/SK-dotnet.sln
+9
View File
@@ -74,6 +74,8 @@ Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Connectors.Memory.Postgres"
EndProject
Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Connectors.Memory.Redis", "src\Connectors\Connectors.Memory.Redis\Connectors.Memory.Redis.csproj", "{3720F5ED-FB4D-485E-8A93-CDE60DEF0805}"
EndProject
Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "Connectors.Memory.Chroma", "src\Connectors\Connectors.Memory.Chroma\Connectors.Memory.Chroma.csproj", "{185E0CE8-C2DA-4E4C-A491-E8EB40316315}"
EndProject
Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Connectors.AI.OpenAI", "src\Connectors\Connectors.AI.OpenAI\Connectors.AI.OpenAI.csproj", "{AFA81EB7-F869-467D-8A90-744305D80AAC}"
EndProject
Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "SemanticKernel.Abstractions", "src\SemanticKernel.Abstractions\SemanticKernel.Abstractions.csproj", "{627742DB-1E52-468A-99BD-6FF1A542D25B}"
@@ -355,6 +357,12 @@ Global
{50FAE231-6F24-4779-9D02-12ABBC9A49E2}.Publish|Any CPU.Build.0 = Publish|Any CPU
{50FAE231-6F24-4779-9D02-12ABBC9A49E2}.Release|Any CPU.ActiveCfg = Release|Any CPU
{50FAE231-6F24-4779-9D02-12ABBC9A49E2}.Release|Any CPU.Build.0 = Release|Any CPU
{185E0CE8-C2DA-4E4C-A491-E8EB40316315}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
{185E0CE8-C2DA-4E4C-A491-E8EB40316315}.Debug|Any CPU.Build.0 = Debug|Any CPU
{185E0CE8-C2DA-4E4C-A491-E8EB40316315}.Publish|Any CPU.ActiveCfg = Debug|Any CPU
{185E0CE8-C2DA-4E4C-A491-E8EB40316315}.Publish|Any CPU.Build.0 = Debug|Any CPU
{185E0CE8-C2DA-4E4C-A491-E8EB40316315}.Release|Any CPU.ActiveCfg = Release|Any CPU
{185E0CE8-C2DA-4E4C-A491-E8EB40316315}.Release|Any CPU.Build.0 = Release|Any CPU
{0D0C4DAD-E6BC-4504-AE3A-EEA4E35920C1}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
{0D0C4DAD-E6BC-4504-AE3A-EEA4E35920C1}.Debug|Any CPU.Build.0 = Debug|Any CPU
{0D0C4DAD-E6BC-4504-AE3A-EEA4E35920C1}.Publish|Any CPU.ActiveCfg = Debug|Any CPU
@@ -418,6 +426,7 @@ Global
{B00AD427-0047-4850-BEF9-BA8237EA9D8B} = {958AD708-F048-4FAF-94ED-D2F2B92748B9}
{DB950192-30F1-48B1-88D7-F43FECCA1A1C} = {958AD708-F048-4FAF-94ED-D2F2B92748B9}
{1C19D805-3573-4477-BF07-40180FCDE1BD} = {958AD708-F048-4FAF-94ED-D2F2B92748B9}
{185E0CE8-C2DA-4E4C-A491-E8EB40316315} = {0247C2C9-86C3-45BA-8873-28B0948EDC0C}
{0D0C4DAD-E6BC-4504-AE3A-EEA4E35920C1} = {9ECD1AA0-75B3-4E25-B0B5-9F0945B64974}
{E6EDAB8F-3406-4DBF-9AAB-DF40DC2CA0FA} = {FA3720F1-C99A-49B2-9577-A940257098BF}
EndGlobalSection
dotnet/samples/KernelSyntaxExamples/Example50_Chroma.cs
+70
View File
@@ -0,0 +1,70 @@
// Copyright (c) Microsoft. All rights reserved.
using System;
using System.Threading.Tasks;
using Microsoft.SemanticKernel;
using Microsoft.SemanticKernel.Connectors.Memory.Chroma;
using Microsoft.SemanticKernel.Memory;
using RepoUtils;
// ReSharper disable once InconsistentNaming
public static class Example50_Chroma
{
private const string MemoryCollectionName = "chroma-test";
public static async Task RunAsync()
{
string endpoint = Env.Var("CHROMA_ENDPOINT");
var memoryStore = new ChromaMemoryStore(endpoint);
IKernel kernel = Kernel.Builder
.WithLogger(ConsoleLogger.Log)
.WithOpenAITextCompletionService("text-davinci-003", Env.Var("OPENAI_API_KEY"))
.WithOpenAITextEmbeddingGenerationService("text-embedding-ada-002", Env.Var("OPENAI_API_KEY"))
.WithMemoryStorage(memoryStore)
//.WithChromaMemoryStore(endpoint) // This method offers an alternative approach to registering Chroma memory store.
.Build();
Console.WriteLine("== Printing Collections in DB ==");
var collections = memoryStore.GetCollectionsAsync();
await foreach (var collection in collections)
{
Console.WriteLine(collection);
}
Console.WriteLine("== Adding Memories ==");
var key1 = await kernel.Memory.SaveInformationAsync(MemoryCollectionName, id: Guid.NewGuid().ToString(), text: "british short hair");
var key2 = await kernel.Memory.SaveInformationAsync(MemoryCollectionName, id: Guid.NewGuid().ToString(), text: "orange tabby");
var key3 = await kernel.Memory.SaveInformationAsync(MemoryCollectionName, id: Guid.NewGuid().ToString(), text: "norwegian forest cat");
Console.WriteLine("== Printing Collections in DB ==");
collections = memoryStore.GetCollectionsAsync();
await foreach (var collection in collections)
{
Console.WriteLine(collection);
}
Console.WriteLine("== Retrieving Memories Through the Kernel ==");
MemoryQueryResult? lookup = await kernel.Memory.GetAsync(MemoryCollectionName, key1);
Console.WriteLine(lookup != null ? lookup.Metadata.Text : "ERROR: memory not found");
Console.WriteLine("== Similarity Searching Memories: My favorite color is orange ==");
var searchResults = kernel.Memory.SearchAsync(MemoryCollectionName, "My favorite color is orange", limit: 3, minRelevanceScore: 0.6);
await foreach (var item in searchResults)
{
Console.WriteLine(item.Metadata.Text + " : " + item.Relevance);
}
Console.WriteLine("== Removing Collection {0} ==", MemoryCollectionName);
await memoryStore.DeleteCollectionAsync(MemoryCollectionName);
Console.WriteLine("== Printing Collections in DB ==");
await foreach (var collection in collections)
{
Console.WriteLine(collection);
}
}
}
dotnet/samples/KernelSyntaxExamples/KernelSyntaxExamples.csproj
+1
View File
@@ -29,6 +29,7 @@
<ProjectReference Include="..\..\src\Connectors\Connectors.AI.OpenAI\Connectors.AI.OpenAI.csproj" />
<ProjectReference Include="..\..\src\Connectors\Connectors.AI.HuggingFace\Connectors.AI.HuggingFace.csproj" />
<ProjectReference Include="..\..\src\Connectors\Connectors.Memory.AzureCognitiveSearch\Connectors.Memory.AzureCognitiveSearch.csproj" />
<ProjectReference Include="..\..\src\Connectors\Connectors.Memory.Chroma\Connectors.Memory.Chroma.csproj" />
<ProjectReference Include="..\..\src\Connectors\Connectors.Memory.Postgres\Connectors.Memory.Postgres.csproj" />
<ProjectReference Include="..\..\src\Connectors\Connectors.Memory.Weaviate\Connectors.Memory.Weaviate.csproj" />
<ProjectReference Include="..\..\src\Connectors\Connectors.Memory.Redis\Connectors.Memory.Redis.csproj" />
dotnet/samples/KernelSyntaxExamples/Program.cs
+3
View File
@@ -154,5 +154,8 @@ public static class Program
await Example49_LogitBias.RunAsync();
Console.WriteLine("== DONE ==");
await Example50_Chroma.RunAsync();
Console.WriteLine("== DONE ==");
}
}
dotnet/src/Connectors/Connectors.Memory.Chroma/ChromaClient.cs
+207
View File
@@ -0,0 +1,207 @@
// Copyright (c) Microsoft. All rights reserved.
using System;
using System.Collections.Generic;
using System.Net.Http;
using System.Runtime.CompilerServices;
using System.Text;
using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Logging.Abstractions;
using Microsoft.SemanticKernel.Connectors.Memory.Chroma.Http.ApiSchema;
using Microsoft.SemanticKernel.Connectors.Memory.Chroma.Http.ApiSchema.Internal;
namespace Microsoft.SemanticKernel.Connectors.Memory.Chroma;
/// <summary>
/// An implementation of a client for the Chroma Vector DB. This class is used to
/// create, delete, and get embeddings data from Chroma Vector DB instance.
/// </summary>
#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.
public class ChromaClient : IChromaClient
#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.
{
/// <summary>
/// Initializes a new instance of the <see cref="ChromaClient"/> class.
/// </summary>
/// <param name="endpoint">Chroma server endpoint URL.</param>
/// <param name="logger">Optional logger instance.</param>
public ChromaClient(string endpoint, ILogger? logger = null)
{
this._httpClient = new HttpClient(NonDisposableHttpClientHandler.Instance, disposeHandler: false);
this._endpoint = endpoint;
this._logger = logger ?? NullLogger<ChromaClient>.Instance;
}
/// <summary>
/// Initializes a new instance of the <see cref="ChromaClient"/> class.
/// </summary>
/// <param name="httpClient">The <see cref="HttpClient"/> instance used for making HTTP requests.</param>
/// <param name="endpoint">Chroma server endpoint URL.</param>
/// <param name="logger">Optional logger instance.</param>
/// <exception cref="ChromaClientException">Occurs when <see cref="HttpClient"/> doesn't have base address and endpoint parameter is not provided.</exception>
public ChromaClient(HttpClient httpClient, string? endpoint = null, ILogger? logger = null)
{
if (string.IsNullOrEmpty(httpClient.BaseAddress?.AbsoluteUri) && string.IsNullOrEmpty(endpoint))
{
throw new ChromaClientException("The HttpClient BaseAddress and endpoint are both null or empty. Please ensure at least one is provided.");
}
this._httpClient = httpClient;
this._endpoint = endpoint;
this._logger = logger ?? NullLogger<ChromaClient>.Instance;
}
/// <inheritdoc />
public async Task CreateCollectionAsync(string collectionName, CancellationToken cancellationToken = default)
{
this._logger.LogDebug("Creating collection {0}", collectionName);
using var request = CreateCollectionRequest.Create(collectionName).Build();
await this.ExecuteHttpRequestAsync(request, cancellationToken).ConfigureAwait(false);
}
/// <inheritdoc />
public async Task<ChromaCollectionModel?> GetCollectionAsync(string collectionName, CancellationToken cancellationToken = default)
{
this._logger.LogDebug("Getting collection {0}", collectionName);
using var request = GetCollectionRequest.Create(collectionName).Build();
(HttpResponseMessage response, string responseContent) = await this.ExecuteHttpRequestAsync(request, cancellationToken).ConfigureAwait(false);
var collection = JsonSerializer.Deserialize<ChromaCollectionModel>(responseContent);
return collection;
}
/// <inheritdoc />
public async Task DeleteCollectionAsync(string collectionName, CancellationToken cancellationToken = default)
{
this._logger.LogDebug("Deleting collection {0}", collectionName);
using var request = DeleteCollectionRequest.Create(collectionName).Build();
await this.ExecuteHttpRequestAsync(request, cancellationToken).ConfigureAwait(false);
}
/// <inheritdoc />
public async IAsyncEnumerable<string> ListCollectionsAsync([EnumeratorCancellation] CancellationToken cancellationToken = default)
{
this._logger.LogDebug("Listing collections");
using var request = ListCollectionsRequest.Create().Build();
(HttpResponseMessage response, string responseContent) = await this.ExecuteHttpRequestAsync(request, cancellationToken).ConfigureAwait(false);
var collections = JsonSerializer.Deserialize<List<ChromaCollectionModel>>(responseContent);
foreach (var collection in collections!)
{
yield return collection.Name;
}
}
/// <inheritdoc />
public async Task AddEmbeddingsAsync(string collectionId, string[] ids, float[][] embeddings, object[]? metadatas = null, CancellationToken cancellationToken = default)
{
this._logger.LogDebug("Adding embeddings to collection with id: {0}", collectionId);
using var request = AddEmbeddingsRequest.Create(collectionId, ids, embeddings, metadatas).Build();
await this.ExecuteHttpRequestAsync(request, cancellationToken).ConfigureAwait(false);
}
/// <inheritdoc />
public async Task<ChromaEmbeddingsModel> GetEmbeddingsAsync(string collectionId, string[] ids, string[]? include = null, CancellationToken cancellationToken = default)
{
this._logger.LogDebug("Getting embeddings from collection with id: {0}", collectionId);
using var request = GetEmbeddingsRequest.Create(collectionId, ids, include).Build();
(HttpResponseMessage response, string responseContent) = await this.ExecuteHttpRequestAsync(request, cancellationToken).ConfigureAwait(false);
var embeddings = JsonSerializer.Deserialize<ChromaEmbeddingsModel>(responseContent);
return embeddings ?? new ChromaEmbeddingsModel();
}
/// <inheritdoc />
public async Task DeleteEmbeddingsAsync(string collectionId, string[] ids, CancellationToken cancellationToken = default)
{
this._logger.LogDebug("Deleting embeddings from collection with id: {0}", collectionId);
using var request = DeleteEmbeddingsRequest.Create(collectionId, ids).Build();
await this.ExecuteHttpRequestAsync(request, cancellationToken).ConfigureAwait(false);
}
/// <inheritdoc />
public async Task<ChromaQueryResultModel> QueryEmbeddingsAsync(string collectionId, float[][] queryEmbeddings, int nResults, string[]? include = null, CancellationToken cancellationToken = default)
{
this._logger.LogDebug("Query embeddings in collection with id: {0}", collectionId);
using var request = QueryEmbeddingsRequest.Create(collectionId, queryEmbeddings, nResults, include).Build();
(HttpResponseMessage response, string responseContent) = await this.ExecuteHttpRequestAsync(request, cancellationToken).ConfigureAwait(false);
var queryResult = JsonSerializer.Deserialize<ChromaQueryResultModel>(responseContent);
return queryResult ?? new ChromaQueryResultModel();
}
#region private ================================================================================
private const string ApiRoute = "api/v1/";
private readonly ILogger _logger;
private readonly HttpClient _httpClient;
private readonly string? _endpoint = null;
private async Task<(HttpResponseMessage response, string responseContent)> ExecuteHttpRequestAsync(
HttpRequestMessage request,
CancellationToken cancellationToken = default)
{
string endpoint = this._endpoint ?? this._httpClient.BaseAddress.ToString();
endpoint = this.SanitizeEndpoint(endpoint);
string operationName = request.RequestUri.ToString();
request.RequestUri = new Uri(new Uri(endpoint), operationName);
HttpResponseMessage response = await this._httpClient.SendAsync(request, cancellationToken).ConfigureAwait(false);
string responseContent = await response.Content.ReadAsStringAsync().ConfigureAwait(false);
try
{
response.EnsureSuccessStatusCode();
}
catch (HttpRequestException e)
{
this._logger.LogError(e, "{0} {1} operation failed: {2}, {3}", request.Method.Method, operationName, e.Message, responseContent);
throw new ChromaClientException($"{request.Method.Method} {operationName} operation failed: {e.Message}, {responseContent}", e);
}
return (response, responseContent);
}
private string SanitizeEndpoint(string endpoint)
{
StringBuilder builder = new(endpoint);
if (!endpoint.EndsWith("/", StringComparison.Ordinal))
{
builder.Append('/');
}
builder.Append(ApiRoute);
return builder.ToString();
}
#endregion
}
dotnet/src/Connectors/Connectors.Memory.Chroma/ChromaClientException.cs
+52
View File
@@ -0,0 +1,52 @@
// Copyright (c) Microsoft. All rights reserved.
using System;
using System.Globalization;
namespace Microsoft.SemanticKernel.Connectors.Memory.Chroma;
/// <summary>
/// Exception to identify issues in <see cref="ChromaClient"/> class.
/// </summary>
public class ChromaClientException : Exception
{
private const string CollectionDoesNotExistErrorFormat = "Collection {0} does not exist";
private const string DeleteNonExistentCollectionErrorMessage = "list index out of range";
/// <summary>
/// Initializes a new instance of the <see cref="ChromaClientException"/> class.
/// </summary>
public ChromaClientException() : base()
{
}
/// <summary>
/// Initializes a new instance of the <see cref="ChromaClientException"/> class.
/// </summary>
/// <param name="message">The message that describes the error.</param>
public ChromaClientException(string message) : base(message)
{
}
/// <summary>
/// Initializes a new instance of the <see cref="ChromaClientException"/> class.
/// </summary>
/// <param name="message">The message that describes the error.</param>
/// <param name="innerException">Instance of inner exception.</param>
public ChromaClientException(string message, Exception innerException) : base(message, innerException)
{
}
/// <summary>
/// Checks if Chroma API error means that collection does not exist.
/// </summary>
/// <param name="collectionName">Collection name.</param>
public bool CollectionDoesNotExistException(string collectionName) =>
this.Message.Contains(string.Format(CultureInfo.InvariantCulture, CollectionDoesNotExistErrorFormat, collectionName));
/// <summary>
/// Checks if Chroma API error means that there was an attempt to delete non-existent collection.
/// </summary>
public bool DeleteNonExistentCollectionException() =>
this.Message.Contains(DeleteNonExistentCollectionErrorMessage);
}
dotnet/src/Connectors/Connectors.Memory.Chroma/ChromaKernelBuilderExtensions.cs
+55
View File
@@ -0,0 +1,55 @@
// Copyright (c) Microsoft. All rights reserved.
using System.Net.Http;
using Microsoft.SemanticKernel.Connectors.Memory.Chroma;
#pragma warning disable IDE0130
namespace Microsoft.SemanticKernel;
#pragma warning restore IDE0130
/// <summary>
/// Provides extension methods for the <see cref="KernelBuilder"/> class to configure Chroma memory connector.
/// </summary>
public static class ChromaKernelBuilderExtensions
{
/// <summary>
/// Registers Chroma memory connector.
/// </summary>
/// <param name="builder">The <see cref="KernelBuilder"/> instance.</param>
/// <param name="endpoint">Chroma server endpoint URL.</param>
/// <returns>Self instance.</returns>
public static KernelBuilder WithChromaMemoryStore(this KernelBuilder builder, string endpoint)
{
builder.WithMemoryStorage((parameters) =>
{
return new ChromaMemoryStore(
HttpClientProvider.GetHttpClient(parameters.Config, null, parameters.Logger),
endpoint,
parameters.Logger);
});
return builder;
}
/// <summary>
/// Registers Chroma memory connector.
/// </summary>
/// <param name="builder">The <see cref="KernelBuilder"/> instance.</param>
/// <param name="httpClient">The <see cref="HttpClient"/> instance used for making HTTP requests.</param>
/// <param name="endpoint">Chroma server endpoint URL. If not specified, the base address of the HTTP client is used.</param>
/// <returns>Self instance.</returns>
public static KernelBuilder WithChromaMemoryStore(this KernelBuilder builder,
HttpClient httpClient,
string? endpoint = null)
{
builder.WithMemoryStorage((parameters) =>
{
return new ChromaMemoryStore(
HttpClientProvider.GetHttpClient(parameters.Config, httpClient, parameters.Logger),
endpoint,
parameters.Logger);
});
return builder;
}
}
dotnet/src/Connectors/Connectors.Memory.Chroma/ChromaMemoryStore.cs
+324
View File
@@ -0,0 +1,324 @@
// Copyright (c) Microsoft. All rights reserved.
using System.Collections.Generic;
using System.Linq;
using System.Net.Http;
using System.Runtime.CompilerServices;
using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Logging.Abstractions;
using Microsoft.SemanticKernel.AI.Embeddings;
using Microsoft.SemanticKernel.Connectors.Memory.Chroma.Http.ApiSchema;
using Microsoft.SemanticKernel.Diagnostics;
using Microsoft.SemanticKernel.Memory;
namespace Microsoft.SemanticKernel.Connectors.Memory.Chroma;
/// <summary>
/// An implementation of <see cref="IMemoryStore" /> for Chroma.
/// </summary>
public class ChromaMemoryStore : IMemoryStore
{
/// <summary>
/// Initializes a new instance of the <see cref="ChromaMemoryStore"/> class.
/// </summary>
/// <param name="endpoint">Chroma server endpoint URL.</param>
/// <param name="logger">Optional logger instance.</param>
public ChromaMemoryStore(string endpoint, ILogger? logger = null)
: this(new ChromaClient(endpoint, logger), logger)
{
}
/// <summary>
/// Initializes a new instance of the <see cref="ChromaMemoryStore"/> class.
/// </summary>
/// <param name="httpClient">The <see cref="HttpClient"/> instance used for making HTTP requests.</param>
/// <param name="endpoint">Chroma server endpoint URL.</param>
/// <param name="logger">Optional logger instance.</param>
public ChromaMemoryStore(HttpClient httpClient, string? endpoint = null, ILogger? logger = null)
: this(new ChromaClient(httpClient, endpoint, logger), logger)
{
}
/// <summary>
/// Initializes a new instance of the <see cref="ChromaMemoryStore"/> class.
/// </summary>
/// <param name="client">Instance of <see cref="IChromaClient"/> implementation.</param>
/// <param name="logger">Optional logger instance.</param>
public ChromaMemoryStore(IChromaClient client, ILogger? logger = null)
{
this._chromaClient = client;
this._logger = logger ?? NullLogger<ChromaMemoryStore>.Instance;
}
/// <inheritdoc />
public async Task CreateCollectionAsync(string collectionName, CancellationToken cancellationToken = default)
{
Verify.NotNullOrWhiteSpace(collectionName);
await this._chromaClient.CreateCollectionAsync(collectionName, cancellationToken).ConfigureAwait(false);
}
/// <inheritdoc />
public async Task DeleteCollectionAsync(string collectionName, CancellationToken cancellationToken = default)
{
Verify.NotNullOrWhiteSpace(collectionName);
try
{
await this._chromaClient.DeleteCollectionAsync(collectionName, cancellationToken).ConfigureAwait(false);
}
catch (ChromaClientException e) when (e.DeleteNonExistentCollectionException())
{
this._logger.LogError("Cannot delete non-existent collection {0}", collectionName);
throw new ChromaMemoryStoreException($"Cannot delete non-existent collection {collectionName}", e);
}
}
/// <inheritdoc />
public async Task<bool> DoesCollectionExistAsync(string collectionName, CancellationToken cancellationToken = default)
{
Verify.NotNullOrWhiteSpace(collectionName);
var collection = await this.GetCollectionAsync(collectionName, cancellationToken).ConfigureAwait(false);
return collection != null;
}
/// <inheritdoc />
public async Task<MemoryRecord?> GetAsync(string collectionName, string key, bool withEmbedding = false, CancellationToken cancellationToken = default)
{
return await this.GetBatchAsync(collectionName, new[] { key }, withEmbedding, cancellationToken)
.FirstOrDefaultAsync(cancellationToken)
.ConfigureAwait(false);
}
/// <inheritdoc />
public async IAsyncEnumerable<MemoryRecord> GetBatchAsync(string collectionName, IEnumerable<string> keys, bool withEmbeddings = false, [EnumeratorCancellation] CancellationToken cancellationToken = default)
{
Verify.NotNullOrWhiteSpace(collectionName);
var collection = await this.GetCollectionOrThrowAsync(collectionName, cancellationToken).ConfigureAwait(false);
var ids = keys.ToArray();
var include = this.GetEmbeddingIncludeTypes(withEmbeddings: withEmbeddings);
var embeddingsModel = await this._chromaClient.GetEmbeddingsAsync(collection.Id, ids, include, cancellationToken).ConfigureAwait(false);
var recordCount = embeddingsModel.Ids?.Count ?? 0;
for (var recordIndex = 0; recordIndex < recordCount; recordIndex++)
{
yield return this.GetMemoryRecordFromEmbeddingsModel(embeddingsModel, recordIndex);
}
}
/// <inheritdoc />
public IAsyncEnumerable<string> GetCollectionsAsync(CancellationToken cancellationToken = default)
{
return this._chromaClient.ListCollectionsAsync(cancellationToken);
}
/// <inheritdoc />
public async Task<(MemoryRecord, double)?> GetNearestMatchAsync(string collectionName, Embedding<float> embedding, double minRelevanceScore = 0, bool withEmbedding = false, CancellationToken cancellationToken = default)
{
var results = this.GetNearestMatchesAsync(
collectionName,
embedding,
minRelevanceScore: minRelevanceScore,
limit: 1,
withEmbeddings: withEmbedding,
cancellationToken: cancellationToken);
(MemoryRecord memoryRecord, double similarityScore) = await results.FirstOrDefaultAsync(cancellationToken).ConfigureAwait(false);
return (memoryRecord, similarityScore);
}
/// <inheritdoc />
public async IAsyncEnumerable<(MemoryRecord, double)> GetNearestMatchesAsync(string collectionName, Embedding<float> embedding, int limit, double minRelevanceScore = 0, bool withEmbeddings = false, [EnumeratorCancellation] CancellationToken cancellationToken = default)
{
Verify.NotNullOrWhiteSpace(collectionName);
var collection = await this.GetCollectionOrThrowAsync(collectionName, cancellationToken).ConfigureAwait(false);
var queryEmbeddings = new float[][] { embedding.Vector.ToArray() };
var nResults = limit;
var include = this.GetEmbeddingIncludeTypes(withEmbeddings: withEmbeddings, withDistances: true);
var queryResultModel = await this._chromaClient.QueryEmbeddingsAsync(collection.Id, queryEmbeddings, nResults, include, cancellationToken).ConfigureAwait(false);
var recordCount = queryResultModel.Ids?.FirstOrDefault()?.Count ?? 0;
for (var recordIndex = 0; recordIndex < recordCount; recordIndex++)
{
(MemoryRecord memoryRecord, double similarityScore) = this.GetMemoryRecordFromQueryResultModel(queryResultModel, recordIndex);
if (similarityScore >= minRelevanceScore)
{
yield return (memoryRecord, similarityScore);
}
}
}
/// <inheritdoc />
public async Task RemoveAsync(string collectionName, string key, CancellationToken cancellationToken = default)
{
await this.RemoveBatchAsync(collectionName, new[] { key }, cancellationToken).ConfigureAwait(false);
}
/// <inheritdoc />
public async Task RemoveBatchAsync(string collectionName, IEnumerable<string> keys, CancellationToken cancellationToken = default)
{
Verify.NotNullOrWhiteSpace(collectionName);
var collection = await this.GetCollectionOrThrowAsync(collectionName, cancellationToken).ConfigureAwait(false);
await this._chromaClient.DeleteEmbeddingsAsync(collection.Id, keys.ToArray(), cancellationToken).ConfigureAwait(false);
}
/// <inheritdoc />
public async Task<string> UpsertAsync(string collectionName, MemoryRecord record, CancellationToken cancellationToken = default)
{
Verify.NotNullOrWhiteSpace(collectionName);
var key = await this.UpsertBatchAsync(collectionName, new[] { record }, cancellationToken)
.FirstOrDefaultAsync(cancellationToken)
.ConfigureAwait(false);
return key ?? string.Empty;
}
/// <inheritdoc />
public async IAsyncEnumerable<string> UpsertBatchAsync(string collectionName, IEnumerable<MemoryRecord> records, [EnumeratorCancellation] CancellationToken cancellationToken = default)
{
Verify.NotNullOrWhiteSpace(collectionName);
var collection = await this.GetCollectionOrThrowAsync(collectionName, cancellationToken).ConfigureAwait(false);
var recordsArray = records.ToArray();
var recordsLength = recordsArray.Length;
var ids = new string[recordsLength];
var embeddings = new float[recordsLength][];
var metadatas = new object[recordsLength];
for (var i = 0; i < recordsLength; i++)
{
ids[i] = recordsArray[i].Metadata.Id;
embeddings[i] = recordsArray[i].Embedding.Vector.ToArray();
metadatas[i] = recordsArray[i].Metadata;
}
await this._chromaClient.AddEmbeddingsAsync(collection.Id, ids, embeddings, metadatas, cancellationToken).ConfigureAwait(false);
foreach (var record in recordsArray)
{
yield return record.Metadata.Id;
}
}
#region private ================================================================================
private const string IncludeMetadatas = "metadatas";
private const string IncludeEmbeddings = "embeddings";
private const string IncludeDistances = "distances";
private readonly ILogger _logger;
private readonly IChromaClient _chromaClient;
private readonly List<string> _defaultEmbeddingIncludeTypes = new() { IncludeMetadatas };
private async Task<ChromaCollectionModel> GetCollectionOrThrowAsync(string collectionName, CancellationToken cancellationToken)
{
return
await this.GetCollectionAsync(collectionName, cancellationToken).ConfigureAwait(false) ??
throw new ChromaMemoryStoreException($"Collection {collectionName} does not exist");
}
private async Task<ChromaCollectionModel?> GetCollectionAsync(string collectionName, CancellationToken cancellationToken)
{
try
{
return await this._chromaClient.GetCollectionAsync(collectionName, cancellationToken).ConfigureAwait(false);
}
catch (ChromaClientException e) when (e.CollectionDoesNotExistException(collectionName))
{
this._logger.LogError("Collection {0} does not exist", collectionName);
return null;
}
}
private string[] GetEmbeddingIncludeTypes(bool withEmbeddings = false, bool withDistances = false)
{
var includeList = new List<string>(this._defaultEmbeddingIncludeTypes);
if (withEmbeddings)
{
includeList.Add(IncludeEmbeddings);
}
if (withDistances)
{
includeList.Add(IncludeDistances);
}
return includeList.ToArray();
}
private MemoryRecord GetMemoryRecordFromEmbeddingsModel(ChromaEmbeddingsModel embeddingsModel, int recordIndex)
{
return this.GetMemoryRecordFromModel(embeddingsModel.Metadatas, embeddingsModel.Embeddings, embeddingsModel.Ids, recordIndex);
}
private (MemoryRecord, double) GetMemoryRecordFromQueryResultModel(ChromaQueryResultModel queryResultModel, int recordIndex)
{
var ids = queryResultModel.Ids?.FirstOrDefault();
var embeddings = queryResultModel.Embeddings?.FirstOrDefault();
var metadatas = queryResultModel.Metadatas?.FirstOrDefault();
var distances = queryResultModel.Distances?.FirstOrDefault();
var memoryRecord = this.GetMemoryRecordFromModel(metadatas, embeddings, ids, recordIndex);
var similarityScore = this.GetSimilarityScore(distances, recordIndex);
return (memoryRecord, similarityScore);
}
private MemoryRecord GetMemoryRecordFromModel(List<Dictionary<string, object>>? metadatas, List<float[]>? embeddings, List<string>? ids, int recordIndex)
{
var metadata = this.GetMetadataForMemoryRecord(metadatas, recordIndex);
var embeddingsVector = this.GetEmbeddingForMemoryRecord(embeddings, recordIndex);
var key = ids?[recordIndex];
return MemoryRecord.FromJsonMetadata(
json: metadata,
embedding: embeddingsVector,
key: key);
}
private string GetMetadataForMemoryRecord(List<Dictionary<string, object>>? metadatas, int recordIndex)
{
return metadatas != null ? JsonSerializer.Serialize(metadatas[recordIndex]) : string.Empty;
}
private Embedding<float> GetEmbeddingForMemoryRecord(List<float[]>? embeddings, int recordIndex)
{
return embeddings != null ? new Embedding<float>(embeddings[recordIndex]) : Embedding<float>.Empty;
}
private double GetSimilarityScore(List<double>? distances, int recordIndex)
{
var similarityScore = distances != null ? 1 - distances[recordIndex] : default;
if (similarityScore < 0)
{
similarityScore = 0;
}
return similarityScore;
}
#endregion
}
dotnet/src/Connectors/Connectors.Memory.Chroma/ChromaMemoryStoreException.cs
+35
View File
@@ -0,0 +1,35 @@
// Copyright (c) Microsoft. All rights reserved.
using System;
namespace Microsoft.SemanticKernel.Connectors.Memory.Chroma;
/// <summary>
/// Exception to identify issues in <see cref="ChromaMemoryStore"/> class.
/// </summary>
public class ChromaMemoryStoreException : Exception
{
/// <summary>
/// Initializes a new instance of the <see cref="ChromaMemoryStoreException"/> class.
/// </summary>
public ChromaMemoryStoreException() : base()
{
}
/// <summary>
/// Initializes a new instance of the <see cref="ChromaMemoryStoreException"/> class.
/// </summary>
/// <param name="message">The message that describes the error.</param>
public ChromaMemoryStoreException(string message) : base(message)
{
}
/// <summary>
/// Initializes a new instance of the <see cref="ChromaMemoryStoreException"/> class.
/// </summary>
/// <param name="message">The message that describes the error.</param>
/// <param name="innerException">Instance of inner exception.</param>
public ChromaMemoryStoreException(string message, Exception innerException) : base(message, innerException)
{
}
}
dotnet/src/Connectors/Connectors.Memory.Chroma/Connectors.Memory.Chroma.csproj
+28
View File
@@ -0,0 +1,28 @@
<Project Sdk="Microsoft.NET.Sdk">
<PropertyGroup>
<!-- THIS PROPERTY GROUP MUST COME FIRST -->
<AssemblyName>Microsoft.SemanticKernel.Connectors.Memory.Chroma</AssemblyName>
<RootNamespace>$(AssemblyName)</RootNamespace>
<TargetFramework>netstandard2.0</TargetFramework>
</PropertyGroup>
<!-- IMPORT NUGET PACKAGE SHARED PROPERTIES -->
<Import Project="$(RepoRoot)/dotnet/nuget/nuget-package.props" />
<Import Project="$(RepoRoot)/dotnet/src/InternalUtilities/src/InternalUtilities.props" />
<PropertyGroup>
<!-- NuGet Package Settings -->
<Title>Semantic Kernel - Chroma Connector</Title>
<Description>Chroma connector for Semantic Kernel skills and semantic memory</Description>
</PropertyGroup>
<ItemGroup>
<PackageReference Include="System.Text.Json" />
</ItemGroup>
<ItemGroup>
<ProjectReference Include="..\..\SemanticKernel\SemanticKernel.csproj" />
</ItemGroup>
</Project>
dotnet/src/Connectors/Connectors.Memory.Chroma/Http/ApiSchema/ChromaCollectionModel.cs
+23
View File
@@ -0,0 +1,23 @@
// Copyright (c) Microsoft. All rights reserved.
using System.Text.Json.Serialization;
namespace Microsoft.SemanticKernel.Connectors.Memory.Chroma.Http.ApiSchema;
/// <summary>
/// Chroma collection model.
/// </summary>
public class ChromaCollectionModel
{
/// <summary>
/// Identifier of collection generated by Chroma.
/// </summary>
[JsonPropertyName("id")]
public string Id { get; set; } = string.Empty;
/// <summary>
/// Collection name provided by client.
/// </summary>
[JsonPropertyName("name")]
public string Name { get; set; } = string.Empty;
}
dotnet/src/Connectors/Connectors.Memory.Chroma/Http/ApiSchema/ChromaEmbeddingsModel.cs
+30
View File
@@ -0,0 +1,30 @@
// Copyright (c) Microsoft. All rights reserved.
using System.Collections.Generic;
using System.Text.Json.Serialization;
namespace Microsoft.SemanticKernel.Connectors.Memory.Chroma.Http.ApiSchema;
/// <summary>
/// Chroma embeddings model.
/// </summary>
public class ChromaEmbeddingsModel
{
/// <summary>
/// Embedding identifiers.
/// </summary>
[JsonPropertyName("ids")]
public List<string> Ids { get; set; } = new();
/// <summary>
/// Embedding vectors.
/// </summary>
[JsonPropertyName("embeddings")]
public List<float[]> Embeddings { get; set; } = new();
/// <summary>
/// Embedding metadatas.
/// </summary>
[JsonPropertyName("metadatas")]
public List<Dictionary<string, object>> Metadatas { get; set; } = new();
}
dotnet/src/Connectors/Connectors.Memory.Chroma/Http/ApiSchema/ChromaQueryResultModel.cs
+36
View File
@@ -0,0 +1,36 @@
// Copyright (c) Microsoft. All rights reserved.
using System.Collections.Generic;
using System.Text.Json.Serialization;
namespace Microsoft.SemanticKernel.Connectors.Memory.Chroma.Http.ApiSchema;
/// <summary>
/// Chroma query result model. Contains result sets after search operation.
/// </summary>
public class ChromaQueryResultModel
{
/// <summary>
/// List of embedding identifiers.
/// </summary>
[JsonPropertyName("ids")]
public List<List<string>> Ids { get; set; } = new();
/// <summary>
/// List of embedding vectors.
/// </summary>
[JsonPropertyName("embeddings")]
public List<List<float[]>> Embeddings { get; set; } = new();
/// <summary>
/// List of embedding metadatas.
/// </summary>
[JsonPropertyName("metadatas")]
public List<List<Dictionary<string, object>>> Metadatas { get; set; } = new();
/// <summary>
/// List of embedding distances.
/// </summary>
[JsonPropertyName("distances")]
public List<List<double>> Distances { get; set; } = new();
}
dotnet/src/Connectors/Connectors.Memory.Chroma/Http/ApiSchema/Internal/AddEmbeddingsRequest.cs
+43
View File
@@ -0,0 +1,43 @@
// Copyright (c) Microsoft. All rights reserved.
using System.Net.Http;
using System.Text.Json.Serialization;
namespace Microsoft.SemanticKernel.Connectors.Memory.Chroma.Http.ApiSchema.Internal;
internal sealed class AddEmbeddingsRequest
{
[JsonIgnore]
public string CollectionId { get; set; }
[JsonPropertyName("ids")]
public string[] Ids { get; set; }
[JsonPropertyName("embeddings")]
public float[][] Embeddings { get; set; }
[JsonPropertyName("metadatas")]
public object[]? Metadatas { get; set; }
public static AddEmbeddingsRequest Create(string collectionId, string[] ids, float[][] embeddings, object[]? metadatas = null)
{
return new AddEmbeddingsRequest(collectionId, ids, embeddings, metadatas);
}
public HttpRequestMessage Build()
{
return HttpRequest.CreatePostRequest($"collections/{this.CollectionId}/add", this);
}
#region private ================================================================================
private AddEmbeddingsRequest(string collectionId, string[] ids, float[][] embeddings, object[]? metadatas = null)
{
this.CollectionId = collectionId;
this.Ids = ids;
this.Embeddings = embeddings;
this.Metadatas = metadatas;
}
#endregion
}
dotnet/src/Connectors/Connectors.Memory.Chroma/Http/ApiSchema/Internal/CreateCollectionRequest.cs
+34
View File
@@ -0,0 +1,34 @@
// Copyright (c) Microsoft. All rights reserved.
using System.Net.Http;
using System.Text.Json.Serialization;
namespace Microsoft.SemanticKernel.Connectors.Memory.Chroma.Http.ApiSchema.Internal;
internal sealed class CreateCollectionRequest
{
[JsonPropertyName("name")]
public string CollectionName { get; set; }
[JsonPropertyName("get_or_create")]
public bool GetOrCreate => true;
public static CreateCollectionRequest Create(string collectionName)
{
return new CreateCollectionRequest(collectionName);
}
public HttpRequestMessage Build()
{
return HttpRequest.CreatePostRequest("collections", this);
}
#region private ================================================================================
private CreateCollectionRequest(string collectionName)
{
this.CollectionName = collectionName;
}
#endregion
}
dotnet/src/Connectors/Connectors.Memory.Chroma/Http/ApiSchema/Internal/DeleteCollectionRequest.cs
+31
View File
@@ -0,0 +1,31 @@
// Copyright (c) Microsoft. All rights reserved.
using System.Net.Http;
using System.Text.Json.Serialization;
namespace Microsoft.SemanticKernel.Connectors.Memory.Chroma.Http.ApiSchema.Internal;
internal sealed class DeleteCollectionRequest
{
[JsonIgnore]
public string CollectionName { get; set; }
public static DeleteCollectionRequest Create(string collectionName)
{
return new DeleteCollectionRequest(collectionName);
}
public HttpRequestMessage Build()
{
return HttpRequest.CreateDeleteRequest($"collections/{this.CollectionName}");
}
#region private ================================================================================
private DeleteCollectionRequest(string collectionName)
{
this.CollectionName = collectionName;
}
#endregion
}
dotnet/src/Connectors/Connectors.Memory.Chroma/Http/ApiSchema/Internal/DeleteEmbeddingsRequest.cs
+35
View File
@@ -0,0 +1,35 @@
// Copyright (c) Microsoft. All rights reserved.
using System.Net.Http;
using System.Text.Json.Serialization;
namespace Microsoft.SemanticKernel.Connectors.Memory.Chroma.Http.ApiSchema.Internal;
internal sealed class DeleteEmbeddingsRequest
{
[JsonIgnore]
public string CollectionId { get; set; }
[JsonPropertyName("ids")]
public string[] Ids { get; set; }
public static DeleteEmbeddingsRequest Create(string collectionId, string[] ids)
{
return new DeleteEmbeddingsRequest(collectionId, ids);
}
public HttpRequestMessage Build()
{
return HttpRequest.CreatePostRequest($"collections/{this.CollectionId}/delete", this);
}
#region private ================================================================================
private DeleteEmbeddingsRequest(string collectionId, string[] ids)
{
this.CollectionId = collectionId;
this.Ids = ids;
}
#endregion
}
dotnet/src/Connectors/Connectors.Memory.Chroma/Http/ApiSchema/Internal/GetCollectionRequest.cs
+31
View File
@@ -0,0 +1,31 @@
// Copyright (c) Microsoft. All rights reserved.
using System.Net.Http;
using System.Text.Json.Serialization;
namespace Microsoft.SemanticKernel.Connectors.Memory.Chroma.Http.ApiSchema.Internal;
internal sealed class GetCollectionRequest
{
[JsonIgnore]
public string CollectionName { get; set; }
public static GetCollectionRequest Create(string collectionName)
{
return new GetCollectionRequest(collectionName);
}
public HttpRequestMessage Build()
{
return HttpRequest.CreateGetRequest($"collections/{this.CollectionName}");
}
#region private ================================================================================
private GetCollectionRequest(string collectionName)
{
this.CollectionName = collectionName;
}
#endregion
}
dotnet/src/Connectors/Connectors.Memory.Chroma/Http/ApiSchema/Internal/GetEmbeddingsRequest.cs
+39
View File
@@ -0,0 +1,39 @@
// Copyright (c) Microsoft. All rights reserved.
using System.Net.Http;
using System.Text.Json.Serialization;
namespace Microsoft.SemanticKernel.Connectors.Memory.Chroma.Http.ApiSchema.Internal;
internal sealed class GetEmbeddingsRequest
{
[JsonIgnore]
public string CollectionId { get; set; }
[JsonPropertyName("ids")]
public string[] Ids { get; set; }
[JsonPropertyName("include")]
public string[]? Include { get; set; }
public static GetEmbeddingsRequest Create(string collectionId, string[] ids, string[]? include = null)
{
return new GetEmbeddingsRequest(collectionId, ids, include);
}
public HttpRequestMessage Build()
{
return HttpRequest.CreatePostRequest($"collections/{this.CollectionId}/get", this);
}
#region private ================================================================================
private GetEmbeddingsRequest(string collectionId, string[] ids, string[]? include = null)
{
this.CollectionId = collectionId;
this.Ids = ids;
this.Include = include;
}
#endregion
}
dotnet/src/Connectors/Connectors.Memory.Chroma/Http/ApiSchema/Internal/ListCollectionsRequest.cs
+25
View File
@@ -0,0 +1,25 @@
// Copyright (c) Microsoft. All rights reserved.
using System.Net.Http;
namespace Microsoft.SemanticKernel.Connectors.Memory.Chroma.Http.ApiSchema.Internal;
internal sealed class ListCollectionsRequest
{
public static ListCollectionsRequest Create()
{
return new ListCollectionsRequest();
}
public HttpRequestMessage Build()
{
return HttpRequest.CreateGetRequest("collections");
}
#region private ================================================================================
private ListCollectionsRequest()
{ }
#endregion
}
dotnet/src/Connectors/Connectors.Memory.Chroma/Http/ApiSchema/Internal/QueryEmbeddingsRequest.cs
+43
View File
@@ -0,0 +1,43 @@
// Copyright (c) Microsoft. All rights reserved.
using System.Net.Http;
using System.Text.Json.Serialization;
namespace Microsoft.SemanticKernel.Connectors.Memory.Chroma.Http.ApiSchema.Internal;
internal sealed class QueryEmbeddingsRequest
{
[JsonIgnore]
public string CollectionId { get; set; }
[JsonPropertyName("query_embeddings")]
public float[][] QueryEmbeddings { get; set; }
[JsonPropertyName("n_results")]
public int NResults { get; set; }
[JsonPropertyName("include")]
public string[]? Include { get; set; }
public static QueryEmbeddingsRequest Create(string collectionId, float[][] queryEmbeddings, int nResults, string[]? include = null)
{
return new QueryEmbeddingsRequest(collectionId, queryEmbeddings, nResults, include);
}
public HttpRequestMessage Build()
{
return HttpRequest.CreatePostRequest($"collections/{this.CollectionId}/query", this);
}
#region private ================================================================================
private QueryEmbeddingsRequest(string collectionId, float[][] queryEmbeddings, int nResults, string[]? include = null)
{
this.CollectionId = collectionId;
this.QueryEmbeddings = queryEmbeddings;
this.NResults = nResults;
this.Include = include;
}
#endregion
}
dotnet/src/Connectors/Connectors.Memory.Chroma/Http/HttpRequest.cs
+59
View File
@@ -0,0 +1,59 @@
// Copyright (c) Microsoft. All rights reserved.
using System.Net.Http;
using System.Text;
using System.Text.Json;
namespace Microsoft.SemanticKernel.Connectors.Memory.Chroma.Http;
internal static class HttpRequest
{
public static HttpRequestMessage CreateGetRequest(string url, object? payload = null)
{
return new HttpRequestMessage(HttpMethod.Get, url)
{
Content = GetJsonContent(payload)
};
}
public static HttpRequestMessage CreatePostRequest(string url, object? payload = null)
{
return new HttpRequestMessage(HttpMethod.Post, url)
{
Content = GetJsonContent(payload)
};
}
public static HttpRequestMessage CreatePutRequest(string url, object? payload = null)
{
return new HttpRequestMessage(HttpMethod.Put, url)
{
Content = GetJsonContent(payload)
};
}
public static HttpRequestMessage CreatePatchRequest(string url, object? payload = null)
{
return new HttpRequestMessage(new HttpMethod("PATCH"), url)
{
Content = GetJsonContent(payload)
};
}
public static HttpRequestMessage CreateDeleteRequest(string url)
{
return new HttpRequestMessage(HttpMethod.Delete, url);
}
private static StringContent? GetJsonContent(object? payload)
{
if (payload == null)
{
return null;
}
string strPayload = payload is string s ? s : JsonSerializer.Serialize(payload);
return new StringContent(strPayload, Encoding.UTF8, "application/json");
}
}
dotnet/src/Connectors/Connectors.Memory.Chroma/IChromaClient.cs
+82
View File
@@ -0,0 +1,82 @@
// Copyright (c) Microsoft. All rights reserved.
using System.Collections.Generic;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.SemanticKernel.Connectors.Memory.Chroma.Http.ApiSchema;
namespace Microsoft.SemanticKernel.Connectors.Memory.Chroma;
/// <summary>
/// Interface for client to make requests to Chroma API.
/// </summary>
public interface IChromaClient
{
/// <summary>
/// Creates Chroma collection.
/// </summary>
/// <param name="collectionName">Collection name.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
Task CreateCollectionAsync(string collectionName, CancellationToken cancellationToken = default);
/// <summary>
/// Returns collection model instance by name.
/// </summary>
/// <param name="collectionName">Collection name.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>Instance of <see cref="ChromaCollectionModel"/> model.</returns>
Task<ChromaCollectionModel?> GetCollectionAsync(string collectionName, CancellationToken cancellationToken = default);
/// <summary>
/// Removes collection by name.
/// </summary>
/// <param name="collectionName">Collection name.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
Task DeleteCollectionAsync(string collectionName, CancellationToken cancellationToken = default);
/// <summary>
/// Returns all collection names.
/// </summary>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>An asynchronous list of collection names.</returns>
IAsyncEnumerable<string> ListCollectionsAsync(CancellationToken cancellationToken = default);
/// <summary>
/// Adds embedding to specified collection.
/// </summary>
/// <param name="collectionId">Collection identifier generated by Chroma.</param>
/// <param name="ids">Array of embedding identifiers.</param>
/// <param name="embeddings">Array of embedding vectors.</param>
/// <param name="metadatas">Array of embedding metadatas.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
Task AddEmbeddingsAsync(string collectionId, string[] ids, float[][] embeddings, object[]? metadatas = null, CancellationToken cancellationToken = default);
/// <summary>
/// Returns embeddings from specified collection.
/// </summary>
/// <param name="collectionId">Collection identifier generated by Chroma.</param>
/// <param name="ids">Array of embedding identifiers.</param>
/// <param name="include">Array of entities to include in response (e.g. "embeddings", "metadatas" "documents"). For more information see: https://github.com/chroma-core/chroma/blob/main/chromadb/api/types.py</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>Instance of <see cref="ChromaEmbeddingsModel"/> model.</returns>
Task<ChromaEmbeddingsModel> GetEmbeddingsAsync(string collectionId, string[] ids, string[]? include = null, CancellationToken cancellationToken = default);
/// <summary>
/// Removes embeddings from specified collection.
/// </summary>
/// <param name="collectionId">Collection identifier generated by Chroma.</param>
/// <param name="ids">Array of embedding identifiers.</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
Task DeleteEmbeddingsAsync(string collectionId, string[] ids, CancellationToken cancellationToken = default);
/// <summary>
/// Searches nearest embeddings by distance in specified collection.
/// </summary>
/// <param name="collectionId">Collection identifier generated by Chroma.</param>
/// <param name="queryEmbeddings">Embeddings to search for.</param>
/// <param name="nResults">Number of results to return.</param>
/// <param name="include">Array of entities to include in response (e.g. "embeddings", "metadatas" "documents"). For more information see: https://github.com/chroma-core/chroma/blob/main/chromadb/api/types.py</param>
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
/// <returns>Instance of <see cref="ChromaQueryResultModel"/> model.</returns>
Task<ChromaQueryResultModel> QueryEmbeddingsAsync(string collectionId, float[][] queryEmbeddings, int nResults, string[]? include = null, CancellationToken cancellationToken = default);
}
dotnet/src/Connectors/Connectors.Memory.Chroma/README.md
+32
View File
@@ -0,0 +1,32 @@
# Microsoft.SemanticKernel.Connectors.Memory.Chroma
This assembly contains implementation of Semantic Kernel Memory Store using [Chroma](https://docs.trychroma.com/), open-source embedding database.
## Quickstart using local Chroma server
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();
```
dotnet/src/Connectors/Connectors.UnitTests/Connectors.UnitTests.csproj
+1
View File
@@ -29,6 +29,7 @@
<ProjectReference Include="..\Connectors.AI.HuggingFace\Connectors.AI.HuggingFace.csproj" />
<ProjectReference Include="..\Connectors.AI.OpenAI\Connectors.AI.OpenAI.csproj" />
<ProjectReference Include="..\Connectors.Memory.AzureCognitiveSearch\Connectors.Memory.AzureCognitiveSearch.csproj" />
<ProjectReference Include="..\Connectors.Memory.Chroma\Connectors.Memory.Chroma.csproj" />
<ProjectReference Include="..\Connectors.Memory.Pinecone\Connectors.Memory.Pinecone.csproj" />
<ProjectReference Include="..\Connectors.Memory.DuckDB\Connectors.Memory.DuckDB.csproj" />
<ProjectReference Include="..\Connectors.Memory.Qdrant\Connectors.Memory.Qdrant.csproj" />
dotnet/src/Connectors/Connectors.UnitTests/Memory/Chroma/ChromaMemoryStoreTests.cs
+334
View File
@@ -0,0 +1,334 @@
// Copyright (c) Microsoft. All rights reserved.
using System;
using System.Collections.Generic;
using System.Linq;
using System.Net.Http;
using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.SemanticKernel.AI.Embeddings;
using Microsoft.SemanticKernel.Connectors.Memory.Chroma;
using Microsoft.SemanticKernel.Connectors.Memory.Chroma.Http.ApiSchema;
using Microsoft.SemanticKernel.Memory;
using Moq;
using Xunit;
namespace SemanticKernel.Connectors.UnitTests.Memory.Chroma;
/// <summary>
/// Unit tests for <see cref="ChromaMemoryStore"/> class.
/// </summary>
public sealed class ChromaMemoryStoreTests : IDisposable
{
private const string CollectionId = "fake-collection-id";
private const string CollectionName = "fake-collection-name";
private readonly HttpMessageHandlerStub _messageHandlerStub;
private readonly HttpClient _httpClient;
private readonly Mock<IChromaClient> _chromaClientMock;
public ChromaMemoryStoreTests()
{
this._messageHandlerStub = new HttpMessageHandlerStub();
this._httpClient = this.GetHttpClientStub();
this._chromaClientMock = new Mock<IChromaClient>();
this._chromaClientMock
.Setup(client => client.GetCollectionAsync(CollectionName, CancellationToken.None))
.ReturnsAsync(new ChromaCollectionModel { Id = CollectionId, Name = CollectionName });
}
[Fact]
public async Task ItUsesProvidedEndpointFromConstructorAsync()
{
// Arrange
const string endpoint = "https://fake-random-test-host/fake-path/";
var store = new ChromaMemoryStore(this._httpClient, endpoint);
// Act
await store.GetAsync("fake-collection", "fake-key");
// Assert
Assert.StartsWith(endpoint, this._messageHandlerStub.RequestUri?.AbsoluteUri, StringComparison.OrdinalIgnoreCase);
}
[Fact]
public async Task ItUsesBaseAddressFromHttpClientAsync()
{
// Arrange
const string baseAddress = "https://fake-random-test-host/fake-path/";
using var httpClient = this.GetHttpClientStub();
httpClient.BaseAddress = new Uri(baseAddress);
var store = new ChromaMemoryStore(httpClient);
// Act
await store.GetAsync("fake-collection", "fake-key");
// Assert
Assert.StartsWith(baseAddress, this._messageHandlerStub.RequestUri?.AbsoluteUri, StringComparison.OrdinalIgnoreCase);
}
[Fact]
public async Task ItCanCreateCollectionAsync()
{
// Arrange
var store = new ChromaMemoryStore(this._chromaClientMock.Object);
// Act
await store.CreateCollectionAsync(CollectionName);
// Assert
this._chromaClientMock.Verify(client => client.CreateCollectionAsync(CollectionName, CancellationToken.None), Times.Once());
}
[Fact]
public async Task ItCanDeleteCollectionAsync()
{
// Arrange
var store = new ChromaMemoryStore(this._chromaClientMock.Object);
// Act
await store.DeleteCollectionAsync(CollectionName);
// Assert
this._chromaClientMock.Verify(client => client.DeleteCollectionAsync(CollectionName, CancellationToken.None), Times.Once());
}
[Fact]
public async Task ItThrowsExceptionOnNonExistentCollectionDeletionAsync()
{
// Arrange
const string collectionName = "non-existent-collection";
const string deleteNonExistentCollectionErrorMessage = "list index out of range";
const string expectedExceptionMessage = $"Cannot delete non-existent collection {collectionName}";
this._chromaClientMock
.Setup(client => client.DeleteCollectionAsync(collectionName, CancellationToken.None))
.Throws(new ChromaClientException(deleteNonExistentCollectionErrorMessage));
var store = new ChromaMemoryStore(this._chromaClientMock.Object);
// Act
var exception = await Record.ExceptionAsync(() => store.DeleteCollectionAsync(collectionName));
// Assert
Assert.IsType<ChromaMemoryStoreException>(exception);
Assert.Equal(expectedExceptionMessage, exception.Message);
}
[Fact]
public async Task ItReturnsTrueWhenCollectionExistsAsync()
{
// Arrange
var store = new ChromaMemoryStore(this._chromaClientMock.Object);
// Act
var doesCollectionExist = await store.DoesCollectionExistAsync(CollectionName);
// Assert
Assert.True(doesCollectionExist);
}
[Fact]
public async Task ItReturnsFalseWhenCollectionDoesNotExistAsync()
{
// Arrange
const string collectionName = "non-existent-collection";
const string collectionDoesNotExistErrorMessage = $"Collection {collectionName} does not exist";
this._chromaClientMock
.Setup(client => client.GetCollectionAsync(collectionName, CancellationToken.None))
.Throws(new ChromaClientException(collectionDoesNotExistErrorMessage));
var store = new ChromaMemoryStore(this._chromaClientMock.Object);
// Act
var doesCollectionExist = await store.DoesCollectionExistAsync(collectionName);
// Assert
Assert.False(doesCollectionExist);
}
[Fact]
public async Task ItCanGetMemoryRecordFromCollectionAsync()
{
// Arrange
var expectedMemoryRecord = this.GetRandomMemoryRecord();
var embeddingsModel = this.GetEmbeddingsModelFromMemoryRecord(expectedMemoryRecord);
this._chromaClientMock
.Setup(client => client.GetEmbeddingsAsync(CollectionId, new[] { expectedMemoryRecord.Key }, It.IsAny<string[]>(), CancellationToken.None))
.ReturnsAsync(embeddingsModel);
var store = new ChromaMemoryStore(this._chromaClientMock.Object);
// Act
var actualMemoryRecord = await store.GetAsync(CollectionName, expectedMemoryRecord.Key, withEmbedding: true);
// Assert
Assert.NotNull(actualMemoryRecord);
this.AssertMemoryRecordEqual(expectedMemoryRecord, actualMemoryRecord);
}
[Fact]
public async Task ItReturnsNullWhenMemoryRecordDoesNotExistAsync()
{
// Arrange
const string memoryRecordKey = "fake-record-key";
this._chromaClientMock
.Setup(client => client.GetEmbeddingsAsync(CollectionId, new[] { memoryRecordKey }, It.IsAny<string[]>(), CancellationToken.None))
.ReturnsAsync(new ChromaEmbeddingsModel());
var store = new ChromaMemoryStore(this._chromaClientMock.Object);
// Act
var actualMemoryRecord = await store.GetAsync(CollectionName, memoryRecordKey, withEmbedding: true);
// Assert
Assert.Null(actualMemoryRecord);
}
[Fact]
public async Task ItThrowsExceptionOnGettingMemoryRecordFromNonExistingCollectionAsync()
{
// Arrange
const string collectionName = "non-existent-collection";
const string memoryRecordKey = "fake-record-key";
const string collectionDoesNotExistErrorMessage = $"Collection {collectionName} does not exist";
this._chromaClientMock
.Setup(client => client.GetCollectionAsync(collectionName, CancellationToken.None))
.Throws(new ChromaClientException(collectionDoesNotExistErrorMessage));
var store = new ChromaMemoryStore(this._chromaClientMock.Object);
// Act
var exception = await Record.ExceptionAsync(() => store.GetAsync(collectionName, memoryRecordKey, withEmbedding: true));
// Assert
Assert.IsType<ChromaMemoryStoreException>(exception);
Assert.Equal(collectionDoesNotExistErrorMessage, exception.Message);
}
[Fact]
public async Task ItCanGetMemoryRecordBatchFromCollectionAsync()
{
// Arrange
var memoryRecord1 = this.GetRandomMemoryRecord();
var memoryRecord2 = this.GetRandomMemoryRecord();
var memoryRecord3 = this.GetRandomMemoryRecord();
var expectedMemoryRecords = new[] { memoryRecord1, memoryRecord2, memoryRecord3 };
var memoryRecordKeys = expectedMemoryRecords.Select(l => l.Key).ToArray();
var embeddingsModel = this.GetEmbeddingsModelFromMemoryRecords(expectedMemoryRecords);
this._chromaClientMock
.Setup(client => client.GetEmbeddingsAsync(CollectionId, memoryRecordKeys, It.IsAny<string[]>(), CancellationToken.None))
.ReturnsAsync(embeddingsModel);
var store = new ChromaMemoryStore(this._chromaClientMock.Object);
// Act
var actualMemoryRecords = await store.GetBatchAsync(CollectionName, memoryRecordKeys, withEmbeddings: true).ToListAsync();
// Assert
Assert.Equal(expectedMemoryRecords.Length, actualMemoryRecords.Count);
for (var i = 0; i < expectedMemoryRecords.Length; i++)
{
this.AssertMemoryRecordEqual(expectedMemoryRecords[i], actualMemoryRecords[i]);
}
}
[Fact]
public async Task ItCanReturnCollectionsAsync()
{
// Arrange
var expectedCollections = new List<string> { "fake-collection-1", "fake-collection-2", "fake-collection-3" };
this._chromaClientMock
.Setup(client => client.ListCollectionsAsync(CancellationToken.None))
.Returns(expectedCollections.ToAsyncEnumerable());
var store = new ChromaMemoryStore(this._chromaClientMock.Object);
// Act
var actualCollections = await store.GetCollectionsAsync().ToListAsync();
// Assert
Assert.Equal(expectedCollections.Count, actualCollections.Count);
for (var i = 0; i < expectedCollections.Count; i++)
{
Assert.Equal(expectedCollections[i], actualCollections[i]);
}
}
public void Dispose()
{
this._httpClient.Dispose();
this._messageHandlerStub.Dispose();
}
#region private ================================================================================
private void AssertMemoryRecordEqual(MemoryRecord expectedRecord, MemoryRecord actualRecord)
{
Assert.Equal(expectedRecord.Key, actualRecord.Key);
Assert.Equal(expectedRecord.Embedding.Vector, actualRecord.Embedding.Vector);
Assert.Equal(expectedRecord.Metadata.Id, actualRecord.Metadata.Id);
Assert.Equal(expectedRecord.Metadata.Text, actualRecord.Metadata.Text);
Assert.Equal(expectedRecord.Metadata.Description, actualRecord.Metadata.Description);
Assert.Equal(expectedRecord.Metadata.AdditionalMetadata, actualRecord.Metadata.AdditionalMetadata);
Assert.Equal(expectedRecord.Metadata.IsReference, actualRecord.Metadata.IsReference);
Assert.Equal(expectedRecord.Metadata.ExternalSourceName, actualRecord.Metadata.ExternalSourceName);
}
private HttpClient GetHttpClientStub()
{
return new HttpClient(this._messageHandlerStub, false);
}
private MemoryRecord GetRandomMemoryRecord(Embedding<float>? embedding = null)
{
var id = Guid.NewGuid().ToString();
var memoryEmbedding = embedding ?? new Embedding<float>(new[] { 1f, 3f, 5f });
return MemoryRecord.LocalRecord(
id: id,
text: "text-" + Guid.NewGuid().ToString(),
description: "description-" + Guid.NewGuid().ToString(),
embedding: memoryEmbedding,
additionalMetadata: "metadata-" + Guid.NewGuid().ToString(),
key: id);
}
private Dictionary<string, object> GetEmbeddingMetadataFromMemoryRecord(MemoryRecord memoryRecord)
{
var serialized = JsonSerializer.Serialize(memoryRecord.Metadata);
return JsonSerializer.Deserialize<Dictionary<string, object>>(serialized)!;
}
private ChromaEmbeddingsModel GetEmbeddingsModelFromMemoryRecords(MemoryRecord[] memoryRecords)
{
var embeddingsModel = new ChromaEmbeddingsModel();
embeddingsModel.Ids.AddRange(memoryRecords.Select(l => l.Key));
embeddingsModel.Embeddings.AddRange(memoryRecords.Select(l => l.Embedding.Vector.ToArray()));
embeddingsModel.Metadatas.AddRange(memoryRecords.Select(this.GetEmbeddingMetadataFromMemoryRecord));
return embeddingsModel;
}
private ChromaEmbeddingsModel GetEmbeddingsModelFromMemoryRecord(MemoryRecord memoryRecord)
{
return this.GetEmbeddingsModelFromMemoryRecords(new[] { memoryRecord });
}
#endregion
}
dotnet/src/IntegrationTests/Connectors/Memory/Chroma/ChromaMemoryStoreTests.cs
+384
View File
@@ -0,0 +1,384 @@
// Copyright (c) Microsoft. All rights reserved.
using System;
using System.Collections.Generic;
using System.Linq;
using System.Net.Http;
using System.Threading.Tasks;
using Microsoft.SemanticKernel.AI.Embeddings;
using Microsoft.SemanticKernel.Connectors.Memory.Chroma;
using Microsoft.SemanticKernel.Memory;
using Xunit;
namespace SemanticKernel.IntegrationTests.Connectors.Memory.Chroma;
/// <summary>
/// Integration tests for <see cref="ChromaMemoryStore"/> class.
/// Tests work with local Chroma server. To setup the server, see dotnet/src/Connectors/Connectors.Memory.Chroma/README.md.
/// </summary>
public sealed class ChromaMemoryStoreTests : IDisposable
{
// If null, all tests will be enabled
private const string SkipReason = "Requires Chroma server up and running";
private const string BaseAddress = "http://localhost:8000";
public ChromaMemoryStoreTests()
{
this._httpClient = new();
this._httpClient.BaseAddress = new Uri(BaseAddress);
this._chromaMemoryStore = new(this._httpClient);
}
[Fact(Skip = SkipReason)]
public async Task ItCanCreateCollectionsAsync()
{
// Arrange
var collectionName1 = this.GetRandomCollectionName();
var collectionName2 = this.GetRandomCollectionName();
var collectionName3 = this.GetRandomCollectionName();
// Act
await this._chromaMemoryStore.CreateCollectionAsync(collectionName1);
await this._chromaMemoryStore.CreateCollectionAsync(collectionName2);
await this._chromaMemoryStore.CreateCollectionAsync(collectionName3);
// Assert
var collections = await this._chromaMemoryStore.GetCollectionsAsync().ToListAsync();
Assert.Contains(collectionName1, collections);
Assert.Contains(collectionName2, collections);
Assert.Contains(collectionName3, collections);
}
[Fact(Skip = SkipReason)]
public async Task ItCanHandleDuplicateNameDuringCollectionCreationAsync()
{
// Arrange
const int expectedCollectionCount = 1;
var collectionName = this.GetRandomCollectionName();
// Act
await this._chromaMemoryStore.CreateCollectionAsync(collectionName);
await this._chromaMemoryStore.CreateCollectionAsync(collectionName);
// Assert
var collections = await this._chromaMemoryStore.GetCollectionsAsync().ToListAsync();
var filteredCollections = collections.Where(collection => collection.Equals(collectionName, StringComparison.Ordinal)).ToList();
Assert.Equal(expectedCollectionCount, filteredCollections.Count);
}
[Theory(Skip = SkipReason)]
[InlineData(true)]
[InlineData(false)]
public async Task ItCanCheckIfCollectionExistsAsync(bool createCollection)
{
// Arrange
var collectionName = this.GetRandomCollectionName();
if (createCollection)
{
await this._chromaMemoryStore.CreateCollectionAsync(collectionName);
}
// Act
bool doesCollectionExist = await this._chromaMemoryStore.DoesCollectionExistAsync(collectionName);
// Assert
Assert.Equal(createCollection, doesCollectionExist);
}
[Fact(Skip = SkipReason)]
public async Task ItCanDeleteExistingCollectionAsync()
{
// Arrange
var collectionName = this.GetRandomCollectionName();
await this._chromaMemoryStore.CreateCollectionAsync(collectionName);
var collectionsBeforeDeletion = await this._chromaMemoryStore.GetCollectionsAsync().ToListAsync();
Assert.Contains(collectionName, collectionsBeforeDeletion);
// Act
await this._chromaMemoryStore.DeleteCollectionAsync(collectionName);
// Assert
var collectionsAfterDeletion = await this._chromaMemoryStore.GetCollectionsAsync().ToListAsync();
Assert.DoesNotContain(collectionName, collectionsAfterDeletion);
}
[Fact(Skip = SkipReason)]
public async Task ItThrowsExceptionOnNonExistentCollectionDeletionAsync()
{
// Arrange
var collectionName = this.GetRandomCollectionName();
// Act
var exception = await Record.ExceptionAsync(() => this._chromaMemoryStore.DeleteCollectionAsync(collectionName));
// Assert
Assert.IsType<ChromaMemoryStoreException>(exception);
Assert.Contains(
$"Cannot delete non-existent collection {collectionName}",
exception.Message,
StringComparison.InvariantCulture);
}
[Fact(Skip = SkipReason)]
public async Task ItReturnsNullOnNonExistentRecordRetrieval()
{
// Arrange
var collectionName = this.GetRandomCollectionName();
var key = Guid.NewGuid().ToString();
await this._chromaMemoryStore.CreateCollectionAsync(collectionName);
// Act
var record = await this._chromaMemoryStore.GetAsync(collectionName, key, true);
// Assert
Assert.Null(record);
}
[Fact(Skip = SkipReason)]
public async Task ItCanUpsertMemoryRecordAsync()
{
// Arrange
var collectionName = this.GetRandomCollectionName();
var expectedRecord = this.GetRandomMemoryRecord();
await this._chromaMemoryStore.CreateCollectionAsync(collectionName);
// Act
var createdRecordKey = await this._chromaMemoryStore.UpsertAsync(collectionName, expectedRecord);
// Assert
Assert.Equal(expectedRecord.Key, createdRecordKey);
var actualRecord = await this._chromaMemoryStore.GetAsync(collectionName, expectedRecord.Key, true);
Assert.NotNull(actualRecord);
this.AssertMemoryRecordEqual(expectedRecord, actualRecord);
}
[Fact(Skip = SkipReason)]
public async Task ItCanUpsertMemoryRecordBatchAsync()
{
// Arrange
var collectionName = this.GetRandomCollectionName();
var expectedRecord1 = this.GetRandomMemoryRecord();
var expectedRecord2 = this.GetRandomMemoryRecord();
var expectedRecord3 = this.GetRandomMemoryRecord();
var batch = new List<MemoryRecord> { expectedRecord1, expectedRecord2, expectedRecord3 };
await this._chromaMemoryStore.CreateCollectionAsync(collectionName);
// Act
var createdRecordKeys = await this._chromaMemoryStore.UpsertBatchAsync(collectionName, batch).ToListAsync();
// Assert
Assert.Equal(expectedRecord1.Key, createdRecordKeys[0]);
Assert.Equal(expectedRecord2.Key, createdRecordKeys[1]);
Assert.Equal(expectedRecord3.Key, createdRecordKeys[2]);
var actualRecords = await this._chromaMemoryStore.GetBatchAsync(collectionName, batch.Select(l => l.Key), true).ToListAsync();
actualRecords.ForEach(Assert.NotNull);
this.AssertMemoryRecordEqual(expectedRecord1, actualRecords[0]);
this.AssertMemoryRecordEqual(expectedRecord2, actualRecords[1]);
this.AssertMemoryRecordEqual(expectedRecord3, actualRecords[2]);
}
[Fact(Skip = SkipReason)]
public async Task ItCanRemoveMemoryRecordAsync()
{
// Arrange
var collectionName = this.GetRandomCollectionName();
var expectedRecord = this.GetRandomMemoryRecord();
await this._chromaMemoryStore.CreateCollectionAsync(collectionName);
await this._chromaMemoryStore.UpsertAsync(collectionName, expectedRecord);
var recordBeforeDeletion = await this._chromaMemoryStore.GetAsync(collectionName, expectedRecord.Key);
Assert.NotNull(recordBeforeDeletion);
// Act
await this._chromaMemoryStore.RemoveAsync(collectionName, expectedRecord.Key);
// Assert
var recordAfterDeletion = await this._chromaMemoryStore.GetAsync(collectionName, expectedRecord.Key);
Assert.Null(recordAfterDeletion);
}
[Fact(Skip = SkipReason)]
public async Task ItCanRemoveMemoryRecordBatchAsync()
{
// Arrange
var collectionName = this.GetRandomCollectionName();
var expectedRecord1 = this.GetRandomMemoryRecord();
var expectedRecord2 = this.GetRandomMemoryRecord();
var expectedRecord3 = this.GetRandomMemoryRecord();
var batch = new List<MemoryRecord> { expectedRecord1, expectedRecord2, expectedRecord3 };
var keys = batch.Select(l => l.Key);
await this._chromaMemoryStore.CreateCollectionAsync(collectionName);
await this._chromaMemoryStore.UpsertBatchAsync(collectionName, batch).ToListAsync();
var recordsBeforeDeletion = await this._chromaMemoryStore.GetBatchAsync(collectionName, keys).ToListAsync();
Assert.Equal(batch.Count, recordsBeforeDeletion.Count);
recordsBeforeDeletion.ForEach(Assert.NotNull);
// Act
await this._chromaMemoryStore.RemoveBatchAsync(collectionName, keys);
// Assert
var recordsAfterDeletion = await this._chromaMemoryStore.GetBatchAsync(collectionName, keys).ToListAsync();
Assert.Empty(recordsAfterDeletion);
}
[Fact(Skip = SkipReason)]
public async Task ItCanGetNearestMatchAsync()
{
// Arrange
var collectionName = this.GetRandomCollectionName();
var expectedRecord1 = this.GetRandomMemoryRecord(new Embedding<float>(new[] { 10f, 10f, 10f }));
var expectedRecord2 = this.GetRandomMemoryRecord(new Embedding<float>(new[] { 5f, 5f, 5f }));
var expectedRecord3 = this.GetRandomMemoryRecord(new Embedding<float>(new[] { 1f, 1f, 1f }));
var searchEmbedding = new Embedding<float>(new[] { 2f, 2f, 2f });
var batch = new List<MemoryRecord> { expectedRecord1, expectedRecord2, expectedRecord3 };
var keys = batch.Select(l => l.Key);
await this._chromaMemoryStore.CreateCollectionAsync(collectionName);
await this._chromaMemoryStore.UpsertBatchAsync(collectionName, batch).ToListAsync();
// Act
var nearestMatch = await this._chromaMemoryStore.GetNearestMatchAsync(collectionName, searchEmbedding, withEmbedding: true);
// Assert
Assert.True(nearestMatch.HasValue);
var actualRecord = nearestMatch.Value.Item1;
Assert.NotNull(actualRecord);
this.AssertMemoryRecordEqual(expectedRecord3, actualRecord);
}
[Fact(Skip = SkipReason)]
public async Task ItCanGetNearestMatchesAsync()
{
// Arrange
var collectionName = this.GetRandomCollectionName();
var expectedRecord1 = this.GetRandomMemoryRecord(new Embedding<float>(new[] { 10f, 10f, 10f }));
var expectedRecord2 = this.GetRandomMemoryRecord(new Embedding<float>(new[] { 5f, 5f, 5f }));
var expectedRecord3 = this.GetRandomMemoryRecord(new Embedding<float>(new[] { 1f, 1f, 1f }));
var searchEmbedding = new Embedding<float>(new[] { 2f, 2f, 2f });
var batch = new List<MemoryRecord> { expectedRecord1, expectedRecord2, expectedRecord3 };
var keys = batch.Select(l => l.Key);
await this._chromaMemoryStore.CreateCollectionAsync(collectionName);
await this._chromaMemoryStore.UpsertBatchAsync(collectionName, batch).ToListAsync();
// Act
var nearestMatches = await this._chromaMemoryStore
.GetNearestMatchesAsync(collectionName, searchEmbedding, batch.Count, withEmbeddings: true)
.ToListAsync();
// Assert
Assert.NotNull(nearestMatches);
Assert.Equal(batch.Count, nearestMatches.Count);
nearestMatches.ForEach(match => Assert.NotNull(match.Item1));
var actualRecord1 = nearestMatches[0].Item1;
var actualRecord2 = nearestMatches[1].Item1;
var actualRecord3 = nearestMatches[2].Item1;
this.AssertMemoryRecordEqual(expectedRecord3, actualRecord1);
this.AssertMemoryRecordEqual(expectedRecord2, actualRecord2);
this.AssertMemoryRecordEqual(expectedRecord1, actualRecord3);
}
[Fact(Skip = SkipReason)]
public async Task ItReturnsNoMatchesFromEmptyCollection()
{
// Arrange
var collectionName = this.GetRandomCollectionName();
var searchEmbedding = new Embedding<float>(new[] { 2f, 2f, 2f });
await this._chromaMemoryStore.CreateCollectionAsync(collectionName);
// Act
var nearestMatch = await this._chromaMemoryStore.GetNearestMatchAsync(collectionName, searchEmbedding, withEmbedding: true);
// Assert
Assert.Null(nearestMatch.Value.Item1);
}
public void Dispose()
{
this.Dispose(true);
GC.SuppressFinalize(this);
}
#region private ================================================================================
private readonly HttpClient _httpClient;
private readonly ChromaMemoryStore _chromaMemoryStore;
private void Dispose(bool disposing)
{
if (disposing)
{
this._httpClient.Dispose();
}
}
private void AssertMemoryRecordEqual(MemoryRecord expectedRecord, MemoryRecord actualRecord)
{
Assert.Equal(expectedRecord.Key, actualRecord.Key);
Assert.Equal(expectedRecord.Embedding.Vector, actualRecord.Embedding.Vector);
Assert.Equal(expectedRecord.Metadata.Id, actualRecord.Metadata.Id);
Assert.Equal(expectedRecord.Metadata.Text, actualRecord.Metadata.Text);
Assert.Equal(expectedRecord.Metadata.Description, actualRecord.Metadata.Description);
Assert.Equal(expectedRecord.Metadata.AdditionalMetadata, actualRecord.Metadata.AdditionalMetadata);
Assert.Equal(expectedRecord.Metadata.IsReference, actualRecord.Metadata.IsReference);
Assert.Equal(expectedRecord.Metadata.ExternalSourceName, actualRecord.Metadata.ExternalSourceName);
}
private string GetRandomCollectionName()
{
return "sk-test-" + Guid.NewGuid();
}
private MemoryRecord GetRandomMemoryRecord(Embedding<float>? embedding = null)
{
var id = Guid.NewGuid().ToString();
var memoryEmbedding = embedding ?? new Embedding<float>(new[] { 1f, 3f, 5f });
return MemoryRecord.LocalRecord(
id: id,
text: "text-" + Guid.NewGuid().ToString(),
description: "description-" + Guid.NewGuid().ToString(),
embedding: memoryEmbedding,
additionalMetadata: "metadata-" + Guid.NewGuid().ToString(),
key: id);
}
#endregion
}
dotnet/src/IntegrationTests/IntegrationTests.csproj
+1
View File
@@ -34,6 +34,7 @@
<ProjectReference Include="..\Connectors\Connectors.AI.HuggingFace\Connectors.AI.HuggingFace.csproj" />
<ProjectReference Include="..\Connectors\Connectors.Memory.Postgres\Connectors.Memory.Postgres.csproj" />
<ProjectReference Include="..\Connectors\Connectors.Memory.Weaviate\Connectors.Memory.Weaviate.csproj" />
<ProjectReference Include="..\Connectors\Connectors.Memory.Chroma\Connectors.Memory.Chroma.csproj" />
<ProjectReference Include="..\Extensions\Planning.ActionPlanner\Planning.ActionPlanner.csproj" />
<ProjectReference Include="..\Extensions\Planning.SequentialPlanner\Planning.SequentialPlanner.csproj" />
<ProjectReference Include="..\Skills\Skills.Web\Skills.Web.csproj" />