ok this should do the embedding

LlmLaraHub · Apr 9, 2024 · 2461af1 · 2461af1
1 parent 84a3c46
commit 2461af1
Show file tree

Hide file tree

Showing 5 changed files with 41 additions and 10 deletions.
diff --git a/app/Domains/Messages/SearchOrSummarizeChatRepo.php b/app/Domains/Messages/SearchOrSummarizeChatRepo.php
@@ -26,10 +26,12 @@ public function search(Chat $chat, string $input): string
             $chat->chatable->getEmbeddingDriver()
         )->embedData($input);
 
+        $embeddingSize = get_embedding_size($chat->chatable->getEmbeddingDriver());
+
         $results = DocumentChunk::query()
             ->join('documents', 'documents.id', '=', 'document_chunks.document_id')
             ->selectRaw(
-                'document_chunks.embedding <-> ? as distance, document_chunks.content, document_chunks.embedding as embedding, document_chunks.id as id',
+                "document_chunks.{$embeddingSize} <-> ? as distance, document_chunks.content, document_chunks.{$embeddingSize} as embedding, document_chunks.id as id",
                 [$embedding->embedding]
             )
             ->where('documents.collection_id', $chat->chatable->id)

diff --git a/app/Models/DocumentChunk.php b/app/Models/DocumentChunk.php
@@ -58,14 +58,7 @@ public function getDriver(): string
     public function getEmbeddingColumn(): string
     {
 
-        $embeddingModel = driverHelper($this->getDriver(), 'embedding_model');
+        return get_embedding_size($this->getEmbeddingDriver());
 
-        $size = config('llmdriver.embedding_sizes.'.$embeddingModel);
-
-        if ($size) {
-            return 'embedding_'.$size;
-        }
-
-        return 'embeding_3072';
     }
 }
diff --git a/app/helpers.php b/app/helpers.php
@@ -58,3 +58,18 @@ function driverHelper(string $driver, string $key): string
         return config("llmdriver.drivers.{$driver}.{$key}");
     }
 }
+
+if (! function_exists('get_embedding_size')) {
+    function get_embedding_size(string $ebmedding_driver): string
+    {
+        $embeddingModel = driverHelper($ebmedding_driver, 'embedding_model');
+
+        $size = config('llmdriver.embedding_sizes.'.$embeddingModel);
+
+        if ($size) {
+            return 'embedding_'.$size;
+        }
+
+        return 'embeding_3072';
+    }
+}
diff --git a/database/factories/DocumentChunkFactory.php b/database/factories/DocumentChunkFactory.php
@@ -44,6 +44,7 @@ public function openAi(): Factory
         return $this->state(function (array $attributes) {
             $collection = Collection::factory()->create([
                 'driver' => DriversEnum::OpenAi,
+                'embedding_driver' => DriversEnum::OpenAi,
             ]);
             $document = Document::factory()->create([
                 'collection_id' => $collection->id,

diff --git a/tests/Feature/HelpersTest.php b/tests/Feature/HelpersTest.php
@@ -2,6 +2,7 @@
 
 namespace Tests\Feature;
 
+use App\Models\DocumentChunk;
 use Tests\TestCase;
 
 class HelpersTest extends TestCase
@@ -12,7 +13,26 @@ class HelpersTest extends TestCase
     public function test_config_helper(): void
     {
 
-        $this->assertEquals(4096, driverHelper('mock', 'embedding_size.mock'));
+        $this->assertEquals('mock', driverHelper('mock', 'embedding_model'));
+
+    }
+
+    public function test_get_embedding_size(): void
+    {
+
+        $model = DocumentChunk::factory()->create();
+
+        $embedding_column = get_embedding_size($model->getEmbeddingDriver());
+
+        $this->assertEquals('embedding_4096', $embedding_column);
+
+        $model = DocumentChunk::factory()
+            ->openAi()
+            ->create();
+
+        $embedding_column = get_embedding_size($model->getEmbeddingDriver());
+
+        $this->assertEquals('embedding_3072', $embedding_column);
 
     }
 }