chromadb-experimental 0.1.0

data/lib/chromadb/embedding_functions.rb

@@ -0,0 +1,121 @@
+# frozen_string_literal: true
+
+require "json_schemer"
+require "porter2stemmer"
+
+module Chroma
+  module EmbeddingFunctions
+    KNOWN_DENSE = {}
+    KNOWN_SPARSE = {}
+
+    module_function
+
+    def register_embedding_function(name, klass)
+      raise ArgumentError, "Embedding function #{name} already registered" if KNOWN_DENSE.key?(name)
+      KNOWN_DENSE[name] = klass
+    end
+
+    def register_sparse_embedding_function(name, klass)
+      raise ArgumentError, "Sparse embedding function #{name} already registered" if KNOWN_SPARSE.key?(name)
+      KNOWN_SPARSE[name] = klass
+    end
+
+    def resolve_name(fn)
+      return nil unless fn
+      if fn.respond_to?(:name)
+        value = fn.name
+        return value if value.is_a?(String)
+      end
+      nil
+    end
+
+    def prepare_embedding_function_config(fn)
+      return { "type" => "legacy" } if fn.nil?
+
+      name = resolve_name(fn)
+      return { "type" => "legacy" } if name.nil?
+      return { "type" => "legacy" } unless fn.respond_to?(:get_config) && fn.class.respond_to?(:build_from_config)
+
+      config = fn.get_config
+      if fn.respond_to?(:validate_config)
+        fn.validate_config(config)
+      elsif fn.class.respond_to?(:validate_config)
+        fn.class.validate_config(config)
+      end
+
+      { "type" => "known", "name" => name, "config" => config }
+    end
+
+    def build_embedding_function(ef_config, client: nil)
+      return nil if ef_config.nil?
+      return nil unless ef_config.is_a?(Hash)
+      return nil unless ef_config["type"] == "known"
+
+      name = ef_config["name"]
+      config = ef_config["config"] || {}
+
+      klass = KNOWN_DENSE[name]
+      return nil unless klass
+
+      if klass.respond_to?(:build_from_config)
+        return klass.build_from_config(config, client: client)
+      end
+
+      nil
+    end
+
+    def build_sparse_embedding_function(ef_config, client: nil)
+      return nil if ef_config.nil?
+      return nil unless ef_config.is_a?(Hash)
+      return nil unless ef_config["type"] == "known"
+
+      name = ef_config["name"]
+      config = ef_config["config"] || {}
+
+      klass = KNOWN_SPARSE[name]
+      return nil unless klass
+
+      if klass.respond_to?(:build_from_config)
+        return klass.build_from_config(config, client: client)
+      end
+
+      nil
+    end
+
+    def validate_config_schema(config, schema_name)
+      schemer = schema_for(schema_name)
+      errors = schemer.validate(config).to_a
+      return if errors.empty?
+
+      messages = errors.map { |err| err["message"] }.uniq
+      raise ArgumentError, "Invalid configuration for #{schema_name}: #{messages.join('; ')}"
+    end
+
+    def schema_for(schema_name)
+      @schemers ||= {}
+      return @schemers[schema_name] if @schemers[schema_name]
+
+      schema_path = File.expand_path("schemas/#{schema_name}.json", __dir__)
+      schema = JSON.parse(File.read(schema_path))
+      @schemers[schema_name] = JSONSchemer.schema(schema)
+    end
+  end
+
+  module SharedState
+    @cloud_api_key = nil
+
+    module_function
+
+    def register_cloud_api_key(key)
+      @cloud_api_key = key if key && !key.to_s.empty?
+    end
+
+    def cloud_api_key
+      @cloud_api_key
+    end
+  end
+end
+
+require_relative "embedding_functions/chroma_cloud_qwen"
+require_relative "embedding_functions/chroma_cloud_splade"
+require_relative "embedding_functions/chroma_bm25"

data/lib/chromadb/errors.rb

@@ -0,0 +1,120 @@
+# frozen_string_literal: true
+
+module Chroma
+  class ChromaError < StandardError
+    attr_accessor :trace_id
+
+    def self.name
+      "ChromaError"
+    end
+  end
+
+  class InvalidDimensionError < ChromaError
+    def self.name
+      "InvalidDimension"
+    end
+  end
+
+  class IDAlreadyExistsError < ChromaError
+    def self.name
+      "IDAlreadyExists"
+    end
+  end
+
+  class ChromaAuthError < ChromaError
+    def self.name
+      "AuthError"
+    end
+  end
+
+  class DuplicateIDError < ChromaError
+    def self.name
+      "DuplicateID"
+    end
+  end
+
+  class InvalidArgumentError < ChromaError
+    def self.name
+      "InvalidArgument"
+    end
+  end
+
+  class InvalidUUIDError < ChromaError
+    def self.name
+      "InvalidUUID"
+    end
+  end
+
+  class InvalidHTTPVersionError < ChromaError
+    def self.name
+      "InvalidHTTPVersion"
+    end
+  end
+
+  class AuthorizationError < ChromaError
+    def self.name
+      "AuthorizationError"
+    end
+  end
+
+  class NotFoundError < ChromaError
+    def self.name
+      "NotFoundError"
+    end
+  end
+
+  class UniqueConstraintError < ChromaError
+    def self.name
+      "UniqueConstraintError"
+    end
+  end
+
+  class BatchSizeExceededError < ChromaError
+    def self.name
+      "BatchSizeExceededError"
+    end
+  end
+
+  class VersionMismatchError < ChromaError
+    def self.name
+      "VersionMismatchError"
+    end
+  end
+
+  class InternalError < ChromaError
+    def self.name
+      "InternalError"
+    end
+  end
+
+  class RateLimitError < ChromaError
+    def self.name
+      "RateLimitError"
+    end
+  end
+
+  class QuotaError < ChromaError
+    def self.name
+      "QuotaError"
+    end
+  end
+
+  ERROR_TYPES = {
+    "InvalidDimension" => InvalidDimensionError,
+    "InvalidArgumentError" => InvalidArgumentError,
+    "IDAlreadyExists" => IDAlreadyExistsError,
+    "DuplicateID" => DuplicateIDError,
+    "InvalidUUID" => InvalidUUIDError,
+    "InvalidHTTPVersion" => InvalidHTTPVersionError,
+    "AuthorizationError" => AuthorizationError,
+    "NotFoundError" => NotFoundError,
+    "BatchSizeExceededError" => BatchSizeExceededError,
+    "VersionMismatchError" => VersionMismatchError,
+    "RateLimitError" => RateLimitError,
+    "AuthError" => ChromaAuthError,
+    "UniqueConstraintError" => UniqueConstraintError,
+    "QuotaError" => QuotaError,
+    "InternalError" => InternalError,
+    "ChromaError" => ChromaError
+  }.freeze
+end

data/lib/chromadb/http_client.rb

@@ -0,0 +1,142 @@
+# frozen_string_literal: true
+
+require "faraday"
+require "uri"
+
+module Chroma
+  class HttpTransport
+    DEFAULT_API_PATH = "/api/v2"
+
+    attr_reader :base_url, :headers
+
+    def initialize(host:, port: 8000, ssl: false, path: DEFAULT_API_PATH, headers: nil, ssl_verify: true, timeout: nil)
+      @base_url = self.class.resolve_url(
+        host: host,
+        port: port,
+        ssl: ssl,
+        default_api_path: path,
+      )
+      @headers = (headers || {}).dup
+      @headers["Content-Type"] ||= "application/json"
+      @headers["User-Agent"] ||= "Chroma Ruby Client v#{Chroma::VERSION} (https://github.com/chroma-core/chroma)"
+      @preflight_checks = nil
+
+      @connection = Faraday.new(url: @base_url, ssl: { verify: ssl_verify }) do |builder|
+        if timeout
+          builder.options.timeout = timeout
+          builder.options.open_timeout = timeout
+        end
+      end
+    end
+
+    def request(method, path, params: nil, json: nil)
+      normalized_path = path.start_with?("/") ? path[1..] : path
+      response = @connection.run_request(method.to_sym, normalized_path, nil, @headers) do |req|
+        req.params.update(clean_params(params)) if params
+        req.body = JSON.generate(json) if json
+      end
+
+      raise_chroma_error(response) unless response.success?
+
+      parse_response(response.body)
+    end
+
+    def get_pre_flight_checks
+      @preflight_checks ||= request(:get, "/pre-flight-checks")
+    end
+
+    def supports_base64_encoding?
+      checks = get_pre_flight_checks
+      checks.is_a?(Hash) && checks["supports_base64_encoding"] == true
+    end
+
+    def max_batch_size
+      checks = get_pre_flight_checks
+      return -1 unless checks.is_a?(Hash)
+      checks.fetch("max_batch_size", -1)
+    end
+
+    def request_headers
+      @headers.dup
+    end
+
+    def self.resolve_url(host:, port:, ssl:, default_api_path: DEFAULT_API_PATH)
+      validate_host(host)
+
+      skip_port = host.start_with?("http://", "https://")
+      parsed = URI.parse(host.start_with?("http") ? host : "http://#{host}")
+
+      scheme = if host.start_with?("http")
+        parsed.scheme
+      else
+        ssl ? "https" : "http"
+      end
+
+      netloc = parsed.host || parsed.path
+      if skip_port && parsed.port
+        netloc = "#{netloc}:#{parsed.port}"
+      end
+      port_value = skip_port ? parsed.port : port
+      path = parsed.path
+      path = default_api_path if path.nil? || path.empty? || path == netloc
+
+      trimmed_path = path.end_with?("/") ? path.chomp("/") : path
+      if default_api_path && !default_api_path.empty? && !trimmed_path.end_with?(default_api_path)
+        path = trimmed_path + default_api_path
+      else
+        path = trimmed_path
+      end
+
+      normalized_path = path.start_with?("/") ? path : "/#{path}"
+      normalized_path = normalized_path.gsub(%r{//+}, "/")
+
+      port_segment = skip_port ? "" : ":#{port_value}"
+      "#{scheme}://#{netloc}#{port_segment}#{normalized_path}"
+    end
+
+    def self.validate_host(host)
+      parsed = URI.parse(host)
+      if host.include?("/") && parsed.scheme.nil?
+        raise ArgumentError,
+              "Invalid URL. Seems that you are trying to pass URL as a host but without specifying the protocol. Please add http:// or https:// to the host."
+      end
+      if host.include?("/") && !%w[http https].include?(parsed.scheme)
+        raise ArgumentError, "Invalid URL. Unrecognized protocol - #{parsed.scheme}."
+      end
+    end
+
+    private
+
+    def clean_params(params)
+      params.reject { |_k, v| v.nil? }
+    end
+
+    def parse_response(body)
+      return nil if body.nil? || body.to_s.strip.empty?
+      JSON.parse(body)
+    rescue JSON::ParserError
+      body
+    end
+
+    def raise_chroma_error(response)
+      body = response.body
+      trace_id = response.headers["chroma-trace-id"]
+
+      begin
+        data = JSON.parse(body)
+        if data.is_a?(Hash) && data["error"]
+          error_class = Chroma::ERROR_TYPES[data["error"]] || Chroma::ChromaError
+          error = error_class.new(data["message"])
+          error.trace_id = trace_id if error.respond_to?(:trace_id=)
+          raise error
+        end
+      rescue JSON::ParserError
+        # fall through
+      end
+
+      message = body.to_s
+      message = "#{message} (trace ID: #{trace_id})" if trace_id
+      raise Chroma::ChromaError, message
+    end
+  end
+end