Add support for model commands

lib/discord_bot/bot.rb

@@ -46,6 +46,14 @@ def command_callback(name, &block)
         instance.application_command(name, &block)
       end
 
+      def reset_model(channel_id:)
+        instance.reset_model(channel_id: channel_id)
+      end
+
+      def set_model(channel_id:, model:)
+        instance.set_model(channel_id: channel_id, model: model)
+      end
+
       def reset_system_prompt(channel_id:)
         instance.reset_system_prompt(channel_id: channel_id)
       end
@@ -164,6 +172,23 @@ def set_channel_system_prompt(channel_id:,system_prompt:)
       conversation(channel_id).set_system_prompt(system_prompt: system_prompt)
     end
 
+    def reset_model(channel_id:)
+      set_model(channel_id: channel_id, model: DiscordBot::LLM::Model.new)
+    end
+
+    def set_model(channel_id:, model:)
+      unless model.is_a?(DiscordBot::LLM::Model)
+        raise ArgumentError,
+          "Tried to set model using invalid class type #{model.class.name}"
+      end
+
+      unless model.available?
+        raise ArgumentError, "Tried to set model, but it's unavailable"
+      end
+
+      @channel_models[channel_id] = model
+    end
+
     private
 
     def conversation(channel_id)
@@ -177,7 +202,7 @@ def model(channel_id)
     def pull_default_model
       Logger.info 'Pulling default LLM model'
       begin
-        DiscordBot::LLM::Model.pull(model_name: DiscordBot::LLM::DEFAULT_MODEL)
+        DiscordBot::LLM::Model.new(model_name: DiscordBot::LLM::DEFAULT_MODEL).pull
       rescue DiscordBot::Errors::FailedToPullModel => error
         Logger.fatal error.message
         shutdown

lib/discord_bot/commands/model.rb

@@ -15,6 +15,7 @@ def register
             command.subcommand(:set, 'Set the current model') do |subcommand|
               subcommand.string(:model, 'What model to use', required: true)
             end
+            command.subcommand(:list, 'List the available models')
           end
         end
 
@@ -34,31 +35,63 @@ def handle
           Bot.command_callback(command_name).subcommand(:set) do |event|
             set_model(DiscordBot::Events::Command.new(event))
           end
+
+          Bot.command_callback(command_name).subcommand(:list) do |event|
+            list_available_models(DiscordBot::Events::Command.new(event))
+          end
         end
 
         def run(command)
           command.respond_with('Please use one of the subcommands instead')
         end
 
         def reset_model(command)
-          # Bot.reset_model(channel_id: command.channel_id)
-          command.respond_with('Pending implementation')
+          Logger.info "#{command.user_id} has reset the LLM model to default for #{command.channel_name}"
+          Bot.reset_model(channel_id: command.channel_id)
+          command.respond_with('Reset the LLM model to default')
         end
 
         def pull_model(command)
           unless command.ran_by_admin?
-            raise DiscordBot::Errors::PermissionDenied,
-              "#{command.user.name} tried running the pull model command without permission"
+            Logger.info "#{command.user.name} tried running the pull model command without permission"
+            command.respond_with('Due to the large size of models, this command is restricted to admins')
+            return
           end
           requested_model = command.options['model']
-          Logger.info "#{command.user.name} has requested the LLM model #{requested_model}"
-          command.respond_with('Pending implementation')
+          Logger.info "#{command.user_id} has requested the LLM model #{requested_model}"
+          command.respond_with("Pulling LLM model \"#{requested_model}\"...")
+          model = DiscordBot::LLM::Model.new(model_name: requested_model)
+          if model.available?
+            command.update_response("\"#{requested_model}\" is already available")
+          else
+            begin
+              model.pull
+              command.update_response("Pulled LLM model \"#{requested_model}\"")
+            rescue DiscordBot::Errors::FailedToPullModel
+              command.update_response("Failed to pull LLM model \"#{requested_model}\"")
+            end
+          end
         rescue DiscordBot::Errors::PermissionDenied => error
           Logger.warn error.message
         end
 
         def set_model(command)
-          command.respond_with('Pending implementation')
+          requested_model = command.options['model']
+          model = DiscordBot::LLM::Model.new(model_name: requested_model)
+          if model.available?
+            Logger.info "#{command.user_id} has set the LLM model to #{requested_model} for \##{command.channel_name}"
+            Bot.set_model(channel_id: command.channel_id, model: model)
+            command.respond_with("Set LLM model to \"#{requested_model}\"")
+          else
+            command.respond_with("That model is currently unavailable. Try running `/model pull #{requested_model}` first")
+          end
+        end
+
+        def list_available_models(command)
+          models = DiscordBot::LLM::Model.available_models
+          # TODO: Also show file size, parameter size, and quantization level
+          formatted_list = models.map{ |model| "- `#{model.name}`" }.join("\n")
+          command.respond_with("The currently available models are:\n#{formatted_list}")
         end
       end
     end

lib/discord_bot/errors.rb

@@ -6,7 +6,6 @@ class Error < StandardError; end
 
   module Errors
     class NotImplementedError < DiscordBot::Error; end
-    class ApiRequestPostFailed < DiscordBot::Error; end
     class FailedToPullModel < DiscordBot::Error; end
     class PermissionDenied < DiscordBot::Error; end
   end

lib/discord_bot/events/base.rb

@@ -10,6 +10,14 @@ def initialize(event)
       def channel_id
         @event.channel.id
       end
+
+      def channel_name
+        @event.channel.name
+      end
+
+      def user_id
+        @event.user.id
+      end
     end
   end
 end

lib/discord_bot/llm/api_request.rb

@@ -6,16 +6,48 @@ class ApiRequest
         accept: :json
       }
 
-      def self.post(endpoint, payload, headers = DEFAULT_HEADERS, timeout = nil)
-        RestClient::Request.execute(
-          method: :post,
-          url: "#{DiscordBot::LLM::API_URL}#{endpoint}",
-          payload: payload,
-          headers: headers,
-          timeout: timeout
-        )
-      rescue StandardError => error
-        raise error.message, DiscordBot::Errors::ApiRequestPostFailed
+      class << self
+        def get(endpoint, headers = DEFAULT_HEADERS, timeout = nil)
+          RestClient::Request.execute(
+            method: :get,
+            url: "#{DiscordBot::LLM::API_URL}#{endpoint}",
+            headers: headers,
+            timeout: timeout
+          )
+        end
+
+        def post(endpoint, payload, headers = DEFAULT_HEADERS, timeout = nil)
+          RestClient::Request.execute(
+            method: :post,
+            url: "#{DiscordBot::LLM::API_URL}#{endpoint}",
+            payload: payload,
+            headers: headers,
+            timeout: timeout
+          )
+        end
+
+        def list_local_models
+          get('/api/tags')
+        end
+
+        def chat(messages:, model_name:)
+          payload = {
+            messages: messages,
+            model: model_name,
+            stream: false
+          }.to_json
+          post('/api/chat', payload)
+        end
+
+        def model_info(model_name:)
+          payload = { name: model_name }.to_json
+          post('/api/show', payload)
+        end
+
+        def pull_model(model_name:)
+          payload = { name: model_name, stream: false }.to_json
+          post('/api/pull', payload)
+        end
       end
     end
   end

lib/discord_bot/llm/model.rb

@@ -3,24 +3,38 @@ module LLM
     class Model
       attr_reader :name
 
-      # TODO: Add automatic retry with backoff
-      def self.pull(model_name:)
-        endpoint = '/api/pull'
-        payload = { name: model_name, stream: false }.to_json
-        DiscordBot::LLM::ApiRequest.post(endpoint, payload)
-      rescue DiscordBot::Error => error
-        raise "Model pull failed due to: \"#{error.message}\"",
-          DiscordBot::Errors::FailedToPullModel
+      def self.available_models
+        response = DiscordBot::LLM::ApiRequest.list_local_models
+        models = JSON.parse(response.body)['models']
+        models.map do |model|
+          DiscordBot::LLM::Model.new(model_name: model['name'])
+        end
       end
 
-      def initialize
-        @name = DiscordBot::LLM::DEFAULT_MODEL
+      def initialize(model_name: nil)
+        if model_name.nil?
+          @name = DiscordBot::LLM::DEFAULT_MODEL
+        else
+          @name = model_name
+        end
       end
 
       # TODO: Use ollama status check to see if model is ready for use
       def available?
+        model_info = DiscordBot::LLM::ApiRequest.model_info(model_name: name)
+        true
+      rescue RestClient::NotFound
         false
       end
+
+      # TODO: Add automatic retry with backoff
+      def pull
+        return if available?
+        DiscordBot::LLM::ApiRequest.pull_model(model_name: name)
+      rescue RestClient::InternalServerError => error
+        raise DiscordBot::Errors::FailedToPullModel,
+          "Model pull failed due to: \"#{error.message}\""
+      end
     end
   end
 end

lib/discord_bot/llm/response.rb

@@ -9,15 +9,12 @@ def initialize(conversation_history:, user_message:, model:)
           message: adjusted_user_message(user_message)
         )
 
-        payload = {
-          messages: conversation_history.messages,
-          model: model.name,
-          stream: false
-        }.to_json
-
         Logger.info 'Requesting LLM Response'
-        @response = DiscordBot::LLM::ApiRequest.post('/api/chat', payload)
-        @body = JSON.parse(@response.body)
+        response = DiscordBot::LLM::ApiRequest.chat(
+          messages: conversation_history.messages,
+          model_name: model.name
+        )
+        @body = JSON.parse(response.body)
 
         conversation_history.append(
           role: 'assistant',