Skip to content

Commit 48db6ae

Browse files
OskarStarkOskarStark
committed
minor #342 [Store][ChromaDB] Add tests for query filtering (OskarStark)
This PR was merged into the main branch. Discussion ---------- [Store][ChromaDB] Add tests for query filtering | Q | A | ------------- | --- | Bug fix? | no | New feature? | no | Docs? | no | Issues | Follows #340 | License | MIT fyi `@valx76` Commits ------- 02a62db Add tests for ChromaDB query filtering feature
2 parents b6cae5c + 02a62db commit 48db6ae

File tree

1 file changed

+343
-0
lines changed

1 file changed

+343
-0
lines changed

‎src/store/tests/Bridge/ChromaDb/StoreTest.php

Lines changed: 343 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
namespace Symfony\AI\Store\Tests\Bridge\ChromaDb;
1313

1414
use Codewithkyrian\ChromaDB\Client;
15+
use Codewithkyrian\ChromaDB\Generated\Responses\QueryItemsResponse;
1516
use Codewithkyrian\ChromaDB\Resources\CollectionResource;
1617
use PHPUnit\Framework\Attributes\CoversClass;
1718
use PHPUnit\Framework\Attributes\DataProvider;
@@ -136,4 +137,346 @@ public static function addDocumentsProvider(): \Iterator
136137
'expectedOriginalDocuments' => [''],
137138
];
138139
}
140+
141+
public function testQueryWithoutFilters()
142+
{
143+
$queryVector = new Vector([0.15, 0.25, 0.35]);
144+
$queryResponse = new QueryItemsResponse(
145+
ids: [['01234567-89ab-cdef-0123-456789abcdef', 'fedcba98-7654-3210-fedc-ba9876543210']],
146+
embeddings: [[[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]],
147+
metadatas: [[['title' => 'Doc 1'], ['title' => 'Doc 2']]],
148+
documents: null,
149+
data: null,
150+
uris: null,
151+
distances: null
152+
);
153+
154+
$collection = $this->createMock(CollectionResource::class);
155+
$client = $this->createMock(Client::class);
156+
157+
$client->expects($this->once())
158+
->method('getOrCreateCollection')
159+
->with('test-collection')
160+
->willReturn($collection);
161+
162+
$collection->expects($this->once())
163+
->method('query')
164+
->with(
165+
[[0.15, 0.25, 0.35]], // queryEmbeddings
166+
null, // queryTexts
167+
null, // queryImages
168+
4, // nResults
169+
null, // where
170+
null, // whereDocument
171+
null // include
172+
)
173+
->willReturn($queryResponse);
174+
175+
$store = new Store($client, 'test-collection');
176+
$documents = $store->query($queryVector);
177+
178+
$this->assertCount(2, $documents);
179+
$this->assertSame('01234567-89ab-cdef-0123-456789abcdef', (string) $documents[0]->id);
180+
$this->assertSame([0.1, 0.2, 0.3], $documents[0]->vector->getData());
181+
$this->assertSame(['title' => 'Doc 1'], $documents[0]->metadata->getArrayCopy());
182+
}
183+
184+
public function testQueryWithWhereFilter()
185+
{
186+
$queryVector = new Vector([0.15, 0.25, 0.35]);
187+
$whereFilter = ['category' => 'technology'];
188+
189+
$queryResponse = new QueryItemsResponse(
190+
ids: [['01234567-89ab-cdef-0123-456789abcdef']],
191+
embeddings: [[[0.1, 0.2, 0.3]]],
192+
metadatas: [[['title' => 'Tech Doc', 'category' => 'technology']]],
193+
documents: null,
194+
data: null,
195+
uris: null,
196+
distances: null
197+
);
198+
199+
$collection = $this->createMock(CollectionResource::class);
200+
$client = $this->createMock(Client::class);
201+
202+
$client->expects($this->once())
203+
->method('getOrCreateCollection')
204+
->with('test-collection')
205+
->willReturn($collection);
206+
207+
$collection->expects($this->once())
208+
->method('query')
209+
->with(
210+
[[0.15, 0.25, 0.35]], // queryEmbeddings
211+
null, // queryTexts
212+
null, // queryImages
213+
4, // nResults
214+
['category' => 'technology'], // where
215+
null, // whereDocument
216+
null // include
217+
)
218+
->willReturn($queryResponse);
219+
220+
$store = new Store($client, 'test-collection');
221+
$documents = $store->query($queryVector, ['where' => $whereFilter]);
222+
223+
$this->assertCount(1, $documents);
224+
$this->assertSame('01234567-89ab-cdef-0123-456789abcdef', (string) $documents[0]->id);
225+
$this->assertSame(['title' => 'Tech Doc', 'category' => 'technology'], $documents[0]->metadata->getArrayCopy());
226+
}
227+
228+
public function testQueryWithWhereDocumentFilter()
229+
{
230+
$queryVector = new Vector([0.15, 0.25, 0.35]);
231+
$whereDocumentFilter = ['$contains' => 'machine learning'];
232+
233+
$queryResponse = new QueryItemsResponse(
234+
ids: [['01234567-89ab-cdef-0123-456789abcdef', 'fedcba98-7654-3210-fedc-ba9876543210']],
235+
embeddings: [[[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]],
236+
metadatas: [[['title' => 'ML Doc 1'], ['title' => 'ML Doc 2']]],
237+
documents: null,
238+
data: null,
239+
uris: null,
240+
distances: null
241+
);
242+
243+
$collection = $this->createMock(CollectionResource::class);
244+
$client = $this->createMock(Client::class);
245+
246+
$client->expects($this->once())
247+
->method('getOrCreateCollection')
248+
->with('test-collection')
249+
->willReturn($collection);
250+
251+
$collection->expects($this->once())
252+
->method('query')
253+
->with(
254+
[[0.15, 0.25, 0.35]], // queryEmbeddings
255+
null, // queryTexts
256+
null, // queryImages
257+
4, // nResults
258+
null, // where
259+
['$contains' => 'machine learning'], // whereDocument
260+
null // include
261+
)
262+
->willReturn($queryResponse);
263+
264+
$store = new Store($client, 'test-collection');
265+
$documents = $store->query($queryVector, ['whereDocument' => $whereDocumentFilter]);
266+
267+
$this->assertCount(2, $documents);
268+
$this->assertSame('01234567-89ab-cdef-0123-456789abcdef', (string) $documents[0]->id);
269+
$this->assertSame('fedcba98-7654-3210-fedc-ba9876543210', (string) $documents[1]->id);
270+
}
271+
272+
public function testQueryWithBothFilters()
273+
{
274+
$queryVector = new Vector([0.15, 0.25, 0.35]);
275+
$whereFilter = ['category' => 'AI', 'status' => 'published'];
276+
$whereDocumentFilter = ['$contains' => 'neural networks'];
277+
278+
$queryResponse = new QueryItemsResponse(
279+
ids: [['01234567-89ab-cdef-0123-456789abcdef']],
280+
embeddings: [[[0.1, 0.2, 0.3]]],
281+
metadatas: [[['title' => 'AI Neural Networks', 'category' => 'AI', 'status' => 'published']]],
282+
documents: null,
283+
data: null,
284+
uris: null,
285+
distances: null
286+
);
287+
288+
$collection = $this->createMock(CollectionResource::class);
289+
$client = $this->createMock(Client::class);
290+
291+
$client->expects($this->once())
292+
->method('getOrCreateCollection')
293+
->with('test-collection')
294+
->willReturn($collection);
295+
296+
$collection->expects($this->once())
297+
->method('query')
298+
->with(
299+
[[0.15, 0.25, 0.35]], // queryEmbeddings
300+
null, // queryTexts
301+
null, // queryImages
302+
4, // nResults
303+
['category' => 'AI', 'status' => 'published'], // where
304+
['$contains' => 'neural networks'], // whereDocument
305+
null // include
306+
)
307+
->willReturn($queryResponse);
308+
309+
$store = new Store($client, 'test-collection');
310+
$documents = $store->query($queryVector, [
311+
'where' => $whereFilter,
312+
'whereDocument' => $whereDocumentFilter,
313+
]);
314+
315+
$this->assertCount(1, $documents);
316+
$this->assertSame('01234567-89ab-cdef-0123-456789abcdef', (string) $documents[0]->id);
317+
$this->assertSame(['title' => 'AI Neural Networks', 'category' => 'AI', 'status' => 'published'], $documents[0]->metadata->getArrayCopy());
318+
}
319+
320+
public function testQueryWithEmptyResults()
321+
{
322+
$queryVector = new Vector([0.15, 0.25, 0.35]);
323+
$whereFilter = ['category' => 'nonexistent'];
324+
325+
$queryResponse = new QueryItemsResponse(
326+
ids: [[]],
327+
embeddings: [[]],
328+
metadatas: [[]],
329+
documents: null,
330+
data: null,
331+
uris: null,
332+
distances: null
333+
);
334+
335+
$collection = $this->createMock(CollectionResource::class);
336+
$client = $this->createMock(Client::class);
337+
338+
$client->expects($this->once())
339+
->method('getOrCreateCollection')
340+
->with('test-collection')
341+
->willReturn($collection);
342+
343+
$collection->expects($this->once())
344+
->method('query')
345+
->with(
346+
[[0.15, 0.25, 0.35]], // queryEmbeddings
347+
null, // queryTexts
348+
null, // queryImages
349+
4, // nResults
350+
['category' => 'nonexistent'], // where
351+
null, // whereDocument
352+
null // include
353+
)
354+
->willReturn($queryResponse);
355+
356+
$store = new Store($client, 'test-collection');
357+
$documents = $store->query($queryVector, ['where' => $whereFilter]);
358+
359+
$this->assertCount(0, $documents);
360+
}
361+
362+
/**
363+
* @param array{where?: array<string, string>, whereDocument?: array<string, mixed>} $options
364+
* @param array<string, mixed>|null $expectedWhere
365+
* @param array<string, mixed>|null $expectedWhereDocument
366+
*/
367+
#[DataProvider('queryFilterProvider')]
368+
public function testQueryWithVariousFilterCombinations(
369+
array $options,
370+
?array $expectedWhere,
371+
?array $expectedWhereDocument,
372+
): void {
373+
$queryVector = new Vector([0.1, 0.2, 0.3]);
374+
375+
$queryResponse = new QueryItemsResponse(
376+
ids: [['01234567-89ab-cdef-0123-456789abcdef']],
377+
embeddings: [[[0.1, 0.2, 0.3]]],
378+
metadatas: [[['title' => 'Test Doc']]],
379+
documents: null,
380+
data: null,
381+
uris: null,
382+
distances: null
383+
);
384+
385+
$collection = $this->createMock(CollectionResource::class);
386+
$client = $this->createMock(Client::class);
387+
388+
$client->expects($this->once())
389+
->method('getOrCreateCollection')
390+
->with('test-collection')
391+
->willReturn($collection);
392+
393+
$collection->expects($this->once())
394+
->method('query')
395+
->with(
396+
[[0.1, 0.2, 0.3]], // queryEmbeddings
397+
null, // queryTexts
398+
null, // queryImages
399+
4, // nResults
400+
$expectedWhere, // where
401+
$expectedWhereDocument,// whereDocument
402+
null // include
403+
)
404+
->willReturn($queryResponse);
405+
406+
$store = new Store($client, 'test-collection');
407+
$documents = $store->query($queryVector, $options);
408+
409+
$this->assertCount(1, $documents);
410+
}
411+
412+
/**
413+
* @return \Iterator<string, array{
414+
* options: array{where?: array<string, string>, whereDocument?: array<string, mixed>},
415+
* expectedWhere: array<string, mixed>|null,
416+
* expectedWhereDocument: array<string, mixed>|null
417+
* }>
418+
*/
419+
public static function queryFilterProvider(): \Iterator
420+
{
421+
yield 'empty options' => [
422+
'options' => [],
423+
'expectedWhere' => null,
424+
'expectedWhereDocument' => null,
425+
];
426+
427+
yield 'only where filter' => [
428+
'options' => ['where' => ['type' => 'article']],
429+
'expectedWhere' => ['type' => 'article'],
430+
'expectedWhereDocument' => null,
431+
];
432+
433+
yield 'only whereDocument filter' => [
434+
'options' => ['whereDocument' => ['$contains' => 'search term']],
435+
'expectedWhere' => null,
436+
'expectedWhereDocument' => ['$contains' => 'search term'],
437+
];
438+
439+
yield 'both filters' => [
440+
'options' => [
441+
'where' => ['status' => 'active'],
442+
'whereDocument' => ['$not_contains' => 'draft'],
443+
],
444+
'expectedWhere' => ['status' => 'active'],
445+
'expectedWhereDocument' => ['$not_contains' => 'draft'],
446+
];
447+
448+
yield 'complex where filter' => [
449+
'options' => [
450+
'where' => [
451+
'category' => 'technology',
452+
'author' => 'john.doe',
453+
'status' => 'published',
454+
],
455+
],
456+
'expectedWhere' => [
457+
'category' => 'technology',
458+
'author' => 'john.doe',
459+
'status' => 'published',
460+
],
461+
'expectedWhereDocument' => null,
462+
];
463+
464+
yield 'complex whereDocument filter' => [
465+
'options' => [
466+
'whereDocument' => [
467+
'$and' => [
468+
['$contains' => 'AI'],
469+
['$contains' => 'machine learning'],
470+
],
471+
],
472+
],
473+
'expectedWhere' => null,
474+
'expectedWhereDocument' => [
475+
'$and' => [
476+
['$contains' => 'AI'],
477+
['$contains' => 'machine learning'],
478+
],
479+
],
480+
];
481+
}
139482
}

0 commit comments

Comments
 (0)