vectorsearch 0.1.0 → 0.1.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 4c9913804bc8aaadc08a60c0a250e2923f2a4ceb28633fdd99e4bf86544203b4
-  data.tar.gz: 8db3a77121d948f6ed709618da5bd4411a87a01aba26feb7942e74e0dd18f207
+  metadata.gz: 160b98d1553c63fae2e50c07dae83e80376eade573874e0fcb28a0d1f7f476ea
+  data.tar.gz: 21fbcbb750cd878ceedec2646ef8db116a839b422f751ce3959aa6da8da78ff1
 SHA512:
-  metadata.gz: d76d13ee23c7219483eac27a37ae61cad335bbd2d4169a76723d19a0e334ab6ac01037923cb2c07a64c853280d2449e54c0107fba7f36a231a624efaf1b68b46
-  data.tar.gz: 43612526795e54138ec0c891e5bbdef2ef712ac287a4776680d26d94e8cabcb9d5976774b5144cbeb3ae95b367e26460230931041f011f1a69671e014060684f
+  metadata.gz: 40991aab084c3eb16d8029b598f70227946edc72411c74de0906289d7e23b16a83fbbd7aa9325c44b6ebc42ee0cac3146d55cf9bf5b1eeb23e862f0f66da66c2
+  data.tar.gz: 1c23318876143377b826c8e619978b83e61da68b10dbed34c203ae0f35a10efeb6fa48843fada3782b3fd7c1785802b0cda1332281ebcc39230bb7ed10c0905d

data/Gemfile.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    vectorsearch (0.1.0)
+    vectorsearch (0.1.2)
       cohere-ruby (~> 0.9.1)
       milvus (~> 0.9.0)
       pinecone (~> 0.1.6)

data/README.md

@@ -1,6 +1,7 @@
 # Vectorsearch
+![Tests status](https://github.com/andreibondarev/vectorsearch/actions/workflows/ci.yml/badge.svg) [![Gem Version](https://badge.fury.io/rb/vectorsearch.svg)](https://badge.fury.io/rb/vectorsearch)
 
-Vectorsearch library is an abstraction layer on top of many popular vector search databases. It is a modern ORM that allows developers to easily chunk, generate embeddings, store, search, query and retrieve data from vector search databases. Vectorsearch offers a straight-forward DSL and abstract developers away from overly complex machine learning/data science-specific configurations.
+Vectorsearch library is an abstraction layer on top of many popular vector search databases. It is a modern ORM that allows developers to easily chunk data, generate embeddings, store, search, query and retrieve data from vector search databases. Vectorsearch offers a straight-forward DSL and abstracts away overly complicated machine learning/data science-specific configurations and concepts
 
 ## Installation
 
@@ -20,12 +21,12 @@ require "vectorsearch"
 
 List of currently supported vector search databases and features:
 
-| Database | Querying | Storage | 
-| -------------------------------------- |
-| Weaviate | :white_check_mark: | WIP |
-| Qdrant   | :white_check_mark: | WIP |
-| Milvus   | :white_check_mark: | WIP |
-| Pinecone | :white_check_mark: | WIP |
+| Database | Querying           | Storage | Schema Management | Backups | Rails Integration | ??? |
+| -------- |:------------------:| -------:| -----------------:| -------:| -----------------:| ---:|
+| Weaviate | :white_check_mark: | WIP     | WIP               | WIP     |                   |     |
+| Qdrant   | :white_check_mark: | WIP     | WIP               | WIP     |                   |     |
+| Milvus   | :white_check_mark: | WIP     | WIP               | WIP     |                   |     |
+| Pinecone | :white_check_mark: | WIP     | WIP               | WIP     |                   |     |
 
 ### Create an instance
 
@@ -40,16 +41,24 @@ client = Vectorsearch::Weaviate.new(
     llm_api_key: ENV["OPENAI_API_KEY"]
 )
 
-# You instantiate any other supported vector search database:
+# You can instantiate any other supported vector search database:
 client = Vectorsearch::Milvus.new(...)
 client = Vectorsearch::Qdrant.new(...)
 client = Vectorsearch::Pinecone.new(...)
 ```
 
+```ruby
+# Creating the default schema
+client.create_default_schema
+```
+
 ```ruby
 # Store your documents in your vector search database
-client.add_documents(
-    documents: []
+client.add_texts(
+    texts: [
+        "Begin by preheating your oven to 375°F (190°C). Prepare four boneless, skinless chicken breasts by cutting a pocket into the side of each breast, being careful not to cut all the way through. Season the chicken with salt and pepper to taste. In a large skillet, melt 2 tablespoons of unsalted butter over medium heat. Add 1 small diced onion and 2 minced garlic cloves, and cook until softened, about 3-4 minutes. Add 8 ounces of fresh spinach and cook until wilted, about 3 minutes. Remove the skillet from heat and let the mixture cool slightly.",
+        "In a bowl, combine the spinach mixture with 4 ounces of softened cream cheese, 1/4 cup of grated Parmesan cheese, 1/4 cup of shredded mozzarella cheese, and 1/4 teaspoon of red pepper flakes. Mix until well combined. Stuff each chicken breast pocket with an equal amount of the spinach mixture. Seal the pocket with a toothpick if necessary. In the same skillet, heat 1 tablespoon of olive oil over medium-high heat. Add the stuffed chicken breasts and sear on each side for 3-4 minutes, or until golden brown."
+    ]
 )
 ```
 

data/lib/vectorsearch/base.rb

@@ -7,8 +7,16 @@ module Vectorsearch
   class Base
     attr_reader :client, :index_name, :llm, :llm_api_key
 
+    DEFAULT_METRIC = "cosine".freeze
+    DEFAULT_COHERE_DIMENSION = 1024
+    DEFAULT_OPENAI_DIMENSION = 1536
+
+    # Currently supported LLMs
+    # TODO: Add support for HuggingFace
     LLMS = %i[openai cohere].freeze
 
+    # @param llm [Symbol] The LLM to use
+    # @param llm_api_key [String] The API key for the LLM
     def initialize(llm:, llm_api_key:)
       validate_llm!(llm: llm)
 
@@ -16,27 +24,96 @@ module Vectorsearch
       @llm_api_key = llm_api_key
     end
 
+    def create_default_schema
+      raise NotImplementedError
+    end
+
+    # TODO
+    def add_texts(texts:)
+      raise NotImplementedError
+    end
+
+    # NotImplementedError will be raised if the subclass does not implement the `ask()` method
+    def ask(question:)
+      raise NotImplementedError
+    end
+
+    # Generate an embedding for a given text
+    # Currently supports OpenAI and Cohere
+    # The LLM-related method will most likely need to be abstracted out into a separate class
+    # @param text [String] The text to generate an embedding for
+    # @return [String] The embedding
     def generate_embedding(text:)
       case llm
       when :openai
-        response = OpenAI::Client.new(access_token: llm_api_key)
-          .embeddings(
-            parameters: {
-              model: "text-embedding-ada-002",
-              input: text
-            }
-          )
+        response = openai_client.embeddings(
+          parameters: {
+            model: "text-embedding-ada-002",
+            input: text
+          }
+        )
         response.dig("data").first.dig("embedding")
       when :cohere
-        response = Cohere::Client.new(api_key: llm_api_key)
-          .embed(
-            texts: [text],
-            model: "small"
-          )
+        response = cohere_client.embed(
+          texts: [text],
+          model: "small"
+        )
         response.dig("embeddings").first
       end
     end
 
+    # Generate a completion for a given prompt
+    # Currently supports OpenAI and Cohere
+    # The LLM-related method will most likely need to be abstracted out into a separate class
+    # @param prompt [String] The prompt to generate a completion for
+    # @return [String] The completion
+    def generate_completion(prompt:)
+      case llm
+      when :openai
+        response = openai_client.completions(
+          parameters: {
+            model: "text-davinci-003",
+            temperature: 0.0,
+            prompt: prompt
+          }
+        )
+        response.dig("choices").first.dig("text")
+      when :cohere
+        response = cohere_client.generate(
+          prompt: prompt,
+          temperature: 0.0
+        )
+        response.dig("generations").first.dig("text")
+      end
+    end
+
+    def generate_prompt(question:, context:)
+      "Context:\n" +
+      "#{context}\n" +
+      "---\n" +
+      "Question: #{question}\n" +
+      "---\n" +
+      "Answer:"
+    end
+
+    private
+
+    def default_dimension
+      if llm == :openai
+        DEFAULT_OPENAI_DIMENSION
+      elsif llm == :cohere
+        DEFAULT_COHERE_DIMENSION
+      end
+    end
+
+    def openai_client
+      @openai_client ||= OpenAI::Client.new(access_token: llm_api_key)
+    end
+
+    def cohere_client
+      @cohere_client ||= Cohere::Client.new(api_key: llm_api_key)
+    end
+
     def validate_llm!(llm:)
       raise ArgumentError, "LLM must be one of #{LLMS}" unless LLMS.include?(llm)
     end

data/lib/vectorsearch/milvus.rb

@@ -19,6 +19,64 @@ module Vectorsearch
       super(llm: llm, llm_api_key: llm_api_key)
     end
 
+    def add_texts(
+      texts:
+    )
+      client.entities.insert(
+        collection_name: index_name,
+        num_rows: texts.count,
+        fields_data: [
+          {
+            field_name: "content",
+            type: ::Milvus::DATA_TYPES["varchar"],
+            field: texts
+          }, {
+            field_name: "vectors",
+            type: ::Milvus::DATA_TYPES["binary_vector"],
+            field: texts.map { |text| generate_embedding(text: text) }
+          }
+        ]
+      )
+    end
+
+    # Create default schema
+    # @return [Hash] The response from the server
+    def create_default_schema
+      client.collections.create(
+        auto_id: true,
+        collection_name: index_name,
+        description: "Default schema created by Vectorsearch",
+        fields: [
+          {
+            name: "id",
+            is_primary_key: true,
+            autoID: true,
+            data_type: ::Milvus::DATA_TYPES["int64"]
+          }, {
+            name: "content",
+            is_primary_key: false,
+            data_type: ::Milvus::DATA_TYPES["varchar"],
+            type_params: [
+              {
+                key: "max_length",
+                value: "32768" # Largest allowed value
+              }
+            ]
+          }, {
+            name: "vectors",
+            data_type: ::Milvus::DATA_TYPES["binary_vector"],
+            is_primary_key: false,
+            type_params: [
+              {
+                key: "dim",
+                value: default_dimension.to_s
+              }
+            ]
+          }
+        ]
+      )
+    end
+
     def similarity_search(
       query:,
       k: 4
@@ -41,7 +99,7 @@ module Vectorsearch
         vectors: [ embedding ],
         dsl_type: 1,
         params: "{\"nprobe\": 10}",
-        anns_field: "book_intro", # Should it get all abstracted away to "content" field?
+        anns_field: "content",
         metric_type: "L2"
       )
     end

data/lib/vectorsearch/pinecone.rb

@@ -22,6 +22,36 @@ module Vectorsearch
       super(llm: llm, llm_api_key: llm_api_key)
     end
 
+    # Add a list of texts to the index
+    # @param texts [Array] The list of texts to add
+    # @return [Hash] The response from the server
+    def add_texts(
+      texts:
+    )
+      vectors = texts.map do |text|
+        {
+          # TODO: Allows passing in your own IDs
+          id: SecureRandom.uuid,
+          metadata: { content: text },
+          values: generate_embedding(text: text)
+        }
+      end
+
+      index = client.index(index_name)
+
+      index.upsert(vectors: vectors)
+    end
+
+    # Create the index with the default schema
+    # @return [Hash] The response from the server
+    def create_default_schema
+      client.create_index(
+        metric: DEFAULT_METRIC,
+        name: index_name,
+        dimension: default_dimension
+      )
+    end
+
     def similarity_search(
       query:,
       k: 4
@@ -40,16 +70,26 @@ module Vectorsearch
     )
       index = client.index(index_name)
 
-      index.query(
+      response = index.query(
         vector: embedding,
         top_k: k,
         include_values: true,
         include_metadata: true
       )
+      response.dig("matches")
     end
 
     def ask(question:)
-      raise NotImplementedError
+      search_results = similarity_search(query: question)
+
+      context = search_results.dig("matches").map do |result|
+        result.dig("metadata").to_s
+      end
+      context = context.join("\n---\n")
+
+      prompt = generate_prompt(question: question, context: context)
+
+      generate_completion(prompt: prompt)
     end
   end
 end

data/lib/vectorsearch/qdrant.rb

@@ -20,6 +20,38 @@ module Vectorsearch
       super(llm: llm, llm_api_key: llm_api_key)
     end
 
+    # Add a list of texts to the index
+    # @param texts [Array] The list of texts to add
+    # @return [Hash] The response from the server
+    def add_texts(
+      texts:
+    )
+      batch = { ids: [], vectors: [], payloads: [] }
+
+      texts.each do |text|
+        batch[:ids].push(SecureRandom.uuid)
+        batch[:vectors].push(generate_embedding(text: text))
+        batch[:payloads].push({ content: text })
+      end
+
+      client.points.upsert(
+        collection_name: index_name,
+        batch: batch
+      )
+    end
+
+    # Create the index with the default schema
+    # @return [Hash] The response from the server
+    def create_default_schema
+      client.collections.create(
+        collection_name: index_name,
+        vectors: {
+          distance: DEFAULT_METRIC.capitalize,
+          size: default_dimension
+        }
+      )
+    end
+
     def similarity_search(
       query:,
       k: 4
@@ -45,7 +77,16 @@ module Vectorsearch
     end
 
     def ask(question:)
-      raise NotImplementedError
+      search_results = similarity_search(query: question)
+
+      context = search_results.dig("result").map do |result|
+        result.dig("payload").to_s
+      end
+      context = context.join("\n---\n")
+
+      prompt = generate_prompt(question: question, context: context)
+
+      generate_completion(prompt: prompt)
     end
   end
 end

data/lib/vectorsearch/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Vectorsearch
-  VERSION = "0.1.0"
+  VERSION = "0.1.2"
 end

data/lib/vectorsearch/weaviate.rb

@@ -25,32 +25,56 @@ module Vectorsearch
     def add_texts(
       texts:
     )
+      objects = []  
       texts.each do |text|
-        text['class'] = index_name
+        objects.push({
+          class_name: index_name,
+          properties: {
+            content: text
+          }
+        })
       end
 
-      client.batch_create(
-        objects: texts     
+      client.objects.batch_create(
+        objects: objects
+      )
+    end
+
+    def create_default_schema
+      client.schema.create(
+        class_name: index_name,
+        vectorizer: "text2vec-#{llm.to_s}",
+        properties: [
+          {
+            dataType: ["text"],
+            name: "content"
+          }
+        ]
       )
     end
 
     # Return documents similar to the query
+    # @param query [String] The query to search for
+    # @param k [Integer|String] The number of results to return
+    # @return [Hash] The search results
     def similarity_search(
       query:,
       k: 4
     )
-      near_text = "{
-        concepts: [\"#{query}\"],
-      }"
+      near_text = "{ concepts: [\"#{query}\"] }"
 
       client.query.get(
         class_name: index_name,
         near_text: near_text,
         limit: k.to_s,
-        fields: "content recipe_id"
+        fields: "content _additional { id }"
       )
     end
 
+    # Return documents similar to the vector
+    # @param embedding [Array] The vector to search for
+    # @param k [Integer|String] The number of results to return
+    # @return [Hash] The search results
     def similarity_search_by_vector(
       embedding:,
       k: 4
@@ -65,17 +89,34 @@ module Vectorsearch
       )
     end
 
+    # Ask a question and return the answer
+    # @param question [String] The question to ask
+    # @return [Hash] The answer
     def ask(
       question:
     )
-      ask_object = "{ question: \"#{question}\" }"
+      # Weaviate currently supports the `ask:` parameter only for the OpenAI LLM (with `qna-openai` module enabled).
+      if llm == :openai
+        ask_object = "{ question: \"#{question}\" }"
 
-      client.query.get(
-        class_name: index_name,
-        ask: ask_object,
-        limit: "1",
-        fields: "_additional { answer { result } }"
-      )
+        client.query.get(
+          class_name: index_name,
+          ask: ask_object,
+          limit: "1",
+          fields: "_additional { answer { result } }"
+        )
+      elsif llm == :cohere
+        search_results = similarity_search(query: question)
+
+        context = search_results.map do |result|
+          result.dig("content").to_s
+        end
+        context = context.join("\n---\n")
+
+        prompt = generate_prompt(question: question, context: context)
+
+        generate_completion(prompt: prompt)
+      end
     end
   end
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: vectorsearch
 version: !ruby/object:Gem::Version
-  version: 0.1.0
+  version: 0.1.2
 platform: ruby
 authors:
 - Andrei Bondarev
 autorequire:
 bindir: exe
 cert_chain: []
-date: 2023-04-30 00:00:00.000000000 Z
+date: 2023-05-01 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: pry-byebug