add ApiCalls for streamed responses

Author: Nicholas Newberry

app/lib/chromadb/client.rb

@@ -108,6 +108,8 @@ def get(path)
         raise ResponseError, @last_response.body
       end
 
+      store_api_call("chromadb", @last_response)
+
       @last_response.parsed_response
     end
 
@@ -119,6 +121,8 @@ def post(path, body = {})
         raise ResponseError, @last_response.body
       end
 
+      store_api_call("chromadb", @last_response)
+
       @last_response.parsed_response
     end
 
@@ -130,7 +134,13 @@ def delete(path)
         raise ResponseError, @last_response.body
       end
 
+      store_api_call("chromadb", @last_response)
+
       @last_response.parsed_response
     end
+
+    def store_api_call(service_name, response)
+      ApiCall.from_httparty(service_name, response).save!
+    end
   end
 end

app/lib/llm_clients/client.rb

@@ -127,7 +127,8 @@ def with_retries
       raise exception
     end
 
-    def store_api_call(service_name, request, response)
+    def store_api_call(service_name, request, response, response_body = nil)
+      response.body = response_body if response_body
       ApiCall.from_net_http(service_name, request, response, @traceable).save!
     end
   end

app/lib/llm_clients/ollama/client.rb

@@ -31,9 +31,12 @@ def helper
       def stream(request, &block)
         @stats = new_stats
 
+        response_contents = ""
+
         Rails.logger.info("Sending request body:\n#{request.body}")
         # TODO: switch to HTTParty for this?
-        Net::HTTP.start(@uri.hostname, @uri.port, use_ssl: @uri.scheme == "https") do |http|
+        # TODO: create ApiCall early and update response body after streaming is done
+        full_response = Net::HTTP.start(@uri.hostname, @uri.port, use_ssl: @uri.scheme == "https") do |http|
           stats[:start_time] = current_time
           http.request(request) do |response|
             raise response_error_for(response) unless response.is_a?(Net::HTTPSuccess)
@@ -57,6 +60,7 @@ def stream(request, &block)
 
               if current_batch_size == @batch_size
                 Rails.logger.debug "==> #{current_batch}"
+                response_contents << current_batch
                 yield current_batch
 
                 current_batch_size = 0
@@ -66,12 +70,17 @@ def stream(request, &block)
 
             if current_batch_size > 0
               Rails.logger.debug "==> #{current_batch}"
+              response_contents << current_batch
               yield current_batch
             end
 
             calculate_stats
           end
         end
+
+        store_api_call("ollama", request, full_response, response_contents)
+
+        full_response
       end
 
       def request(request)

app/models/api_call.rb

@@ -4,9 +4,28 @@ class ApiCall < ApplicationRecord
   enum http_method: { get: 0, post: 1, put: 2, delete: 3, patch: 4, head: 5 }, _prefix: :http_method
 
   class << self
-    # there's no way to resolve AbcSize that doesn't make the code ugly
     # rubocop: disable Metrics/AbcSize
-    def from_net_http(service_name, request, response, traceable)
+    def from_httparty(service_name, response, traceable = nil)
+      request = response.request
+
+      api_call = new(
+        service_name:,
+        http_method: http_method_from_httparty(request),
+        url: request.uri.to_s,
+        headers: request.options[:headers],
+        body: json_body(request.raw_body),
+        body_length: request.raw_body&.length,
+        response_code: response.code,
+        response_headers: response.headers,
+        response_body: json_body(response.body),
+        response_length: response.body.length
+      )
+
+      api_call.traceable = traceable if traceable
+      api_call
+    end
+
+    def from_net_http(service_name, request, response, traceable = nil)
       api_call = new(
         service_name:,
         http_method: request.method.downcase,
@@ -28,9 +47,15 @@ def from_net_http(service_name, request, response, traceable)
     private
 
     def json_body(body)
+      return if body.nil?
+
       JSON.parse(body)
-    rescue JSON::ParserError
+    rescue JSON::ParserError, TypeError
       body
     end
+
+    def http_method_from_httparty(request)
+      request.http_method.name.split("::").last.downcase
+    end
   end
 end