DEV: use a proper object for tool definition (#1337)

* DEV: use a proper object for tool definition

This moves away from using a loose hash to define tools, which
is error prone.

Instead given a proper object we will also be able to coerce the
return values to match tool definition correctly

* fix xml tools

* fix anthropic tools

* fix specs... a few more to go

* specs are passing

* FIX: coerce values for XML tool calls

* Update spec/lib/completions/tool_definition_spec.rb

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: SamSaffron

lib/completions/dialects/claude_tools.rb

@@ -9,33 +9,8 @@ def initialize(tools)
         end
 
         def translated_tools
-          # Transform the raw tools into the required Anthropic Claude API format
           raw_tools.map do |t|
-            properties = {}
-            required = []
-
-            if t[:parameters]
-              properties = {}
-
-              t[:parameters].each do |param|
-                mapped = { type: param[:type], description: param[:description] }
-                mapped[:items] = { type: param[:item_type] } if param[:item_type]
-                mapped[:enum] = param[:enum] if param[:enum]
-                properties[param[:name]] = mapped
-              end
-              required =
-                t[:parameters].select { |param| param[:required] }.map { |param| param[:name] }
-            end
-
-            {
-              name: t[:name],
-              description: t[:description],
-              input_schema: {
-                type: "object",
-                properties: properties,
-                required: required,
-              },
-            }
+            { name: t.name, description: t.description, input_schema: t.parameters_json_schema }
           end
         end
 

lib/completions/dialects/cohere_tools.rb

@@ -38,27 +38,24 @@ def tool_results(messages)
         end
 
         def translated_tools
-          raw_tools.map do |t|
-            tool = t.dup
+          raw_tools.map do |tool|
+            defs = {}
 
-            tool[:parameter_definitions] = t[:parameters]
-              .to_a
-              .reduce({}) do |memo, p|
-                name = p[:name]
-                memo[name] = {
-                  description: p[:description],
-                  type: cohere_type(p[:type], p[:item_type]),
-                  required: p[:required],
-                }
+            tool.parameters.each do |p|
+              name = p.name
+              defs[name] = {
+                description: p.description,
+                type: cohere_type(p.type, p.item_type),
+                required: p.required,
+              }
 
-                memo[name][:default] = p[:default] if p[:default]
-                memo
-              end
+              #defs[name][:default] = p.default if p.default
+            end
 
             {
-              name: tool[:name] == "search" ? "search_local" : tool[:name],
-              description: tool[:description],
-              parameter_definitions: tool[:parameter_definitions],
+              name: tool.name == "search" ? "search_local" : tool.name,
+              description: tool.description,
+              parameter_definitions: defs,
             }
           end
         end
@@ -72,6 +69,7 @@ def instructions
         attr_reader :raw_tools
 
         def cohere_type(type, item_type)
+          type = type.to_s
           case type
           when "string"
             "str"

lib/completions/dialects/gemini.rb

@@ -47,22 +47,8 @@ def tools
 
           translated_tools =
             prompt.tools.map do |t|
-              tool = t.slice(:name, :description)
-
-              if t[:parameters]
-                tool[:parameters] = t[:parameters].reduce(
-                  { type: "object", required: [], properties: {} },
-                ) do |memo, p|
-                  name = p[:name]
-                  memo[:required] << name if p[:required]
-
-                  memo[:properties][name] = p.except(:name, :required, :item_type)
-
-                  memo[:properties][name][:items] = { type: p[:item_type] } if p[:item_type]
-                  memo
-                end
-              end
-
+              tool = { name: t.name, description: t.description }
+              tool[:parameters] = t.parameters_json_schema if t.parameters
               tool
             end
 

lib/completions/dialects/nova_tools.rb

@@ -17,10 +17,10 @@ def translated_tools
               @raw_tools.map do |tool|
                 {
                   toolSpec: {
-                    name: tool[:name],
-                    description: tool[:description],
+                    name: tool.name,
+                    description: tool.description,
                     inputSchema: {
-                      json: convert_tool_to_input_schema(tool),
+                      json: tool.parameters_json_schema,
                     },
                   },
                 }
@@ -51,32 +51,6 @@ def from_raw_tool(raw_message)
             },
           }
         end
-
-        private
-
-        def convert_tool_to_input_schema(tool)
-          tool = tool.transform_keys(&:to_sym)
-          properties = {}
-          tool[:parameters].each do |param|
-            schema = {}
-            type = param[:type]
-            type = "string" if !%w[string number boolean integer array].include?(type)
-
-            schema[:type] = type
-
-            if enum = param[:enum]
-              schema[:enum] = enum
-            end
-
-            schema[:items] = { type: param[:item_type] } if type == "array"
-
-            schema[:required] = true if param[:required]
-
-            properties[param[:name]] = schema
-          end
-
-          { type: "object", properties: properties }
-        end
       end
     end
   end

lib/completions/dialects/ollama_tools.rb

@@ -14,23 +14,15 @@ def instructions
         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 }
+          raw_tools.map do |tool|
+            {
+              type: "function",
+              function: {
+                name: tool.name,
+                description: tool.description,
+                parameters: tool.parameters_json_schema,
+              },
+            }
           end
         end
 

lib/completions/dialects/open_ai_tools.rb

@@ -9,25 +9,15 @@ def initialize(tools)
         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[:properties][name][:items] = { type: p[:item_type] } if p[:item_type]
-                memo
-              end
-
-            { type: "function", function: tool }
+          raw_tools.map do |tool|
+            {
+              type: "function",
+              function: {
+                name: tool.name,
+                description: tool.description,
+                parameters: tool.parameters_json_schema,
+              },
+            }
           end
         end
 

lib/completions/dialects/xml_tools.rb

@@ -9,42 +9,45 @@ def initialize(tools)
         end
 
         def translated_tools
-          raw_tools.reduce(+"") do |tools, function|
+          result = +""
+
+          raw_tools.each do |tool|
             parameters = +""
-            if function[:parameters].present?
-              function[:parameters].each do |parameter|
+            if tool.parameters.present?
+              tool.parameters.each do |parameter|
                 parameters << <<~PARAMETER
                   <parameter>
-                  <name>#{parameter[:name]}</name>
-                  <type>#{parameter[:type]}</type>
-                  <description>#{parameter[:description]}</description>
-                  <required>#{parameter[:required]}</required>
+                  <name>#{parameter.name}</name>
+                  <type>#{parameter.type}</type>
+                  <description>#{parameter.description}</description>
+                  <required>#{parameter.required}</required>
                 PARAMETER
-                if parameter[:enum]
-                  parameters << "<options>#{parameter[:enum].join(",")}</options>\n"
+                if parameter.item_type
+                  parameters << "<array_item_type>#{parameter.item_type}</array_item_type>\n"
                 end
+                parameters << "<options>#{parameter.enum.join(",")}</options>\n" if parameter.enum
                 parameters << "</parameter>\n"
               end
             end
 
-            tools << <<~TOOLS
+            result << <<~TOOLS
               <tool_description>
-              <tool_name>#{function[:name]}</tool_name>
-              <description>#{function[:description]}</description>
+              <tool_name>#{tool.name}</tool_name>
+              <description>#{tool.description}</description>
               <parameters>
               #{parameters}</parameters>
               </tool_description>
             TOOLS
           end
+          result
         end
 
         def instructions
           return "" if raw_tools.blank?
 
           @instructions ||=
             begin
-              has_arrays =
-                raw_tools.any? { |tool| tool[:parameters]&.any? { |p| p[:type] == "array" } }
+              has_arrays = raw_tools.any? { |tool| tool.parameters&.any? { |p| p.type == "array" } }
 
               (<<~TEXT).strip
               #{tool_preamble(include_array_tip: has_arrays)}

lib/completions/endpoints/base.rb

@@ -28,7 +28,7 @@ def endpoint_for(provider_name)
               DiscourseAi::Completions::Endpoints::OpenRouter,
             ]
 
-            endpoints << DiscourseAi::Completions::Endpoints::Ollama if Rails.env.development?
+            endpoints << DiscourseAi::Completions::Endpoints::Ollama if !Rails.env.production?
 
             if Rails.env.test? || Rails.env.development?
               endpoints << DiscourseAi::Completions::Endpoints::Fake
@@ -166,6 +166,7 @@ def perform_completion!(
               xml_tool_processor =
                 XmlToolProcessor.new(
                   partial_tool_calls: partial_tool_calls,
+                  tool_definitions: dialect.prompt.tools,
                 ) if xml_tools_enabled? && dialect.prompt.has_tools?
 
               to_strip = xml_tags_to_strip(dialect)

lib/completions/prompt.rb

@@ -5,8 +5,8 @@ module Completions
     class Prompt
       INVALID_TURN = Class.new(StandardError)
 
-      attr_reader :messages
-      attr_accessor :tools, :topic_id, :post_id, :max_pixels, :tool_choice
+      attr_reader :messages, :tools
+      attr_accessor :topic_id, :post_id, :max_pixels, :tool_choice
 
       def initialize(
         system_message_text = nil,
@@ -37,10 +37,25 @@ def initialize(
         @messages.each { |message| validate_message(message) }
         @messages.each_cons(2) { |last_turn, new_turn| validate_turn(last_turn, new_turn) }
 
-        @tools = tools
+        self.tools = tools
         @tool_choice = tool_choice
       end
 
+      def tools=(tools)
+        raise ArgumentError, "tools must be an array" if !tools.is_a?(Array) && !tools.nil?
+
+        @tools =
+          tools.map do |tool|
+            if tool.is_a?(Hash)
+              ToolDefinition.from_hash(tool)
+            elsif tool.is_a?(ToolDefinition)
+              tool
+            else
+              raise ArgumentError, "tool must be a hash or a ToolDefinition was #{tool.class}"
+            end
+          end
+      end
+
       # this new api tries to create symmetry between responses and prompts
       # this means anything we get back from the model via endpoint can be easily appended
       def push_model_response(response)