Skip to content

Commit 14f747f

Browse files
cincuranetcincuranet
authored
Implement strongly typed Where filter (#65)
1 parent a7c7172 commit 14f747f

6 files changed

+370
-77
lines changed

‎ChromaDB.Client.Tests/CollectionClientCUDTests.cs

+1-4
Original file line numberDiff line numberDiff line change
@@ -223,10 +223,7 @@ await client.Add([id1, id2],
223223
{ "key", "value" },
224224
}]);
225225
await client.Delete([id1, id2],
226-
where: new Dictionary<string, object>
227-
{
228-
{ "key", "value" },
229-
});
226+
where: ChromaWhere.Equal("key", "value"));
230227
}
231228

232229
[Test]

‎ChromaDB.Client.Tests/CollectionClientGetTests.cs

+104-20
Original file line numberDiff line numberDiff line change
@@ -193,11 +193,11 @@ public async Task GetLimitOffsetIncludeAll()
193193
}
194194

195195
[Test]
196-
public async Task GetWhereIncludeAll()
196+
public async Task GetWhereEqualIncludeAll()
197197
{
198198
var client = await Init();
199199
var result = await client.Get(
200-
where: new Dictionary<string, object> { { MetadataKey2, Metadata2[MetadataKey2] } },
200+
where: ChromaWhere.Equal(MetadataKey2, Metadata2[MetadataKey2]),
201201
include: ChromaGetInclude.Embeddings | ChromaGetInclude.Metadatas | ChromaGetInclude.Documents);
202202
Assert.That(result, Has.Count.EqualTo(1));
203203
Assert.That(result[0].Id, Is.EqualTo(Id2));
@@ -207,11 +207,67 @@ public async Task GetWhereIncludeAll()
207207
}
208208

209209
[Test]
210-
public async Task GetWhereOperatorIncludeAll()
210+
public async Task GetWhereNotEqualIncludeAll()
211211
{
212212
var client = await Init();
213213
var result = await client.Get(
214-
where: new Dictionary<string, object> { { MetadataKey2, new Dictionary<string, object> { { "$lt", Metadata2[MetadataKey2] } } } },
214+
where: ChromaWhere.NotEqual(MetadataKey2, Metadata1[MetadataKey2]),
215+
include: ChromaGetInclude.Embeddings | ChromaGetInclude.Metadatas | ChromaGetInclude.Documents);
216+
Assert.That(result, Has.Count.EqualTo(1));
217+
Assert.That(result[0].Id, Is.EqualTo(Id2));
218+
Assert.That(result[0].Embeddings, Is.EqualTo(Embeddings2).Using(EmbeddingsComparer.Instance));
219+
Assert.That(result[0].Metadata, Is.EqualTo(Metadata2));
220+
Assert.That(result[0].Document, Is.EqualTo(Doc2));
221+
}
222+
223+
[Test]
224+
public async Task GetWhereInIncludeAll()
225+
{
226+
var client = await Init();
227+
var result = await client.Get(
228+
where: ChromaWhere.In(MetadataKey2, Metadata2[MetadataKey2]),
229+
include: ChromaGetInclude.Embeddings | ChromaGetInclude.Metadatas | ChromaGetInclude.Documents);
230+
Assert.That(result, Has.Count.EqualTo(1));
231+
Assert.That(result[0].Id, Is.EqualTo(Id2));
232+
Assert.That(result[0].Embeddings, Is.EqualTo(Embeddings2).Using(EmbeddingsComparer.Instance));
233+
Assert.That(result[0].Metadata, Is.EqualTo(Metadata2));
234+
Assert.That(result[0].Document, Is.EqualTo(Doc2));
235+
}
236+
237+
[Test]
238+
public async Task GetWhereNotInIncludeAll()
239+
{
240+
var client = await Init();
241+
var result = await client.Get(
242+
where: ChromaWhere.NotIn(MetadataKey2, Metadata1[MetadataKey2]),
243+
include: ChromaGetInclude.Embeddings | ChromaGetInclude.Metadatas | ChromaGetInclude.Documents);
244+
Assert.That(result, Has.Count.EqualTo(1));
245+
Assert.That(result[0].Id, Is.EqualTo(Id2));
246+
Assert.That(result[0].Embeddings, Is.EqualTo(Embeddings2).Using(EmbeddingsComparer.Instance));
247+
Assert.That(result[0].Metadata, Is.EqualTo(Metadata2));
248+
Assert.That(result[0].Document, Is.EqualTo(Doc2));
249+
}
250+
251+
[Test]
252+
public async Task GetWhereGreaterThanIncludeAll()
253+
{
254+
var client = await Init();
255+
var result = await client.Get(
256+
where: ChromaWhere.GreaterThan(MetadataKey2, Metadata1[MetadataKey2]),
257+
include: ChromaGetInclude.Embeddings | ChromaGetInclude.Metadatas | ChromaGetInclude.Documents);
258+
Assert.That(result, Has.Count.EqualTo(1));
259+
Assert.That(result[0].Id, Is.EqualTo(Id2));
260+
Assert.That(result[0].Embeddings, Is.EqualTo(Embeddings2).Using(EmbeddingsComparer.Instance));
261+
Assert.That(result[0].Metadata, Is.EqualTo(Metadata2));
262+
Assert.That(result[0].Document, Is.EqualTo(Doc2));
263+
}
264+
265+
[Test]
266+
public async Task GetWhereLessThanIncludeAll()
267+
{
268+
var client = await Init();
269+
var result = await client.Get(
270+
where: ChromaWhere.LessThan(MetadataKey2, Metadata2[MetadataKey2]),
215271
include: ChromaGetInclude.Embeddings | ChromaGetInclude.Metadatas | ChromaGetInclude.Documents);
216272
Assert.That(result, Has.Count.EqualTo(1));
217273
Assert.That(result[0].Id, Is.EqualTo(Id1));
@@ -221,39 +277,67 @@ public async Task GetWhereOperatorIncludeAll()
221277
}
222278

223279
[Test]
224-
public async Task GetWhereDocumentContainsIncludeDocuments()
280+
public async Task GetWhereGreaterThanOrEqualIncludeAll()
225281
{
226282
var client = await Init();
227283
var result = await client.Get(
228-
whereDocument: ChromaWhereDocument.Contains(Doc1[^1]),
229-
include: ChromaGetInclude.Documents);
284+
where: ChromaWhere.GreaterThanOrEqual(MetadataKey2, Metadata2[MetadataKey2]),
285+
include: ChromaGetInclude.Embeddings | ChromaGetInclude.Metadatas | ChromaGetInclude.Documents);
230286
Assert.That(result, Has.Count.EqualTo(1));
231-
Assert.That(result[0].Id, Is.EqualTo(Id1));
232-
Assert.That(result[0].Embeddings, Is.Null);
233-
Assert.That(result[0].Metadata, Is.Null);
234-
Assert.That(result[0].Document, Is.EqualTo(Doc1));
287+
Assert.That(result[0].Id, Is.EqualTo(Id2));
288+
Assert.That(result[0].Embeddings, Is.EqualTo(Embeddings2).Using(EmbeddingsComparer.Instance));
289+
Assert.That(result[0].Metadata, Is.EqualTo(Metadata2));
290+
Assert.That(result[0].Document, Is.EqualTo(Doc2));
235291
}
236292

237293
[Test]
238-
public async Task GetWhereDocumentNotContainsIncludeDocuments()
294+
public async Task GetWhereLessThanOrEqualIncludeAll()
239295
{
240296
var client = await Init();
241297
var result = await client.Get(
242-
whereDocument: ChromaWhereDocument.NotContains(Doc2[^1]),
243-
include: ChromaGetInclude.Documents);
298+
where: ChromaWhere.LessThanOrEqual(MetadataKey2, Metadata1[MetadataKey2]),
299+
include: ChromaGetInclude.Embeddings | ChromaGetInclude.Metadatas | ChromaGetInclude.Documents);
244300
Assert.That(result, Has.Count.EqualTo(1));
245301
Assert.That(result[0].Id, Is.EqualTo(Id1));
246-
Assert.That(result[0].Embeddings, Is.Null);
247-
Assert.That(result[0].Metadata, Is.Null);
302+
Assert.That(result[0].Embeddings, Is.EqualTo(Embeddings1).Using(EmbeddingsComparer.Instance));
303+
Assert.That(result[0].Metadata, Is.EqualTo(Metadata1));
248304
Assert.That(result[0].Document, Is.EqualTo(Doc1));
249305
}
250306

251307
[Test]
252-
public async Task GetWhereDocumentAndIncludeDocuments()
308+
public async Task GetWhereAndIncludeAll()
309+
{
310+
var client = await Init();
311+
var result = await client.Get(
312+
where: ChromaWhere.Equal(MetadataKey2, Metadata2[MetadataKey2]) && ChromaWhere.NotEqual(MetadataKey2, Metadata1[MetadataKey2]),
313+
include: ChromaGetInclude.Embeddings | ChromaGetInclude.Metadatas | ChromaGetInclude.Documents);
314+
Assert.That(result, Has.Count.EqualTo(1));
315+
Assert.That(result[0].Id, Is.EqualTo(Id2));
316+
Assert.That(result[0].Embeddings, Is.EqualTo(Embeddings2).Using(EmbeddingsComparer.Instance));
317+
Assert.That(result[0].Metadata, Is.EqualTo(Metadata2));
318+
Assert.That(result[0].Document, Is.EqualTo(Doc2));
319+
}
320+
321+
[Test]
322+
public async Task GetWhereOrIncludeAll()
323+
{
324+
var client = await Init();
325+
var result = await client.Get(
326+
where: ChromaWhere.Equal(MetadataKey2, Metadata2[MetadataKey2]) || ChromaWhere.NotEqual(MetadataKey2, Metadata1[MetadataKey2]),
327+
include: ChromaGetInclude.Embeddings | ChromaGetInclude.Metadatas | ChromaGetInclude.Documents);
328+
Assert.That(result, Has.Count.EqualTo(1));
329+
Assert.That(result[0].Id, Is.EqualTo(Id2));
330+
Assert.That(result[0].Embeddings, Is.EqualTo(Embeddings2).Using(EmbeddingsComparer.Instance));
331+
Assert.That(result[0].Metadata, Is.EqualTo(Metadata2));
332+
Assert.That(result[0].Document, Is.EqualTo(Doc2));
333+
}
334+
335+
[Test]
336+
public async Task GetWhereDocumentContainsIncludeDocuments()
253337
{
254338
var client = await Init();
255339
var result = await client.Get(
256-
whereDocument: ChromaWhereDocument.Contains(Doc1[^1]) && ChromaWhereDocument.NotContains(Doc2[^1]),
340+
whereDocument: ChromaWhereDocument.Contains(Doc1[^1]),
257341
include: ChromaGetInclude.Documents);
258342
Assert.That(result, Has.Count.EqualTo(1));
259343
Assert.That(result[0].Id, Is.EqualTo(Id1));
@@ -263,11 +347,11 @@ public async Task GetWhereDocumentAndIncludeDocuments()
263347
}
264348

265349
[Test]
266-
public async Task GetWhereDocumentOrIncludeDocuments()
350+
public async Task GetWhereDocumentNotContainsIncludeDocuments()
267351
{
268352
var client = await Init();
269353
var result = await client.Get(
270-
whereDocument: ChromaWhereDocument.Contains(Doc1[^1]) || ChromaWhereDocument.NotContains(Doc2[^1]),
354+
whereDocument: ChromaWhereDocument.NotContains(Doc2[^1]),
271355
include: ChromaGetInclude.Documents);
272356
Assert.That(result, Has.Count.EqualTo(1));
273357
Assert.That(result[0].Id, Is.EqualTo(Id1));

0 commit comments

Comments
 (0)