chroma-db 0.1.0 → 0.3.0

data/lib/.rbnext/3.1/chroma/resources/collection.rb

@@ -0,0 +1,385 @@
+# frozen_string_literal: true
+
+module Chroma
+    using RubyNext
+  module Resources
+    # A Collection class represents a store for your embeddings, documents, and any additional metadata.
+    # This class can be instantiated by receiving the collection's name and metadata hash.
+    class Collection
+      include Chroma::APIOperations::Request
+
+      attr_reader :name
+      attr_reader :metadata
+
+      def initialize(name:, metadata: nil)
+        @name = name
+        @metadata = metadata
+      end
+
+      # Query the collection and return an array of embeddings.
+      #
+      # query_embeddings - An array of the embeddings to use for querying the collection.
+      # results - The maximum number of results to return. 10 by default.
+      # where - A Hash of additional conditions to filter the query results (optional).
+      # where_document - A Hash of additional conditions to filter the associated documents (optional).
+      # include  - An Array of the additional information to include in the query results (optional). Metadatas,
+      #   documents, and distances by default.
+      #
+      # Examples
+      #
+      #   collection = Chroma::Resource::Collection.get("ruby-documentation")
+      #   embeddings = collection.query(query_embeddings: [[1.5, 2.9, 3.3]], results: 5)
+      #
+      # Return an Array of Embedding with query results.
+      def query(query_embeddings:, results: 10, where: {}, where_document: {}, include: %w[metadatas documents distances])
+        payload = {
+          query_embeddings: query_embeddings,
+          n_results: results,
+          where: where,
+          where_document: where_document,
+          include: include
+        }
+
+        result = self.class.execute_request(:post, "#{Chroma.api_url}/collections/#{name}/query", payload)
+
+        if result.success?
+          build_embeddings_response(result.success.body)
+        else
+          raise_failure_error(result)
+        end
+      end
+
+      # Get embeddings from the collection.
+      #
+      # ids - An Array of the specific embedding IDs to retrieve (optional).
+      # where - A Hash of additional conditions to filter the query results (optional).
+      # sort - The sorting criteria for the query results (optional).
+      # limit - The maximum number of embeddings to retrieve (optional).
+      # offset - The offset for pagination (optional).
+      # page - The page number for pagination (optional).
+      # page_size - The page size for pagination (optional).
+      # where_document - A Hash of additional conditions to filter the associated documents (optional).
+      # include - An Array of the additional information to include in the query results (optional). Metadatas,
+      #   and documents by default.
+      #
+      # Examples
+      #
+      #   collection = Chroma::Resource::Collection.get("ruby-documentation")
+      #   embeddings = collection.get([Array#sort, "Array#each"])
+      #
+      # Returns an Array of Embeddings
+      def get(ids: nil, where: {}, sort: nil, limit: nil, offset: nil, page: nil, page_size: nil, where_document: {}, include: %w[metadatas documents])
+        if !page.nil? && !page_size.nil?
+          offset = (page - 1) * page_size
+          limit = page_size
+        end
+
+        payload = {
+          ids: ids,
+          where: where,
+          sort: sort,
+          limit: limit,
+          offset: offset,
+          where_document: where_document,
+          include: include
+        }
+
+        result = self.class.execute_request(:post, "#{Chroma.api_url}/collections/#{name}/get", payload)
+
+        if result.success?
+          build_embeddings_response(result.success.body)
+        else
+          raise_failure_error(result)
+        end
+      end
+
+      # Add one or many embeddings to the collection.
+      #
+      # embeddings - An Array of Embeddings or one Embedding to add.
+      #
+      # Examples
+      #
+      #   collection = Chroma::Resource::Collection.get("ruby-documentation")
+      #   collection.add(Embedding.new(id: "Array#fetch", embeddings[9.8, 2.3, 2.9], metadata: {url: "https://..."}))
+      #
+      # Returns true with success or raises a Chroma::Error on failure.
+      def add(embeddings = [])
+        embeddings_array = Array(embeddings)
+        return false if embeddings_array.size == 0
+
+        payload = build_embeddings_payload(embeddings_array)
+
+        result = self.class.execute_request(:post, "#{Chroma.api_url}/collections/#{name}/add", payload)
+
+        return true if result.success?
+
+        raise_failure_error(result)
+      end
+
+      # Delete embeddings from the collection.
+      #
+      # ids [Array<Integer>, nil] The specific embedding IDs to delete (optional).
+      # where [Hash] Additional conditions to filter the embeddings to delete (optional).
+      # where_document [Hash] Additional conditions to filter the associated documents (optional).
+      #
+      # Examples
+      #
+      #   collection = Chroma::Resource::Collection.get("ruby-documentation")
+      #   collection.delete(["Array#fetch", "Array#sort"])
+      #
+      # Returns an Array of deleted global ids.
+      def delete(ids: nil, where: {}, where_document: {})
+        payload = {
+          ids: ids,
+          where: where,
+          where_document: where_document
+        }
+
+        result = self.class.execute_request(:post, "#{Chroma.api_url}/collections/#{name}/delete", payload)
+
+        return result.success.body if result.success?
+
+        raise_failure_error(result)
+      end
+
+      # Update one or many embeddings to the collection.
+      #
+      # embeddings - An Array of Embeddings or one Embedding to add.
+      #
+      # Examples
+      #
+      #   collection = Chroma::Resource::Collection.get("ruby-documentation")
+      #   collection.update(Embedding.new(id: "Array#fetch", embeddings[9.8, 2.3, 2.9], metadata: {url: "https://..."}))
+      #
+      # Returns true with success or raises a Chroma::Error on failure.
+      def update(embeddings = [])
+        embeddings_array = Array(embeddings)
+        return false if embeddings_array.size == 0
+
+        payload = build_embeddings_payload(embeddings_array)
+        payload.delete(:increment_index)
+
+        result = self.class.execute_request(:post, "#{Chroma.api_url}/collections/#{name}/update", payload)
+
+        return true if result.success?
+
+        raise_failure_error(result)
+      end
+
+      # Upsert (insert or update) one or many embeddings to the collection.
+      #
+      # embeddings - An Array of Embeddings or one Embedding to add.
+      #
+      # Examples
+      #
+      #   collection = Chroma::Resource::Collection.get("ruby-documentation")
+      #   embeddings = [
+      #     Embedding.new(id: "Array#fetch", embeddings[9.8, 2.3, 2.9], metadata: {url: "https://..."}),
+      #     Embedding.new(id: "Array#select", embeddings[5.6, 3.1, 4.7], metadata: {url: "https://..."})
+      #   ]
+      #   collection.upsert()
+      #
+      # Returns true with success or raises a Chroma::Error on failure.
+      def upsert(embeddings = [])
+        embeddings_array = Array(embeddings)
+        return false if embeddings_array.size == 0
+
+        payload = build_embeddings_payload(embeddings_array)
+
+        result = self.class.execute_request(:post, "#{Chroma.api_url}/collections/#{name}/upsert", payload)
+
+        return true if result.success?
+
+        raise_failure_error(result)
+      end
+
+      # Count the number of embeddings in a collection.
+      #
+      # Examples
+      #
+      #   collection = Chroma::Resource::Collection.get("ruby-documentation")
+      #   collection.count
+      #
+      # Returns the count of embeddings in the collection.
+      def count
+        result = self.class.execute_request(:get, "#{Chroma.api_url}/collections/#{name}/count")
+
+        return result.success.body if result.success?
+
+        raise_failure_error(result)
+      end
+
+      # Modify the name and metadata of the current collection.
+      #
+      # new_name - The new name for the collection.
+      # new_metadata - The new metadata hash for the collection.
+      #
+      # Examples:
+      #
+      #   collection = Chroma::Resource::Collection.get("ruby-documentation")
+      #   collection.modify("ruby-3.2-documentation")
+      #
+      # Returns nothing.
+      def modify(new_name, new_metadata: {})
+        payload = {new_name: new_name}
+        payload[:new_metadata] = new_metadata if new_metadata.any?
+
+        result = self.class.execute_request(:put, "#{Chroma.api_url}/collections/#{name}", payload)
+
+        if result.success?
+          @name = new_name
+          @metadata = new_metadata
+        else
+          raise_failure_error(result)
+        end
+      end
+
+      # Creates an index for the collection.
+      #
+      # Examples:
+      #
+      #   collection = Chroma::Resource::Collection.get("ruby-documentation")
+      #   collection.create_index
+      #
+      # Returns true on success or raise a Chroma::Error on failure.
+      def create_index
+        result = self.class.execute_request(:post, "#{Chroma.api_url}/collections/#{name}/create_index")
+
+        return true if result.success?
+
+        raise_failure_error(result)
+      end
+
+      # Create a new collection on the database.
+      #
+      # name - The name of the collection. Name needs to be between 3-63 characters, starts and ends
+      #   with an alphanumeric character, contains only alphanumeric characters, underscores or hyphens (-), and
+      #   contains no two consecutive periods
+      # metadata - A hash of additional metadata associated with the collection.
+      #
+      # Examples
+      #
+      #   collection = Chorma::Resources::Collection.create("ruby-documentation", {source: "Ruby lang website"})
+      #
+      # Returns the created collection object.
+      def self.create(name, metadata = nil)
+        payload = {name: name, metadata: metadata, get_or_create: false}
+
+        result = execute_request(:post, "#{Chroma.api_url}/collections", payload)
+
+        if result.success?
+          data = result.success.body
+          new(name: data["name"], metadata: data["metadata"])
+        else
+          raise_failure_error(result)
+        end
+      end
+
+      # Retrieves a collection from the database.
+      #
+      # name - The name of the collection to retrieve.
+      #
+      # Examples
+      #
+      #   collection = Chroma::Resources::Colection.get("ruby-documentation")
+      #
+      # Returns The retrieved collection object. Raises Chroma::InvalidRequestError if not found.
+      def self.get(name)
+        result = execute_request(:get, "#{Chroma.api_url}/collections/#{name}")
+
+        if result.success?
+          data = result.success.body
+          new(name: data["name"], metadata: data["metadata"])
+        else
+          raise_failure_error(result)
+        end
+      end
+
+      # Retrieves all collections in the database.
+      #
+      # Examples
+      #
+      #   collections = Chroma::Resources::Collection.list
+      #
+      # Returns An array of all collections in the database.
+      def self.list
+        result = execute_request(:get, "#{Chroma.api_url}/collections")
+
+        if result.success?
+          data = result.success.body
+          data.map { |item| new(name: item["name"], metadata: item["metadata"]) }
+        else
+          raise_failure_error(result)
+        end
+      end
+
+      # Deletes a collection from the database.
+      #
+      # name - The name of the collection to delete.
+      #
+      # Examples
+      #
+      #   Chroma::Resources::Collection.delete("ruby-documentation")
+      #
+      # Returns true if the collection was successfully deleted, raise Chroma::InvalidRequestError otherwise.
+      def self.delete(name)
+        result = execute_request(:delete, "#{Chroma.api_url}/collections/#{name}")
+
+        return true if result.success?
+
+        raise_failure_error(result)
+      end
+
+      def self.raise_failure_error(result)
+        case result.failure.error
+        in Exception => exception
+          raise Chroma::APIConnectionError.new(exception.message)
+        in Net::HTTPInternalServerError => response
+          if response.body.is_a?(String) && (response.body.include?("ValueError") || response.body.include?("IndexError") || response.body.include?("TypeError"))
+            raise Chroma::InvalidRequestError.new(result.failure.body, status: result.failure.status, body: result.failure.body)
+          else
+            raise Chroma::APIConnectionError.new(result.failure.body, status: result.failure.status, body: result.failure.body)
+          end
+        else
+          raise Chroma::APIError.new(result.failure.body, status: result.failure.status, body: result.failure.body)
+        end
+      end
+      private_class_method :raise_failure_error
+
+      private
+
+      def build_embeddings_payload(embeddings, increment_index = true)
+        payload = {ids: [], embeddings: [], metadatas: [], documents: [], increment_index: increment_index}
+
+        embeddings.each do |embedding|
+          payload[:ids] << embedding.id
+          payload[:embeddings] << embedding.embedding
+          payload[:metadatas] << embedding.metadata
+          payload[:documents] << embedding.document
+        end
+
+        payload
+      end
+
+      def build_embeddings_response(result)
+        Chroma::Util.log_debug("Building embeddings from #{result.inspect}")
+
+        result_ids = result.fetch("ids", []).flatten
+        result_embeddings = (result.dig("embeddings") || []).flatten
+        result_documents = (result.dig("documents") || []).flatten
+        result_metadatas = (result.dig("metadatas") || []).flatten
+        result_distances = (result.dig("distances") || []).flatten
+
+        Chroma::Util.log_debug("Ids #{result_ids.inspect}")
+        Chroma::Util.log_debug("Embeddings #{result_embeddings.inspect}")
+        Chroma::Util.log_debug("Documents #{result_documents.inspect}")
+        Chroma::Util.log_debug("Metadatas #{result_metadatas.inspect}")
+        Chroma::Util.log_debug("distances #{result_distances.inspect}")
+
+        result_ids.map.with_index do |id, index|
+          Chroma::Resources::Embedding.new(id: id, embedding: result_embeddings[index], document: result_documents[index], metadata: result_metadatas[index], distance: result_distances[index])
+        end
+      end
+    end
+  end
+end

data/lib/chroma/api_operations/request.rb

@@ -0,0 +1,130 @@
+# frozen_string_literal: true
+
+module Chroma
+  module APIOperations
+    using RubyNext
+
+    # Request's response Data object.
+    #
+    # status - HTTP status code. It is zero when a request fails due to network error.
+    # body - Parsed JSON object or response body.
+    # headers - HTTP response headers.
+    # error - Exception or Net::HTTPResponse object if the response is not Net::HTTPSuccess
+    #
+    # NOTE: Not supported yet by Ruby Next
+    # Response = Data.define(:status, :body, :headers, :error)
+    Response = Struct.new("Response", :status, :body, :headers, :error)
+
+    # Request module provides functionality to perform HTTP requests.
+    module Request
+      module ClassMethods
+        include Dry::Monads[:result]
+
+        # Execute an HTTP request and return a monad wrapping the response.
+        #
+        # method - The HTTP method to use (e.g. 'GET', 'POST'). Method must be a `Symbol`.
+        # url - The URL to send the request to.
+        # params - The query parameters or request body. Params needs to be in a form of a Hash.
+        # options - Additional options to pass to the request.
+        #
+        # A `Dry::Monads::Result` monad wrapping the response, either a success or failure.
+        # The response is a `Chroma::APIOperations::Response` Data object.
+        #
+        # Examples
+        #
+        #   result = execute_request(:get, "https://example.com", {name: "test request"})
+        #   if result.success?
+        #     puts "Response status: #{result.success.status}"
+        #     puts "Response body: #{result.success.body}"
+        #   else
+        #     puts "Request failed with error: #{result.failure.error}"
+        #   end
+        def execute_request(method, url, params = {}, options = {})
+          uri = URI.parse(url)
+
+          request = build_request(method, uri, params)
+
+          use_ssl = options.delete(:use_ssl) || false
+          response = Net::HTTP.start(uri.hostname, uri.port, use_ssl:) do |http|
+            Chroma::Util.log_debug("Sending a request", {method:, uri:, params:})
+            http.request(request)
+          end
+
+          build_response(response)
+        rescue => ex
+          build_response(ex)
+        end
+
+        private def build_response(response)
+          case response
+          in Net::HTTPSuccess => success_response
+            Chroma::Util.log_info("Successful response", code: success_response.code)
+
+            build_response_details(success_response)
+          in Net::HTTPRedirection => redirect_response
+            Chroma::Util.log_info("Server redirect response", code: redirect_response.code, location: redirect_response["location"])
+
+            build_response_details(redirect_response)
+          in Net::HTTPClientError => client_error_response
+            Chroma::Util.log_error("Client error response", code: client_error_response.code, body: client_error_response.body)
+
+            build_response_details(client_error_response)
+          in Net::HTTPServerError => server_error_response
+            Chroma::Util.log_error("Server error response", code: server_error_response.code)
+
+            build_response_details(server_error_response, parse_body: false)
+          else
+            Chroma::Util.log_error("An error happened", error: response.to_s)
+
+            build_response_details(response, exception: true, parse_body: false)
+          end
+        end
+
+        private def build_response_details(response, exception: false, parse_body: true)
+          response_data = Chroma::APIOperations::Response.new(
+            exception ? 0 : response.code.to_i,
+            if exception
+              exception.to_s
+            else
+              (parse_body ? body_to_json(response.body) : response.body)
+            end,
+            exception ? {} : response.each_header.to_h,
+            response.is_a?(Net::HTTPSuccess) ? nil : response
+          )
+
+          case response
+          in Net::HTTPSuccess
+            return Success(response_data)
+          else
+            return Failure(response_data)
+          end
+        end
+
+        private def body_to_json(content)
+          JSON.parse(content, symbolize_keys: true)
+        rescue JSON::ParserError, TypeError
+          content
+        end
+
+        private def build_request(method, uri, params)
+          request = case method
+          when :post then Net::HTTP::Post.new(uri)
+          when :put then Net::HTTP::Put.new(uri)
+          when :delete then Net::HTTP::Delete.new(uri)
+          else
+            Net::HTTP::Get.new(uri)
+          end
+
+          request.content_type = "application/json"
+          request.body = params.to_json if params.size > 0
+
+          request
+        end
+      end
+
+      def self.included(base)
+        base.extend(ClassMethods)
+      end
+    end
+  end
+end

data/lib/chroma/chroma.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+module Chroma
+  # map to the same values as the standard library's logger
+  LEVEL_DEBUG = Logger::DEBUG
+  LEVEL_ERROR = Logger::ERROR
+  LEVEL_INFO = Logger::INFO
+
+  @config = Chroma::ChromaConfiguration.setup
+
+  class << self
+    extend Forwardable
+
+    attr_reader :config
+
+    # User configuration options
+    def_delegators :@config, :connect_host, :connect_host=
+    def_delegators :@config, :api_base, :api_base=
+    def_delegators :@config, :api_version, :api_version=
+    def_delegators :@config, :log_level, :log_level=
+    def_delegators :@config, :logger, :logger=
+  end
+
+  def self.api_url
+    "#{connect_host}/#{api_base}/#{api_version}"
+  end
+
+  Chroma.log_level = ENV["CHROMA_LOG"].to_i unless ENV["CHROMA_LOG"].nil?
+end

data/lib/chroma/chroma_configuration.rb

@@ -0,0 +1,75 @@
+# frozen_string_literal: true
+
+module Chroma
+  # The ChromaConfiguration class allows you to configure the connection settings for a Chroma service client.
+  # It can be initialized directly calling #new method.
+  #
+  # Examples
+  #
+  #   configuration = Chroma::ChromaConfiguration.new
+  #   configurarion.connect_host = "https://chroma.example.com"
+  #
+  # Or via a setup block.
+  #
+  # Examples
+  #
+  #   Chroma::ChromaConfiguration.setup do |config|
+  #     config.connect_host = "https://chroma.example.com"
+  #   end
+  class ChromaConfiguration
+    # Sets the host name for the Chroma service.
+    #
+    # Examples
+    #
+    #   config.connect_host = "chroma.example.com"
+    #
+    # Returns the String host name.
+    attr_accessor :connect_host
+
+    # Sets the base path for the Chroma API.
+    #
+    # Examples
+    #
+    #   config.api_base = "/api"
+    #
+    # Returns the String API base path.
+    attr_accessor :api_base
+    # Sets the version of the Chroma API.
+    #
+    # Examples
+    #
+    #   config.api_version = "v1"
+    #
+    # Returns the String API version.
+    attr_accessor :api_version
+    # Sets the logger for the Chroma service client.
+    #
+    # Examples
+    #
+    #   config.logger = Logger.new(STDOUT)
+    #
+    # Returns the Logger instance.
+    attr_accessor :logger
+    # Sets the logger's log level for the Chroma service client.
+    #
+    # Examples
+    #
+    #   config.log_level = Chroma::LEVEL_INFO
+    #
+    # Returns the log level constant
+    attr_accessor :log_level
+
+    def self.setup
+      new.tap do |instance|
+        yield(instance) if block_given?
+      end
+    end
+
+    def initialize
+      @api_base = "api"
+      @api_version = "v1"
+
+      @log_level = Chroma::LEVEL_INFO
+    end
+  end
+end

data/lib/chroma/errors.rb

@@ -0,0 +1,71 @@
+# frozen_string_literal: true
+
+module Chroma
+  # ChromaError is the base error from which all other more specific Chrome
+  # errors derive.
+  class ChromaError < StandardError
+    attr_accessor :response
+
+    attr_reader :message
+    attr_reader :error
+    attr_reader :body
+    attr_reader :status
+
+    def initialize(message = nil, status: nil, body: nil, response: nil)
+      @message = message
+      @status = status
+      @body = body
+      @response = response
+    end
+
+    def to_s
+      status_value = @status.nil? ? "" : "(Status #{@status}) "
+      "#{status_value}#{@message}"
+    end
+  end
+
+  # APIConnectionError is raised in the event that the SDK can't connect to
+  # Chroma's database. That can be for a variety of different reasons from a
+  # downed network to a bad TLS certificate.
+  class APIConnectionError < ChromaError
+  end
+
+  # InvalidRequestError is raised when a request is initiated with invalid
+  # parameters.
+  class InvalidRequestError < ChromaError
+    def initialize(message = nil, status: nil, body: nil, response: nil)
+      super
+
+      decode_message
+    end
+
+    private
+
+    def decode_message
+      if message.include?("ValueError")
+        error = JSON.parse(message)
+        value = error.fetch("error", "").sub("ValueError('", "").sub("')", "")
+        @message = value if !value.nil?
+      end
+    rescue JSON::ParserError, TypeError
+    end
+  end
+
+  # APIError is a generic error that may be raised in cases where none of the
+  # other named errors cover the problem.
+  class APIError < ChromaError
+    def initialize(message = nil, status: nil, body: nil, response: nil)
+      super
+
+      decode_message
+    end
+
+    private
+
+    def decode_message
+      if message.is_a?(Hash) && message.has_key?("error")
+        @message = message.fetch("error")
+      end
+    end
+  end
+end