Class: Pikuri::VectorDb::Backend::Chroma

Inherits:
Object
  • Object
show all
Defined in:
lib/pikuri/vector_db/backend/chroma.rb

Overview

Thin Faraday HTTP client against a self-hosted Chroma server (v2 API), targeting Chroma 0.5.x+. The persistent backend, behind the same duck-typed Pikuri::VectorDb::Backend protocol as InMemory (same method names, return shapes, and ArgumentError contract on empty input / non-positive top_k; the vector-dim contract diverges — see below). Hand-rolled rather than a chroma-db gem dep so the wire protocol stays auditable in one readable file; the cost is tracking the v2 API by hand.

Two ways to get one: bring your own — +Chroma.new(host:, port:, collection:)+ against an existing deployment (this class is purely the HTTP client); or let pikuri manage it — Server::Chroma.ensure_running's #client(collection:) returns a Chroma pointed at a supervised container. Docker lifecycle and HTTP wire protocol share nothing, so each is its own class.

v2 endpoints used: collections (get-or-create), .../upsert, .../query, .../count, DELETE .../{id}, .../delete (where-filtered), .../get (where-filtered projection).

BYO embeddings (Chroma's embedder is never invoked)

Chroma collections can carry a server-side embedding function; we always send pre-computed embeddings and never use it. pikuri's Embedder is the single source of truth — a parallel Chroma-side embedder would split it invisibly (local embedder in pikuri + OpenAIEmbeddingFunction in Chroma ⇒ every indexed document silently lands at OpenAI).

Contracts

  • Vector-dim diverges from InMemory: enforced server-side (first upsert sets the dim; a mismatch is HTTP 4xx → RuntimeError, not the ArgumentError InMemory raises). Same loud-failure shape; not worth parsing Chroma's error envelope to coerce the class.
  • Lazy collection resolution: new doesn't touch the server; the first +#upsert+/+#query+/+#count+ resolves (creating if missing) the collection by name and caches the id. #delete_all clears it; the next #upsert re-creates.
  • Cosine: collection created hnsw.space: 'cosine'; Chroma returns distance [0,2], #query converts via 1 - distance so the Result score matches InMemory's similarity scale.
  • Metadata keys: JSON round-trips Symbol keys to Strings, so #upsert stringifies and #query re-symbolizes — a queried Chunk looks identical to one stored in InMemory. source rides as a reserved metadata key (Chroma has no native source).

Constant Summary collapse

MANIFEST_PAGE_SIZE =

Rows per /get page in #sources_with_hashes — caps the parse working set of the boot manifest on a large corpus (one row per file, so a 50k-file corpus is ~50 localhost round trips, not one multi-MB response). Small corpora finish in one page.

1_000

Instance Method Summary collapse

Constructor Details

#initialize(host:, port:, collection:, tenant: 'default_tenant', database: 'default_database', connection: nil) ⇒ Chroma

Parameters:

  • host (String)
  • port (Integer)
  • collection (String) —

    Chroma collection name — a Chroma-specific identifier, so it lives here, not on VectorDb::Extension (a no-op for InMemory).

  • tenant (String) (defaults to: 'default_tenant') —

    Chroma v2 tenant; defaults to Chroma's default.

  • database (String) (defaults to: 'default_database') —

    Chroma v2 database; defaults to Chroma's default.

  • connection (Faraday::Connection, nil) (defaults to: nil) —

    DI point for tests.

Raises:

  • (ArgumentError) —

    on empty host or empty collection.



70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
 
# File 'lib/pikuri/vector_db/backend/chroma.rb', line 70

def initialize(host:, port:, collection:,
               tenant: 'default_tenant',
               database: 'default_database',
               connection: nil)
  raise ArgumentError, 'host must be non-empty' if host.nil? || host.to_s.empty?
  raise ArgumentError, 'collection must be non-empty' if collection.nil? || collection.to_s.empty?

  @host = host
  @port = port
  @collection_name = collection
  @tenant = tenant
  @database = database
  @collection_id = nil
  @connection = connection || Faraday.new(url: "http://#{host}:#{port}") do |f|
    f.request :json
    f.response :json
    f.adapter Faraday.default_adapter
  end
end

Instance Method Details

#count ⇒ Integer

Returns current chunk count. Zero before the first #upsert.

Returns:

  • (Integer) —

    current chunk count. Zero before the first #upsert.



192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
 
# File 'lib/pikuri/vector_db/backend/chroma.rb', line 192

def count
  return 0 if @collection_id.nil? && !collection_exists?

  response = @connection.get("#{collection_path}/count")
  unless response.status == 200
    raise "Backend::Chroma: GET #{collection_path}/count returned " \
          "HTTP #{response.status}: #{response.body.inspect}"
  end

  body = response.body
  # Chroma v2 returns the count as a bare integer.
  return body if body.is_a?(Integer)
  return body['count'] if body.is_a?(Hash) && body['count'].is_a?(Integer)

  raise "Backend::Chroma: count response was not an Integer (got #{body.inspect})"
end

#delete_all ⇒ void

This method returns an undefined value.

Drop the collection (the nuke-and-reload reindex path the Indexer drives); the next #upsert re-creates. No-op if none was created; a 404 on the DELETE is treated as "already gone" — idempotent.



178
179
180
181
182
183
184
185
186
187
188
 
# File 'lib/pikuri/vector_db/backend/chroma.rb', line 178

def delete_all
  return nil if @collection_id.nil? && !collection_exists?

  response = @connection.delete(collection_path)
  unless [200, 204, 404].include?(response.status)
    raise "Backend::Chroma: DELETE #{collection_path} returned " \
          "HTTP #{response.status}: #{response.body.inspect}"
  end
  @collection_id = nil
  nil
end

#delete_by_source(source) ⇒ void

This method returns an undefined value.

Remove every chunk whose source matches, via a metadata-filtered POST .../delete (+source+ is the reserved metadata key #upsert writes). The scoped counterpart to #delete_all. No-op when the collection doesn't exist yet.

Parameters:

Raises:

  • (RuntimeError) —

    on HTTP failure.



218
219
220
221
222
223
 
# File 'lib/pikuri/vector_db/backend/chroma.rb', line 218

def delete_by_source(source)
  return nil if @collection_id.nil? && !collection_exists?

  post_json("#{collection_path}/delete", { where: { 'source' => source } })
  nil
end

#query(vector:, top_k:) ⇒ Array<Backend::Result>

k-NN query by cosine similarity. Returns at most top_k Results descending by score. +score+ is +1 - cosine_distance+ so the value matches InMemory's cosine-similarity scale.

Parameters:

  • vector (Array<Float>)
  • top_k (Integer)

Returns:

Raises:

  • (ArgumentError) —

    on non-positive top_k.

  • (RuntimeError) —

    on HTTP failure.



137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
 
# File 'lib/pikuri/vector_db/backend/chroma.rb', line 137

def query(vector:, top_k:)
  raise ArgumentError, "top_k must be positive (got #{top_k})" if top_k <= 0

  # If we've never upserted, the collection doesn't
  # exist yet — semantic answer is "no hits."
  return [] if @collection_id.nil? && !collection_exists?

  response_body = post_json("#{collection_path}/query", {
                              query_embeddings: [vector],
                              n_results: top_k,
                              include: %w[documents metadatas distances]
                            })

  ids = (response_body['ids']       || [[]]).first || []
  docs = (response_body['documents'] || [[]]).first || []
  metas = (response_body['metadatas'] || [[]]).first || []
  dists = (response_body['distances'] || [[]]).first || []

  ids.each_with_index.map do |id, i|
    meta = metas[i] || {}
    # Pull +source+ back out of the metadata blob;
    # symbolize the remaining keys for round-trip
    # consistency with InMemory.
    source = meta['source'] || ''
    chunk_meta = {}
    meta.each do |k, v|
      next if k == 'source'

      chunk_meta[k.to_sym] = v
    end

    chunk = Chunk.new(id: id, source: source, text: docs[i] || '', metadata: chunk_meta)
    Result.new(chunk: chunk, score: 1.0 - dists[i].to_f)
  end
end

#replace_source(source:, chunks:, vectors:) ⇒ void

This method returns an undefined value.

Replace all chunks for one source: delete the old set, then upsert the new (the incremental-reindex unit, Indexer#reindex_file!).

Not transactional (the InMemory divergence): two HTTP calls, so a #query between them can see the source with zero chunks — a window InMemory#replace_source closes with its monitor but Chroma can't. The Indexer mitigates the common failure by embedding before calling here, so an embedder outage never reaches this and the old chunks stay. Delete-then-upsert, not the reverse — the reverse would delete the just-written chunks.

Parameters:

  • source (String) —

    the Chunk#source being replaced.

  • chunks (Array<Chunk>) —

    the new chunk set.

  • vectors (Array<Array<Float>>) —

    parallel to chunks.

Raises:

  • (ArgumentError) —

    on empty input or length mismatch.

  • (RuntimeError) —

    on HTTP failure.



242
243
244
245
246
 
# File 'lib/pikuri/vector_db/backend/chroma.rb', line 242

def replace_source(source:, chunks:, vectors:)
  delete_by_source(source)
  upsert(chunks: chunks, vectors: vectors)
  nil
end

#source_indexed?(source) ⇒ Boolean

Is source in the corpus? Scoped existence check for Tools::Read's membership gate: a +where+-filtered /get capped at one row, include: [] so the response carries only ids — O(1) transport regardless of corpus size, never the full #sources_with_hashes manifest. See the Backend protocol yardoc.

Parameters:

Returns:

  • (Boolean) —

    true if at least one chunk has this source.

Raises:

  • (RuntimeError) —

    on HTTP failure.



298
299
300
301
302
303
304
305
306
307
308
 
# File 'lib/pikuri/vector_db/backend/chroma.rb', line 298

def source_indexed?(source)
  return false if @collection_id.nil? && !collection_exists?

  body = post_json("#{collection_path}/get", {
                     where: { 'source' => source },
                     include: [],
                     limit: 1
                   })
  ids = body.is_a?(Hash) ? (body['ids'] || []) : []
  !ids.empty?
end

#sources_with_hashes ⇒ Hash{String => String, nil}

The boot-sweep reference: source → stored content hash, one metadata row per file (not per chunk) via where: { offset: 0 } (every file has one chunk at offset 0), include: ['metadatas'] (drops the heavy embeddings/documents), and MANIFEST_PAGE_SIZE +limit+/+offset+ paging so a large corpus never materializes one multi-MB response. (Two +offset+s collide in wording: the where one is a chunk metadata field, the top-level one is the pagination cursor.)

Pagination assumes the manifest isn't mutating mid-read; the Watcher drives this from its single worker thread, so no reindex runs concurrently.

Returns:

  • (Hash{String => String, nil}) —

    source → content hash. Empty when the collection doesn't exist yet.

Raises:

  • (RuntimeError) —

    on HTTP failure.



263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
 
# File 'lib/pikuri/vector_db/backend/chroma.rb', line 263

def sources_with_hashes
  return {} if @collection_id.nil? && !collection_exists?

  result = {}
  cursor = 0
  loop do
    body = post_json("#{collection_path}/get", {
                       where: { 'offset' => 0 },
                       include: ['metadatas'],
                       limit: MANIFEST_PAGE_SIZE,
                       offset: cursor
                     })
    metas = body.is_a?(Hash) ? (body['metadatas'] || []) : []
    metas.each do |meta|
      next unless meta.is_a?(Hash) && meta['source']

      result[meta['source']] = meta['hash']
    end
    break if metas.size < MANIFEST_PAGE_SIZE

    cursor += metas.size
  end
  result
end

#upsert(chunks:, vectors:) ⇒ void

This method returns an undefined value.

Insert-or-replace by chunk.id (parallel equal-length arrays). Vector-dim mismatch surfaces as RuntimeError server-side (InMemory raises ArgumentError — see the class header).

Parameters:

  • chunks (Array<Chunk>)
  • vectors (Array<Array<Float>>)

Raises:

  • (ArgumentError) —

    on empty input or length mismatch.

  • (RuntimeError) —

    on HTTP failure.



99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
 
# File 'lib/pikuri/vector_db/backend/chroma.rb', line 99

def upsert(chunks:, vectors:)
  raise ArgumentError, 'upsert called with empty chunks/vectors' if chunks.empty?
  if chunks.size != vectors.size
    raise ArgumentError, "size mismatch: #{chunks.size} chunks vs #{vectors.size} vectors"
  end

  ensure_collection!

  metadatas = chunks.map do |c|
    # Serialize +source+ as a reserved key in Chroma's
    # +metadata+; merge in the user's metadata Hash with
    # keys stringified for JSON round-trip stability.
    base = { 'source' => c.source }
    c.metadata.each { |k, v| base[k.to_s] = v }
    base
  end

  body = {
    ids: chunks.map(&:id),
    embeddings: vectors,
    documents: chunks.map(&:text),
    metadatas: metadatas
  }

  post_json("#{collection_path}/upsert", body)
  nil
end