FEATURE: Tools for models from Ollama provider (#819)

 Adds support for Ollama function calling

Author: nvh0412

app/models/llm_model.rb

@@ -31,6 +31,7 @@ def self.provider_params
       },
       ollama: {
         disable_system_prompt: :checkbox,
+        enable_native_tool: :checkbox,
       },
     }
   end

config/locales/client.en.yml

@@ -312,6 +312,7 @@ en:
           region: "AWS Bedrock Region"
           organization: "Optional OpenAI Organization ID"
           disable_system_prompt: "Disable system message in prompts"
+          enable_native_tool: "Enable native tool support"
 
       related_topics:
         title: "Related Topics"

lib/completions/dialects/ollama.rb

@@ -10,14 +10,24 @@ def can_translate?(model_provider)
           end
         end
 
-        # TODO: Add tool suppport
+        def native_tool_support?
+          enable_native_tool?
+        end
 
         def max_prompt_tokens
           llm_model.max_prompt_tokens
         end
 
         private
 
+        def tools_dialect
+          if enable_native_tool?
+            @tools_dialect ||= DiscourseAi::Completions::Dialects::OllamaTools.new(prompt.tools)
+          else
+            super
+          end
+        end
+
         def tokenizer
           llm_model.tokenizer_class
         end
@@ -26,8 +36,28 @@ def model_msg(msg)
           { role: "assistant", content: msg[:content] }
         end
 
+        def tool_call_msg(msg)
+          tools_dialect.from_raw_tool_call(msg)
+        end
+
+        def tool_msg(msg)
+          tools_dialect.from_raw_tool(msg)
+        end
+
         def system_msg(msg)
-          { role: "system", content: msg[:content] }
+          msg = { role: "system", content: msg[:content] }
+
+          if tools_dialect.instructions.present?
+            msg[:content] = msg[:content].dup << "\n\n#{tools_dialect.instructions}"
+          end
+
+          msg
+        end
+
+        def enable_native_tool?
+          return @enable_native_tool if defined?(@enable_native_tool)
+
+          @enable_native_tool = llm_model.lookup_custom_param("enable_native_tool")
         end
 
         def user_msg(msg)

lib/completions/dialects/ollama_tools.rb

@@ -0,0 +1,58 @@
+# frozen_string_literal: true
+
+module DiscourseAi
+  module Completions
+    module Dialects
+      # TODO: Define the Tool class to be inherited by all tools.
+      class OllamaTools
+        def initialize(tools)
+          @raw_tools = tools
+        end
+
+        def instructions
+          "" # Noop. Tools are listed separate.
+        end
+
+        def translated_tools
+          raw_tools.map do |t|
+            tool = t.dup
+
+            tool[:parameters] = t[:parameters]
+              .to_a
+              .reduce({ type: "object", properties: {}, required: [] }) do |memo, p|
+                name = p[:name]
+                memo[:required] << name if p[:required]
+
+                except = %i[name required item_type]
+                except << :enum if p[:enum].blank?
+
+                memo[:properties][name] = p.except(*except)
+                memo
+              end
+
+            { type: "function", function: tool }
+          end
+        end
+
+        def from_raw_tool_call(raw_message)
+          call_details = JSON.parse(raw_message[:content], symbolize_names: true)
+          call_details[:name] = raw_message[:name]
+
+          {
+            role: "assistant",
+            content: nil,
+            tool_calls: [{ type: "function", function: call_details }],
+          }
+        end
+
+        def from_raw_tool(raw_message)
+          { role: "tool", content: raw_message[:content], name: raw_message[:name] }
+        end
+
+        private
+
+        attr_reader :raw_tools
+      end
+    end
+  end
+end

lib/completions/endpoints/ollama.rb

@@ -37,11 +37,28 @@ def model_uri
           URI(llm_model.url)
         end
 
-        def prepare_payload(prompt, model_params, _dialect)
+        def native_tool_support?
+          @native_tool_support
+        end
+
+        def has_tool?(_response_data)
+          @has_function_call
+        end
+
+        def prepare_payload(prompt, model_params, dialect)
+          @native_tool_support = dialect.native_tool_support?
+
+          # https://github.com/ollama/ollama/blob/main/docs/api.md#parameters-1
+          # Due to ollama enforce a 'stream: false' for tool calls, instead of complicating the code,
+          # we will just disable streaming for all ollama calls if native tool support is enabled
+
           default_options
             .merge(model_params)
             .merge(messages: prompt)
-            .tap { |payload| payload[:stream] = false if !@streaming_mode }
+            .tap { |payload| payload[:stream] = false if @native_tool_support || !@streaming_mode }
+            .tap do |payload|
+              payload[:tools] = dialect.tools if @native_tool_support && dialect.tools.present?
+            end
         end
 
         def prepare_request(payload)
@@ -58,7 +75,66 @@ def extract_completion_from(response_raw)
           parsed = JSON.parse(response_raw, symbolize_names: true)
           return if !parsed
 
-          parsed.dig(:message, :content)
+          response_h = parsed.dig(:message)
+
+          @has_function_call ||= response_h.dig(:tool_calls).present?
+          @has_function_call ? response_h.dig(:tool_calls, 0) : response_h.dig(:content)
+        end
+
+        def add_to_function_buffer(function_buffer, payload: nil, partial: nil)
+          @args_buffer ||= +""
+
+          if @streaming_mode
+            return function_buffer if !partial
+          else
+            partial = payload
+          end
+
+          f_name = partial.dig(:function, :name)
+
+          @current_function ||= function_buffer.at("invoke")
+
+          if f_name
+            current_name = function_buffer.at("tool_name").content
+
+            if current_name.blank?
+              # first call
+            else
+              # we have a previous function, so we need to add a noop
+              @args_buffer = +""
+              @current_function =
+                function_buffer.at("function_calls").add_child(
+                  Nokogiri::HTML5::DocumentFragment.parse(noop_function_call_text + "\n"),
+                )
+            end
+          end
+
+          @current_function.at("tool_name").content = f_name if f_name
+          @current_function.at("tool_id").content = partial[:id] if partial[:id]
+
+          args = partial.dig(:function, :arguments)
+
+          # allow for SPACE within arguments
+          if args && args != ""
+            @args_buffer << args.to_json
+
+            begin
+              json_args = JSON.parse(@args_buffer, symbolize_names: true)
+
+              argument_fragments =
+                json_args.reduce(+"") do |memo, (arg_name, value)|
+                  memo << "\n<#{arg_name}>#{value}</#{arg_name}>"
+                end
+              argument_fragments << "\n"
+
+              @current_function.at("parameters").children =
+                Nokogiri::HTML5::DocumentFragment.parse(argument_fragments)
+            rescue JSON::ParserError
+              return function_buffer
+            end
+          end
+
+          function_buffer
         end
       end
     end

spec/fabricators/llm_model_fabricator.rb

@@ -87,4 +87,5 @@
   api_key "ABC"
   tokenizer "DiscourseAi::Tokenizer::Llama3Tokenizer"
   url "http://api.ollama.ai/api/chat"
+  provider_params { { enable_native_tool: true } }
 end

spec/lib/completions/dialects/ollama_spec.rb

@@ -7,15 +7,37 @@
   let(:context) { DialectContext.new(described_class, model) }
 
   describe "#translate" do
-    it "translates a prompt written in our generic format to the Ollama format" do
-      ollama_version = [
-        { role: "system", content: context.system_insts },
-        { role: "user", content: context.simple_user_input },
-      ]
+    context "when native tool support is enabled" do
+      it "translates a prompt written in our generic format to the Ollama format" do
+        ollama_version = [
+          { role: "system", content: context.system_insts },
+          { role: "user", content: context.simple_user_input },
+        ]
 
-      translated = context.system_user_scenario
+        translated = context.system_user_scenario
 
-      expect(translated).to eq(ollama_version)
+        expect(translated).to eq(ollama_version)
+      end
+    end
+
+    context "when native tool support is disabled - XML tools" do
+      it "includes the instructions in the system message" do
+        allow(model).to receive(:lookup_custom_param).with("enable_native_tool").and_return(false)
+
+        DiscourseAi::Completions::Dialects::XmlTools
+          .any_instance
+          .stubs(:instructions)
+          .returns("Instructions")
+
+        ollama_version = [
+          { role: "system", content: "#{context.system_insts}\n\nInstructions" },
+          { role: "user", content: context.simple_user_input },
+        ]
+
+        translated = context.system_user_scenario
+
+        expect(translated).to eq(ollama_version)
+      end
     end
 
     it "trims content if it's getting too long" do
@@ -33,4 +55,40 @@
       expect(context.dialect(nil).max_prompt_tokens).to eq(10_000)
     end
   end
+
+  describe "#tools" do
+    context "when native tools are enabled" do
+      it "returns the translated tools from the OllamaTools class" do
+        tool = instance_double(DiscourseAi::Completions::Dialects::OllamaTools)
+
+        allow(model).to receive(:lookup_custom_param).with("enable_native_tool").and_return(true)
+        allow(tool).to receive(:translated_tools)
+        allow(DiscourseAi::Completions::Dialects::OllamaTools).to receive(:new).and_return(tool)
+
+        context.dialect_tools
+
+        expect(DiscourseAi::Completions::Dialects::OllamaTools).to have_received(:new).with(
+          context.prompt.tools,
+        )
+        expect(tool).to have_received(:translated_tools)
+      end
+    end
+
+    context "when native tools are disabled" do
+      it "returns the translated tools from the XmlTools class" do
+        tool = instance_double(DiscourseAi::Completions::Dialects::XmlTools)
+
+        allow(model).to receive(:lookup_custom_param).with("enable_native_tool").and_return(false)
+        allow(tool).to receive(:translated_tools)
+        allow(DiscourseAi::Completions::Dialects::XmlTools).to receive(:new).and_return(tool)
+
+        context.dialect_tools
+
+        expect(DiscourseAi::Completions::Dialects::XmlTools).to have_received(:new).with(
+          context.prompt.tools,
+        )
+        expect(tool).to have_received(:translated_tools)
+      end
+    end
+  end
 end