link ApiCalls with Conversations

Author: Nicholas Newberry

app/lib/chromadb/client.rb

@@ -4,9 +4,11 @@ module Chromadb
   class Client
     attr_reader :last_response
 
-    def initialize(host = nil, port = nil)
+    def initialize(host = nil, port = nil, traceable: nil)
       @host = host || ENV.fetch("CHROMADB_HOST") { "localhost" }
       @port = port || ENV.fetch("CHROMADB_PORT") { 8000 }
+      @traceable = traceable
+
       @url = "http://#{@host}:#{@port}"
     end
 
@@ -137,7 +139,7 @@ def delete(path)
     end
 
     def store_api_call(service_name, response)
-      ApiCall.from_httparty(service_name, response).save!
+      ApiCall.from_httparty(service_name, response, @traceable).save!
     end
   end
 end

app/models/api_call.rb

@@ -1,5 +1,5 @@
 class ApiCall < ApplicationRecord
-  belongs_to :traceable, optional: true
+  belongs_to :traceable, polymorphic: true, optional: true
 
   enum http_method: { get: 0, post: 1, put: 2, delete: 3, patch: 4, head: 5 }, _prefix: :http_method
 

app/models/conversation.rb

@@ -4,6 +4,7 @@ class Conversation < ApplicationRecord
   has_many :messages, dependent: :destroy
   has_many :conversation_collections, dependent: :destroy
   has_many :collections, through: :conversation_collections, dependent: :destroy
+  has_many :api_calls, as: :traceable, dependent: :destroy
 
   include Turbo::Broadcastable
 

app/models/message.rb

@@ -2,6 +2,7 @@ class Message < ApplicationRecord
   belongs_to :conversation
   belongs_to :author, polymorphic: true
   has_one :user, through: :conversation
+  has_many :api_calls, as: :traceable, dependent: :destroy
 
   include Turbo::Broadcastable
 
@@ -20,10 +21,6 @@ class Message < ApplicationRecord
       target: "message_#{id}",
       partial: "messages/message"
     )
-    broadcast_action_to(
-      user_dom_id("conversations"),
-      action: "scroll_to_bottom"
-    )
   }
 
   def previous(count = 1)

app/services/api/search.rb

@@ -4,10 +4,11 @@ class SearchError < StandardError; end
 
     include Rails.application.routes.url_helpers
 
-    def initialize(collection, base_url: nil, browser_base_url: nil)
+    def initialize(collection, base_url: nil, browser_base_url: nil, traceable: nil)
       @collection = collection
       @base_url = base_url
       @browser_base_url = URI(browser_base_url)
+      @traceable = traceable
     end
 
     def search(query)
@@ -103,7 +104,7 @@ def collection_id
     end
 
     def embedder
-      @embedder ||= Embedder.new(@collection.embedding_model)
+      @embedder ||= Embedder.new(@collection.embedding_model, traceable: @traceable)
     end
 
     def chroma

app/services/embedder.rb

@@ -1,6 +1,7 @@
 class Embedder
-  def initialize(embedding_model)
+  def initialize(embedding_model, traceable: nil)
     @embedding_model = embedding_model
+    @traceable = traceable
   end
 
   def embed(text)
@@ -15,7 +16,8 @@ def client
     @client ||= LlmClients::Ollama::Client.new(
       endpoint:,
       api_key: "todo",
-      embedding_model: @embedding_model.model
+      embedding_model: @embedding_model.model,
+      traceable: @traceable
     )
   end
 

app/services/reply_to_message.rb

@@ -87,14 +87,16 @@ def streamer
     @streamer ||= ResponseStreamer.new(
       endpoint: ModelServer.active_server.url,
       model: model_config.model,
-      provider: ModelServer.active_server.provider
+      provider: ModelServer.active_server.provider,
+      traceable: @conversation
     )
   end
 
   def searcher
     @searcher ||= Search::SearchMultiple.new(
       collections_to_search,
-      num_results: Setting.get(:num_chunks_to_include)
+      num_results: Setting.get(:num_chunks_to_include),
+      traceable: @conversation
     )
   end
 

app/services/response_streamer.rb

@@ -13,11 +13,12 @@ class UnsupportedServerError < ResponseStreamerError; end
 
   BATCH_SIZE = 16
 
-  def initialize(endpoint:, model:, provider:, api_key: nil)
+  def initialize(endpoint:, model:, provider:, api_key: nil, traceable: nil)
     @endpoint = endpoint
     @model = model
     @provider = provider
     @api_key = api_key
+    @traceable = traceable
 
     @chunk = ""
   end
@@ -63,7 +64,8 @@ def client
       endpoint: @endpoint,
       api_key: @api_key,
       model: @model,
-      batch_size: BATCH_SIZE
+      batch_size: BATCH_SIZE,
+      traceable: @traceable
     )
   end
 end

app/services/search/search.rb

@@ -1,7 +1,8 @@
 module Search
   class Search < Base
-    def initialize(collection)
+    def initialize(collection, traceable: nil)
       @collection = collection
+      @traceable = traceable
 
       super()
     end
@@ -62,11 +63,11 @@ def collection_id
     end
 
     def embedder
-      @embedder ||= Embedder.new(@collection.embedding_model)
+      @embedder ||= Embedder.new(@collection.embedding_model, traceable: @traceable)
     end
 
     def chroma
-      @chroma ||= Chromadb::Client.new
+      @chroma ||= Chromadb::Client.new(traceable: @traceable)
     end
   end
 end

app/services/search/search_multiple.rb

@@ -1,8 +1,9 @@
 module Search
   class SearchMultiple < Base
-    def initialize(collections, num_results: 10)
+    def initialize(collections, num_results: 10, traceable: nil)
       @collections = collections
       @num_results = num_results
+      @traceable = traceable
 
       super()
     end
@@ -19,7 +20,7 @@ def search(query, &)
 
     def searchers
       @searchers ||= @collections.map do |collection|
-        Search.new(collection)
+        Search.new(collection, traceable: @traceable)
       end
     end
   end