Sets headers and query params correctly for hosted version (#28)

.gitignore

@@ -7,6 +7,8 @@
 /spec/reports/
 /tmp/
 
+/test/hosted.rb
+
 .rbnext/
 
 .rubocop.yml

CHANGELOG.md

@@ -1,5 +1,9 @@
 ## [Unreleased]
 
+## 0.8.1 2024-10-16
+
+- Adds support for Chroma hosted service. See README for more details how to set your `api_key`, `tenant`, and `database`.
+
 ## [0.6.0] 2023-06-05
 
 - Fix failure to rescue from API call exceptions.

Gemfile.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    chroma-db (0.8.0)
+    chroma-db (0.8.1)
       dry-monads (~> 1.6)
       ruby-next (~> 1.0, >= 1.0.3)
       zeitwerk (~> 2.6.0)

README.md

@@ -54,7 +54,7 @@ Chroma.database = "my_database" # Optional
 
 ## Requirements
 
-- Ruby 3.0.6 or newer
+- Ruby 3.1.4 or newer
 - Chroma Database 0.4.24 or later running as a client/server model.
 
 For Chroma database 0.3.22 or older, please use version 0.3.0 of this gem.

lib/chroma/api_operations/request.rb

@@ -94,9 +94,9 @@ def execute_request(method, url, params = {}, options = {})
 
           case response
           in Net::HTTPSuccess
-            return Success(response_data)
+            Success(response_data)
           else
-            return Failure(response_data)
+            Failure(response_data)
           end
         end
 
@@ -107,6 +107,8 @@ def execute_request(method, url, params = {}, options = {})
         end
 
         private def build_request(method, uri, params)
+          uri.query = URI.encode_www_form(tenant: Chroma.tenant, database: Chroma.database) unless Chroma.api_key.nil?
+
           request = case method
           when :post then Net::HTTP::Post.new(uri)
           when :put then Net::HTTP::Put.new(uri)
@@ -115,10 +117,11 @@ def execute_request(method, url, params = {}, options = {})
             Net::HTTP::Get.new(uri)
           end
 
+          api_key = ENV.fetch("CHROMA_SERVER_AUTHN_CREDENTIALS", Chroma.api_key)
           request.content_type = "application/json"
           request.body = params.to_json if params.size > 0
           request.basic_auth(uri.user, uri.password) if !uri.user.nil?
-          request['X-Chroma-Token'] = ENV.fetch('CHROMA_SERVER_AUTHN_CREDENTIALS', nil) if ENV.fetch('CHROMA_SERVER_AUTHN_CREDENTIALS', nil)
+          request["X-Chroma-Token"] = api_key unless api_key.nil?
           request
         end
       end

lib/chroma/chroma.rb

@@ -25,18 +25,7 @@ class << self
   end
 
   def self.api_url
-    base_url = "#{connect_host}/#{api_base}/#{api_version}"
-    uri = URI(base_url)
-
-    unless api_key.nil?
-      query_params = {
-        tenant: tenant,
-        database: database
-      }
-      query_params["x-chroma-token"] = api_key
-      uri.query = URI.encode_www_form(query_params)
-    end
-    uri.to_s
+    "#{connect_host}/#{api_base}/#{api_version}"
   end
 
   Chroma.log_level = ENV["CHROMA_LOG"].to_i unless ENV["CHROMA_LOG"].nil?

lib/chroma/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Chroma
-  VERSION = "0.8.0"
+  VERSION = "0.8.1"
 end

test/chroma/api_operations/test_request.rb

@@ -76,4 +76,16 @@ def test_it_handles_successuful_request
     assert response.success?
     assert_equal(200, response.success.status)
   end
+
+  def test_it_adds_api_key_to_request
+    Chroma.api_key = "1234abcd"
+
+    url = "#{@url}/collection"
+    stub_successful_request(url, {tenant: Chroma.tenant, database: Chroma.database, "X-Chroma-Token": Chroma.api_key})
+    subject = Class.new.include(Chroma::APIOperations::Request)
+
+    response = subject.execute_request(:get, url)
+
+    assert response.success?
+  end
 end

test/chroma/test_chroma.rb

@@ -38,6 +38,27 @@ def test_it_build_an_url_with_api_url
     assert_equal("https://chroma-server.org/other/api/v2", Chroma.api_url)
   end
 
+  def test_it_allows_to_set_api_key
+    api_key = "1234abcd"
+    Chroma.api_key = api_key
+
+    assert_equal(api_key, Chroma.api_key)
+  end
+
+  def test_it_allows_to_set_tenant_and_database
+    assert_equal("default_tenant", Chroma.tenant)
+    assert_equal("default_database", Chroma.database)
+
+    tenant = "test_tenant"
+    database = "test_database"
+
+    Chroma.tenant = tenant
+    Chroma.database = database
+
+    assert_equal(tenant, Chroma.tenant)
+    assert_equal(database, Chroma.database)
+  end
+
   def test_it_allows_log_level_configuration
     Chroma.log_level = 0
     assert_equal ::Logger::DEBUG, Chroma.log_level
@@ -49,13 +70,6 @@ def test_it_allows_logger_configuration
     assert_equal logger, Chroma.logger
   end
 
-  def test_it_sets_default_database_and_tenant_when_api_key_is_set
-    Chroma.connect_host = "https://chroma-server.org"
-    Chroma.api_key = "1234abcd"
-
-    assert_equal("https://chroma-server.org/api/v1?tenant=default_tenant&database=default_database&x-chroma-token=1234abcd", Chroma.api_url)
-  end
-
   def test_it_sets_default_database_and_tenant_when_api_key_is_not_set
     Chroma.connect_host = "https://chroma-server.org"
     Chroma.api_key = nil

test/support/stub_request.rb

@@ -27,7 +27,15 @@ def stub_network_error(url)
     @stubs << stub_request(:any, url).to_timeout
   end
 
-  def stub_successful_request(url)
-    @stubs << stub_request(:any, url).to_return(status: 200, body: %({"name": "data-index"}), headers: {"Content-Type": "application/json"})
+  def stub_successful_request(url, query = {})
+    @stubs << if query.empty?
+      stub_request(:any, url)
+        .to_return(status: 200, body: %({"name": "data-index"}), headers: {"Content-Type": "application/json"})
+    else
+      api_key = query.delete(:"X-Chroma-Token")
+      stub_request(:any, url)
+        .with(query: hash_including(query), headers: {"X-Chroma-Token": api_key})
+        .to_return(status: 200, body: %({"name": "data-index"}), headers: {"Content-Type": "application/json"})
+    end
   end
 end