switch to chat api

Author: Nicholas Newberry

app/lib/llm_clients/client.rb

@@ -41,17 +41,6 @@ def initialize(
 
     private
 
-    def uri(path)
-      URI("#{@endpoint}#{"/" unless path.starts_with?("/") || @endpoint.ends_with?("/")}#{path}")
-    end
-
-    def headers
-      headers = { "Content-Type": "application/json" }
-      headers[:Authorization] = "Bearer #{@api_key}" if @api_key
-
-      headers
-    end
-
     def context(prompt, model)
       template = template_for(model)
 

app/lib/llm_clients/ollama/chat.rb

@@ -1,33 +1,42 @@
 module LlmClients
   module Ollama
     class Chat
-      def initialize(conversation)
-        @conversation = conversation
+      def initialize(messages)
+        @messages = messages
       end
 
-      def prompt
-        {
-          model: @conversation.model_config.model,
-          messages:,
-        }
+      def chat_history
+        @chat_history ||= generate_chat_history
       end
 
       private
 
-      def messages
-        @conversation.messages.map do |message|
+      def generate_chat_history
+        # use message content only for all messages but the last
+        messages = conversation_messages[..-2].map do |message|
           {
             role: role_for_message(message),
             content: message.content,
           }
         end
+
+        # use the augmented prompt for the last message
+        last_message = conversation_messages.last
+        messages << {
+          role: role_for_message(last_message),
+          content: last_message.prompt || last_message.content,
+        }
       end
 
       def role_for_message(message)
         return nil if message.author.nil?
 
         message.author.is_a?(User) ? "user" : "assistant"
       end
+
+      def conversation_messages
+        @conversation_messages ||= @messages.sort
+      end
     end
   end
 end

app/lib/llm_clients/ollama/client.rb

@@ -3,11 +3,35 @@ module Ollama
     class Client < LlmClients::Client
       NETWORK_TIMEOUT = 8
 
+      def complete(prompt, &)
+        request = helper.completion_request(prompt)
+
+        stream(request, &)
+      end
+
+      def chat(chat, &)
+        request = helper.chat_request(chat.chat_history)
+
+        stream(request, &)
+      end
+
+      def embed(content)
+        request = helper.embed_request(content)
+
+        request(request)
+      end
+
+      private
+
+      def helper
+        @helper ||= RequestHelper.new(@endpoint, @api_key, @embedding_model, @model, @temperature)
+      end
+
       # rubocop:disable all
-      def complete(prompt, &block)
+      def stream(request, &block)
         @stats = new_stats
-        request = completion_request(prompt)
 
+        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|
           stats[:start_time] = current_time
@@ -49,45 +73,36 @@ def complete(prompt, &block)
           end
         end
       end
-      # rubocop:enable all
 
-      def embed(prompt)
+      def request(request)
         response = with_retries do
-          response = HTTParty.post(uri(embedding_path), headers:, timeout: NETWORK_TIMEOUT, body: {
-            model: @embedding_model,
-            prompt:,
-          }.to_json)
+          response = Net::HTTP.start(@uri.hostname, @uri.port, use_ssl: @uri.scheme == "https") do |http|
+            http.request(request)
+          end
 
-          raise RetryableError if response.code == 500
+          # sometimes ollama just returns a 500 on an embed request when running locally, then is fine
+          raise RetryableError if response.code == "500"
 
           response
         end
 
-        raise response_error_for(response) unless response.success?
+        raise response_error_for(response) unless response.kind_of?(Net::HTTPSuccess)
 
-        response.parsed_response
+        parse_response(response)
       rescue SocketError
         raise ResponseError.new("Unable to connect to the active server", [ModelServer.active_server.url])
       end
+      # rubocop:enable all
 
-      private
-
-      def completion_request(prompt)
-        request = Net::HTTP::Post.new(uri(completion_path), **headers)
-        request.body = {
-          model: @model,
-          prompt:,
-          temperature: @temperature,
-          stream: true,
-          max_tokens: 200,
-        }.to_json
-
-        request
+      def parse_response(http_response)
+        JSON.parse(http_response.body)
+      rescue JSON::ParserError
+        http_response.body
       end
 
       def extract_message(response_string)
         response_hash = JSON.parse(response_string)
-        message = response_hash["response"]
+        message = response_hash["response"] || response_hash["message"]["content"]
         Rails.logger.debug { "<== '#{message}'" }
 
         message
@@ -103,14 +118,6 @@ def template_for(model)
         { prefix: "", suffix: "" }
       end
 
-      def completion_path
-        "api/generate"
-      end
-
-      def embedding_path
-        "api/embeddings"
-      end
-
       def formatter
         @formatter ||= FORMATTERS.find { |key, _value| key.match?(@model) ? value : nil }
       end

app/lib/llm_clients/ollama/request_helper.rb

@@ -0,0 +1,59 @@
+module LlmClients
+  module Ollama
+    class RequestHelper
+      def initialize(endpoint, api_key, embedding_model, model, temperature)
+        @endpoint = endpoint
+        @api_key = api_key
+        @embedding_model = embedding_model
+        @model = model
+        @temperature = temperature
+      end
+
+      def embed_request(content)
+        request = Net::HTTP::Post.new(uri(embedding_path), **headers)
+        request.body = { model: @embedding_model, prompt: content }.to_json
+
+        request
+      end
+
+      def completion_request(prompt)
+        request = Net::HTTP::Post.new(uri(completion_path), **headers)
+        request.body = { model: @model, prompt:, temperature: @temperature, stream: true, max_tokens: 200 }.to_json
+
+        request
+      end
+
+      def chat_request(messages)
+        request = Net::HTTP::Post.new(uri(chat_path), **headers)
+        request.body = { model: @model, messages: }.to_json
+
+        request
+      end
+
+      private
+
+      def uri(path)
+        URI("#{@endpoint}#{"/" unless path.starts_with?("/") || @endpoint.ends_with?("/")}#{path}")
+      end
+
+      def headers
+        headers = { "Content-Type": "application/json" }
+        headers[:Authorization] = "Bearer #{@api_key}" if @api_key
+
+        headers
+      end
+
+      def completion_path
+        "api/generate"
+      end
+
+      def embedding_path
+        "api/embeddings"
+      end
+
+      def chat_path
+        "api/chat"
+      end
+    end
+  end
+end

app/services/prompt_augmentor.rb

@@ -0,0 +1,17 @@
+class PromptAugmentor
+  def initialize(message, search_hits)
+    @message = message
+    @search_hits = search_hits
+  end
+
+  def prompt
+    @prompt ||= begin
+      prompt = "Here is some context that may help you answer the following question:\n\n"
+      @search_hits.each do |hit|
+        prompt << "#{hit.chunk.content}\n\n"
+      end
+
+      prompt << "Question: #{@message.content}"
+    end
+  end
+end

app/services/reply_to_message.rb

@@ -7,20 +7,22 @@ class ReplyToMessage
 
   def initialize(message)
     @message = message
+    @conversation = @message.conversation
   end
 
   def execute
-    search if collections_to_search.any?
+    augment_message_prompt
+    chat.chat_history
 
     @reply = Message.create!(
       content: "",
-      author: @message.conversation.model_config,
-      conversation: @message.conversation
+      author: @conversation.model_config,
+      conversation: @conversation
     )
 
     prepend(reply_id, "messages/spinner")
 
-    streamer.stream do |message|
+    streamer.chat(chat) do |message|
       Rails.logger.info("Got back: #{message}")
       @reply.update!(content: @reply.content + message)
 
@@ -29,22 +31,28 @@ def execute
 
     @reply.update!(statistics:)
   rescue LlmClients::ResponseError, Errno::ECONNREFUSED, EOFError, ResponseStreamer::ResponseStreamerError,
-      ResponseStreamer::NetworkError => e
+         ResponseStreamer::NetworkError => e
     Rails.logger.error("\n#{e.class.name}: #{e.message}#{e.backtrace.join("\n")}")
 
     active_message.update!(error: { message: e })
   end
 
   private
 
+  def augment_message_prompt
+    return unless collections_to_search.any?
+
+    @message.update!(prompt: prompt_augmentor.prompt)
+  end
+
   def active_message
     @reply || @message
   end
 
   def collections_to_search
-    @collections_to_search ||= if @message.conversation.search_collections
-      message_collections = @message.conversation.collections
-      message_collections.any? ? message_collections : @message.conversation.user.collections
+    @collections_to_search ||= if @conversation.search_collections
+      message_collections = @conversation.collections
+      message_collections.any? ? message_collections : @conversation.user.collections
     else
       []
     end
@@ -68,55 +76,44 @@ def prepend(target, partial, locals = {})
   end
 
   def streamer
+    # TODO: add api_key to ModelServer
     @streamer ||= ResponseStreamer.new(
-      {
-        endpoint: ModelServer.active_server.url,
-        model: model_config.model,
-        provider: ModelServer.active_server.provider,
-      },
-      prompt || @message.content
+      endpoint: ModelServer.active_server.url,
+      model: model_config.model,
+      provider: ModelServer.active_server.provider
     )
   end
 
-  def search_hits
-    @search_hits ||= searcher.search(@message.content)
-  end
-
-  def prompt
-    return unless collections_to_search.any?
-
-    @prompt ||= begin
-      prompt = "Here is some context that may help you answer the following question:\n\n"
-      search_hits.each do |hit|
-        prompt << "#{hit.chunk.content}\n\n"
-      end
-
-      prompt << "Question: #{@message.content}"
-    end
-  end
-
-  def search
-    @message.update!(prompt:) if prompt.present?
-  end
-
   def searcher
-    @searcher ||= Search::SearchMultiple.new(collections_to_search, num_results: 10)
+    @searcher ||= Search::SearchMultiple.new(collections_to_search, num_results: 5)
   end
 
   def model_config
-    @message.conversation.model_config
+    @conversation.model_config
   end
 
   def reply_id
     @reply_id ||= "message_#{@reply.id}"
   end
 
   def channel_name
-    "#{@message.conversation.user.id}_conversations"
+    "#{@conversation.user.id}_conversations"
   end
 
   def statistics
     streamer.stats.merge({ server: ModelServer.active_server.name })
   end
+
+  def chat
+    @chat ||= LlmClients::Ollama::Chat.new(@conversation.messages.sort)
+  end
+
+  def prompt_augmentor
+    @prompt_augmentor ||= begin
+      search_hits = searcher.search(@message.content)
+
+      PromptAugmentor.new(@message, search_hits)
+    end
+  end
 end
 # rubocop:enable Metrics/AbcSize