sambanova and vllm

Author: SamSaffron

lib/completions/endpoints/base.rb

@@ -312,10 +312,12 @@ def non_streaming_response(
           response_data.each { |partial| partials_raw << partial.to_s }
 
           if xml_tool_processor
-            processed = (xml_tool_processor << response_data)
-            processed << xml_tool_processor.finish
-            response_data = []
-            processed.flatten.compact.each { |partial| response_data << partial }
+            response_data.each do |partial|
+              processed = (xml_tool_processor << partial)
+              processed << xml_tool_processor.finish
+              response_data = []
+              processed.flatten.compact.each { |inner| response_data << inner }
+            end
           end
 
           if xml_stripper

lib/completions/endpoints/samba_nova.rb

@@ -55,29 +55,31 @@ def final_log_update(log)
           log.response_tokens = @completion_tokens if @completion_tokens
         end
 
-        def extract_completion_from(response_raw)
+        def xml_tools_enabled?
+          true
+        end
+
+        def decode(response_raw)
           json = JSON.parse(response_raw, symbolize_names: true)
+          [json.dig(:choices, 0, :message, :content)]
+        end
+
+        def decode_chunk(chunk)
+          @json_decoder ||= JsonStreamDecoder.new
+          (@json_decoder << chunk).map { |json|
+            text = json.dig(:choices, 0, :delta, :content)
 
-          if @streaming_mode
             @prompt_tokens ||= json.dig(:usage, :prompt_tokens)
             @completion_tokens ||= json.dig(:usage, :completion_tokens)
-          end
 
-          parsed = json.dig(:choices, 0)
-          return if !parsed
-
-          @streaming_mode ? parsed.dig(:delta, :content) : parsed.dig(:message, :content)
-        end
-
-        def partials_from(decoded_chunk)
-          decoded_chunk
-            .split("\n")
-            .map do |line|
-              data = line.split("data: ", 2)[1]
-              data == "[DONE]" ? nil : data
+            if !text.to_s.empty?
+              text
+            else
+              nil
             end
-            .compact
+          }.flatten.compact
         end
+
       end
     end
   end

lib/completions/endpoints/vllm.rb

@@ -56,6 +56,27 @@ def prepare_request(payload)
           Net::HTTP::Post.new(model_uri, headers).tap { |r| r.body = payload }
         end
 
+        def xml_tools_enabled?
+          true
+        end
+
+        def decode(response_raw)
+          json = JSON.parse(response_raw, symbolize_names: true)
+          [json.dig(:choices, 0, :message, :content)]
+        end
+
+        def decode_chunk(chunk)
+          @json_decoder ||= JsonStreamDecoder.new
+          (@json_decoder << chunk).map do |parsed|
+            text = parsed.dig(:choices, 0, :delta, :content)
+            if text.to_s.empty?
+              nil
+            else
+              text
+            end
+          end.compact
+        end
+
         def partials_from(decoded_chunk)
           decoded_chunk
             .split("\n")

spec/lib/completions/endpoints/endpoint_compliance.rb

@@ -201,6 +201,7 @@ def streaming_mode_simple_prompt(mock)
       expect(log.raw_request_payload).to be_present
       expect(log.raw_response_payload).to be_present
       expect(log.request_tokens).to eq(endpoint.prompt_size(dialect.translate))
+
       expect(log.response_tokens).to eq(
         endpoint.llm_model.tokenizer_class.size(mock.streamed_simple_deltas[0...-1].join),
       )

spec/lib/completions/endpoints/samba_nova_spec.rb

@@ -22,10 +22,15 @@
       },
     ).to_return(status: 200, body: body, headers: {})
 
-    response = +""
+    response = []
     llm.generate("who are you?", user: Discourse.system_user) { |partial| response << partial }
 
-    expect(response).to eq("I am a bot")
+    expect(response).to eq(["I am a bot"])
+
+    log = AiApiAuditLog.order(:id).last
+
+    expect(log.request_tokens).to eq(21)
+    expect(log.response_tokens).to eq(41)
   end
 
   it "can perform regular completions" do

spec/lib/completions/endpoints/vllm_spec.rb

@@ -136,18 +136,13 @@ def stub_streamed_response(prompt, deltas, tool_call: false)
 
       result = llm.generate(prompt, user: Discourse.system_user)
 
-      expected = <<~TEXT
-        <function_calls>
-        <invoke>
-        <tool_name>calculate</tool_name>
-        <parameters>
-        <expression>1+1</expression></parameters>
-        <tool_id>tool_0</tool_id>
-        </invoke>
-        </function_calls>
-      TEXT
+      expected = DiscourseAi::Completions::ToolCall.new(
+        name: "calculate",
+        id: "tool_0",
+        parameters: { expression: "1+1" },
+      )
 
-      expect(result.strip).to eq(expected.strip)
+      expect(result).to eq(expected)
     end
   end