extend data loader

Author: ilvalerione

src/Exceptions/DataReaderException.php

@@ -0,0 +1,8 @@
+<?php
+
+namespace NeuronAI\Exceptions;
+
+class DataReaderException extends NeuronException
+{
+
+}

src/Providers/Embeddings/AbstractEmbeddingProvider.php

@@ -0,0 +1,18 @@
+<?php
+
+namespace NeuronAI\Providers\Embeddings;
+
+use NeuronAI\RAG\Document;
+
+abstract class AbstractEmbeddingProvider implements EmbeddingsProviderInterface
+{
+    public function embedDocuments(array $documents): array
+    {
+        /** @var Document $document */
+        foreach ($documents as $index => $document) {
+            $documents[$index] = $this->embedDocument($document);
+        }
+
+        return $documents;
+    }
+}

src/Providers/Embeddings/EmbeddingsProviderInterface.php

@@ -12,4 +12,6 @@ interface EmbeddingsProviderInterface
     public function embedText(string $text): array;
 
     public function embedDocument(Document $document): Document;
+
+    public function embedDocuments(array $documents): array;
 }

src/Providers/Embeddings/VoyageEmbeddingProvider.php

@@ -5,7 +5,7 @@
 use NeuronAI\RAG\Document;
 use GuzzleHttp\Client;
 
-class VoyageEmbeddingProvider implements EmbeddingsProviderInterface
+class VoyageEmbeddingProvider extends AbstractEmbeddingProvider
 {
     protected Client $client;
 

src/RAG/DataLoader/FileDataLoader.php

@@ -6,25 +6,51 @@
 
 class FileDataLoader implements DataLoaderInterface
 {
-    public function __construct(
-        protected string $filePath,
-        protected array $extensions = [],
-    ) {}
+    /**
+     * @var array<string, ReaderInterface>
+     */
+    protected array $readers = [];
+
+    public function __construct(protected string $path, array $readers = [])
+    {
+        $this->setReaders($readers);
+    }
+
+    public static function file(...$args): static
+    {
+        return new static(...$args);
+    }
+
+    public function addReader(string $fileExtension, ReaderInterface $reader): self
+    {
+        $this->readers[$fileExtension] = $reader;
+        return $this;
+    }
+
+    /**
+     * @param array $readers
+     * @return FileDataLoader
+     */
+    public function setReaders(array $readers): self
+    {
+        $this->readers = $readers;
+        return $this;
+    }
 
     public function getDocuments(): array
     {
-        if (! file_exists($this->filePath)) {
+        if (! file_exists($this->path)) {
             return [];
         }
 
         // If it's a directory
-        if (is_dir($this->filePath)) {
-            return $this->getDocumentsFromDirectory($this->filePath);
+        if (is_dir($this->path)) {
+            return $this->getDocumentsFromDirectory($this->path);
         }
 
         // If it's a file
         try {
-            return [$this->getDocument($this->getContentFromFile($this->filePath), $this->filePath)];
+            return [$this->getDocument($this->getContentFromFile($this->path), $this->path)];
         } catch (\Throwable $exception) {
             return [];
         }
@@ -67,10 +93,12 @@ protected function getContentFromFile(string $path): string|false
     {
         $fileExtension = strtolower(pathinfo($path, PATHINFO_EXTENSION));
 
-        return match ($fileExtension) {
-            'pdf' => PdfParser::getText($path),
-            default => file_get_contents($path)
-        };
+        if (\array_key_exists($fileExtension, $this->readers)) {
+            $reader = $this->readers[$fileExtension];
+            return $reader::getText($path);
+        }
+
+        return TextFileReader::getText($path);
     }
 
 

src/RAG/DataLoader/PdfParser.php renamed to src/RAG/DataLoader/PdfReader.php

@@ -3,6 +3,7 @@
 namespace NeuronAI\RAG\DataLoader;
 
 use Closure;
+use NeuronAI\Exceptions\DataReaderException;
 use Symfony\Component\Process\Exception\ProcessFailedException;
 use Symfony\Component\Process\Process;
 
@@ -11,7 +12,7 @@
  *
  * https://en.wikipedia.org/wiki/Pdftotext
  */
-class PdfParser
+class PdfReader implements ReaderInterface
 {
     protected string $pdf;
 
@@ -28,6 +29,12 @@ public function __construct(?string $binPath = null)
         $this->binPath = $binPath ?? $this->findPdfToText();
     }
 
+    public function setBinPath(string $binPath): self
+    {
+        $this->binPath = $binPath;
+        return $this;
+    }
+
     protected function findPdfToText(): string
     {
         $commonPaths = [
@@ -44,13 +51,13 @@ protected function findPdfToText(): string
             }
         }
 
-        throw new \LogicException("The pdftotext binary was not found or is not executable.");
+        throw new DataReaderException("The pdftotext binary was not found or is not executable.");
     }
 
     public function setPdf(string $pdf): self
     {
         if (!is_readable($pdf)) {
-            throw new \Exception("Could not read `{$pdf}`");
+            throw new DataReaderException("Could not read `{$pdf}`");
         }
 
         $this->pdf = $pdf;
@@ -97,11 +104,10 @@ public function setTimeout($timeout): self
         return $this;
     }
 
-    public function text(?Closure $callback = null): string
+    public function text(): string
     {
         $process = new Process(array_merge([$this->binPath], $this->options, [$this->pdf, '-']));
         $process->setTimeout($this->timeout);
-        $process = $callback ? $callback($process) : $process;
         $process->run();
         if (!$process->isSuccessful()) {
             throw new ProcessFailedException($process);
@@ -114,16 +120,23 @@ public function text(?Closure $callback = null): string
      * @throws \Exception
      */
     public static function getText(
-        string $pdf,
-        ?string $binPath = null,
-        array $options = [],
-        $timeout = 60,
-        ?Closure $callback = null
+        string $filePath,
+        array $options = []
     ): string {
-        return (new static($binPath))
-            ->setOptions($options)
-            ->setTimeout($timeout)
-            ->setPdf($pdf)
-            ->text($callback);
+        $instance = new static();
+
+        if (\array_key_exists('binPath', $options)) {
+            $instance->setBinPath($options['binPath']);
+        }
+
+        if (\array_key_exists('options', $options)) {
+            $instance->setOptions($options['options']);
+        }
+
+        if (\array_key_exists('timeout', $options)) {
+            $instance->setTimeout($options['timeout']);
+        }
+
+        return $instance->text();
     }
 }

src/RAG/DataLoader/ReaderInterface.php

@@ -0,0 +1,8 @@
+<?php
+
+namespace NeuronAI\RAG\DataLoader;
+
+interface ReaderInterface
+{
+    public static function getText(string $filePath, array $options = []): string;
+}

src/RAG/DataLoader/TextFileReader.php

@@ -0,0 +1,11 @@
+<?php
+
+namespace NeuronAI\RAG\DataLoader;
+
+class TextFileReader implements ReaderInterface
+{
+    public static function getText(string $filePath, array $options = []): string
+    {
+        return file_get_contents($filePath);
+    }
+}