feat: implement document loader & transformer pipeline for store indexing

Author: chr-hertel

examples/store/document-splitting.php

@@ -0,0 +1,16 @@
+<?php
+
+declare(strict_types=1);
+
+use PhpLlm\LlmChain\Store\Document\Loader\TextFileLoader;
+use PhpLlm\LlmChain\Store\Document\Transformer\TextSplitTransformer;
+
+require_once dirname(__DIR__, 2).'/vendor/autoload.php';
+
+$loader = new TextFileLoader();
+$splitter = new TextSplitTransformer();
+$source = dirname(__DIR__, 2).'/tests/Fixture/lorem.txt';
+
+$documents = iterator_to_array($splitter($loader($source)));
+
+dump($documents);

src/Store/Document/Loader/TextFileLoader.php

@@ -0,0 +1,34 @@
+<?php
+
+declare(strict_types=1);
+
+namespace PhpLlm\LlmChain\Store\Document\Loader;
+
+use PhpLlm\LlmChain\Store\Document\LoaderInterface;
+use PhpLlm\LlmChain\Store\Document\Metadata;
+use PhpLlm\LlmChain\Store\Document\TextDocument;
+use PhpLlm\LlmChain\Store\Exception\RuntimeException;
+use Symfony\Component\Uid\Uuid;
+
+/**
+ * @author Christopher Hertel <[email protected]>
+ */
+final readonly class TextFileLoader implements LoaderInterface
+{
+    public function __invoke(string $source, array $options = []): iterable
+    {
+        if (!is_file($source)) {
+            throw new RuntimeException(\sprintf('File "%s" does not exist.', $source));
+        }
+
+        $content = file_get_contents($source);
+
+        if (false === $content) {
+            throw new RuntimeException(\sprintf('Unable to read file "%s"', $source));
+        }
+
+        yield new TextDocument(Uuid::v4(), trim($content), new Metadata([
+            'source' => $source,
+        ]));
+    }
+}

src/Store/Document/LoaderInterface.php

@@ -0,0 +1,19 @@
+<?php
+
+declare(strict_types=1);
+
+namespace PhpLlm\LlmChain\Store\Document;
+
+/**
+ * @author Christopher Hertel <[email protected]>
+ */
+interface LoaderInterface
+{
+    /**
+     * @param string               $source  Identifier for the loader to load the documents from, e.g. file path, folder, or URL.
+     * @param array<string, mixed> $options loader specific set of options to control the loading process
+     *
+     * @return iterable<TextDocument> iterable of TextDocuments loaded from the source
+     */
+    public function __invoke(string $source, array $options = []): iterable;
+}

src/Store/Document/Transformer/ChainTransformer.php

@@ -0,0 +1,32 @@
+<?php
+
+declare(strict_types=1);
+
+namespace PhpLlm\LlmChain\Store\Document\Transformer;
+
+use PhpLlm\LlmChain\Store\Document\TransformerInterface;
+
+final readonly class ChainTransformer implements TransformerInterface
+{
+    /**
+     * @var TransformerInterface[]
+     */
+    private array $transformers;
+
+    /**
+     * @param iterable<TransformerInterface> $transformers
+     */
+    public function __construct(iterable $transformers)
+    {
+        $this->transformers = $transformers instanceof \Traversable ? iterator_to_array($transformers) : $transformers;
+    }
+
+    public function __invoke(iterable $documents, array $options = []): iterable
+    {
+        foreach ($this->transformers as $transformer) {
+            $documents = $transformer($documents, $options);
+        }
+
+        return $documents;
+    }
+}

src/Store/Document/Transformer/ChunkDelayTransformer.php

@@ -0,0 +1,44 @@
+<?php
+
+declare(strict_types=1);
+
+namespace PhpLlm\LlmChain\Store\Document\Transformer;
+
+use PhpLlm\LlmChain\Store\Document\TransformerInterface;
+use Symfony\Component\Clock\ClockInterface;
+
+/**
+ * This transformer splits the batch of documents into chunks and delays in-between with x seconds, which is useful
+ * when indexing a lot of documents and facing API rate limits.
+ *
+ * @author Christopher Hertel <[email protected]>
+ */
+final readonly class ChunkDelayTransformer implements TransformerInterface
+{
+    public const OPTION_CHUNK_SIZE = 'chunk_size';
+    public const OPTION_DELAY = 'delay';
+
+    public function __construct(
+        private ClockInterface $clock,
+    ) {
+    }
+
+    /**
+     * @param array{chunk_size?: int, delay?: int} $options
+     */
+    public function __invoke(iterable $documents, array $options = []): iterable
+    {
+        $chunkSize = $options[self::OPTION_CHUNK_SIZE] ?? 50;
+        $delay = $options[self::OPTION_DELAY] ?? 10;
+
+        $counter = 0;
+        foreach ($documents as $document) {
+            yield $document;
+            ++$counter;
+
+            if ($chunkSize === $counter && 0 !== $delay) {
+                $this->clock->sleep($delay);
+            }
+        }
+    }
+}

src/Store/Document/Transformer/TextSplitTransformer.php

@@ -0,0 +1,62 @@
+<?php
+
+declare(strict_types=1);
+
+namespace PhpLlm\LlmChain\Store\Document\Transformer;
+
+use PhpLlm\LlmChain\Store\Document\Metadata;
+use PhpLlm\LlmChain\Store\Document\TextDocument;
+use PhpLlm\LlmChain\Store\Document\TransformerInterface;
+use PhpLlm\LlmChain\Store\Exception\InvalidArgumentException;
+use Symfony\Component\Uid\Uuid;
+
+/**
+ * Splits a TextDocument into smaller chunks of specified size with optional overlap.
+ * If the document's content is shorter than the specified chunk size, it returns the original document as a single chunk.
+ * Overlap cannot be negative and must be less than the chunk size.
+ *
+ * @author Christopher Hertel <[email protected]>
+ */
+final readonly class TextSplitTransformer implements TransformerInterface
+{
+    public const OPTION_CHUNK_SIZE = 'chunk_size';
+    public const OPTION_OVERLAP = 'overlap';
+
+    /**
+     * @param array{chunk_size?: int, overlap?: int} $options
+     */
+    public function __invoke(iterable $documents, array $options = []): iterable
+    {
+        $chunkSize = $options[self::OPTION_CHUNK_SIZE] ?? 1000;
+        $overlap = $options[self::OPTION_OVERLAP] ?? 200;
+
+        if ($overlap < 0 || $overlap >= $chunkSize) {
+            throw new InvalidArgumentException('Overlap must be non-negative and less than chunk size.');
+        }
+
+        foreach ($documents as $document) {
+            if (mb_strlen($document->content) <= $chunkSize) {
+                yield $document;
+
+                continue;
+            }
+
+            $text = $document->content;
+            $length = mb_strlen($text);
+            $start = 0;
+
+            while ($start < $length) {
+                $end = min($start + $chunkSize, $length);
+                $chunkText = mb_substr($text, $start, $end - $start);
+
+                yield new TextDocument(Uuid::v4(), $chunkText, new Metadata([
+                    'parent_id' => $document->id,
+                    'text' => $chunkText,
+                    ...$document->metadata,
+                ]));
+
+                $start += ($chunkSize - $overlap);
+            }
+        }
+    }
+}

src/Store/Document/TransformerInterface.php

@@ -0,0 +1,23 @@
+<?php
+
+declare(strict_types=1);
+
+namespace PhpLlm\LlmChain\Store\Document;
+
+/**
+ * A Transformer is designed to mutate a stream of TextDocuments with the purpose of preparing them for indexing.
+ * It can reduce or expand the number of documents, modify their content or metadata.
+ * It should not act blocking, but is expected to iterate over incoming documents and yield prepared ones.
+ *
+ * @author Christopher Hertel <[email protected]>
+ */
+interface TransformerInterface
+{
+    /**
+     * @param iterable<TextDocument> $documents
+     * @param array<string, mixed>   $options
+     *
+     * @return iterable<TextDocument>
+     */
+    public function __invoke(iterable $documents, array $options = []): iterable;
+}

src/Store/Document/Vectorizer.php

@@ -0,0 +1,53 @@
+<?php
+
+declare(strict_types=1);
+
+namespace PhpLlm\LlmChain\Store\Document;
+
+use PhpLlm\LlmChain\Platform\Capability;
+use PhpLlm\LlmChain\Platform\Model;
+use PhpLlm\LlmChain\Platform\PlatformInterface;
+
+/**
+ * The Vectorizer encapsulates the logic to convert a collection of TextDocuments into VectorDocuments. It checks for
+ * the model's capabilities to handle batch processing or handles it with HttpClient's concurrency feature.
+ */
+final readonly class Vectorizer
+{
+    public function __construct(
+        private PlatformInterface $platform,
+        private Model $model,
+    ) {
+    }
+
+    /**
+     * @param TextDocument[] $documents
+     *
+     * @return VectorDocument[]
+     */
+    public function vectorizeDocuments(array $documents): array
+    {
+        if ($this->model->supports(Capability::INPUT_MULTIPLE)) {
+            $response = $this->platform->request($this->model, array_map(fn (TextDocument $document) => $document->content, $documents));
+
+            $vectors = $response->getContent();
+        } else {
+            $responses = [];
+            foreach ($documents as $document) {
+                $responses[] = $this->platform->request($this->model, $document->content);
+            }
+
+            $vectors = [];
+            foreach ($responses as $response) {
+                $vectors = array_merge($vectors, $response->getContent());
+            }
+        }
+
+        $vectorDocuments = [];
+        foreach ($documents as $i => $document) {
+            $vectorDocuments[] = new VectorDocument($document->id, $vectors[$i], $document->metadata);
+        }
+
+        return $vectorDocuments;
+    }
+}

src/Store/Indexer.php

@@ -4,87 +4,48 @@
 
 namespace PhpLlm\LlmChain\Store;
 
-use PhpLlm\LlmChain\Platform\Capability;
-use PhpLlm\LlmChain\Platform\Model;
-use PhpLlm\LlmChain\Platform\PlatformInterface;
 use PhpLlm\LlmChain\Store\Document\TextDocument;
-use PhpLlm\LlmChain\Store\Document\VectorDocument;
+use PhpLlm\LlmChain\Store\Document\Vectorizer;
 use Psr\Log\LoggerInterface;
 use Psr\Log\NullLogger;
-use Symfony\Component\Clock\Clock;
-use Symfony\Component\Clock\ClockInterface;
 
 /**
+ * Converts a collection of TextDocuments into VectorDocuments and pushes them to a store implementation.
+ *
  * @author Christopher Hertel <[email protected]>
  */
 final readonly class Indexer
 {
-    private ClockInterface $clock;
-
     public function __construct(
-        private PlatformInterface $platform,
-        private Model $model,
+        private Vectorizer $vectorizer,
         private StoreInterface $store,
-        ?ClockInterface $clock = null,
         private LoggerInterface $logger = new NullLogger(),
     ) {
-        $this->clock = $clock ?? Clock::get();
     }
 
     /**
      * @param TextDocument|iterable<TextDocument> $documents
+     * @param int                                 $chunkSize number of documents to vectorize and store in one batch
      */
-    public function index(TextDocument|iterable $documents, int $chunkSize = 0, int $sleep = 0): void
+    public function index(TextDocument|iterable $documents, int $chunkSize = 50): void
     {
         if ($documents instanceof TextDocument) {
             $documents = [$documents];
         }
 
-        if ([] === $documents) {
-            $this->logger->debug('No documents to index');
-
-            return;
-        }
-
-        $chunks = 0 !== $chunkSize ? array_chunk($documents, $chunkSize) : [$documents];
-
-        foreach ($chunks as $chunk) {
-            $this->store->add(...$this->createVectorDocuments($chunk));
-
-            if (0 !== $sleep) {
-                $this->clock->sleep($sleep);
-            }
-        }
-    }
-
-    /**
-     * @param TextDocument[] $documents
-     *
-     * @return VectorDocument[]
-     */
-    private function createVectorDocuments(array $documents): array
-    {
-        if ($this->model->supports(Capability::INPUT_MULTIPLE)) {
-            $response = $this->platform->request($this->model, array_map(fn (TextDocument $document) => $document->content, $documents));
-
-            $vectors = $response->getContent();
-        } else {
-            $responses = [];
-            foreach ($documents as $document) {
-                $responses[] = $this->platform->request($this->model, $document->content);
-            }
+        $counter = 0;
+        $chunk = [];
+        foreach ($documents as $document) {
+            $chunk[] = $document;
 
-            $vectors = [];
-            foreach ($responses as $response) {
-                $vectors = array_merge($vectors, $response->getContent());
+            if ($chunkSize === \count($chunk)) {
+                $this->store->add(...$this->vectorizer->vectorizeDocuments($chunk));
+                $chunk = [];
             }
-        }
 
-        $vectorDocuments = [];
-        foreach ($documents as $i => $document) {
-            $vectorDocuments[] = new VectorDocument($document->id, $vectors[$i], $document->metadata);
+            ++$counter;
         }
 
-        return $vectorDocuments;
+        $this->logger->debug(0 === $counter ? 'No documents to index' : \sprintf('Indexed %d documents', $counter));
     }
 }