Skip to content

Commit 65023b2

Browse files
committed
feature #1029 [Store][ChromaDB] Add score to query results (OskarStark)
This PR was merged into the main branch. Discussion ---------- [Store][ChromaDB] Add score to query results | Q | A | ------------- | --- | Bug fix? | no | New feature? | yes | Docs? | no | Issues | -- | License | MIT Commits ------- 7dc5458 [Store] Add score to ChromaDb query results
2 parents bf5a432 + 7dc5458 commit 65023b2

File tree

2 files changed

+66
-0
lines changed

2 files changed

+66
-0
lines changed

src/store/src/Bridge/ChromaDb/Store.php

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@ public function query(Vector $vector, array $options = []): iterable
7070
id: Uuid::fromString($queryResponse->ids[0][$i]),
7171
vector: new Vector($queryResponse->embeddings[0][$i]),
7272
metadata: new Metadata($queryResponse->metadatas[0][$i]),
73+
score: $queryResponse->distances[0][$i] ?? null,
7374
);
7475
}
7576
}

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

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -354,6 +354,71 @@ public function testQueryWithEmptyResults()
354354
$this->assertCount(0, $documents);
355355
}
356356

357+
public function testQueryReturnsDistancesAsScore()
358+
{
359+
$queryVector = new Vector([0.15, 0.25, 0.35]);
360+
$queryResponse = new QueryItemsResponse(
361+
ids: [['01234567-89ab-cdef-0123-456789abcdef', 'fedcba98-7654-3210-fedc-ba9876543210']],
362+
embeddings: [[[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]],
363+
metadatas: [[['title' => 'Doc 1'], ['title' => 'Doc 2']]],
364+
documents: null,
365+
data: null,
366+
uris: null,
367+
distances: [[0.123, 0.456]]
368+
);
369+
370+
$collection = $this->createMock(CollectionResource::class);
371+
$client = $this->createMock(Client::class);
372+
373+
$client->expects($this->once())
374+
->method('getOrCreateCollection')
375+
->with('test-collection')
376+
->willReturn($collection);
377+
378+
$collection->expects($this->once())
379+
->method('query')
380+
->willReturn($queryResponse);
381+
382+
$store = new Store($client, 'test-collection');
383+
$documents = iterator_to_array($store->query($queryVector));
384+
385+
$this->assertCount(2, $documents);
386+
$this->assertSame(0.123, $documents[0]->score);
387+
$this->assertSame(0.456, $documents[1]->score);
388+
}
389+
390+
public function testQueryReturnsNullScoreWhenDistancesNotAvailable()
391+
{
392+
$queryVector = new Vector([0.15, 0.25, 0.35]);
393+
$queryResponse = new QueryItemsResponse(
394+
ids: [['01234567-89ab-cdef-0123-456789abcdef']],
395+
embeddings: [[[0.1, 0.2, 0.3]]],
396+
metadatas: [[['title' => 'Doc 1']]],
397+
documents: null,
398+
data: null,
399+
uris: null,
400+
distances: null
401+
);
402+
403+
$collection = $this->createMock(CollectionResource::class);
404+
$client = $this->createMock(Client::class);
405+
406+
$client->expects($this->once())
407+
->method('getOrCreateCollection')
408+
->with('test-collection')
409+
->willReturn($collection);
410+
411+
$collection->expects($this->once())
412+
->method('query')
413+
->willReturn($queryResponse);
414+
415+
$store = new Store($client, 'test-collection');
416+
$documents = iterator_to_array($store->query($queryVector));
417+
418+
$this->assertCount(1, $documents);
419+
$this->assertNull($documents[0]->score);
420+
}
421+
357422
/**
358423
* @param array{where?: array<string, string>, whereDocument?: array<string, mixed>} $options
359424
* @param array<string, mixed>|null $expectedWhere

0 commit comments

Comments
 (0)