refactor: split Document into TextDocument and VectorDocument

php-llm · Oct 3, 2024 · 5aab0d7 · 5aab0d7
1 parent 2552463
commit 5aab0d7
Show file tree

Hide file tree

Showing 15 changed files with 85 additions and 116 deletions.
diff --git a/examples/store-mongodb-similarity-search.php b/examples/store-mongodb-similarity-search.php
@@ -2,8 +2,8 @@
 
 use MongoDB\Client as MongoDBClient;
 use PhpLlm\LlmChain\Chain;
-use PhpLlm\LlmChain\Document\Document;
 use PhpLlm\LlmChain\Document\Metadata;
+use PhpLlm\LlmChain\Document\TextDocument;
 use PhpLlm\LlmChain\DocumentEmbedder;
 use PhpLlm\LlmChain\Message\Message;
 use PhpLlm\LlmChain\Message\MessageBag;
@@ -46,7 +46,7 @@
 
 // create embeddings and documents
 foreach ($movies as $movie) {
-    $documents[] = new Document(
+    $documents[] = new TextDocument(
         id: Uuid::v4(),
         text: $movie['title'].' '.$movie['description'],
         metadata: new Metadata($movie),

diff --git a/examples/store-pinecone-similarity-search.php b/examples/store-pinecone-similarity-search.php
@@ -1,8 +1,8 @@
 <?php
 
 use PhpLlm\LlmChain\Chain;
-use PhpLlm\LlmChain\Document\Document;
 use PhpLlm\LlmChain\Document\Metadata;
+use PhpLlm\LlmChain\Document\TextDocument;
 use PhpLlm\LlmChain\DocumentEmbedder;
 use PhpLlm\LlmChain\Message\Message;
 use PhpLlm\LlmChain\Message\MessageBag;
@@ -40,7 +40,7 @@
 
 // create embeddings and documents
 foreach ($movies as $movie) {
-    $documents[] = new Document(
+    $documents[] = new TextDocument(
         id: Uuid::v4(),
         text: 'Title: '.$movie['title'].PHP_EOL.'Director: '.$movie['director'].PHP_EOL.$movie['description'],
         metadata: new Metadata($movie),

diff --git a/src/Document/EmbeddedDocument.php b/src/Document/EmbeddedDocument.php
diff --git a/src/Document/Document.php → src/Document/TextDocument.php b/src/Document/Document.php → src/Document/TextDocument.php
@@ -7,7 +7,7 @@
 use Symfony\Component\Uid\Uuid;
 use Webmozart\Assert\Assert;
 
-readonly class Document
+readonly class TextDocument
 {
     public function __construct(
         public Uuid $id,
@@ -16,14 +16,4 @@ public function __construct(
     ) {
         Assert::stringNotEmpty(trim($this->text));
     }
-
-    public function withVector(Vector $vector): EmbeddedDocument
-    {
-        return new EmbeddedDocument(
-            $this->id,
-            $this->text,
-            $vector,
-            $this->metadata,
-        );
-    }
 }
diff --git a/src/Document/VectorDocument.php b/src/Document/VectorDocument.php
@@ -0,0 +1,17 @@
+<?php
+
+declare(strict_types=1);
+
+namespace PhpLlm\LlmChain\Document;
+
+use Symfony\Component\Uid\Uuid;
+
+final readonly class VectorDocument
+{
+    public function __construct(
+        public Uuid $id,
+        public Vector $vector,
+        public Metadata $metadata = new Metadata([]),
+    ) {
+    }
+}
diff --git a/src/DocumentEmbedder.php b/src/DocumentEmbedder.php
@@ -4,7 +4,8 @@
 
 namespace PhpLlm\LlmChain;
 
-use PhpLlm\LlmChain\Document\Document;
+use PhpLlm\LlmChain\Document\TextDocument;
+use PhpLlm\LlmChain\Document\VectorDocument;
 use PhpLlm\LlmChain\Store\StoreInterface;
 use Psr\Log\LoggerInterface;
 use Psr\Log\NullLogger;
@@ -25,11 +26,11 @@ public function __construct(
     }
 
     /**
-     * @param Document|list<Document> $documents
+     * @param TextDocument|TextDocument[] $documents
      */
-    public function embed(Document|array $documents, int $chunkSize = 0, int $sleep = 0): void
+    public function embed(TextDocument|array $documents, int $chunkSize = 0, int $sleep = 0): void
     {
-        if ($documents instanceof Document) {
+        if ($documents instanceof TextDocument) {
             $documents = [$documents];
         }
 
@@ -42,14 +43,14 @@ public function embed(Document|array $documents, int $chunkSize = 0, int $sleep
         $chunks = 0 !== $chunkSize ? array_chunk($documents, $chunkSize) : [$documents];
 
         foreach ($chunks as $chunk) {
-            $vectors = $this->embeddings->multiCreate(array_map(fn (Document $document) => $document->text, $chunk));
+            $vectors = $this->embeddings->multiCreate(array_map(fn (TextDocument $document) => $document->text, $chunk));
 
-            $embeddedDocuments = [];
+            $vectorDocuments = [];
             foreach ($chunk as $i => $document) {
-                $embeddedDocuments[] = $document->withVector($vectors[$i]);
+                $vectorDocuments[] = new VectorDocument($document->id, $vectors[$i], $document->metadata);
             }
 
-            $this->store->addDocuments($embeddedDocuments);
+            $this->store->add($vectorDocuments);
 
             if (0 !== $sleep) {
                 $this->clock->sleep($sleep);

diff --git a/src/Store/Azure/SearchStore.php b/src/Store/Azure/SearchStore.php
@@ -4,9 +4,9 @@
 
 namespace PhpLlm\LlmChain\Store\Azure;
 
-use PhpLlm\LlmChain\Document\EmbeddedDocument;
 use PhpLlm\LlmChain\Document\Metadata;
 use PhpLlm\LlmChain\Document\Vector;
+use PhpLlm\LlmChain\Document\VectorDocument;
 use PhpLlm\LlmChain\Store\VectorStoreInterface;
 use Symfony\Component\Uid\Uuid;
 use Symfony\Contracts\HttpClient\HttpClientInterface;
@@ -26,21 +26,17 @@ public function __construct(
     ) {
     }
 
-    public function addDocument(EmbeddedDocument $document): void
+    public function add(array|VectorDocument $documents): void
     {
-        $this->addDocuments([$document]);
-    }
+        if ($documents instanceof VectorDocument) {
+            $documents = [$documents];
+        }
 
-    public function addDocuments(array $documents): void
-    {
         $this->request('index', [
             'value' => array_map([$this, 'convertDocumentToIndexableArray'], $documents),
         ]);
     }
 
-    /**
-     * @return list<EmbeddedDocument>
-     */
     public function query(Vector $vector, array $options = []): array
     {
         $result = $this->request('search', [
@@ -73,23 +69,21 @@ private function request(string $endpoint, array $payload): array
     /**
      * @return array<string, mixed>
      */
-    private function convertDocumentToIndexableArray(EmbeddedDocument $document): array
+    private function convertDocumentToIndexableArray(VectorDocument $document): array
     {
         return array_merge([
             'id' => $document->id,
             $this->vectorFieldName => $document->vector->getData(),
-            'text' => $document->text,
         ], $document->metadata->getArrayCopy());
     }
 
     /**
      * @param array<string, mixed> $data
      */
-    private function convertArrayToDocument(array $data): EmbeddedDocument
+    private function convertArrayToDocument(array $data): VectorDocument
     {
-        return new EmbeddedDocument(
+        return new VectorDocument(
             id: Uuid::fromString($data['id']),
-            text: $data['text'],
             vector: $data[$this->vectorFieldName] ? new Vector($data[$this->vectorFieldName]) : null,
             metadata: new Metadata($data),
         );

diff --git a/src/Store/ChromaDB/Store.php b/src/Store/ChromaDB/Store.php
@@ -5,9 +5,9 @@
 namespace PhpLlm\LlmChain\Store\ChromaDB;
 
 use Codewithkyrian\ChromaDB\Client;
-use PhpLlm\LlmChain\Document\EmbeddedDocument;
 use PhpLlm\LlmChain\Document\Metadata;
 use PhpLlm\LlmChain\Document\Vector;
+use PhpLlm\LlmChain\Document\VectorDocument;
 use PhpLlm\LlmChain\Store\VectorStoreInterface;
 use Symfony\Component\Uid\Uuid;
 
@@ -19,13 +19,12 @@ public function __construct(
     ) {
     }
 
-    public function addDocument(EmbeddedDocument $document): void
+    public function add(array|VectorDocument $documents): void
     {
-        $this->addDocuments([$document]);
-    }
+        if ($documents instanceof VectorDocument) {
+            $documents = [$documents];
+        }
 
-    public function addDocuments(array $documents): void
-    {
         $ids = [];
         $vectors = [];
         $metadata = [];
@@ -49,9 +48,8 @@ public function query(Vector $vector, array $options = []): array
 
         $documents = [];
         for ($i = 0; $i < count($queryResponse->metadatas[0]); ++$i) {
-            $documents[] = new EmbeddedDocument(
+            $documents[] = new VectorDocument(
                 id: Uuid::fromString($queryResponse->ids[0][$i]),
-                text: '???',
                 vector: new Vector($queryResponse->embeddings[0][$i]),
                 metadata: new Metadata($queryResponse->metadatas[0][$i]),
             );

diff --git a/src/Store/MongoDB/Store.php b/src/Store/MongoDB/Store.php
@@ -8,9 +8,9 @@
 use MongoDB\Client;
 use MongoDB\Collection;
 use MongoDB\Driver\Exception\CommandException;
-use PhpLlm\LlmChain\Document\EmbeddedDocument;
 use PhpLlm\LlmChain\Document\Metadata;
 use PhpLlm\LlmChain\Document\Vector;
+use PhpLlm\LlmChain\Document\VectorDocument;
 use PhpLlm\LlmChain\Exception\InvalidArgumentException;
 use PhpLlm\LlmChain\Store\InitializableStoreInterface;
 use PhpLlm\LlmChain\Store\VectorStoreInterface;
@@ -61,13 +61,12 @@ public function __construct(
     ) {
     }
 
-    public function addDocument(EmbeddedDocument $document): void
+    public function add(array|VectorDocument $documents): void
     {
-        $this->addDocuments([$document]);
-    }
+        if ($documents instanceof VectorDocument) {
+            $documents = [$documents];
+        }
 
-    public function addDocuments(array $documents): void
-    {
         $operations = [];
 
         foreach ($documents as $document) {
@@ -76,7 +75,6 @@ public function addDocuments(array $documents): void
                 array_filter([
                     'metadata' => $document->metadata->getArrayCopy(),
                     $this->vectorFieldName => $document->vector->getData(),
-                    'text' => $document->text,
                 ]),
                 ['upsert' => true], // insert if not exists
             ];
@@ -100,8 +98,6 @@ public function addDocuments(array $documents): void
      *     numCandidates?: positive-int,
      *     filter?: array<mixed>
      * } $options
-     *
-     * @return EmbeddedDocument[]
      */
     public function query(Vector $vector, array $options = []): array
     {
@@ -125,9 +121,8 @@ public function query(Vector $vector, array $options = []): array
         $documents = [];
 
         foreach ($results as $result) {
-            $documents[] = new EmbeddedDocument(
+            $documents[] = new VectorDocument(
                 id: $this->toUuid($result['_id']),
-                text: $result['text'],
                 vector: new Vector($result[$this->vectorFieldName]),
                 metadata: new Metadata($result['metadata'] ?? []),
             );

diff --git a/src/Store/Pinecone/Store.php b/src/Store/Pinecone/Store.php
@@ -4,9 +4,9 @@
 
 namespace PhpLlm\LlmChain\Store\Pinecone;
 
-use PhpLlm\LlmChain\Document\EmbeddedDocument;
 use PhpLlm\LlmChain\Document\Metadata;
 use PhpLlm\LlmChain\Document\Vector;
+use PhpLlm\LlmChain\Document\VectorDocument;
 use PhpLlm\LlmChain\Store\VectorStoreInterface;
 use Probots\Pinecone\Client;
 use Probots\Pinecone\Resources\Data\VectorResource;
@@ -25,19 +25,17 @@ public function __construct(
     ) {
     }
 
-    public function addDocument(EmbeddedDocument $document): void
+    public function add(array|VectorDocument $documents): void
     {
-        $this->addDocuments([$document]);
-    }
+        if ($documents instanceof VectorDocument) {
+            $documents = [$documents];
+        }
 
-    public function addDocuments(array $documents): void
-    {
         $vectors = [];
         foreach ($documents as $document) {
             $vectors[] = [
                 'id' => (string) $document->id,
                 'values' => $document->vector->getData(),
-                'text' => $document->text,
                 'metadata' => $document->metadata->getArrayCopy(),
             ];
         }
@@ -61,9 +59,8 @@ public function query(Vector $vector, array $options = []): array
 
         $documents = [];
         foreach ($response->json()['matches'] as $match) {
-            $documents[] = new EmbeddedDocument(
+            $documents[] = new VectorDocument(
                 id: Uuid::fromString($match['id']),
-                text: $match['text'],
                 vector: new Vector($match['values']),
                 metadata: new Metadata($match['metadata']),
             );

diff --git a/src/Store/StoreInterface.php b/src/Store/StoreInterface.php
@@ -4,14 +4,12 @@
 
 namespace PhpLlm\LlmChain\Store;
 
-use PhpLlm\LlmChain\Document\EmbeddedDocument;
+use PhpLlm\LlmChain\Document\VectorDocument;
 
 interface StoreInterface
 {
-    public function addDocument(EmbeddedDocument $document): void;
-
     /**
-     * @param EmbeddedDocument[] $documents
+     * @param VectorDocument|VectorDocument[] $documents
      */
-    public function addDocuments(array $documents): void;
+    public function add(VectorDocument|array $documents): void;
 }
diff --git a/src/Store/VectorStoreInterface.php b/src/Store/VectorStoreInterface.php
@@ -4,15 +4,15 @@
 
 namespace PhpLlm\LlmChain\Store;
 
-use PhpLlm\LlmChain\Document\EmbeddedDocument;
 use PhpLlm\LlmChain\Document\Vector;
+use PhpLlm\LlmChain\Document\VectorDocument;
 
 interface VectorStoreInterface extends StoreInterface
 {
     /**
      * @param array<string, mixed> $options
      *
-     * @return EmbeddedDocument[]
+     * @return VectorDocument[]
      */
     public function query(Vector $vector, array $options = []): array;
 }