[Store] Add score to ChromaDb query results

Author: OskarStark

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

@@ -70,6 +70,7 @@ public function query(Vector $vector, array $options = []): iterable
                 id: Uuid::fromString($queryResponse->ids[0][$i]),
                 vector: new Vector($queryResponse->embeddings[0][$i]),
                 metadata: new Metadata($queryResponse->metadatas[0][$i]),
+                score: $queryResponse->distances[0][$i] ?? null,
             );
         }
     }

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

@@ -354,6 +354,71 @@ public function testQueryWithEmptyResults()
         $this->assertCount(0, $documents);
     }
 
+    public function testQueryReturnsDistancesAsScore()
+    {
+        $queryVector = new Vector([0.15, 0.25, 0.35]);
+        $queryResponse = new QueryItemsResponse(
+            ids: [['01234567-89ab-cdef-0123-456789abcdef', 'fedcba98-7654-3210-fedc-ba9876543210']],
+            embeddings: [[[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]],
+            metadatas: [[['title' => 'Doc 1'], ['title' => 'Doc 2']]],
+            documents: null,
+            data: null,
+            uris: null,
+            distances: [[0.123, 0.456]]
+        );
+
+        $collection = $this->createMock(CollectionResource::class);
+        $client = $this->createMock(Client::class);
+
+        $client->expects($this->once())
+            ->method('getOrCreateCollection')
+            ->with('test-collection')
+            ->willReturn($collection);
+
+        $collection->expects($this->once())
+            ->method('query')
+            ->willReturn($queryResponse);
+
+        $store = new Store($client, 'test-collection');
+        $documents = iterator_to_array($store->query($queryVector));
+
+        $this->assertCount(2, $documents);
+        $this->assertSame(0.123, $documents[0]->score);
+        $this->assertSame(0.456, $documents[1]->score);
+    }
+
+    public function testQueryReturnsNullScoreWhenDistancesNotAvailable()
+    {
+        $queryVector = new Vector([0.15, 0.25, 0.35]);
+        $queryResponse = new QueryItemsResponse(
+            ids: [['01234567-89ab-cdef-0123-456789abcdef']],
+            embeddings: [[[0.1, 0.2, 0.3]]],
+            metadatas: [[['title' => 'Doc 1']]],
+            documents: null,
+            data: null,
+            uris: null,
+            distances: null
+        );
+
+        $collection = $this->createMock(CollectionResource::class);
+        $client = $this->createMock(Client::class);
+
+        $client->expects($this->once())
+            ->method('getOrCreateCollection')
+            ->with('test-collection')
+            ->willReturn($collection);
+
+        $collection->expects($this->once())
+            ->method('query')
+            ->willReturn($queryResponse);
+
+        $store = new Store($client, 'test-collection');
+        $documents = iterator_to_array($store->query($queryVector));
+
+        $this->assertCount(1, $documents);
+        $this->assertNull($documents[0]->score);
+    }
+
     /**
      * @param array{where?: array<string, string>, whereDocument?: array<string, mixed>} $options
      * @param array<string, mixed>|null                                                  $expectedWhere