[Store][InMemory] Fix cosine similarity sorting

Author: RamyHakam

src/store/src/InMemoryStore.php

@@ -20,7 +20,7 @@
  */
 final class InMemoryStore implements VectorStoreInterface
 {
-    public const COSINE_SIMILARITY = 'cosine';
+    public const COSINE_DISTANCE = 'cosine';
     public const ANGULAR_DISTANCE = 'angular';
     public const EUCLIDEAN_DISTANCE = 'euclidean';
     public const MANHATTAN_DISTANCE = 'manhattan';
@@ -32,7 +32,7 @@ final class InMemoryStore implements VectorStoreInterface
     private array $documents = [];
 
     public function __construct(
-        private readonly string $similarity = self::COSINE_SIMILARITY,
+        private readonly string $distance = self::COSINE_DISTANCE,
     ) {
     }
 
@@ -48,13 +48,13 @@ public function add(VectorDocument ...$documents): void
      */
     public function query(Vector $vector, array $options = [], ?float $minScore = null): array
     {
-        $strategy = match ($this->similarity) {
-            self::COSINE_SIMILARITY => $this->cosineSimilarity(...),
+        $strategy = match ($this->distance) {
+            self::COSINE_DISTANCE => $this->cosineDistance(...),
             self::ANGULAR_DISTANCE => $this->angularDistance(...),
             self::EUCLIDEAN_DISTANCE => $this->euclideanDistance(...),
             self::MANHATTAN_DISTANCE => $this->manhattanDistance(...),
             self::CHEBYSHEV_DISTANCE => $this->chebyshevDistance(...),
-            default => throw new InvalidArgumentException(\sprintf('Unsupported similarity strategy "%s"', $this->similarity)),
+            default => throw new InvalidArgumentException(\sprintf('Unsupported distance metric "%s"', $this->distance)),
         };
 
         $currentEmbeddings = array_map(
@@ -80,6 +80,11 @@ public function query(Vector $vector, array $options = [], ?float $minScore = nu
         );
     }
 
+    private function cosineDistance(VectorDocument $embedding, Vector $against): float
+    {
+        return 1 - $this->cosineSimilarity($embedding, $against);
+    }
+
     private function cosineSimilarity(VectorDocument $embedding, Vector $against): float
     {
         $currentEmbeddingVectors = $embedding->vector->getData();

src/store/tests/InMemoryStoreTest.php

@@ -23,7 +23,7 @@
 final class InMemoryStoreTest extends TestCase
 {
     #[Test]
-    public function storeCanSearchUsingCosineSimilarity(): void
+    public function storeCanSearchUsingCosineDistance(): void
     {
         $store = new InMemoryStore();
         $store->add(
@@ -32,19 +32,44 @@ public function storeCanSearchUsingCosineSimilarity(): void
             new VectorDocument(Uuid::v4(), new Vector([0.3, 0.7, 0.1])),
         );
 
-        self::assertCount(3, $store->query(new Vector([0.0, 0.1, 0.6])));
+        $result = $store->query(new Vector([0.0, 0.1, 0.6]));
+        self::assertCount(3, $result);
+        self::assertSame([0.1, 0.1, 0.5], $result[0]->vector->getData());
 
         $store->add(
             new VectorDocument(Uuid::v4(), new Vector([0.1, 0.1, 0.5])),
             new VectorDocument(Uuid::v4(), new Vector([0.7, -0.3, 0.0])),
             new VectorDocument(Uuid::v4(), new Vector([0.3, 0.7, 0.1])),
         );
 
-        self::assertCount(6, $store->query(new Vector([0.0, 0.1, 0.6])));
+        $result = $store->query(new Vector([0.0, 0.1, 0.6]));
+        self::assertCount(6, $result);
+        self::assertSame([0.1, 0.1, 0.5], $result[0]->vector->getData());
     }
 
     #[Test]
-    public function storeCanSearchUsingCosineSimilarityWithMaxItems(): void
+    public function storeCanSearchUsingCosineDistanceAndReturnCorrectOrder(): void
+    {
+        $store = new InMemoryStore();
+        $store->add(
+            new VectorDocument(Uuid::v4(), new Vector([0.1, 0.1, 0.5])),
+            new VectorDocument(Uuid::v4(), new Vector([0.7, -0.3, 0.0])),
+            new VectorDocument(Uuid::v4(), new Vector([0.3, 0.7, 0.1])),
+            new VectorDocument(Uuid::v4(), new Vector([0.3, 0.1, 0.6])),
+            new VectorDocument(Uuid::v4(), new Vector([0.0, 0.1, 0.6])),
+        );
+
+        $result = $store->query(new Vector([0.0, 0.1, 0.6]));
+        self::assertCount(5, $result);
+        self::assertSame([0.0, 0.1, 0.6], $result[0]->vector->getData());
+        self::assertSame([0.1, 0.1, 0.5], $result[1]->vector->getData());
+        self::assertSame([0.3, 0.1, 0.6], $result[2]->vector->getData());
+        self::assertSame([0.3, 0.7, 0.1], $result[3]->vector->getData());
+        self::assertSame([0.7, -0.3, 0.0], $result[4]->vector->getData());
+    }
+
+    #[Test]
+    public function storeCanSearchUsingCosineDistanceWithMaxItems(): void
     {
         $store = new InMemoryStore();
         $store->add(