chromadb-experimental 0.1.0

data/lib/chromadb/search/search.rb

@@ -0,0 +1,73 @@
+# frozen_string_literal: true
+
+require "set"
+module Chroma
+  module Search
+    class Search
+      attr_reader :where_clause, :rank_expression, :limit_config, :select_config
+
+      def initialize(where: nil, rank: nil, limit: nil, select: nil)
+        @where_clause = WhereExpression.from(where) if where
+        @rank_expression = RankExpression.from(rank) if rank
+        @limit_config = Limit.from(limit)
+        @select_config = Select.from(select)
+      end
+
+      def where(where = nil)
+        clone_with(where: WhereExpression.from(where))
+      end
+
+      def rank(rank = nil)
+        clone_with(rank: RankExpression.from(rank))
+      end
+
+      def limit(limit = nil, offset = nil)
+        if limit.is_a?(Numeric)
+          clone_with(limit: Limit.from(limit.to_i, offset))
+        else
+          clone_with(limit: Limit.from(limit))
+        end
+      end
+
+      def select(*keys)
+        if keys.length == 1 && (keys[0].is_a?(Array) || keys[0].is_a?(Set))
+          return clone_with(select: Select.from(keys[0]))
+        end
+        if keys.length == 1 && keys[0].is_a?(Select)
+          return clone_with(select: Select.from(keys[0]))
+        end
+        if keys.length == 1 && keys[0].is_a?(Hash) && (keys[0].key?(:keys) || keys[0].key?("keys"))
+          return clone_with(select: Select.from(keys[0]))
+        end
+
+        clone_with(select: Select.from(keys))
+      end
+
+      def select_all
+        clone_with(select: Select.all)
+      end
+
+      def to_h
+        payload = {
+          "limit" => @limit_config.to_h,
+          "select" => @select_config.to_h
+        }
+
+        payload["filter"] = @where_clause.to_h if @where_clause
+        payload["rank"] = @rank_expression.to_h if @rank_expression
+        payload
+      end
+
+      private
+
+      def clone_with(where: @where_clause, rank: @rank_expression, limit: @limit_config, select: @select_config)
+        instance = self.class.allocate
+        instance.instance_variable_set(:@where_clause, where)
+        instance.instance_variable_set(:@rank_expression, rank)
+        instance.instance_variable_set(:@limit_config, limit)
+        instance.instance_variable_set(:@select_config, select)
+        instance
+      end
+    end
+  end
+end

data/lib/chromadb/search/select.rb

@@ -0,0 +1,54 @@
+# frozen_string_literal: true
+
+module Chroma
+  module Search
+    class Select
+      def initialize(keys = [])
+        unique = []
+        seen = {}
+        Array(keys).each do |key|
+          normalized = key.respond_to?(:name) ? key.name : key
+          unless normalized.is_a?(String)
+            raise TypeError, "Select keys must be strings or Key instances"
+          end
+          next if seen[normalized]
+          seen[normalized] = true
+          unique << normalized
+        end
+        @keys = unique
+      end
+
+      def self.from(input)
+        return Select.new(input.values) if input.is_a?(Select)
+        return Select.new if input.nil?
+
+        if input.is_a?(Hash) && (input.key?(:keys) || input.key?("keys"))
+          keys = input[:keys] || input["keys"] || []
+          return Select.new(keys)
+        end
+
+        if input.is_a?(String)
+          return Select.new([ input ])
+        end
+
+        if input.respond_to?(:each)
+          return Select.new(input)
+        end
+
+        raise TypeError, "Unsupported select input"
+      end
+
+      def self.all
+        Select.new([ K::DOCUMENT, K::EMBEDDING, K::METADATA, K::SCORE ])
+      end
+
+      def values
+        @keys.dup
+      end
+
+      def to_h
+        { "keys" => values }
+      end
+    end
+  end
+end

data/lib/chromadb/search/where.rb

@@ -0,0 +1,157 @@
+# frozen_string_literal: true
+
+module Chroma
+  module Search
+    class WhereExpression
+      def and(other)
+        target = self.class.from(other)
+        return self unless target
+        AndWhere.combine(self, target)
+      end
+
+      def or(other)
+        target = self.class.from(other)
+        return self unless target
+        OrWhere.combine(self, target)
+      end
+
+      def to_h
+        raise NotImplementedError
+      end
+
+      def self.from(input)
+        return input if input.is_a?(WhereExpression)
+        return nil if input.nil?
+        unless input.is_a?(Hash)
+          raise TypeError, "Where input must be a WhereExpression or Hash"
+        end
+        parse_where_hash(input)
+      end
+
+      def self.create_comparison(key, operator, value)
+        ComparisonWhere.new(key, operator, value)
+      end
+
+      def self.parse_where_hash(data)
+        if data.key?("$and")
+          raise ArgumentError, "$and cannot be combined with other keys" if data.length != 1
+          raw = data["$and"]
+          unless raw.is_a?(Array) && !raw.empty?
+            raise TypeError, "$and must be a non-empty array"
+          end
+          conditions = raw.map.with_index do |item, index|
+            expr = from(item)
+            raise TypeError, "Invalid where clause at index #{index}" if expr.nil?
+            expr
+          end
+          return conditions[0] if conditions.length == 1
+
+          return conditions.drop(1).reduce(conditions[0]) { |acc, cond| AndWhere.combine(acc, cond) }
+        end
+
+        if data.key?("$or")
+          raise ArgumentError, "$or cannot be combined with other keys" if data.length != 1
+          raw = data["$or"]
+          unless raw.is_a?(Array) && !raw.empty?
+            raise TypeError, "$or must be a non-empty array"
+          end
+          conditions = raw.map.with_index do |item, index|
+            expr = from(item)
+            raise TypeError, "Invalid where clause at index #{index}" if expr.nil?
+            expr
+          end
+          return conditions[0] if conditions.length == 1
+
+          return conditions.drop(1).reduce(conditions[0]) { |acc, cond| OrWhere.combine(acc, cond) }
+        end
+
+        entries = data.to_a
+        if entries.length != 1
+          raise ArgumentError, "Where hash must contain exactly one field"
+        end
+        field, value = entries[0]
+        unless value.is_a?(Hash)
+          return ComparisonWhere.new(field, "$eq", value)
+        end
+
+        operator_entries = value.to_a
+        if operator_entries.length != 1
+          raise ArgumentError, "Operator hash for field '#{field}' must contain exactly one operator"
+        end
+
+        operator, operand = operator_entries[0]
+        unless %w[$eq $ne $gt $gte $lt $lte $in $nin $contains $not_contains $regex $not_regex].include?(operator)
+          raise ArgumentError, "Unsupported where operator: #{operator}"
+        end
+
+        ComparisonWhere.new(field, operator, operand)
+      end
+    end
+
+    class AndWhere < WhereExpression
+      def initialize(conditions)
+        @conditions = conditions
+      end
+
+      def to_h
+        { "$and" => @conditions.map(&:to_h) }
+      end
+
+      def operands
+        @conditions.dup
+      end
+
+      def self.combine(left, right)
+        flattened = []
+        [ left, right ].each do |expr|
+          if expr.is_a?(AndWhere)
+            flattened.concat(expr.operands)
+          else
+            flattened << expr
+          end
+        end
+        return flattened[0] if flattened.length == 1
+        AndWhere.new(flattened)
+      end
+    end
+
+    class OrWhere < WhereExpression
+      def initialize(conditions)
+        @conditions = conditions
+      end
+
+      def to_h
+        { "$or" => @conditions.map(&:to_h) }
+      end
+
+      def operands
+        @conditions.dup
+      end
+
+      def self.combine(left, right)
+        flattened = []
+        [ left, right ].each do |expr|
+          if expr.is_a?(OrWhere)
+            flattened.concat(expr.operands)
+          else
+            flattened << expr
+          end
+        end
+        return flattened[0] if flattened.length == 1
+        OrWhere.new(flattened)
+      end
+    end
+
+    class ComparisonWhere < WhereExpression
+      def initialize(key, operator, value)
+        @key = key
+        @operator = operator
+        @value = value
+      end
+
+      def to_h
+        { @key => { @operator => @value } }
+      end
+    end
+  end
+end

data/lib/chromadb/search.rb

@@ -0,0 +1,8 @@
+# frozen_string_literal: true
+
+require_relative "search/key"
+require_relative "search/where"
+require_relative "search/rank"
+require_relative "search/limit"
+require_relative "search/select"
+require_relative "search/search"

data/lib/chromadb/types/results.rb

@@ -0,0 +1,96 @@
+# frozen_string_literal: true
+
+module Chroma
+  module Types
+    class GetResult
+      attr_reader :ids, :embeddings, :metadatas, :documents, :uris, :data, :included
+
+      def initialize(ids:, embeddings:, metadatas:, documents:, uris:, data:, included: nil)
+        @ids = ids
+        @embeddings = embeddings
+        @metadatas = metadatas
+        @documents = documents
+        @uris = uris
+        @data = data
+        @included = included
+      end
+
+      def to_h
+        {
+          ids: @ids,
+          embeddings: @embeddings,
+          metadatas: @metadatas,
+          documents: @documents,
+          uris: @uris,
+          data: @data,
+          included: @included
+        }
+      end
+    end
+
+    class QueryResult < GetResult
+      attr_reader :distances
+
+      def initialize(ids:, embeddings:, metadatas:, documents:, uris:, data:, distances:, included: nil)
+        super(ids: ids, embeddings: embeddings, metadatas: metadatas, documents: documents, uris: uris, data: data, included: included)
+        @distances = distances
+      end
+
+      def to_h
+        super.merge(distances: @distances)
+      end
+    end
+
+    class SearchResult
+      attr_reader :ids, :documents, :embeddings, :metadatas, :scores, :select
+
+      def initialize(response)
+        @ids = response.fetch("ids")
+        payload_count = @ids.length
+
+        @documents = normalize_payload_array(response["documents"], payload_count)
+        @embeddings = normalize_payload_array(response["embeddings"], payload_count)
+        raw_metadatas = normalize_payload_array(response["metadatas"], payload_count)
+        @metadatas = raw_metadatas.map { |payload| payload ? Types::Validation.deserialize_metadatas(payload) : nil }
+        @scores = normalize_payload_array(response["scores"], payload_count)
+        @select = response["select"] || []
+      end
+
+      def rows
+        results = []
+        @ids.each_with_index do |id_batch, batch_index|
+          docs = @documents[batch_index] || []
+          embeds = @embeddings[batch_index] || []
+          metas = @metadatas[batch_index] || []
+          scores = @scores[batch_index] || []
+
+          batch_rows = id_batch.map.with_index do |id, row_index|
+            row = { id: id }
+            doc = docs[row_index]
+            row[:document] = doc unless doc.nil?
+            emb = embeds[row_index]
+            row[:embedding] = emb unless emb.nil?
+            meta = metas[row_index]
+            row[:metadata] = meta unless meta.nil?
+            score = scores[row_index]
+            row[:score] = score unless score.nil?
+            row
+          end
+          results << batch_rows
+        end
+        results
+      end
+
+      private
+
+      def normalize_payload_array(payload, count)
+        return Array.new(count) { nil } if payload.nil?
+        return payload.map { |item| item ? item.dup : nil } if payload.length == count
+
+        result = payload.map { |item| item ? item.dup : nil }
+        result.fill(nil, result.length...count)
+        result
+      end
+    end
+  end
+end

data/lib/chromadb/types/sparse_vector.rb

@@ -0,0 +1,86 @@
+# frozen_string_literal: true
+
+module Chroma
+  module Types
+    class SparseVector
+      attr_reader :indices, :values, :labels
+
+      def initialize(indices:, values:, labels: nil)
+        @indices = indices
+        @values = values
+        @labels = labels
+        validate!
+      end
+
+      def to_h
+        result = {
+          TYPE_KEY => SPARSE_VECTOR_TYPE_VALUE,
+          "indices" => @indices,
+          "values" => @values
+        }
+        result["tokens"] = @labels if @labels
+        result
+      end
+
+      def self.from_h(data)
+        unless data.is_a?(Hash) && data[TYPE_KEY] == SPARSE_VECTOR_TYPE_VALUE
+          raise ArgumentError,
+                "Expected #{TYPE_KEY}='#{SPARSE_VECTOR_TYPE_VALUE}', got #{data[TYPE_KEY].inspect}"
+        end
+        new(indices: data.fetch("indices"), values: data.fetch("values"), labels: data["tokens"])
+      end
+
+      private
+
+      def validate!
+        unless @indices.is_a?(Array)
+          raise ArgumentError, "Expected SparseVector indices to be an Array, got #{@indices.class}"
+        end
+        unless @values.is_a?(Array)
+          raise ArgumentError, "Expected SparseVector values to be an Array, got #{@values.class}"
+        end
+        if @indices.length != @values.length
+          raise ArgumentError,
+                "SparseVector indices and values must have the same length, got #{@indices.length} indices and #{@values.length} values"
+        end
+
+        if @labels
+          unless @labels.is_a?(Array)
+            raise ArgumentError, "Expected SparseVector labels to be an Array, got #{@labels.class}"
+          end
+          if @labels.length != @indices.length
+            raise ArgumentError,
+                  "SparseVector labels must match indices length, got #{@labels.length} labels and #{@indices.length} indices"
+          end
+        end
+
+        @indices.each_with_index do |idx, i|
+          unless idx.is_a?(Integer)
+            raise ArgumentError,
+                  "SparseVector indices must be integers, got #{idx.inspect} at position #{i}"
+          end
+          if idx.negative?
+            raise ArgumentError,
+                  "SparseVector indices must be non-negative, got #{idx} at position #{i}"
+          end
+        end
+
+        @values.each_with_index do |val, i|
+          unless val.is_a?(Numeric)
+            raise ArgumentError,
+                  "SparseVector values must be numeric, got #{val.inspect} at position #{i}"
+          end
+        end
+
+        if @indices.length > 1
+          @indices.each_cons(2).with_index do |(prev, curr), i|
+            if curr <= prev
+              raise ArgumentError,
+                    "SparseVector indices must be strictly ascending, found indices[#{i + 1}]=#{curr} <= indices[#{i}]=#{prev}"
+            end
+          end
+        end
+      end
+    end
+  end
+end