chromadb-experimental 0.1.0

data/lib/chromadb/schemas/chroma-cloud-qwen.json

@@ -0,0 +1,61 @@
+{
+    "$schema": "http://json-schema.org/draft-07/schema#",
+    "title": "Chroma Cloud Qwen Embedding Function Schema",
+    "description": "Schema for the Chroma Cloud Qwen embedding function configuration",
+    "version": "1.0.0",
+    "type": "object",
+    "properties": {
+        "model": {
+            "type": "string",
+            "enum": [
+                "Qwen/Qwen3-Embedding-0.6B"
+            ],
+            "description": "The specific Qwen model to use for embeddings"
+        },
+        "task": {
+            "type": [
+                "string",
+                "null"
+            ],
+            "description": "The task for which embeddings are being generated. If null or empty, empty instructions will be used."
+        },
+        "instructions": {
+            "type": "object",
+            "description": "A mapping of tasks to instructions for targets (documents/queries)",
+            "properties": {
+                "nl_to_code": {
+                    "type": "object",
+                    "properties": {
+                        "documents": {
+                            "type": "string",
+                            "description": "Instructions for embedding documents"
+                        },
+                        "query": {
+                            "type": "string",
+                            "description": "Instructions for embedding queries"
+                        }
+                    },
+                    "required": [
+                        "documents",
+                        "query"
+                    ],
+                    "additionalProperties": false
+                }
+            },
+            "required": [
+                "nl_to_code"
+            ],
+            "additionalProperties": false
+        },
+        "api_key_env_var": {
+            "type": "string",
+            "description": "Environment variable name that contains your API key for the Chroma Embedding API",
+            "default": "CHROMA_API_KEY"
+        }
+    },
+    "required": [
+        "model",
+        "task"
+    ],
+    "additionalProperties": false
+}

data/lib/chromadb/schemas/chroma-cloud-splade.json

@@ -0,0 +1,31 @@
+{
+    "$schema": "http://json-schema.org/draft-07/schema#",
+    "title": "Chroma Cloud Splade Embedding Function Schema",
+    "description": "Schema for the Chroma Cloud Splade sparse embedding function configuration",
+    "version": "1.0.0",
+    "type": "object",
+    "properties": {
+        "model": {
+            "type": "string",
+            "enum": [
+                "prithivida/Splade_PP_en_v1"
+            ],
+            "description": "The specific Splade model to use for sparse embeddings"
+        },
+        "api_key_env_var": {
+            "type": "string",
+            "description": "Environment variable name that contains your API key for the Chroma Embedding API",
+            "default": "CHROMA_API_KEY"
+        },
+        "include_tokens": {
+            "type": "boolean",
+            "description": "Whether to store token labels in the sparse vector output",
+            "default": false
+        }
+    },
+    "required": [
+        "api_key_env_var",
+        "model"
+    ],
+    "additionalProperties": false
+}

data/lib/chromadb/schemas/chroma_bm25.json

@@ -0,0 +1,37 @@
+{
+    "$schema": "http://json-schema.org/draft-07/schema#",
+    "title": "Chroma BM25 Embedding Function Schema",
+    "description": "Schema for the Chroma BM25 sparse embedding function configuration",
+    "version": "1.0.0",
+    "type": "object",
+    "properties": {
+        "k": {
+            "type": "number",
+            "description": "BM25 saturation parameter controlling term frequency scaling"
+        },
+        "b": {
+            "type": "number",
+            "description": "BM25 length normalization parameter"
+        },
+        "avg_doc_length": {
+            "type": "number",
+            "description": "Average document length in tokens used for normalization"
+        },
+        "token_max_length": {
+            "type": "number",
+            "description": "Maximum token length allowed before filtering"
+        },
+        "stopwords": {
+            "type": "array",
+            "description": "Optional custom stopword list (in lowercase) to override the defaults",
+            "items": {
+                "type": "string"
+            }
+        },
+        "include_tokens": {
+            "type": "boolean",
+            "description": "Whether to store token strings in the sparse vectors (default: true)"
+        }
+    },
+    "additionalProperties": false
+}

data/lib/chromadb/search/key.rb

@@ -0,0 +1,94 @@
+# frozen_string_literal: true
+
+module Chroma
+  module Search
+    class Key
+      attr_reader :name
+
+      def initialize(name)
+        @name = name
+      end
+
+      def eq(value)
+        WhereExpression.create_comparison(@name, "$eq", value)
+      end
+
+      def ne(value)
+        WhereExpression.create_comparison(@name, "$ne", value)
+      end
+
+      def gt(value)
+        WhereExpression.create_comparison(@name, "$gt", value)
+      end
+
+      def gte(value)
+        WhereExpression.create_comparison(@name, "$gte", value)
+      end
+
+      def lt(value)
+        WhereExpression.create_comparison(@name, "$lt", value)
+      end
+
+      def lte(value)
+        WhereExpression.create_comparison(@name, "$lte", value)
+      end
+
+      def is_in(values)
+        array = iterable_to_array(values)
+        assert_non_empty(array, "$in requires at least one value")
+        WhereExpression.create_comparison(@name, "$in", array)
+      end
+
+      def not_in(values)
+        array = iterable_to_array(values)
+        assert_non_empty(array, "$nin requires at least one value")
+        WhereExpression.create_comparison(@name, "$nin", array)
+      end
+
+      def contains(value)
+        raise TypeError, "$contains requires a string value" unless value.is_a?(String)
+        WhereExpression.create_comparison(@name, "$contains", value)
+      end
+
+      def not_contains(value)
+        raise TypeError, "$not_contains requires a string value" unless value.is_a?(String)
+        WhereExpression.create_comparison(@name, "$not_contains", value)
+      end
+
+      def regex(pattern)
+        raise TypeError, "$regex requires a string pattern" unless pattern.is_a?(String)
+        WhereExpression.create_comparison(@name, "$regex", pattern)
+      end
+
+      def not_regex(pattern)
+        raise TypeError, "$not_regex requires a string pattern" unless pattern.is_a?(String)
+        WhereExpression.create_comparison(@name, "$not_regex", pattern)
+      end
+
+      private
+
+      def iterable_to_array(values)
+        return values.to_a if values.respond_to?(:to_a)
+        Array(values)
+      end
+
+      def assert_non_empty(array, message)
+        raise ArgumentError, message if array.empty?
+      end
+    end
+
+    module K
+      ID = Key.new("#id")
+      DOCUMENT = Key.new("#document")
+      EMBEDDING = Key.new("#embedding")
+      METADATA = Key.new("#metadata")
+      SCORE = Key.new("#score")
+
+      module_function
+
+      def [](name)
+        Key.new(name)
+      end
+    end
+  end
+end

data/lib/chromadb/search/limit.rb

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+module Chroma
+  module Search
+    class Limit
+      attr_reader :offset, :limit
+
+      def initialize(offset: 0, limit: nil)
+        unless offset.is_a?(Integer) && offset >= 0
+          raise TypeError, "Limit offset must be a non-negative integer"
+        end
+        if !limit.nil?
+          unless limit.is_a?(Integer) && limit > 0
+            raise TypeError, "Limit must be a positive integer when provided"
+          end
+        end
+        @offset = offset
+        @limit = limit
+      end
+
+      def self.from(input, offset_override = nil)
+        return Limit.new(offset: input.offset, limit: input.limit) if input.is_a?(Limit)
+        if input.is_a?(Numeric)
+          return Limit.new(limit: input.to_i, offset: offset_override || 0)
+        end
+        return Limit.new if input.nil?
+        if input.is_a?(Hash)
+          return Limit.new(offset: input[:offset] || input["offset"] || 0,
+                           limit: input[:limit] || input["limit"])
+        end
+        raise TypeError, "Invalid limit input"
+      end
+
+      def to_h
+        result = { "offset" => @offset }
+        result["limit"] = @limit if @limit
+        result
+      end
+    end
+  end
+end

data/lib/chromadb/search/rank.rb

@@ -0,0 +1,425 @@
+# frozen_string_literal: true
+
+module Chroma
+  module Search
+    class RankExpression
+      def to_h
+        raise NotImplementedError
+      end
+
+      def add(*others)
+        return self if others.empty?
+        expressions = [ self ] + others.map.with_index { |item, idx| ::Chroma::Search.require_rank(item, "add operand #{idx}") }
+        SumRankExpression.create(expressions)
+      end
+
+      def subtract(other)
+        SubRankExpression.new(self, ::Chroma::Search.require_rank(other, "subtract operand"))
+      end
+
+      def multiply(*others)
+        return self if others.empty?
+        expressions = [ self ] + others.map.with_index { |item, idx| ::Chroma::Search.require_rank(item, "multiply operand #{idx}") }
+        MulRankExpression.create(expressions)
+      end
+
+      def divide(other)
+        DivRankExpression.new(self, ::Chroma::Search.require_rank(other, "divide operand"))
+      end
+
+      def negate
+        multiply(-1)
+      end
+
+      def abs
+        AbsRankExpression.new(self)
+      end
+
+      def exp
+        ExpRankExpression.new(self)
+      end
+
+      def log
+        LogRankExpression.new(self)
+      end
+
+      def max(*others)
+        return self if others.empty?
+        expressions = [ self ] + others.map.with_index { |item, idx| ::Chroma::Search.require_rank(item, "max operand #{idx}") }
+        MaxRankExpression.create(expressions)
+      end
+
+      def min(*others)
+        return self if others.empty?
+        expressions = [ self ] + others.map.with_index { |item, idx| ::Chroma::Search.require_rank(item, "min operand #{idx}") }
+        MinRankExpression.create(expressions)
+      end
+
+      def self.from(input)
+        return input if input.is_a?(RankExpression)
+        return nil if input.nil?
+        if input.is_a?(Numeric)
+          return ValueRankExpression.new(::Chroma::Search.require_number(input, "Val requires a numeric value"))
+        end
+        if input.is_a?(Hash)
+          return RawRankExpression.new(input)
+        end
+        raise TypeError, "Rank input must be a RankExpression, number, or Hash"
+      end
+    end
+
+    class RawRankExpression < RankExpression
+      def initialize(raw)
+        @raw = raw
+      end
+
+      def to_h
+        deep_clone(@raw)
+      end
+
+      private
+
+      def deep_clone(value)
+        Marshal.load(Marshal.dump(value))
+      end
+    end
+
+    class ValueRankExpression < RankExpression
+      def initialize(value)
+        @value = value
+      end
+
+      def to_h
+        { "$val" => @value }
+      end
+    end
+
+    class SumRankExpression < RankExpression
+      def initialize(ranks)
+        @ranks = ranks
+      end
+
+      def self.create(ranks)
+        flattened = []
+        ranks.each do |rank|
+          if rank.is_a?(SumRankExpression)
+            flattened.concat(rank.operands)
+          else
+            flattened << rank
+          end
+        end
+        return flattened[0] if flattened.length == 1
+        SumRankExpression.new(flattened)
+      end
+
+      def operands
+        @ranks.dup
+      end
+
+      def to_h
+        { "$sum" => @ranks.map(&:to_h) }
+      end
+    end
+
+    class SubRankExpression < RankExpression
+      def initialize(left, right)
+        @left = left
+        @right = right
+      end
+
+      def to_h
+        { "$sub" => { "left" => @left.to_h, "right" => @right.to_h } }
+      end
+    end
+
+    class MulRankExpression < RankExpression
+      def initialize(ranks)
+        @ranks = ranks
+      end
+
+      def self.create(ranks)
+        flattened = []
+        ranks.each do |rank|
+          if rank.is_a?(MulRankExpression)
+            flattened.concat(rank.operands)
+          else
+            flattened << rank
+          end
+        end
+        return flattened[0] if flattened.length == 1
+        MulRankExpression.new(flattened)
+      end
+
+      def operands
+        @ranks.dup
+      end
+
+      def to_h
+        { "$mul" => @ranks.map(&:to_h) }
+      end
+    end
+
+    class DivRankExpression < RankExpression
+      def initialize(left, right)
+        @left = left
+        @right = right
+      end
+
+      def to_h
+        { "$div" => { "left" => @left.to_h, "right" => @right.to_h } }
+      end
+    end
+
+    class AbsRankExpression < RankExpression
+      def initialize(operand)
+        @operand = operand
+      end
+
+      def to_h
+        { "$abs" => @operand.to_h }
+      end
+    end
+
+    class ExpRankExpression < RankExpression
+      def initialize(operand)
+        @operand = operand
+      end
+
+      def to_h
+        { "$exp" => @operand.to_h }
+      end
+    end
+
+    class LogRankExpression < RankExpression
+      def initialize(operand)
+        @operand = operand
+      end
+
+      def to_h
+        { "$log" => @operand.to_h }
+      end
+    end
+
+    class MaxRankExpression < RankExpression
+      def initialize(ranks)
+        @ranks = ranks
+      end
+
+      def self.create(ranks)
+        flattened = []
+        ranks.each do |rank|
+          if rank.is_a?(MaxRankExpression)
+            flattened.concat(rank.operands)
+          else
+            flattened << rank
+          end
+        end
+        return flattened[0] if flattened.length == 1
+        MaxRankExpression.new(flattened)
+      end
+
+      def operands
+        @ranks.dup
+      end
+
+      def to_h
+        { "$max" => @ranks.map(&:to_h) }
+      end
+    end
+
+    class MinRankExpression < RankExpression
+      def initialize(ranks)
+        @ranks = ranks
+      end
+
+      def self.create(ranks)
+        flattened = []
+        ranks.each do |rank|
+          if rank.is_a?(MinRankExpression)
+            flattened.concat(rank.operands)
+          else
+            flattened << rank
+          end
+        end
+        return flattened[0] if flattened.length == 1
+        MinRankExpression.new(flattened)
+      end
+
+      def operands
+        @ranks.dup
+      end
+
+      def to_h
+        { "$min" => @ranks.map(&:to_h) }
+      end
+    end
+
+    class KnnRankExpression < RankExpression
+      def initialize(config)
+        @config = config
+      end
+
+      def to_h
+        base = {
+          "query" => @config[:query],
+          "key" => @config[:key],
+          "limit" => @config[:limit]
+        }
+        base["default"] = @config[:default_value] if @config.key?(:default_value)
+        base["return_rank"] = true if @config[:return_rank]
+        { "$knn" => base }
+      end
+    end
+
+    def self.Val(value)
+      ValueRankExpression.new(require_number(value, "Val requires a numeric value"))
+    end
+
+    def self.Knn(options)
+      KnnRankExpression.new(normalize_knn_options(options))
+    end
+
+    def self.Rrf(ranks:, k: 60, weights: nil, normalize: false)
+      unless k.is_a?(Integer) && k > 0
+        raise TypeError, "Rrf k must be a positive integer"
+      end
+      unless ranks.is_a?(Array) && !ranks.empty?
+        raise TypeError, "Rrf requires at least one rank expression"
+      end
+
+      expressions = ranks.map.with_index { |rank, index| require_rank(rank, "ranks[#{index}]") }
+
+      weight_values = weights ? weights.dup : Array.new(expressions.length, 1)
+      if weight_values.length != expressions.length
+        raise ArgumentError, "Number of weights must match number of ranks"
+      end
+      if weight_values.any? { |value| !value.is_a?(Numeric) || value.negative? }
+        raise TypeError, "Weights must be non-negative numbers"
+      end
+
+      if normalize
+        total = weight_values.reduce(0.0) { |sum, value| sum + value }
+        raise ArgumentError, "Weights must sum to a positive value when normalize=true" if total <= 0
+        weight_values = weight_values.map { |value| value / total }
+      end
+
+      terms = expressions.map.with_index do |rank, index|
+        weight = weight_values[index]
+        numerator = Val(weight)
+        denominator = rank.add(k)
+        numerator.divide(denominator)
+      end
+
+      fused = terms.reduce { |acc, term| acc.add(term) }
+      fused.negate
+    end
+
+    def self.Sum(*inputs)
+      raise ArgumentError, "Sum requires at least one rank expression" if inputs.empty?
+      expressions = inputs.map.with_index { |rank, index| require_rank(rank, "Sum operand #{index}") }
+      SumRankExpression.create(expressions)
+    end
+
+    def self.Sub(left, right)
+      SubRankExpression.new(require_rank(left, "Sub left"), require_rank(right, "Sub right"))
+    end
+
+    def self.Mul(*inputs)
+      raise ArgumentError, "Mul requires at least one rank expression" if inputs.empty?
+      expressions = inputs.map.with_index { |rank, index| require_rank(rank, "Mul operand #{index}") }
+      MulRankExpression.create(expressions)
+    end
+
+    def self.Div(left, right)
+      DivRankExpression.new(require_rank(left, "Div left"), require_rank(right, "Div right"))
+    end
+
+    def self.Abs(input)
+      require_rank(input, "Abs").abs
+    end
+
+    def self.Exp(input)
+      require_rank(input, "Exp").exp
+    end
+
+    def self.Log(input)
+      require_rank(input, "Log").log
+    end
+
+    def self.Max(*inputs)
+      raise ArgumentError, "Max requires at least one rank expression" if inputs.empty?
+      expressions = inputs.map.with_index { |rank, index| require_rank(rank, "Max operand #{index}") }
+      MaxRankExpression.create(expressions)
+    end
+
+    def self.Min(*inputs)
+      raise ArgumentError, "Min requires at least one rank expression" if inputs.empty?
+      expressions = inputs.map.with_index { |rank, index| require_rank(rank, "Min operand #{index}") }
+      MinRankExpression.create(expressions)
+    end
+
+    def self.require_rank(input, context)
+      result = RankExpression.from(input)
+      raise TypeError, "#{context} must be a rank expression" unless result
+      result
+    end
+
+    def self.require_number(value, message)
+      unless value.is_a?(Numeric) && value.finite?
+        raise TypeError, message
+      end
+      value.to_f
+    end
+
+    def self.normalize_knn_options(options)
+      limit = options[:limit] || options["limit"] || 128
+      unless limit.is_a?(Integer) && limit > 0
+        raise TypeError, "Knn limit must be a positive integer"
+      end
+
+      query_input = options[:query] || options["query"]
+      query = if query_input.is_a?(String)
+        query_input
+      elsif query_input.is_a?(Chroma::Types::SparseVector)
+        { "indices" => query_input.indices.dup, "values" => query_input.values.dup }
+      elsif query_input.is_a?(Hash) && query_input.key?(:indices) && query_input.key?(:values)
+        { "indices" => query_input[:indices].dup, "values" => query_input[:values].dup }
+      else
+        normalize_dense_vector(query_input)
+      end
+
+      key_input = options[:key] || options["key"]
+      key = key_input.respond_to?(:name) ? key_input.name : key_input
+      key ||= "#embedding"
+      raise TypeError, "Knn key must be a String or Key instance" unless key.is_a?(String)
+
+      default_value = options.key?(:default) ? options[:default] : options["default"]
+      if !default_value.nil?
+        default_value = require_number(default_value, "Knn default must be a number")
+        raise TypeError, "Knn default must be a finite number" unless default_value.finite?
+      end
+
+      return_rank = options.key?(:returnRank) ? options[:returnRank] : options["returnRank"]
+      return_rank = options.key?(:return_rank) ? options[:return_rank] : return_rank
+
+      {
+        query: query.is_a?(Hash) ? Marshal.load(Marshal.dump(query)) : query,
+        key: key,
+        limit: limit,
+        return_rank: !!return_rank
+      }.tap do |config|
+        config[:default_value] = default_value unless default_value.nil?
+      end
+    end
+
+    def self.normalize_dense_vector(vector)
+      array = vector.is_a?(Array) ? vector.dup : vector.to_a
+      array.map do |value|
+        unless value.is_a?(Numeric) && value.finite?
+          raise TypeError, "Dense query vector values must be finite numbers"
+        end
+        value.to_f
+      end
+    end
+  end
+end