From 184c45d26f813ef01b7dfec276803848ead5b150 Mon Sep 17 00:00:00 2001 From: Jiri Cincura Date: Mon, 24 Jun 2024 13:35:08 +0200 Subject: [PATCH] Implement create collection. --- ChromaDB.Client.Tests/ChromaDBAuthTests.cs | 2 +- .../ChromaDBCollectionTests.cs | 55 +++++++++++++++++++ ChromaDB.Client/Models/Collection.cs | 1 + .../Requests/DBCreateCollectionRequest.cs | 12 ++++ .../Implementations/ChromaDBClient.cs | 10 ++++ .../Services/Interfaces/IChromaDBClient.cs | 2 + README.md | 3 +- Samples/ChromaDB.Client.Sample/Program.cs | 7 ++- 8 files changed, 89 insertions(+), 3 deletions(-) create mode 100644 ChromaDB.Client.Tests/ChromaDBCollectionTests.cs create mode 100644 ChromaDB.Client/Models/Requests/DBCreateCollectionRequest.cs diff --git a/ChromaDB.Client.Tests/ChromaDBAuthTests.cs b/ChromaDB.Client.Tests/ChromaDBAuthTests.cs index 88f6207..6ae7b26 100644 --- a/ChromaDB.Client.Tests/ChromaDBAuthTests.cs +++ b/ChromaDB.Client.Tests/ChromaDBAuthTests.cs @@ -11,6 +11,6 @@ public async Task NoAuth() { using var httpClient = new ChromaDBHttpClient(ConfigurationOptions); var client = new ChromaDBClient(ConfigurationOptions, httpClient); - await Assert.ThatAsync(client.Heartbeat, Throws.Nothing); + await Assert.ThatAsync(client.Heartbeat, Throws.Nothing); } } diff --git a/ChromaDB.Client.Tests/ChromaDBCollectionTests.cs b/ChromaDB.Client.Tests/ChromaDBCollectionTests.cs new file mode 100644 index 0000000..5d2db4f --- /dev/null +++ b/ChromaDB.Client.Tests/ChromaDBCollectionTests.cs @@ -0,0 +1,55 @@ +using ChromaDB.Client.Models.Requests; +using ChromaDB.Client.Services.Implementations; +using NUnit.Framework; + +namespace ChromaDB.Client.Tests; + +[TestFixture] +public class ChromaDBCollectionTests : ChromaDBTestsBase +{ + [Test] + public async Task CreateCollectionWithoutMetadata() + { + using var httpClient = new ChromaDBHttpClient(ConfigurationOptions); + var client = new ChromaDBClient(ConfigurationOptions, httpClient); + var result = await client.CreateCollection(new DBCreateCollectionRequest() + { + Name = $"collection{Random.Shared.Next()}", + }); + Assert.That(result.Success, Is.True); + } + + [Test] + public async Task CreateCollectionWithMetadata() + { + using var httpClient = new ChromaDBHttpClient(ConfigurationOptions); + var client = new ChromaDBClient(ConfigurationOptions, httpClient); + var result = await client.CreateCollection(new DBCreateCollectionRequest() + { + Name = $"collection{Random.Shared.Next()}", + Metadata = new Dictionary() + { + { "test", "foo" }, + { "test2", 10 }, + }, + }); + Assert.That(result.Success, Is.True); + } + + [Test] + public async Task CreateCollectionAlreadyExists() + { + using var httpClient = new ChromaDBHttpClient(ConfigurationOptions); + var client = new ChromaDBClient(ConfigurationOptions, httpClient); + await client.CreateCollection(new DBCreateCollectionRequest() + { + Name = $"collection_exists", + }); + var result = await client.CreateCollection(new DBCreateCollectionRequest() + { + Name = $"collection_exists", + }); + Assert.That(result.Success, Is.False); + Assert.That(result.ReasonPhrase, Is.Not.Null.And.Not.Empty); + } +} diff --git a/ChromaDB.Client/Models/Collection.cs b/ChromaDB.Client/Models/Collection.cs index 2160ea1..031de32 100644 --- a/ChromaDB.Client/Models/Collection.cs +++ b/ChromaDB.Client/Models/Collection.cs @@ -9,6 +9,7 @@ namespace ChromaDB.Client.Models; [ChromaGetRoute(Endpoint = "collections?tenant={tenant}&database={database}", Source = typeof(Collection), ResponseType = typeof(List))] [ChromaPostRoute(Endpoint = "collections/{collection_id}/get", Source = typeof(Collection), RequestType = typeof(CollectionGetRequest), ResponseType = typeof(CollectionEntriesResponse))] [ChromaPostRoute(Endpoint = "collections/{collection_id}/query", Source = typeof(Collection), RequestType = typeof(CollectionQueryRequest), ResponseType = typeof(CollectionEntriesQueryResponse))] +[ChromaPostRoute(Endpoint = "collections", Source = typeof(Collection), RequestType = typeof(DBCreateCollectionRequest), ResponseType = typeof(Collection))] public class Collection { [JsonPropertyName("id")] diff --git a/ChromaDB.Client/Models/Requests/DBCreateCollectionRequest.cs b/ChromaDB.Client/Models/Requests/DBCreateCollectionRequest.cs new file mode 100644 index 0000000..809e0c8 --- /dev/null +++ b/ChromaDB.Client/Models/Requests/DBCreateCollectionRequest.cs @@ -0,0 +1,12 @@ +using System.Text.Json.Serialization; + +namespace ChromaDB.Client.Models.Requests; + +public class DBCreateCollectionRequest +{ + [JsonPropertyName("name")] + public required string Name { get; init; } + + [JsonPropertyName("metadata")] + public IDictionary? Metadata { get; init; } +} diff --git a/ChromaDB.Client/Services/Implementations/ChromaDBClient.cs b/ChromaDB.Client/Services/Implementations/ChromaDBClient.cs index 556f677..fa94361 100644 --- a/ChromaDB.Client/Services/Implementations/ChromaDBClient.cs +++ b/ChromaDB.Client/Services/Implementations/ChromaDBClient.cs @@ -51,4 +51,14 @@ public async Task> Heartbeat() { return await _httpClient.Get(new RequestQueryParams()); } + + public async Task> CreateCollection(DBCreateCollectionRequest request, string? tenant = null, string? database = null) + { + tenant = tenant is not null and not [] ? tenant : _currentTenant.Name; + database = database is not null and not [] ? database : _currentDatabase.Name; + RequestQueryParams requestParams = new RequestQueryParams() + .Add("{tenant}", tenant) + .Add("{database}", database); + return await _httpClient.Post(request, requestParams); + } } diff --git a/ChromaDB.Client/Services/Interfaces/IChromaDBClient.cs b/ChromaDB.Client/Services/Interfaces/IChromaDBClient.cs index cacd057..f03ca99 100644 --- a/ChromaDB.Client/Services/Interfaces/IChromaDBClient.cs +++ b/ChromaDB.Client/Services/Interfaces/IChromaDBClient.cs @@ -1,4 +1,5 @@ using ChromaDB.Client.Models; +using ChromaDB.Client.Models.Requests; namespace ChromaDB.Client.Services.Interfaces; @@ -7,4 +8,5 @@ public interface IChromaDBClient Task> GetCollectionByName(string name, string? tenant = null, string? database = null); Task>> GetCollections(string? tenant = null, string? database = null); Task> Heartbeat(); + Task> CreateCollection(DBCreateCollectionRequest request, string? tenant = null, string? database = null); } diff --git a/README.md b/README.md index a43132c..c3b982f 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,8 @@ With ChromaDB.Client, you can easily connect to a ChromaDB instance, create and ## Features - [x] Basic connection and authentication (partially done) -- [ ] Collection creation and deletion +- [x] Collection creation (partially done) +- [ ] Collection deletion - [x] Collection retrieval and modification (partially done) - [ ] Document insertion, deletion, and update - [ ] Document retrieval by ID or filter diff --git a/Samples/ChromaDB.Client.Sample/Program.cs b/Samples/ChromaDB.Client.Sample/Program.cs index bf4c25c..1a34d8d 100644 --- a/Samples/ChromaDB.Client.Sample/Program.cs +++ b/Samples/ChromaDB.Client.Sample/Program.cs @@ -9,6 +9,11 @@ using IChromaDBHttpClient httpClient = new ChromaDBHttpClient(configOptions); IChromaDBClient client = new ChromaDBClient(configOptions, httpClient); +BaseResponse createCollectionResponse = await client.CreateCollection(new DBCreateCollectionRequest() +{ + Name = "string5", +}, database: "test", tenant: "nedeljko"); + BaseResponse> collections = await client.GetCollections(database: "test", tenant: "nedeljko"); BaseResponse collection1 = await client.GetCollectionByName("string5", database: "test", tenant: "nedeljko"); @@ -18,7 +23,7 @@ BaseResponse> getResponse = await string5Client.Get(new CollectionGetRequest() { Ids = ["340a36ad-c38a-406c-be38-250174aee5a4"], - Include = ["metadatas", "documents", "embeddings"] + Include = ["metadatas", "documents", "embeddings"], }); BaseResponse queryResponse = await string5Client.Query(new CollectionQueryRequest()