formatting and spec fixes

Author: SamSaffron

lib/completions/dialects/nova.rb

@@ -33,16 +33,11 @@ def to_payload(options = nil)
             stop_sequences = options[:stop_sequences]
             max_tokens = options[:max_tokens]
 
-            inference_config =
-              options&.slice(:temperature, :top_p, :top_k)
+            inference_config = options&.slice(:temperature, :top_p, :top_k)
 
-            if stop_sequences.present?
-              inference_config[:stopSequences] = stop_sequences
-            end
+            inference_config[:stopSequences] = stop_sequences if stop_sequences.present?
 
-            if max_tokens.present?
-              inference_config[:max_new_tokens] = max_tokens
-            end
+            inference_config[:max_new_tokens] = max_tokens if max_tokens.present?
 
             result = { system: system, messages: messages }
             result[:inferenceConfig] = inference_config if inference_config.present?

lib/completions/endpoints/aws_bedrock.rb

@@ -167,9 +167,7 @@ def bedrock_decode(chunk)
           while decoded
             parsed = JSON.parse(decoded.payload.string)
             if exception = decoded.headers[":exception-type"]
-              Rails.logger.error(
-                "#{self.class.name}: #{exception}: #{parsed}",
-              )
+              Rails.logger.error("#{self.class.name}: #{exception}: #{parsed}")
               # TODO based on how often this happens, we may want to raise so we
               # can retry, this may catch rate limits for example
             end
@@ -203,8 +201,10 @@ def final_log_update(log)
 
         def processor
           if dialect.is_a?(DiscourseAi::Completions::Dialects::Claude)
-          @processor ||=
-            DiscourseAi::Completions::AnthropicMessageProcessor.new(streaming_mode: @streaming_mode)
+            @processor ||=
+              DiscourseAi::Completions::AnthropicMessageProcessor.new(
+                streaming_mode: @streaming_mode,
+              )
           else
             @processor ||=
               DiscourseAi::Completions::NovaMessageProcessor.new(streaming_mode: @streaming_mode)

lib/completions/nova_message_processor.rb

@@ -69,12 +69,10 @@ def process_streamed_message(parsed)
       @current_tool_call.append(tool_progress)
     end
 
-    if parsed[:contentBlockStop] && @current_tool_call
-      result = @current_tool_call.to_tool_call
-    end
+    result = @current_tool_call.to_tool_call if parsed[:contentBlockStop] && @current_tool_call
 
     if metadata = parsed[:metadata]
-      @input_tokens = metadata.dig(:usage,:inputTokens)
+      @input_tokens = metadata.dig(:usage, :inputTokens)
       @output_tokens = metadata.dig(:usage, :outputTokens)
     end
 

spec/lib/completions/dialects/claude_spec.rb

@@ -1,7 +1,6 @@
 # frozen_string_literal: true
 
 RSpec.describe DiscourseAi::Completions::Dialects::Claude do
-
   fab!(:llm_model) { Fabricate(:anthropic_model, name: "claude-3-opus") }
 
   let :opus_dialect_klass do

spec/lib/completions/dialects/nova_spec.rb

@@ -62,16 +62,9 @@
           {
             name: "get_weather",
             description: "Get the weather in a city",
-            input_schema: {
-              type: "object",
-              properties: {
-                location: {
-                  type: "string",
-                  description: "the city name",
-                },
-              },
-              required: ["location"],
-            },
+            parameters: [
+              { name: "location", type: "string", description: "the city name", required: true },
+            ],
           },
         ]
 
@@ -98,21 +91,16 @@
                     json: {
                       type: "object",
                       properties: {
-                        location: {
+                        "location" => {
                           type: "string",
-                          description: "the city name",
+                          required: true,
                         },
                       },
-                      required: ["location"],
                     },
                   },
                 },
               },
             ],
-            toolChoice: {
-              auto: {
-              },
-            },
           },
         )
       end
@@ -126,18 +114,18 @@
 
         dialect = nova_dialect_klass.new(prompt, llm_model)
 
-        options = { temperature: 0.7, top_p: 0.9, max_new_tokens: 100, stop_sequences: ["STOP"] }
+        options = { temperature: 0.7, top_p: 0.9, max_tokens: 100, stop_sequences: ["STOP"] }
 
         translated = dialect.translate
 
         expected = {
-          system: { text: "You are a helpful bot" },
+          system: [{ text: "You are a helpful bot" }],
           messages: [{ role: "user", content: [{ text: "Hello" }] }],
           inferenceConfig: {
             temperature: 0.7,
             top_p: 0.9,
+            stopSequences: ["STOP"],
             max_new_tokens: 100,
-            stop_sequences: ["STOP"],
           },
         }
 

spec/lib/completions/endpoints/nova_spec.rb

@@ -223,7 +223,7 @@ def encode_message(message)
     )
 
     # lets continue and ensure all messages are mapped correctly
-    prompt.push(type: :tool_call, name: "time", content: {timezone: "EST"}.to_json, id: "111")
+    prompt.push(type: :tool_call, name: "time", content: { timezone: "EST" }.to_json, id: "111")
     prompt.push(type: :tool, name: "time", content: "1pm".to_json, id: "111")
 
     # lets just return the tool call again, this is about ensuring we encode the prompt right
@@ -239,9 +239,45 @@ def encode_message(message)
       proxy.generate(prompt, user: user, max_tokens: 200) { |partial| response << partial }
     end
 
-    expected = {:system=>[{:text=>"You are a helpful assistant."}], :messages=>[{:role=>"user", :content=>[{:text=>"what is the time in EST"}]}, {:role=>"assistant", :content=>[{:toolUse=>{:toolUseId=>"111", :name=>"time", :input=>nil}}]}, {:role=>"user", :content=>[{:toolResult=>{:toolUseId=>"111", :content=>[{:json=>"1pm"}]}}]}], :inferenceConfig=>{:max_new_tokens=>200}, :toolConfig=>{:tools=>[{:toolSpec=>{:name=>"time", :description=>"Will look up the current time", :inputSchema=>{:json=>{:type=>"object", :properties=>{:timezone=>{:type=>"string", :required=>true}}}}}}]}}
+    expected = {
+      system: [{ text: "You are a helpful assistant." }],
+      messages: [
+        { role: "user", content: [{ text: "what is the time in EST" }] },
+        {
+          role: "assistant",
+          content: [{ toolUse: { toolUseId: "111", name: "time", input: nil } }],
+        },
+        {
+          role: "user",
+          content: [{ toolResult: { toolUseId: "111", content: [{ json: "1pm" }] } }],
+        },
+      ],
+      inferenceConfig: {
+        max_new_tokens: 200,
+      },
+      toolConfig: {
+        tools: [
+          {
+            toolSpec: {
+              name: "time",
+              description: "Will look up the current time",
+              inputSchema: {
+                json: {
+                  type: "object",
+                  properties: {
+                    timezone: {
+                      type: "string",
+                      required: true,
+                    },
+                  },
+                },
+              },
+            },
+          },
+        ],
+      },
+    }
 
     expect(JSON.parse(request.body, symbolize_names: true)).to eq(expected)
-
   end
 end