Chunk documents ✨

Gemfile

@@ -68,6 +68,7 @@ group :test do
 end
 
 # Nosia dependencies
+gem "baran"
 gem "commonmarker"
 gem "faraday"
 gem "langchainrb", github: "cbldev/langchainrb", branch: "main"

Gemfile.lock

@@ -435,6 +435,7 @@ PLATFORMS
   x86_64-linux
 
 DEPENDENCIES
+  baran
   bootsnap
   brakeman
   capybara

app/jobs/add_document_job.rb

@@ -5,6 +5,6 @@ def perform(document_id)
     document = Document.find(document_id)
     document.titlize!
     document.parse!
-    document.vectorize! if document.content.present?
+    document.chunkify! if document.content.present?
   end
 end

app/jobs/get_ai_response_job.rb

@@ -17,7 +17,21 @@ def create_message(chat:)
 
   def call_ollama(chat:)
     message = create_message(chat:)
-    Document.ask(chat.last_question, k: 4) do |stream|
+
+    search_results = Chunk.similarity_search(chat.last_question, k: 4)
+    message.update(similar_document_ids: search_results.pluck(:document_id).uniq)
+
+    question = chat.last_question
+    context = search_results.map(&:content).join("\n---\n")
+
+    prompt = [
+      ["Contexte :", context].join("\n"),
+      ["Question :", question].join("\n"),
+      ["Réponse :"]
+    ].join("\n---\n")
+
+    llm = LangchainrbRails.config.vectorsearch.llm
+    llm.chat(messages: [{role: "user", content: prompt}]) do |stream|
       new_content = stream.raw_response.dig("message", "content")
       message.update(content: message.content + new_content) if new_content
     end

app/models/chunk.rb

@@ -0,0 +1,5 @@
+class Chunk < ApplicationRecord
+  include Vectorizable
+
+  belongs_to :document
+end

app/models/chunk/vectorizable.rb

@@ -0,0 +1,15 @@
+module Chunk::Vectorizable
+  extend ActiveSupport::Concern
+
+  included do
+    vectorsearch
+  end
+
+  def as_vector
+    self.content
+  end
+
+  def vectorize!
+    upsert_to_vectorsearch
+  end
+end

app/models/document.rb

@@ -1,5 +1,5 @@
 class Document < ApplicationRecord
-  include Parsable, Vectorizable
+  include Chunkable, Parsable, Vectorizable
 
   has_one_attached :file
 

app/models/document/chunkable.rb

@@ -0,0 +1,24 @@
+module Document::Chunkable
+  extend ActiveSupport::Concern
+
+  included do
+    has_many :chunks, dependent: :destroy
+  end
+
+  def chunkify!
+    splitter = ::Baran::CharacterTextSplitter.new(
+      chunk_size: 1024,
+      chunk_overlap: 64,
+      separator: "\n\n"
+    )
+
+    new_chunks = splitter.chunks(self.content)
+
+    self.chunks.destroy_all
+
+    new_chunks.each do |new_chunk|
+      chunk = self.chunks.create!(content: new_chunk.dig(:text))
+      chunk.vectorize!
+    end
+  end
+end

app/models/message.rb

@@ -26,6 +26,10 @@ def broadcast_updated
     )
   end
 
+  def similar_documents
+    Document.where(id: similar_document_ids)
+  end
+
   def to_html
     Commonmarker.to_html(content)
   end

app/views/messages/_message.html.erb

@@ -8,6 +8,29 @@
       </div>
     </div>
   <% else %>
+    <% if message.similar_document_ids.any? %>
+      <div class="px-2 py-4 sm:px-4">
+        <h2 class="text-lg font-semibold text-slate-900 dark:text-slate-200">Documents similaires</h2>
+        <ul role="list" class="grid grid-cols-1 gap-6 text-slate-900 dark:text-slate-200 sm:grid-cols-2 lg:grid-cols-3">
+          <% message.similar_documents.each do |document| %>
+            <li class="group col-span-1 rounded-lg bg-slate-50 shadow transition-colors duration-300 hover:bg-blue-600 dark:bg-slate-900 dark:hover:bg-blue-600">
+              <%= link_to document_path(document), class: "flex cursor-pointer items-center justify-between space-x-6 truncate p-6" do %>
+                <div class="flex-1 truncate">
+                  <div class="flex items-center space-x-3">
+                    <h3 class="text-sm font-medium text-slate-900 transition-colors duration-300 group-hover:text-slate-50 dark:text-slate-200">
+                      <%= document.title %>
+                    </h3>
+                  </div>
+                  <p class="mt-1 truncate text-sm text-slate-500 transition-colors duration-300 group-hover:text-slate-300">
+                    <%= document.content.first(50) %>
+                  </p>
+                </div>
+              <% end %>
+            </li>
+          <% end %>
+        </ul>
+      </div>
+    <% end %>
     <div class="mb-4 flex rounded-xl bg-slate-50 px-2 py-6 dark:bg-slate-900 sm:px-4">
       <div class="flex max-w-3xl items-center rounded-xl">
         <div>

db/migrate/20240628203936_create_chunks.rb

@@ -0,0 +1,11 @@
+class CreateChunks < ActiveRecord::Migration[8.0]
+  def change
+    create_table :chunks do |t|
+      t.belongs_to :document, null: false, foreign_key: true
+      t.text :content
+      t.vector :embedding, limit: 768
+
+      t.timestamps
+    end
+  end
+end

db/migrate/20240628210314_add_similar_document_ids_to_messages.rb

@@ -0,0 +1,5 @@
+class AddSimilarDocumentIdsToMessages < ActiveRecord::Migration[8.0]
+  def change
+    add_column :messages, :similar_document_ids, :string, array: true, default: []
+  end
+end