From cfad22e21b78b70ec4b7b8a5c3979d7cff493535 Mon Sep 17 00:00:00 2001 From: Carl Lange Date: Wed, 15 Jan 2025 16:10:02 +0000 Subject: [PATCH 1/3] Add sqlite-vec support --- examples/sqlite_vec_example.rb | 47 ++++ lib/langchain/vectorsearch/sqlite_vec.rb | 154 +++++++++++ .../langchain/vectorsearch/sqlite_vec_spec.rb | 259 ++++++++++++++++++ 3 files changed, 460 insertions(+) create mode 100644 examples/sqlite_vec_example.rb create mode 100644 lib/langchain/vectorsearch/sqlite_vec.rb create mode 100644 spec/langchain/vectorsearch/sqlite_vec_spec.rb diff --git a/examples/sqlite_vec_example.rb b/examples/sqlite_vec_example.rb new file mode 100644 index 000000000..27edeeada --- /dev/null +++ b/examples/sqlite_vec_example.rb @@ -0,0 +1,47 @@ +require "langchain" + +# Initialize the LLM (using OpenAI in this example) +llm = Langchain::LLM::Ollama.new + +# Initialize the SQLite-vec vectorstore +db = Langchain::Vectorsearch::SqliteVec.new( + url: ":memory:", # Use a file-based DB (or ":memory:" for in-memory) + index_name: "documents", + namespace: "test", + llm: llm +) + +# Create the schema +db.create_default_schema + +# Add some sample texts +texts = [ + "Ruby is a dynamic, open source programming language with a focus on simplicity and productivity.", + "Python is a programming language that lets you work quickly and integrate systems more effectively.", + "JavaScript is a lightweight, interpreted programming language with first-class functions.", + "Rust is a multi-paradigm, general-purpose programming language designed for performance and safety." +] + +puts "Adding texts..." +ids = db.add_texts(texts: texts) +puts "Added #{ids.size} texts with IDs: #{ids.join(", ")}" + +# Search for similar texts +query = "What programming language is focused on memory safety?" +puts "\nSearching for: #{query}" +results = db.similarity_search(query: query) + +puts "\nResults:" +results.each do |result| + puts "- #{result[1]}" +end + +# Ask a question +question = "Which programming language emphasizes simplicity?" +puts "\nAsking: #{question}" +response = db.ask(question: question) +puts "Answer: #{response.chat_completion}" + +# Clean up +db.destroy_default_schema +File.delete("test_vectors.sqlite3") if File.exist?("test_vectors.sqlite3") diff --git a/lib/langchain/vectorsearch/sqlite_vec.rb b/lib/langchain/vectorsearch/sqlite_vec.rb new file mode 100644 index 000000000..6cd7d43b5 --- /dev/null +++ b/lib/langchain/vectorsearch/sqlite_vec.rb @@ -0,0 +1,154 @@ +# frozen_string_literal: true + +require "sqlite_vec" +module Langchain::Vectorsearch + class SqliteVec < Base + # + # The SQLite vector search adapter using sqlite-vec + # + # Gem requirements: + # gem "sqlite3", "~> 2.5" + # gem "sqlite_vec", "~> 0.16.0" + # + # Usage: + # sqlite_vec = Langchain::Vectorsearch::SqliteVec.new(url:, index_name:, llm:, namespace: nil) + # + + attr_reader :db, :table_name, :namespace_column, :namespace + + # @param url [String] The path to the SQLite database file (or :memory: for in-memory) + # @param index_name [String] The name of the table to use for the index + # @param llm [Object] The LLM client to use + # @param namespace [String] The namespace to use for the index when inserting/querying + def initialize(url:, index_name:, llm:, namespace: nil) + depends_on "sqlite3" + depends_on "sqlite_vec" + + @db = SQLite3::Database.new(url) + @db.enable_load_extension(true) + ::SqliteVec.load(@db) + @db.enable_load_extension(false) + + @table_name = index_name + @namespace_column = "namespace" + @namespace = namespace + + super(llm: llm) + end + + # Create default schema + def create_default_schema + @db.execute("CREATE VIRTUAL TABLE IF NOT EXISTS #{table_name} USING vec0( + embedding float[#{llm.default_dimensions}], + content TEXT, + #{namespace_column} TEXT + )") + end + + # Destroy default schema + def destroy_default_schema + @db.execute("DROP TABLE IF EXISTS #{table_name}") + end + + # Add a list of texts to the index + # @param texts [Array] The texts to add to the index + # @param ids [Array] The ids to add to the index, in the same order as the texts + # @return [Array] The ids of the added texts + def add_texts(texts:, ids: nil) + if ids.nil? || ids.empty? + max_rowid = @db.execute("SELECT MAX(rowid) FROM #{table_name}").first.first || 0 + ids = texts.map.with_index do |_, i| + max_rowid + i + 1 + end + end + + @db.transaction do + texts.zip(ids).each do |text, id| + embedding = llm.embed(text: text).embedding + @db.execute( + "INSERT INTO #{table_name}(rowid, content, embedding, #{namespace_column}) VALUES (?, ?, ?, ?)", + [id, text, embedding.pack("f*"), namespace] + ) + end + end + + ids + end + + # Update a list of ids and corresponding texts in the index + # @param texts [Array] The texts to update in the index + # @param ids [Array] The ids to update in the index, in the same order as the texts + # @return [Array] The ids of the updated texts + def update_texts(texts:, ids:) + @db.transaction do + texts.zip(ids).each do |text, id| + embedding = llm.embed(text: text).embedding + @db.execute( + "UPDATE #{table_name} SET content = ?, embedding = ? WHERE rowid = ?", + [text, embedding.pack("f*"), id] + ) + end + end + ids + end + + # Remove a list of texts from the index + # @param ids [Array] The ids of the texts to remove from the index + # @return [Integer] The number of texts removed from the index + def remove_texts(ids:) + @db.execute("DELETE FROM #{table_name} WHERE rowid IN (#{ids.join(",")})") + ids.length + end + + # Search for similar texts in the index + # @param query [String] The text to search for + # @param k [Integer] The number of top results to return + # @return [Array] The results of the search + def similarity_search(query:, k: 4) + embedding = llm.embed(text: query).embedding + similarity_search_by_vector(embedding: embedding, k: k) + end + + # Search for similar texts in the index by vector + # @param embedding [Array] The vector to search for + # @param k [Integer] The number of top results to return + # @return [Array] The results of the search + def similarity_search_by_vector(embedding:, k: 4) + namespace_condition = namespace ? "AND #{namespace_column} = ?" : "" + query_params = [embedding.pack("f*")] + query_params << namespace if namespace + + @db.execute(<<-SQL, query_params) + SELECT + rowid, + content, + distance + FROM #{table_name} + WHERE embedding MATCH ? + #{namespace_condition} + ORDER BY distance + LIMIT #{k} + SQL + end + + # Ask a question and return the answer + # @param question [String] The question to ask + # @param k [Integer] The number of results to have in context + # @yield [String] Stream responses back one String at a time + # @return [String] The answer to the question + def ask(question:, k: 4, &) + search_results = similarity_search(query: question, k: k) + + context = search_results.map { |result| result[1].to_s } + context = context.join("\n---\n") + + prompt = generate_rag_prompt(question: question, context: context) + + messages = [{role: "user", content: prompt}] + response = llm.chat(messages: messages, &) + + response.context = context + response + end + end +end diff --git a/spec/langchain/vectorsearch/sqlite_vec_spec.rb b/spec/langchain/vectorsearch/sqlite_vec_spec.rb new file mode 100644 index 000000000..306b20454 --- /dev/null +++ b/spec/langchain/vectorsearch/sqlite_vec_spec.rb @@ -0,0 +1,259 @@ +# frozen_string_literal: true + +RSpec.describe Langchain::Vectorsearch::SqliteVec do + subject { + described_class.new( + url: ":memory:", + index_name: "test_items", + llm: Langchain::LLM::OpenAI.new(api_key: "123") + ) + } + + before { subject.create_default_schema } + after { subject.destroy_default_schema } + + describe "#add_texts" do + before do + allow_any_instance_of( + OpenAI::Client + ).to receive(:embeddings) + .with( + parameters: { + dimensions: 1536, + model: "text-embedding-3-small", + input: "Hello World" + } + ) + .and_return({ + "object" => "list", + "data" => [ + {"embedding" => 1536.times.map { rand }} + ] + }) + end + + it "adds texts" do + result = subject.add_texts(texts: ["Hello World", "Hello World"]) + expect(result.size).to eq(2) + end + + it "adds texts with a namespace" do + allow(subject).to receive(:namespace).and_return("test_namespace") + result = subject.add_texts(texts: ["Hello World", "Hello World"]) + expect(result.size).to eq(2) + + count = subject.db.get_first_value("SELECT COUNT(*) FROM test_items WHERE namespace = 'test_namespace'") + expect(count).to eq(2) + end + end + + describe "#update_texts" do + let(:text_embedding_mapping) do + { + "Hello World" => 1536.times.map { rand }, + "Hello World".reverse => 1536.times.map { rand } + } + end + + before do + text_embedding_mapping.each do |input, embedding| + allow_any_instance_of( + OpenAI::Client + ).to receive(:embeddings) + .with( + parameters: { + dimensions: 1536, + model: "text-embedding-3-small", + input: input + } + ) + .and_return({ + "object" => "list", + "data" => [ + {"embedding" => embedding} + ] + }) + end + end + + it "updates texts" do + values = subject.add_texts(texts: ["Hello World", "Hello World"]) + result = subject.update_texts(texts: ["Hello World", "Hello World".reverse], ids: values) + expect(result.size).to eq(2) + end + end + + describe "#remove_texts" do + before do + allow_any_instance_of( + OpenAI::Client + ).to receive(:embeddings) + .with( + parameters: { + dimensions: 1536, + model: "text-embedding-3-small", + input: "Hello World" + } + ) + .and_return({ + "object" => "list", + "data" => [ + {"embedding" => 1536.times.map { rand }} + ] + }) + end + + it "removes texts" do + values = subject.add_texts(texts: ["Hello World", "Hello World"]) + expect(values.length).to eq(2) + + result = subject.remove_texts(ids: values) + expect(result).to eq(2) + end + end + + describe "#similarity_search" do + before do + allow_any_instance_of( + OpenAI::Client + ).to receive(:embeddings) + .with( + parameters: { + dimensions: 1536, + model: "text-embedding-3-small", + input: "earth" + } + ) + .and_return({ + "object" => "list", + "data" => [ + {"embedding" => 1536.times.map { 0 }} + ] + }) + end + + before do + # Add a document with zero vector (should be closest to our search) + subject.db.execute( + "INSERT INTO test_items(rowid, content, embedding) VALUES (?, ?, ?)", + [1, "something about earth", 1536.times.map { 0 }.pack("f*")] + ) + + # Add some random documents + 2.upto(5) do |i| + subject.db.execute( + "INSERT INTO test_items(rowid, content, embedding) VALUES (?, ?, ?)", + [i, "Hello World", 1536.times.map { rand }.pack("f*")] + ) + end + end + + it "searches for similar texts" do + result = subject.similarity_search(query: "earth") + expect(result.first[1]).to eq("something about earth") + end + + it "searches for similar texts using a namespace" do + namespace = "foo_namespace" + subject.db.execute( + "INSERT INTO test_items(rowid, content, embedding, namespace) VALUES (?, ?, ?, ?)", + [6, "a namespaced chunk of text", 1536.times.map { 0 }.pack("f*"), namespace] + ) + + allow(subject).to receive(:namespace).and_return(namespace) + result = subject.similarity_search(query: "earth") + expect(result.first[1]).to eq("a namespaced chunk of text") + end + end + + describe "#similarity_search_by_vector" do + before do + # Add a document with zero vector (should be closest to our search) + subject.db.execute( + "INSERT INTO test_items(rowid, content, embedding) VALUES (?, ?, ?)", + [1, "Some valuable data", 1536.times.map { 0 }.pack("f*")] + ) + + # Add some random documents + 2.upto(5) do |i| + subject.db.execute( + "INSERT INTO test_items(rowid, content, embedding) VALUES (?, ?, ?)", + [i, "Hello World", 1536.times.map { rand }.pack("f*")] + ) + end + end + + it "searches for similar vectors" do + result = subject.similarity_search_by_vector(embedding: 1536.times.map { 0 }) + expect(result.count).to eq(4) + expect(result.first[1]).to eq("Some valuable data") + end + end + + describe "#ask" do + let(:question) { "How many times is 'lorem' mentioned in this text?" } + let(:text) { "Lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed non risus. Suspendisse lectus tortor, dignissim sit amet, adipiscing nec, ultricies sed, dolor." } + let(:messages) { [{role: "user", content: "Context:\n#{text}\n---\nQuestion: #{question}\n---\nAnswer:"}] } + let(:response) { double(completion: answer) } + let(:answer) { "5 times" } + let(:k) { 4 } + + before do + allow_any_instance_of( + OpenAI::Client + ).to receive(:embeddings) + .with( + parameters: { + dimensions: 1536, + model: "text-embedding-3-small", + input: question + } + ) + .and_return({ + "object" => "list", + "data" => [ + {"embedding" => 1536.times.map { 0 }} + ] + }) + end + + before do + subject.db.execute( + "INSERT INTO test_items(rowid, content, embedding) VALUES (?, ?, ?)", + [1, text, 1536.times.map { 0 }.pack("f*")] + ) + end + + context "without block" do + before do + allow(subject.llm).to receive(:chat).with(messages: messages).and_return(response) + expect(response).to receive(:context=).with(text) + end + + it "asks a question and returns the answer" do + expect(subject.ask(question: question, k: k).completion).to eq(answer) + end + end + + context "with block" do + let(:block) { proc { |chunk| puts "Received chunk: #{chunk}" } } + + before do + allow(subject.llm).to receive(:chat) do |parameters| + if parameters[:prompt] == prompt && parameters[:stream].is_a?(Proc) + parameters[:stream].call("Received chunk from llm.chat") + end + end + end + + it "asks a question and yields the chunk to the block" do + expect do + captured_output = capture(:stdout) do + subject.ask(question: question, &block) + end + expect(captured_output).to match(/Received chunk from llm.chat/) + end + end + end + end +end From c482cead51b3f3d923e8eadfccb33d80eee16d54 Mon Sep 17 00:00:00 2001 From: Carl Lange Date: Wed, 15 Jan 2025 16:15:25 +0000 Subject: [PATCH 2/3] Small fixes to the example script --- examples/sqlite_vec_example.rb | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/examples/sqlite_vec_example.rb b/examples/sqlite_vec_example.rb index 27edeeada..469c2e977 100644 --- a/examples/sqlite_vec_example.rb +++ b/examples/sqlite_vec_example.rb @@ -1,11 +1,11 @@ require "langchain" -# Initialize the LLM (using OpenAI in this example) +# Initialize the LLM (using Ollama in this example) llm = Langchain::LLM::Ollama.new # Initialize the SQLite-vec vectorstore db = Langchain::Vectorsearch::SqliteVec.new( - url: ":memory:", # Use a file-based DB (or ":memory:" for in-memory) + url: ":memory:", # Use a file-based DB by passing a path or ":memory:" for in-memory index_name: "documents", namespace: "test", llm: llm @@ -43,5 +43,4 @@ puts "Answer: #{response.chat_completion}" # Clean up -db.destroy_default_schema -File.delete("test_vectors.sqlite3") if File.exist?("test_vectors.sqlite3") +db.destroy_default_schema \ No newline at end of file From 6818dd118f6d2dd9349519981b04338b3b77dbaf Mon Sep 17 00:00:00 2001 From: Carl Lange Date: Thu, 16 Jan 2025 22:50:04 +0000 Subject: [PATCH 3/3] honestly, not sure what these changes are but they appear to be needed. --- Gemfile.lock | 8 ++++++++ lib/langchain/dependency_helper.rb | 7 ++++++- lib/langchain/llm/ollama.rb | 2 +- 3 files changed, 15 insertions(+), 2 deletions(-) diff --git a/Gemfile.lock b/Gemfile.lock index 89efb2ef0..a8e4efd4e 100644 --- a/Gemfile.lock +++ b/Gemfile.lock @@ -399,6 +399,13 @@ GEM spreadsheet (1.3.1) bigdecimal ruby-ole + sqlite3 (1.7.3) + mini_portile2 (~> 2.8.0) + sqlite3 (1.7.3-aarch64-linux) + sqlite3 (1.7.3-arm-linux) + sqlite3 (1.7.3-arm64-darwin) + sqlite3 (1.7.3-x86_64-darwin) + sqlite3 (1.7.3-x86_64-linux) standard (1.39.1) language_server-protocol (~> 3.17.0.2) lint_roller (~> 1.0) @@ -490,6 +497,7 @@ DEPENDENCIES ruby-openai (~> 7.1.0) safe_ruby (~> 1.0.4) sequel (~> 5.87.0) + sqlite3 (~> 1.7.0) standard (>= 1.35.1) vcr weaviate-ruby (~> 0.9.2) diff --git a/lib/langchain/dependency_helper.rb b/lib/langchain/dependency_helper.rb index 23814f131..78748d50d 100644 --- a/lib/langchain/dependency_helper.rb +++ b/lib/langchain/dependency_helper.rb @@ -15,7 +15,12 @@ class VersionError < ScriptError; end # @raise [VersionError] If the gem is installed, but the version does not meet the requirements # def depends_on(gem_name, req: true) - gem(gem_name) # require the gem + if gem_name == "sqlite_vec" + require "sqlite_vec" + return true + else + gem(gem_name) # require the gem + end return(true) unless defined?(Bundler) # If we're in a non-bundler environment, we're no longer able to determine if we'll meet requirements diff --git a/lib/langchain/llm/ollama.rb b/lib/langchain/llm/ollama.rb index 1fd1fd358..faeb27450 100644 --- a/lib/langchain/llm/ollama.rb +++ b/lib/langchain/llm/ollama.rb @@ -24,7 +24,7 @@ class Ollama < Base llama2: 4_096, llama3: 4_096, "llama3.1": 4_096, - "llama3.2": 4_096, + "llama3.2": 3_072, llava: 4_096, mistral: 4_096, "mistral-openorca": 4_096,