|
| 1 | +# frozen_string_literal: true |
| 2 | + |
| 3 | +module Langchain::Vectorsearch |
| 4 | + class Pgvector < Base |
| 5 | + # |
| 6 | + # The PostgreSQL vector search adapter |
| 7 | + # |
| 8 | + # Gem requirements: |
| 9 | + # gem "pgvector", "~> 0.2" |
| 10 | + # |
| 11 | + # Usage: |
| 12 | + # pgvector = Langchain::Vectorsearch::Pgvector.new(llm:, model_name:) |
| 13 | + # |
| 14 | + |
| 15 | + # The operators supported by the PostgreSQL vector search adapter |
| 16 | + OPERATORS = [ |
| 17 | + "cosine", |
| 18 | + "euclidean", |
| 19 | + "inner_product" |
| 20 | + ] |
| 21 | + DEFAULT_OPERATOR = "cosine" |
| 22 | + |
| 23 | + attr_reader :db, :operator, :llm |
| 24 | + attr_accessor :model |
| 25 | + |
| 26 | + # @param url [String] The URL of the PostgreSQL database |
| 27 | + # @param index_name [String] The name of the table to use for the index |
| 28 | + # @param llm [Object] The LLM client to use |
| 29 | + # @param namespace [String] The namespace to use for the index when inserting/querying |
| 30 | + def initialize(llm:) |
| 31 | + # If the line below is called, the generator fails as calls to |
| 32 | + # LangchainrbRails.config.vectorsearch will generate an exception. |
| 33 | + # These happen in the template files. |
| 34 | + # depends_on "neighbor" |
| 35 | + |
| 36 | + @operator = DEFAULT_OPERATOR |
| 37 | + |
| 38 | + super(llm: llm) |
| 39 | + end |
| 40 | + |
| 41 | + # Add a list of texts to the index |
| 42 | + # @param texts [Array<String>] The texts to add to the index |
| 43 | + # @param ids [Array<String>] The ids to add to the index, in the same order as the texts |
| 44 | + # @return [Array<Integer>] The the ids of the added texts. |
| 45 | + def add_texts(texts:, ids:) |
| 46 | + embeddings = texts.map do |text| |
| 47 | + llm.embed(text: text).embedding |
| 48 | + end |
| 49 | + |
| 50 | + # I believe the records returned by #find must be in the |
| 51 | + # same order as the embeddings. I _think_ this works for uuid ids but didn't test |
| 52 | + # deeply. |
| 53 | + # TODO - implement find_each so we don't load all records into memory |
| 54 | + model.find(ids).each.with_index do |record, i| |
| 55 | + record.update_column(:embedding, embeddings[i]) |
| 56 | + end |
| 57 | + end |
| 58 | + |
| 59 | + def update_texts(texts:, ids:) |
| 60 | + add_texts(texts: texts, ids: ids) |
| 61 | + end |
| 62 | + |
| 63 | + # Invoke a rake task that will create an initializer (`config/initializers/langchain.rb`) file |
| 64 | + # and db/migrations/* files |
| 65 | + def create_default_schema |
| 66 | + Rake::Task["pgvector"].invoke |
| 67 | + end |
| 68 | + |
| 69 | + # Destroy default schema |
| 70 | + def destroy_default_schema |
| 71 | + # Tell the user to rollback the migration |
| 72 | + end |
| 73 | + |
| 74 | + # Search for similar texts in the index |
| 75 | + # @param query [String] The text to search for |
| 76 | + # @param k [Integer] The number of top results to return |
| 77 | + # @return [Array<Hash>] The results of the search |
| 78 | + # TODO - drop the named "query:" param so it is the same interface as #ask? |
| 79 | + def similarity_search(query:, k: 4) |
| 80 | + embedding = llm.embed(text: query).embedding |
| 81 | + |
| 82 | + similarity_search_by_vector( |
| 83 | + embedding: embedding, |
| 84 | + k: k |
| 85 | + ) |
| 86 | + end |
| 87 | + |
| 88 | + # Search for similar texts in the index by the passed in vector. |
| 89 | + # You must generate your own vector using the same LLM that generated the embeddings stored in the Vectorsearch DB. |
| 90 | + # @param embedding [Array<Float>] The vector to search for |
| 91 | + # @param k [Integer] The number of top results to return |
| 92 | + # @return [Array<Hash>] The results of the search |
| 93 | + # TODO - drop the named "embedding:" param so it is the same interface as #ask? |
| 94 | + def similarity_search_by_vector(embedding:, k: 4) |
| 95 | + model |
| 96 | + .nearest_neighbors(:embedding, embedding, distance: operator) |
| 97 | + .limit(k) |
| 98 | + end |
| 99 | + |
| 100 | + # Ask a question and return the answer |
| 101 | + # @param question [String] The question to ask |
| 102 | + # @param k [Integer] The number of results to have in context |
| 103 | + # @yield [String] Stream responses back one String at a time |
| 104 | + # @return [String] The answer to the question |
| 105 | + def ask(question, k: 4, &block) |
| 106 | + # Noisy as the embedding column has a lot of data |
| 107 | + ActiveRecord::Base.logger.silence do |
| 108 | + search_results = similarity_search(query: question, k: k) |
| 109 | + |
| 110 | + context = search_results.map do |result| |
| 111 | + result.as_vector |
| 112 | + end |
| 113 | + context = context.join("\n---\n") |
| 114 | + |
| 115 | + prompt = generate_rag_prompt(question: question, context: context) |
| 116 | + |
| 117 | + llm.chat(prompt: prompt, &block) |
| 118 | + end |
| 119 | + end |
| 120 | + end |
| 121 | +end |
| 122 | + |
| 123 | +# Rails connection when configuring vectorsearch |
| 124 | +# Update READMEs |
| 125 | +# Rails migration to create a migration |
0 commit comments