feature #257 [Store] Add PSR-6 support for vector store (Guikingone)

This PR was squashed before being merged into the main branch.

Discussion
----------

[Store] Add PSR-6 support for vector store

| Q             | A
| ------------- | ---
| Bug fix?      | no
| New feature?  | yes
| Docs?         | yes
| Issues        | None
| License       | MIT

Hi 👋🏻

This PR aims to introduce a `CacheStore` based on PSR-6, it's heavily inspired by the `CacheStore` from `MessageStore` layer, this store introduce a refactoring of the `InMemoryStore` to use common logic for both stores (as we can't rely on the `CacheItemPool` implementation).

Thanks for the feedback.

Commits
-------

3d82bf8 [Store] Add PSR-6 support for vector store

Author: OskarStark

examples/rag/cache.php

@@ -0,0 +1,63 @@
+<?php
+
+/*
+ * This file is part of the Symfony package.
+ *
+ * (c) Fabien Potencier <[email protected]>
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+
+use Symfony\AI\Agent\Agent;
+use Symfony\AI\Agent\Toolbox\AgentProcessor;
+use Symfony\AI\Agent\Toolbox\Tool\SimilaritySearch;
+use Symfony\AI\Agent\Toolbox\Toolbox;
+use Symfony\AI\Fixtures\Movies;
+use Symfony\AI\Platform\Bridge\OpenAi\Embeddings;
+use Symfony\AI\Platform\Bridge\OpenAi\Gpt;
+use Symfony\AI\Platform\Bridge\OpenAi\PlatformFactory;
+use Symfony\AI\Platform\Message\Message;
+use Symfony\AI\Platform\Message\MessageBag;
+use Symfony\AI\Store\CacheStore;
+use Symfony\AI\Store\Document\Metadata;
+use Symfony\AI\Store\Document\TextDocument;
+use Symfony\AI\Store\Document\Vectorizer;
+use Symfony\AI\Store\Indexer;
+use Symfony\Component\Cache\Adapter\ArrayAdapter;
+use Symfony\Component\Uid\Uuid;
+
+require_once dirname(__DIR__).'/bootstrap.php';
+
+// initialize the store
+$store = new CacheStore(new ArrayAdapter());
+
+// create embeddings and documents
+foreach (Movies::all() as $i => $movie) {
+    $documents[] = new TextDocument(
+        id: Uuid::v4(),
+        content: 'Title: '.$movie['title'].\PHP_EOL.'Director: '.$movie['director'].\PHP_EOL.'Description: '.$movie['description'],
+        metadata: new Metadata($movie),
+    );
+}
+
+// create embeddings for documents
+$platform = PlatformFactory::create(env('OPENAI_API_KEY'), http_client());
+$vectorizer = new Vectorizer($platform, $embeddings = new Embeddings());
+$indexer = new Indexer($vectorizer, $store, logger());
+$indexer->index($documents);
+
+$model = new Gpt(Gpt::GPT_4O_MINI);
+
+$similaritySearch = new SimilaritySearch($platform, $embeddings, $store);
+$toolbox = new Toolbox([$similaritySearch], logger: logger());
+$processor = new AgentProcessor($toolbox);
+$agent = new Agent($platform, $model, [$processor], [$processor], logger());
+
+$messages = new MessageBag(
+    Message::forSystem('Please answer all user questions only using SimilaritySearch function.'),
+    Message::ofUser('Which movie fits the theme of the mafia?')
+);
+$result = $agent->call($messages);
+
+echo $result->getContent().\PHP_EOL;

src/ai-bundle/config/options.php

@@ -161,6 +161,15 @@
                             ->end()
                         ->end()
                     ->end()
+                    ->arrayNode('cache')
+                        ->normalizeKeys(false)
+                        ->useAttributeAsKey('name')
+                        ->arrayPrototype()
+                            ->children()
+                                ->scalarNode('service')->cannotBeEmpty()->defaultValue('cache.app')->end()
+                            ->end()
+                        ->end()
+                    ->end()
                     ->arrayNode('chroma_db')
                         ->normalizeKeys(false)
                         ->useAttributeAsKey('name')

src/ai-bundle/src/AiBundle.php

@@ -48,6 +48,7 @@
 use Symfony\AI\Store\Bridge\Qdrant\Store as QdrantStore;
 use Symfony\AI\Store\Bridge\SurrealDb\Store as SurrealDbStore;
 use Symfony\AI\Store\Bridge\Typesense\Store as TypesenseStore;
+use Symfony\AI\Store\CacheStore;
 use Symfony\AI\Store\Document\Vectorizer;
 use Symfony\AI\Store\Indexer;
 use Symfony\AI\Store\InMemoryStore;
@@ -471,6 +472,21 @@ private function processStoreConfig(string $type, array $stores, ContainerBuilde
             }
         }
 
+        if ('cache' === $type) {
+            foreach ($stores as $name => $store) {
+                $arguments = [
+                    new Reference($store['service']),
+                ];
+
+                $definition = new Definition(CacheStore::class);
+                $definition
+                    ->addTag('ai.store')
+                    ->setArguments($arguments);
+
+                $container->setDefinition('ai.store.'.$type.'.'.$name, $definition);
+            }
+        }
+
         if ('chroma_db' === $type) {
             foreach ($stores as $name => $store) {
                 $definition = new Definition(ChromaDbStore::class);

src/ai-bundle/tests/DependencyInjection/AiBundleTest.php

@@ -201,6 +201,11 @@ private function getFullConfig(): array
                             'vector_field' => 'contentVector',
                         ],
                     ],
+                    'cache' => [
+                        'my_cache_store' => [
+                            'service' => 'cache.system',
+                        ],
+                    ],
                     'chroma_db' => [
                         'my_chroma_store' => [
                             'collection' => 'my_collection',

src/store/CHANGELOG.md

@@ -28,7 +28,7 @@ CHANGELOG
    - Orchestrates document processing pipeline
    - Accepts TextDocuments, vectorizes and stores in chunks
    - Configurable batch processing
- * Add `InMemoryStore` implementation with multiple distance algorithms:
+ * Add `InMemoryStore` and `CacheStore` implementations with multiple distance algorithms:
    - Cosine similarity
    - Angular distance
    - Euclidean distance

src/store/composer.json

@@ -40,7 +40,8 @@
         "mongodb/mongodb": "^1.21 || ^2.0",
         "phpstan/phpstan": "^2.0",
         "phpunit/phpunit": "^11.5",
-        "probots-io/pinecone-php": "^1.0"
+        "probots-io/pinecone-php": "^1.0",
+        "symfony/cache": "^7.3"
     },
     "config": {
         "sort-packages": true

src/store/doc/index.rst

@@ -43,10 +43,17 @@ You can find more advanced usage in combination with an Agent using the store fo
 * `Similarity Search with MongoDB (RAG)`_
 * `Similarity Search with Neo4j (RAG)`_
 * `Similarity Search with Pinecone (RAG)`_
+* `Similarity Search with PSR-6 Cache (RAG)`_
 * `Similarity Search with Qdrant (RAG)`_
 * `Similarity Search with SurrealDB (RAG)`_
 * `Similarity Search with Typesense (RAG)`_
 
+.. note::
+
+    Both `InMemory` and `PSR-6 cache` vector stores will load all the data into the
+    memory of the PHP process. They can be used only the amount of data fits in the
+    PHP memory limit, typically for testing.
+
 Supported Stores
 ----------------
 
@@ -59,6 +66,7 @@ Supported Stores
 * `Neo4j`_
 * `Pinecone`_ (requires `probots-io/pinecone-php` as additional dependency)
 * `Postgres`_ (requires `ext-pdo`)
+* `PSR-6 Cache`_
 * `Qdrant`_
 * `SurrealDB`_
 * `Typesense`_
@@ -97,12 +105,13 @@ This leads to a store implementing two methods::
 .. _`Retrieval Augmented Generation`: https://de.wikipedia.org/wiki/Retrieval-Augmented_Generation
 .. _`Similarity Search with MariaDB (RAG)`: https://github.com/symfony/ai/blob/main/examples/rag/mariadb.php
 .. _`Similarity Search with MongoDB (RAG)`: https://github.com/symfony/ai/blob/main/examples/rag/mongodb.php
-.. _`Similarity Search with Pinecone (RAG)`: https://github.com/symfony/ai/blob/main/examples/rag/pinecone.php
 .. _`Similarity Search with Meilisearch (RAG)`: https://github.com/symfony/ai/blob/main/examples/rag/meilisearch.php
-.. _`Similarity Search with SurrealDB (RAG)`: https://github.com/symfony/ai/blob/main/examples/rag/surrealdb.php
 .. _`Similarity Search with memory storage (RAG)`: https://github.com/symfony/ai/blob/main/examples/rag/in-memory.php
-.. _`Similarity Search with Qdrant (RAG)`: https://github.com/symfony/ai/blob/main/examples/rag/qdrant.php
 .. _`Similarity Search with Neo4j (RAG)`: https://github.com/symfony/ai/blob/main/examples/rag/neo4j.php
+.. _`Similarity Search with Pinecone (RAG)`: https://github.com/symfony/ai/blob/main/examples/rag/pinecone.php
+.. _`Similarity Search with PSR-6 Cache (RAG)`: https://github.com/symfony/ai/blob/main/examples/rag/cache.php
+.. _`Similarity Search with Qdrant (RAG)`: https://github.com/symfony/ai/blob/main/examples/rag/qdrant.php
+.. _`Similarity Search with SurrealDB (RAG)`: https://github.com/symfony/ai/blob/main/examples/rag/surrealdb.php
 .. _`Similarity Search with Typesense (RAG)`: https://github.com/symfony/ai/blob/main/examples/rag/typesense.php
 .. _`Azure AI Search`: https://azure.microsoft.com/products/ai-services/ai-search
 .. _`Chroma`: https://www.trychroma.com/
@@ -117,3 +126,4 @@ This leads to a store implementing two methods::
 .. _`Neo4j`: https://neo4j.com/
 .. _`Typesense`: https://typesense.org/
 .. _`GitHub`: https://github.com/symfony/ai/issues/16
+.. _`PSR-6 Cache`: https://www.php-fig.org/psr/psr-6/

src/store/src/CacheStore.php

@@ -0,0 +1,74 @@
+<?php
+
+/*
+ * This file is part of the Symfony package.
+ *
+ * (c) Fabien Potencier <[email protected]>
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+
+namespace Symfony\AI\Store;
+
+use Psr\Cache\CacheItemPoolInterface;
+use Symfony\AI\Platform\Vector\Vector;
+use Symfony\AI\Store\Document\Metadata;
+use Symfony\AI\Store\Document\VectorDocument;
+use Symfony\AI\Store\Exception\RuntimeException;
+use Symfony\Component\Uid\Uuid;
+use Symfony\Contracts\Cache\CacheInterface;
+
+/**
+ * @author Guillaume Loulier <[email protected]>
+ */
+final readonly class CacheStore implements VectorStoreInterface
+{
+    public function __construct(
+        private CacheInterface&CacheItemPoolInterface $cache,
+        private DistanceCalculator $distanceCalculator = new DistanceCalculator(),
+        private string $cacheKey = '_vectors',
+    ) {
+        if (!interface_exists(CacheItemPoolInterface::class)) {
+            throw new RuntimeException('For using the CacheStore as vector store, a PSR-6 cache implementation is required. Try running "composer require symfony/cache" or another PSR-6 compatible cache.');
+        }
+    }
+
+    public function add(VectorDocument ...$documents): void
+    {
+        $existingVectors = $this->cache->get($this->cacheKey, static fn (): array => []);
+
+        $newVectors = array_map(static fn (VectorDocument $document): array => [
+            'id' => $document->id->toRfc4122(),
+            'vector' => $document->vector->getData(),
+            'metadata' => $document->metadata->getArrayCopy(),
+        ], $documents);
+
+        $cacheItem = $this->cache->getItem($this->cacheKey);
+
+        $cacheItem->set([
+            ...$existingVectors,
+            ...$newVectors,
+        ]);
+
+        $this->cache->save($cacheItem);
+    }
+
+    /**
+     * @param array{
+     *     maxItems?: positive-int
+     * } $options If maxItems is provided, only the top N results will be returned
+     */
+    public function query(Vector $vector, array $options = []): array
+    {
+        $documents = $this->cache->getItem($this->cacheKey)->get() ?? [];
+
+        $vectorDocuments = array_map(static fn (array $document): VectorDocument => new VectorDocument(
+            id: Uuid::fromString($document['id']),
+            vector: new Vector($document['vector']),
+            metadata: new Metadata($document['metadata']),
+        ), $documents);
+
+        return $this->distanceCalculator->calculate($vectorDocuments, $vector, $options['maxItems'] ?? null);
+    }
+}

src/store/src/DistanceCalculator.php

@@ -0,0 +1,133 @@
+<?php
+
+/*
+ * This file is part of the Symfony package.
+ *
+ * (c) Fabien Potencier <[email protected]>
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+
+namespace Symfony\AI\Store;
+
+use Symfony\AI\Platform\Vector\Vector;
+use Symfony\AI\Store\Document\VectorDocument;
+
+/**
+ * @author Guillaume Loulier <[email protected]>
+ */
+final readonly class DistanceCalculator
+{
+    public function __construct(
+        private DistanceStrategy $strategy = DistanceStrategy::COSINE_DISTANCE,
+    ) {
+    }
+
+    /**
+     * @param VectorDocument[] $documents
+     * @param ?int             $maxItems  If maxItems is provided, only the top N results will be returned
+     *
+     * @return VectorDocument[]
+     */
+    public function calculate(array $documents, Vector $vector, ?int $maxItems = null): array
+    {
+        $strategy = match ($this->strategy) {
+            DistanceStrategy::COSINE_DISTANCE => $this->cosineDistance(...),
+            DistanceStrategy::ANGULAR_DISTANCE => $this->angularDistance(...),
+            DistanceStrategy::EUCLIDEAN_DISTANCE => $this->euclideanDistance(...),
+            DistanceStrategy::MANHATTAN_DISTANCE => $this->manhattanDistance(...),
+            DistanceStrategy::CHEBYSHEV_DISTANCE => $this->chebyshevDistance(...),
+        };
+
+        $currentEmbeddings = array_map(
+            static fn (VectorDocument $vectorDocument): array => [
+                'distance' => $strategy($vectorDocument, $vector),
+                'document' => $vectorDocument,
+            ],
+            $documents,
+        );
+
+        usort(
+            $currentEmbeddings,
+            static fn (array $embedding, array $nextEmbedding): int => $embedding['distance'] <=> $nextEmbedding['distance'],
+        );
+
+        if (null !== $maxItems && $maxItems < \count($currentEmbeddings)) {
+            $currentEmbeddings = \array_slice($currentEmbeddings, 0, $maxItems);
+        }
+
+        return array_map(
+            static fn (array $embedding): VectorDocument => $embedding['document'],
+            $currentEmbeddings,
+        );
+    }
+
+    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();
+
+        $dotProduct = array_sum(array: array_map(
+            static fn (float $a, float $b): float => $a * $b,
+            $currentEmbeddingVectors,
+            $against->getData(),
+        ));
+
+        $currentEmbeddingLength = sqrt(array_sum(array_map(
+            static fn (float $value): float => $value ** 2,
+            $currentEmbeddingVectors,
+        )));
+
+        $againstLength = sqrt(array_sum(array_map(
+            static fn (float $value): float => $value ** 2,
+            $against->getData(),
+        )));
+
+        return fdiv($dotProduct, $currentEmbeddingLength * $againstLength);
+    }
+
+    private function angularDistance(VectorDocument $embedding, Vector $against): float
+    {
+        $cosineSimilarity = $this->cosineSimilarity($embedding, $against);
+
+        return fdiv(acos($cosineSimilarity), \M_PI);
+    }
+
+    private function euclideanDistance(VectorDocument $embedding, Vector $against): float
+    {
+        return sqrt(array_sum(array_map(
+            static fn (float $a, float $b): float => ($a - $b) ** 2,
+            $embedding->vector->getData(),
+            $against->getData(),
+        )));
+    }
+
+    private function manhattanDistance(VectorDocument $embedding, Vector $against): float
+    {
+        return array_sum(array_map(
+            static fn (float $a, float $b): float => abs($a - $b),
+            $embedding->vector->getData(),
+            $against->getData(),
+        ));
+    }
+
+    private function chebyshevDistance(VectorDocument $embedding, Vector $against): float
+    {
+        $embeddingsAsPower = array_map(
+            static fn (float $currentValue, float $againstValue): float => abs($currentValue - $againstValue),
+            $embedding->vector->getData(),
+            $against->getData(),
+        );
+
+        return array_reduce(
+            array: $embeddingsAsPower,
+            callback: static fn (float $value, float $current): float => max($value, $current),
+            initial: 0.0,
+        );
+    }
+}