wip

RuliLG · RuliLG · commit ebd4f6685dfb · 2024-05-05T21:54:05.000+01:00
diff --git a/config/knowledge_base.php b/config/knowledge_base.php
@@ -6,6 +6,10 @@
         'timeout' => env('KNOWLEDGE_BASE_TIMEOUT', 90),
     ],
     'models' => [
-        'knowledge_base_id' => \Borah\KnowledgeBase\Models\KnowledgeBaseId::class,
+        'knowledge_base_chunk' => \Borah\KnowledgeBase\Models\KnowledgeBaseChunk::class,
+    ],
+    'chunking' => [
+        'size' => env('KNOWLEDGE_BASE_CHUNK_SIZE', 1000),
+        'overlap' => env('KNOWLEDGE_BASE_CHUNK_OVERLAP', 100),
     ],
 ];
diff --git a/database/migrations/2024_01_01_000000_create_knowledge_base_chunks_table.php b/database/migrations/2024_01_01_000000_create_knowledge_base_chunks_table.php
@@ -11,8 +11,10 @@
      */
     public function up(): void
     {
-        Schema::create('knowledge_base_ids', function (Blueprint $table) {
+        Schema::create('knowledge_base_chunks', function (Blueprint $table) {
             $table->uuid('id')->primary();
+            $table->unsignedInteger('order');
+            $table->mediumText('text');
             $table->morphs('model');
         });
     }
diff --git a/src/Client/KnowledgeBaseClient.php b/src/Client/KnowledgeBaseClient.php
@@ -2,6 +2,7 @@
 
 namespace Borah\KnowledgeBase\Client;
 
+use Borah\KnowledgeBase\DTO\KnowledgeBaseChunkItem;
 use Borah\KnowledgeBase\DTO\KnowledgeBaseQueryResponse;
 use Illuminate\Http\Client\PendingRequest;
 use Illuminate\Support\Facades\Http;
@@ -27,6 +28,7 @@ public function __construct()
      */
     public function upsert(array $data): bool
     {
+        logger()->debug('[KnowledgeBaseClient] Upserting data', $data);
         return $this->client->post('/insert', [
             'data' => $data,
         ])
@@ -42,6 +44,7 @@ public function upsert(array $data): bool
      */
     public function destroy(mixed $id): bool
     {
+        logger()->debug('[KnowledgeBaseClient] Deleting record', ['id' => $id]);
         return $this->client->delete('/delete/'.$id)
             ->throw()
             ->json('success') ?: false;
@@ -57,6 +60,12 @@ public function destroy(mixed $id): bool
      */
     public function query(string $text, int $k = 10, ?array $entities = null, ?array $where = null): KnowledgeBaseQueryResponse
     {
+        logger()->debug('[KnowledgeBaseClient] Querying knowledge base', [
+            'text' => $text,
+            'k' => $k,
+            'entities' => $entities,
+            'where' => $where,
+        ]);
         $params = [
             'query' => $text,
             'k' => $k,
@@ -76,4 +85,28 @@ public function query(string $text, int $k = 10, ?array $entities = null, ?array
 
         return KnowledgeBaseQueryResponse::from($response);
     }
+
+    /**
+     * Generates chunks of text for the selected records.
+     *
+     * @param array<\Borah\KnowledgeBase\DTO\KnowledgeBaseChunkItem> $records
+     *
+     * @return array<\Borah\KnowledgeBase\DTO\KnowledgeBaseChunkItem>
+     */
+    public function chunk(array $records, ?int $chunkSize = null, ?int $chunkOverlap = null): array
+    {
+        logger()->debug('[KnowledgeBaseClient] Chunking data');
+        $response = $this->client->post('/chunk', [
+            'data' => $records,
+            'chunk_size' => $chunkSize ?? config('knowledge_base.chunking.size'),
+            'chunk_overlap' => $chunkOverlap ?? config('knowledge_base.chunking.overlap'),
+        ])
+            ->throw()
+            ->collect('chunks');
+
+        return $response->map(fn (array $item) => new KnowledgeBaseChunkItem(
+            id: $item['id'],
+            text: $item['text'],
+        ))->all();
+    }
 }
diff --git a/src/DTO/KnowledgeBaseChunkItem.php b/src/DTO/KnowledgeBaseChunkItem.php
@@ -0,0 +1,23 @@
+<?php
+
+namespace Borah\KnowledgeBase\DTO;
+
+use Illuminate\Contracts\Support\Arrayable;
+
+class KnowledgeBaseChunkItem implements Arrayable
+{
+    public function __construct(
+        public readonly mixed $id,
+        public readonly string $text,
+    ) {
+
+    }
+
+    public function toArray(): array
+    {
+        return [
+            'id' => $this->id,
+            'text' => $this->text,
+        ];
+    }
+}
diff --git a/src/DTO/KnowledgeBaseQueryResponse.php b/src/DTO/KnowledgeBaseQueryResponse.php
@@ -2,7 +2,7 @@
 
 namespace Borah\KnowledgeBase\DTO;
 
-use Borah\KnowledgeBase\Models\KnowledgeBaseId;
+use Borah\KnowledgeBase\Models\KnowledgeBaseChunk;
 use Illuminate\Database\Eloquent\Collection;
 
 class KnowledgeBaseQueryResponse
@@ -25,14 +25,16 @@ public static function from(array $response): static
     public function models(): Collection
     {
         $ids = collect($this->results)->pluck('id');
-        $intermediateRecords = KnowledgeBaseId::query()
+        $intermediateRecords = KnowledgeBaseChunk::query()
             ->with('model')
             ->whereIn('id', $ids)
             ->get();
 
         // sort the records in the same order as the results
         $intermediateRecords = $intermediateRecords->sortBy(fn ($record) => array_search($record->id, $ids->toArray()));
 
-        return $intermediateRecords->map(fn ($record) => $record->model);
+        return $intermediateRecords
+            ->map(fn ($record) => $record->model)
+            ->unique();
     }
 }
diff --git a/src/DTO/KnowledgeBaseQueryResponseItem.php b/src/DTO/KnowledgeBaseQueryResponseItem.php
@@ -8,7 +8,6 @@ public function __construct(
         public readonly mixed $id,
         public readonly string $entity,
         public readonly string $text,
-        public readonly float $vectorScore,
         public readonly float $rerankingScore,
         public readonly ?array $payload = null,
     ) {
@@ -21,7 +20,6 @@ public static function from(array $result): static
             id: $result['id'],
             entity: $result['entity'],
             text: $result['text'],
-            vectorScore: $result['vector_score'],
             rerankingScore: $result['reranking_score'],
             payload: $result['payload'] ?? null,
         );
diff --git a/src/KnowledgeBase.php b/src/KnowledgeBase.php
@@ -32,7 +32,11 @@ public function destroy(Model $model): bool
         $client = new KnowledgeBaseClient();
         $items = $model->knowledgeInsertItems();
 
-        return $client->destroy($items[0]->id);
+        foreach ($items as $item) {
+            $client->destroy($item->id);
+        }
+
+        return true;
     }
 
     public function query(string $text, int $k = 10, ?array $entities = null, ?array $where = null): KnowledgeBaseQueryResponse
@@ -41,4 +45,18 @@ public function query(string $text, int $k = 10, ?array $entities = null, ?array
 
         return $client->query($text, $k, $entities, $where);
     }
+
+    /**
+     * Generates chunks of text for the selected records.
+     *
+     * @param array<\Borah\KnowledgeBase\DTO\KnowledgeBaseChunkItem> $records
+     *
+     * @return array<\Borah\KnowledgeBase\DTO\KnowledgeBaseChunkItem>
+     */
+    public function chunk(array $records, ?int $chunkSize = null, ?int $chunkOverlap = null): array
+    {
+        $client = new KnowledgeBaseClient();
+
+        return $client->chunk($records, $chunkSize, $chunkOverlap);
+    }
 }
diff --git a/src/KnowledgeBaseServiceProvider.php b/src/KnowledgeBaseServiceProvider.php
@@ -18,7 +18,7 @@ public function configurePackage(Package $package): void
         $package
             ->name('knowledge-base-laravel')
             ->hasConfigFile('knowledge_base')
-            ->hasMigration('2024_01_01_000000_create_knowledge_base_ids_table')
+            ->hasMigration('2024_01_01_000000_create_knowledge_base_chunks_table')
             ->hasCommand(ReimportKnowledgeBaseCommand::class);
     }
 }
diff --git a/src/Models/KnowledgeBaseChunk.php b/src/Models/KnowledgeBaseChunk.php
@@ -5,14 +5,15 @@
 use Illuminate\Database\Eloquent\Concerns\HasUuids;
 use Illuminate\Database\Eloquent\Model;
 
-class KnowledgeBaseId extends Model
+class KnowledgeBaseChunk extends Model
 {
     use HasUuids;
 
     protected $fillable = [
         'id',
         'model_type',
         'model_id',
+        'text',
     ];
 
     public $timestamps = false;
diff --git a/src/Traits/BelongsToKnowledgeBase.php b/src/Traits/BelongsToKnowledgeBase.php
@@ -2,12 +2,13 @@
 
 namespace Borah\KnowledgeBase\Traits;
 
+use Borah\KnowledgeBase\Client\KnowledgeBaseClient;
 use Borah\KnowledgeBase\DTO\KnowledgeBaseQueryResponse;
 use Borah\KnowledgeBase\DTO\KnowledgeEmbeddingText;
 use Borah\KnowledgeBase\DTO\KnowledgeInsertItem;
 use Borah\KnowledgeBase\Facades\KnowledgeBase;
 use Illuminate\Database\Eloquent\Model;
-use Illuminate\Database\Eloquent\Relations\MorphOne;
+use Illuminate\Database\Eloquent\Relations\MorphMany;
 use function Illuminate\Events\queueable;
 
 trait BelongsToKnowledgeBase
@@ -27,33 +28,60 @@ public static function bootBelongsToKnowledgeBase()
         }));
     }
 
-    public function knowledgeBaseId(): MorphOne
+    public function knowledgeBaseChunks(): MorphMany
     {
-        return $this->morphOne(config('knowledge_base.models.knowledge_base_id'), 'model');
+        return $this->morphMany(config('knowledge_base.models.knowledge_base_chunk'), 'model');
+    }
+
+    public function knowledgeBasePayload(): array
+    {
+        return [];
     }
 
     /**
      * @return KnowledgeInsertItem[]
      */
     public function knowledgeInsertItems(): array
     {
-        $knowledgeBaseId = $this->knowledgeBaseId ?? $this->knowledgeBaseId()->create();
         $texts = $this->getEmbeddingsTexts();
-        if (! is_array($texts)) {
+        if (!is_array($texts)) {
             $texts = [$texts];
         }
 
-        return collect($texts)
-            ->map(fn (KnowledgeEmbeddingText $text) => new KnowledgeInsertItem(
-                id: $knowledgeBaseId->id,
+        $items = [];
+        $existingChunks = $this->knowledgeBaseChunks()->get()->sortBy('order');
+        foreach ($texts as $i => $text) {
+            $chunk = $existingChunks[$i] ?? $this->knowledgeBaseChunks()->create([
+                'text' => $text->text,
+                'order' => $i,
+            ]);
+
+            if (!$chunk->wasRecentlyCreated && $chunk->text !== $text->text) {
+                $chunk->update(['text' => $text->text]);
+            }
+
+            $items[] = new KnowledgeInsertItem(
+                id: $chunk->id,
                 entity: $text->entity,
                 text: $text->text,
                 payload: [
-                    ...$this->toArray(),
+                    ...$this->knowledgeBasePayload(),
                     'original_record_id' => $this->getKey(),
                 ],
-            ))
-            ->toArray();
+            );
+        }
+
+        $chunksToDelete = $this->knowledgeBaseChunks()
+            ->where('order', '>=', count($texts))
+            ->select('id', 'order')
+            ->get();
+        $client = new KnowledgeBaseClient();
+        foreach ($chunksToDelete as $chunk) {
+            $client->destroy($chunk->id);
+            $chunk->delete();
+        }
+
+        return $items;
     }
 
     public static function searchInKnowledgeBase(string $query, int $k = 10, ?array $where = null): KnowledgeBaseQueryResponse

Original file line number	Diff line number	Diff line change
`@@ -11,8 +11,10 @@`
`11`	`11`	`*/`
`12`	`12`	`public function up(): void`
`13`	`13`	`{`
`14`		`- Schema::create('knowledge_base_ids', function (Blueprint $table) {`
	`14`	`+ Schema::create('knowledge_base_chunks', function (Blueprint $table) {`
`15`	`15`	`$table->uuid('id')->primary();`
	`16`	`+ $table->unsignedInteger('order');`
	`17`	`+ $table->mediumText('text');`
`16`	`18`	`$table->morphs('model');`
`17`	`19`	`});`
`18`	`20`	`}`
Original file line number	Diff line number	Diff line change
`@@ -18,7 +18,7 @@ public function configurePackage(Package $package): void`
`18`	`18`	`$package`
`19`	`19`	`->name('knowledge-base-laravel')`
`20`	`20`	`->hasConfigFile('knowledge_base')`
`21`		`- ->hasMigration('2024_01_01_000000_create_knowledge_base_ids_table')`
	`21`	`+ ->hasMigration('2024_01_01_000000_create_knowledge_base_chunks_table')`
`22`	`22`	`->hasCommand(ReimportKnowledgeBaseCommand::class);`
`23`	`23`	`}`
`24`	`24`	`}`