FEATURE: AI Bot RAG support. (#537)
This PR lets you associate uploads to an AI persona, which we'll split and generate embeddings from. When building the system prompt to get a bot reply, we'll do a similarity search followed by a re-ranking (if available). This will let us find the most relevant fragments from the body of knowledge you associated with the persona, resulting in better, more informed responses.
For now, we'll only allow plain-text files, but this will change in the future.
Commits:
* FEATURE: RAG embeddings for the AI Bot
This first commit introduces a UI where admins can upload text files, which we'll store, split into fragments,
and generate embeddings of. In a next commit, we'll use those to give the bot additional information during
conversations.
* Basic asymmetric similarity search to provide guidance in system prompt
* Fix tests and lint
* Apply reranker to fragments
* Uploads filter, css adjustments and file validations
* Add placeholder for rag fragments
* Update annotations
2024-04-01 12:43:34 -04:00
|
|
|
# frozen_string_literal: true
|
|
|
|
|
|
|
|
RSpec.describe Jobs::DigestRagUpload do
|
|
|
|
fab!(:persona) { Fabricate(:ai_persona) }
|
|
|
|
fab!(:upload)
|
|
|
|
|
|
|
|
let(:document_file) { StringIO.new("some text" * 200) }
|
|
|
|
|
|
|
|
let(:truncation) { DiscourseAi::Embeddings::Strategies::Truncation.new }
|
|
|
|
let(:vector_rep) do
|
|
|
|
DiscourseAi::Embeddings::VectorRepresentations::Base.current_representation(truncation)
|
|
|
|
end
|
|
|
|
|
|
|
|
let(:expected_embedding) { [0.0038493] * vector_rep.dimensions }
|
|
|
|
|
2024-04-04 10:02:16 -04:00
|
|
|
let(:document_with_metadata) { plugin_file_from_fixtures("doc_with_metadata.txt", "rag") }
|
|
|
|
|
|
|
|
let(:parsed_document_with_metadata) do
|
|
|
|
plugin_file_from_fixtures("parsed_doc_with_metadata.txt", "rag")
|
|
|
|
end
|
|
|
|
|
|
|
|
let(:upload_with_metadata) do
|
|
|
|
UploadCreator.new(document_with_metadata, "document.txt").create_for(Discourse.system_user.id)
|
|
|
|
end
|
|
|
|
|
FEATURE: AI Bot RAG support. (#537)
This PR lets you associate uploads to an AI persona, which we'll split and generate embeddings from. When building the system prompt to get a bot reply, we'll do a similarity search followed by a re-ranking (if available). This will let us find the most relevant fragments from the body of knowledge you associated with the persona, resulting in better, more informed responses.
For now, we'll only allow plain-text files, but this will change in the future.
Commits:
* FEATURE: RAG embeddings for the AI Bot
This first commit introduces a UI where admins can upload text files, which we'll store, split into fragments,
and generate embeddings of. In a next commit, we'll use those to give the bot additional information during
conversations.
* Basic asymmetric similarity search to provide guidance in system prompt
* Fix tests and lint
* Apply reranker to fragments
* Uploads filter, css adjustments and file validations
* Add placeholder for rag fragments
* Update annotations
2024-04-01 12:43:34 -04:00
|
|
|
before do
|
|
|
|
SiteSetting.ai_embeddings_enabled = true
|
|
|
|
SiteSetting.ai_embeddings_discourse_service_api_endpoint = "http://test.com"
|
2024-04-04 10:02:16 -04:00
|
|
|
SiteSetting.authorized_extensions = "txt"
|
FEATURE: AI Bot RAG support. (#537)
This PR lets you associate uploads to an AI persona, which we'll split and generate embeddings from. When building the system prompt to get a bot reply, we'll do a similarity search followed by a re-ranking (if available). This will let us find the most relevant fragments from the body of knowledge you associated with the persona, resulting in better, more informed responses.
For now, we'll only allow plain-text files, but this will change in the future.
Commits:
* FEATURE: RAG embeddings for the AI Bot
This first commit introduces a UI where admins can upload text files, which we'll store, split into fragments,
and generate embeddings of. In a next commit, we'll use those to give the bot additional information during
conversations.
* Basic asymmetric similarity search to provide guidance in system prompt
* Fix tests and lint
* Apply reranker to fragments
* Uploads filter, css adjustments and file validations
* Add placeholder for rag fragments
* Update annotations
2024-04-01 12:43:34 -04:00
|
|
|
|
|
|
|
WebMock.stub_request(
|
|
|
|
:post,
|
|
|
|
"#{SiteSetting.ai_embeddings_discourse_service_api_endpoint}/api/v1/classify",
|
|
|
|
).to_return(status: 200, body: JSON.dump(expected_embedding))
|
|
|
|
end
|
|
|
|
|
|
|
|
describe "#execute" do
|
2024-04-04 10:02:16 -04:00
|
|
|
context "when processing an upload containing metadata" do
|
|
|
|
it "correctly splits on metadata boundary" do
|
|
|
|
described_class.new.execute(upload_id: upload_with_metadata.id, ai_persona_id: persona.id)
|
|
|
|
|
|
|
|
parsed = +""
|
|
|
|
first = true
|
|
|
|
RagDocumentFragment
|
|
|
|
.where(upload: upload_with_metadata)
|
|
|
|
.order(:fragment_number)
|
|
|
|
.each do |fragment|
|
|
|
|
parsed << "\n\n" if !first
|
|
|
|
parsed << "metadata: #{fragment.metadata}\n"
|
|
|
|
parsed << "number: #{fragment.fragment_number}\n"
|
|
|
|
parsed << fragment.fragment
|
|
|
|
first = false
|
|
|
|
end
|
|
|
|
|
|
|
|
# to rebuild parsed
|
|
|
|
# File.write("/tmp/testing", parsed)
|
|
|
|
|
|
|
|
expect(parsed).to eq(parsed_document_with_metadata.read)
|
|
|
|
end
|
|
|
|
end
|
FEATURE: AI Bot RAG support. (#537)
This PR lets you associate uploads to an AI persona, which we'll split and generate embeddings from. When building the system prompt to get a bot reply, we'll do a similarity search followed by a re-ranking (if available). This will let us find the most relevant fragments from the body of knowledge you associated with the persona, resulting in better, more informed responses.
For now, we'll only allow plain-text files, but this will change in the future.
Commits:
* FEATURE: RAG embeddings for the AI Bot
This first commit introduces a UI where admins can upload text files, which we'll store, split into fragments,
and generate embeddings of. In a next commit, we'll use those to give the bot additional information during
conversations.
* Basic asymmetric similarity search to provide guidance in system prompt
* Fix tests and lint
* Apply reranker to fragments
* Uploads filter, css adjustments and file validations
* Add placeholder for rag fragments
* Update annotations
2024-04-01 12:43:34 -04:00
|
|
|
context "when processing an upload for the first time" do
|
|
|
|
before { File.expects(:open).returns(document_file) }
|
|
|
|
|
|
|
|
it "splits an upload into chunks" do
|
|
|
|
subject.execute(upload_id: upload.id, ai_persona_id: persona.id)
|
|
|
|
|
|
|
|
created_fragment = RagDocumentFragment.last
|
|
|
|
|
|
|
|
expect(created_fragment).to be_present
|
|
|
|
expect(created_fragment.fragment).to be_present
|
|
|
|
expect(created_fragment.fragment_number).to eq(2)
|
|
|
|
end
|
|
|
|
|
|
|
|
it "queue jobs to generate embeddings for each fragment" do
|
|
|
|
expect { subject.execute(upload_id: upload.id, ai_persona_id: persona.id) }.to change(
|
|
|
|
Jobs::GenerateRagEmbeddings.jobs,
|
|
|
|
:size,
|
|
|
|
).by(1)
|
|
|
|
end
|
|
|
|
end
|
|
|
|
|
|
|
|
it "doesn't generate new fragments if we already processed the upload" do
|
|
|
|
Fabricate(:rag_document_fragment, upload: upload, ai_persona: persona)
|
|
|
|
previous_count = RagDocumentFragment.where(upload: upload, ai_persona: persona).count
|
|
|
|
|
|
|
|
subject.execute(upload_id: upload.id, ai_persona_id: persona.id)
|
|
|
|
updated_count = RagDocumentFragment.where(upload: upload, ai_persona: persona).count
|
|
|
|
|
|
|
|
expect(updated_count).to eq(previous_count)
|
|
|
|
end
|
|
|
|
end
|
|
|
|
end
|