From d3c91a29ca1d47ed258c3cc539fbf3c2bcc1d007 Mon Sep 17 00:00:00 2001 From: aakashb95 Date: Thu, 19 Feb 2026 14:23:45 +0530 Subject: [PATCH] feat(captain): add document chunk schema foundation --- ...n_document_chunks_and_chunking_progress.rb | 65 +++++++++++++++++++ enterprise/app/models/captain/assistant.rb | 1 + enterprise/app/models/captain/document.rb | 16 +++++ .../app/models/captain/document_chunk.rb | 38 +++++++++++ .../app/models/enterprise/concerns/account.rb | 1 + 5 files changed, 121 insertions(+) create mode 100644 db/migrate/20260219143000_add_captain_document_chunks_and_chunking_progress.rb create mode 100644 enterprise/app/models/captain/document_chunk.rb diff --git a/db/migrate/20260219143000_add_captain_document_chunks_and_chunking_progress.rb b/db/migrate/20260219143000_add_captain_document_chunks_and_chunking_progress.rb new file mode 100644 index 000000000..26b9fb207 --- /dev/null +++ b/db/migrate/20260219143000_add_captain_document_chunks_and_chunking_progress.rb @@ -0,0 +1,65 @@ +class AddCaptainDocumentChunksAndChunkingProgress < ActiveRecord::Migration[7.0] + def up + create_document_chunks_table + add_chunking_progress_columns + end + + def down + remove_chunking_progress_columns + drop_table :captain_document_chunks + end + + private + + def create_document_chunks_table + create_table :captain_document_chunks do |t| + t.bigint :document_id, null: false + t.bigint :assistant_id, null: false + t.bigint :account_id, null: false + t.text :content, null: false + t.text :context + t.vector :embedding, limit: 1536 + t.integer :position, null: false + t.integer :token_count + t.column :searchable, :tsvector + + t.timestamps + end + + add_document_chunk_indexes + end + + def add_document_chunk_indexes + add_index :captain_document_chunks, :document_id + add_index :captain_document_chunks, [:document_id, :position], unique: true + add_index :captain_document_chunks, [:account_id, :assistant_id] + add_index :captain_document_chunks, :searchable, using: :gin + add_index :captain_document_chunks, + :embedding, + using: :ivfflat, + name: 'index_captain_document_chunks_on_embedding', + opclass: :vector_l2_ops + end + + def add_chunking_progress_columns + add_column :captain_documents, :chunking_status, :integer, default: 0, null: false + add_column :captain_documents, :expected_chunk_count, :integer, default: 0, null: false + add_column :captain_documents, :indexed_chunk_count, :integer, default: 0, null: false + add_column :captain_documents, :chunks_generated_at, :datetime + add_column :captain_documents, :last_chunk_error, :text + + add_index :captain_documents, :chunking_status + add_index :captain_documents, :chunks_generated_at + end + + def remove_chunking_progress_columns + remove_index :captain_documents, :chunking_status + remove_index :captain_documents, :chunks_generated_at + + remove_column :captain_documents, :chunking_status + remove_column :captain_documents, :expected_chunk_count + remove_column :captain_documents, :indexed_chunk_count + remove_column :captain_documents, :chunks_generated_at + remove_column :captain_documents, :last_chunk_error + end +end diff --git a/enterprise/app/models/captain/assistant.rb b/enterprise/app/models/captain/assistant.rb index 4f039d57a..6e1d8108d 100644 --- a/enterprise/app/models/captain/assistant.rb +++ b/enterprise/app/models/captain/assistant.rb @@ -25,6 +25,7 @@ class Captain::Assistant < ApplicationRecord belongs_to :account has_many :documents, class_name: 'Captain::Document', dependent: :destroy_async + has_many :document_chunks, class_name: 'Captain::DocumentChunk', dependent: :destroy_async has_many :responses, class_name: 'Captain::AssistantResponse', dependent: :destroy_async has_many :captain_inboxes, class_name: 'CaptainInbox', diff --git a/enterprise/app/models/captain/document.rb b/enterprise/app/models/captain/document.rb index 751162b28..99b9d8cb8 100644 --- a/enterprise/app/models/captain/document.rb +++ b/enterprise/app/models/captain/document.rb @@ -3,8 +3,13 @@ # Table name: captain_documents # # id :bigint not null, primary key +# chunking_status :integer default("pending"), not null +# chunks_generated_at :datetime # content :text +# expected_chunk_count :integer default(0), not null # external_link :string not null +# indexed_chunk_count :integer default(0), not null +# last_chunk_error :text # metadata :jsonb # name :string # status :integer default("in_progress"), not null @@ -18,6 +23,8 @@ # index_captain_documents_on_account_id (account_id) # index_captain_documents_on_assistant_id (assistant_id) # index_captain_documents_on_assistant_id_and_external_link (assistant_id,external_link) UNIQUE +# index_captain_documents_on_chunks_generated_at (chunks_generated_at) +# index_captain_documents_on_chunking_status (chunking_status) # index_captain_documents_on_status (status) # class Captain::Document < ApplicationRecord @@ -26,6 +33,7 @@ class Captain::Document < ApplicationRecord belongs_to :assistant, class_name: 'Captain::Assistant' has_many :responses, class_name: 'Captain::AssistantResponse', dependent: :destroy, as: :documentable + has_many :chunks, class_name: 'Captain::DocumentChunk', dependent: :destroy belongs_to :account has_one_attached :pdf_file @@ -44,6 +52,14 @@ class Captain::Document < ApplicationRecord available: 1 } + enum chunking_status: { + pending: 0, + chunking: 1, + indexing: 2, + ready: 3, + failed: 4 + }, _prefix: true + before_create :ensure_within_plan_limit after_create_commit :enqueue_crawl_job after_create_commit :update_document_usage diff --git a/enterprise/app/models/captain/document_chunk.rb b/enterprise/app/models/captain/document_chunk.rb new file mode 100644 index 000000000..0143099a6 --- /dev/null +++ b/enterprise/app/models/captain/document_chunk.rb @@ -0,0 +1,38 @@ +# == Schema Information +# +# Table name: captain_document_chunks +# +# id :bigint not null, primary key +# content :text not null +# context :text +# embedding :vector(1536) +# position :integer not null +# searchable :tsvector +# token_count :integer +# created_at :datetime not null +# updated_at :datetime not null +# account_id :bigint not null +# assistant_id :bigint not null +# document_id :bigint not null +# +# Indexes +# +# index_captain_document_chunks_on_account_id_and_assistant_id (account_id,assistant_id) +# index_captain_document_chunks_on_document_id (document_id) +# index_captain_document_chunks_on_document_id_and_position (document_id,position) UNIQUE +# index_captain_document_chunks_on_embedding (embedding) USING ivfflat +# index_captain_document_chunks_on_searchable (searchable) USING gin +# +class Captain::DocumentChunk < ApplicationRecord + self.table_name = 'captain_document_chunks' + + belongs_to :document, class_name: 'Captain::Document' + belongs_to :assistant, class_name: 'Captain::Assistant' + belongs_to :account + + has_neighbors :embedding, normalize: true + + validates :content, presence: true + validates :position, presence: true + validates :position, uniqueness: { scope: :document_id } +end diff --git a/enterprise/app/models/enterprise/concerns/account.rb b/enterprise/app/models/enterprise/concerns/account.rb index 01693ac79..e523db1f7 100644 --- a/enterprise/app/models/enterprise/concerns/account.rb +++ b/enterprise/app/models/enterprise/concerns/account.rb @@ -12,6 +12,7 @@ module Enterprise::Concerns::Account has_many :captain_assistants, dependent: :destroy_async, class_name: 'Captain::Assistant' has_many :captain_assistant_responses, dependent: :destroy_async, class_name: 'Captain::AssistantResponse' has_many :captain_documents, dependent: :destroy_async, class_name: 'Captain::Document' + has_many :captain_document_chunks, dependent: :destroy_async, class_name: 'Captain::DocumentChunk' has_many :captain_custom_tools, dependent: :destroy_async, class_name: 'Captain::CustomTool' has_many :copilot_threads, dependent: :destroy_async