Discourse/discourse-ai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

FEATURE: improve tool support (#904)

Browse Source 
This re-implements tool support in DiscourseAi::Completions::Llm #generate

Previously tool support was always returned via XML and it would be the responsibility of the caller to parse XML

New implementation has the endpoints return ToolCall objects.

Additionally this simplifies the Llm endpoint interface and gives it more clarity. Llms must implement

decode, decode_chunk (for streaming)

It is the implementers responsibility to figure out how to decode chunks, base no longer implements. To make this easy we ship a flexible json decoder which is easy to wire up.

Also (new)

    Better debugging for PMs, we now have a next / previous button to see all the Llm messages associated with a PM
    Token accounting is fixed for vllm (we were not correctly counting tokens)
This commit is contained in:
Sam
2024-11-12 08:14:30 +11:00
committed by GitHub
parent 644141ff08
commit e817b7dc11
43 changed files with 1685 additions and 1293 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/controllers/discourse_ai/ai_bot/bot_controller.rb
+8
View File
@@ -6,6 +6,14 @@ module DiscourseAi
requires_plugin ::DiscourseAi::PLUGIN_NAME
requires_login
def show_debug_info_by_id
log = AiApiAuditLog.find(params[:id])
raise Discourse::NotFound if !log.topic
guardian.ensure_can_debug_ai_bot_conversation!(log.topic)
render json: AiApiAuditLogSerializer.new(log, root: false), status: 200
end
def show_debug_info
post = Post.find(params[:post_id])
guardian.ensure_can_debug_ai_bot_conversation!(post)
app/models/ai_api_audit_log.rb
+8
View File
@@ -14,6 +14,14 @@ class AiApiAuditLog < ActiveRecord::Base
Ollama = 7
SambaNova = 8
end
def next_log_id
self.class.where("id > ?", id).where(topic_id: topic_id).order(id: :asc).pluck(:id).first
end
def prev_log_id
self.class.where("id < ?", id).where(topic_id: topic_id).order(id: :desc).pluck(:id).first
end
end
# == Schema Information
app/serializers/ai_api_audit_log_serializer.rb
+3 -1
View File
@@ -12,5 +12,7 @@ class AiApiAuditLogSerializer < ApplicationSerializer
:post_id,
:feature_name,
:language_model,
:created_at
:created_at,
:prev_log_id,
:next_log_id
end
assets/javascripts/discourse/components/modal/debug-ai-modal.gjs
+46 -5
View File
@@ -7,6 +7,7 @@ import { htmlSafe } from "@ember/template";
import DButton from "discourse/components/d-button";
import DModal from "discourse/components/d-modal";
import { ajax } from "discourse/lib/ajax";
import { popupAjaxError } from "discourse/lib/ajax-error";
import { clipboardCopy, escapeExpression } from "discourse/lib/utilities";
import i18n from "discourse-common/helpers/i18n";
import discourseLater from "discourse-common/lib/later";
@@ -63,6 +64,28 @@ export default class DebugAiModal extends Component {
this.copy(this.info.raw_response_payload);
}
async loadLog(logId) {
try {
await ajax(`/discourse-ai/ai-bot/show-debug-info/${logId}.json`).then(
(result) => {
this.info = result;
}
);
} catch (e) {
popupAjaxError(e);
}
}
@action
prevLog() {
this.loadLog(this.info.prev_log_id);
}
@action
nextLog() {
this.loadLog(this.info.next_log_id);
}
copy(text) {
clipboardCopy(text);
this.justCopiedText = I18n.t("discourse_ai.ai_bot.conversation_shared");
@@ -73,11 +96,13 @@ export default class DebugAiModal extends Component {
}
loadApiRequestInfo() {
ajax(
`/discourse-ai/ai-bot/post/${this.args.model.id}/show-debug-info.json`
).then((result) => {
this.info = result;
});
ajax(`/discourse-ai/ai-bot/post/${this.args.model.id}/show-debug-info.json`)
.then((result) => {
this.info = result;
})
.catch((e) => {
popupAjaxError(e);
});
}
get requestActive() {
@@ -147,6 +172,22 @@ export default class DebugAiModal extends Component {
@action={{this.copyResponse}}
@label="discourse_ai.ai_bot.debug_ai_modal.copy_response"
/>
{{#if this.info.prev_log_id}}
<DButton
class="btn"
@icon="angles-left"
@action={{this.prevLog}}
@label="discourse_ai.ai_bot.debug_ai_modal.previous_log"
/>
{{/if}}
{{#if this.info.next_log_id}}
<DButton
class="btn"
@icon="angles-right"
@action={{this.nextLog}}
@label="discourse_ai.ai_bot.debug_ai_modal.next_log"
/>
{{/if}}
<span class="ai-debut-modal__just-copied">{{this.justCopiedText}}</span>
</:footer>
</DModal>
config/locales/client.en.yml
+2
View File
@@ -415,6 +415,8 @@ en:
response_tokens: "Response tokens:"
request: "Request"
response: "Response"
next_log: "Next"
previous_log: "Previous"
share_full_topic_modal:
title: "Share Conversation Publicly"
config/routes.rb
+1
View File
@@ -22,6 +22,7 @@ DiscourseAi::Engine.routes.draw do
scope module: :ai_bot, path: "/ai-bot", defaults: { format: :json } do
get "bot-username" => "bot#show_bot_username"
get "post/:post_id/show-debug-info" => "bot#show_debug_info"
get "show-debug-info/:id" => "bot#show_debug_info_by_id"
post "post/:post_id/stop-streaming" => "bot#stop_streaming_response"
end
lib/ai_bot/bot.rb
+12 -7
View File
@@ -100,6 +100,7 @@ module DiscourseAi
llm_kwargs[:top_p] = persona.top_p if persona.top_p
needs_newlines = false
tools_ran = 0
while total_completions <= MAX_COMPLETIONS && ongoing_chain
tool_found = false
@@ -107,9 +108,10 @@ module DiscourseAi
result =
llm.generate(prompt, feature_name: "bot", **llm_kwargs) do |partial, cancel|
tools = persona.find_tools(partial, bot_user: user, llm: llm, context: context)
tool = persona.find_tool(partial, bot_user: user, llm: llm, context: context)
tool = nil if tools_ran >= MAX_TOOLS
if (tools.present?)
if tool.present?
tool_found = true
# a bit hacky, but extra newlines do no harm
if needs_newlines
@@ -117,13 +119,16 @@ module DiscourseAi
needs_newlines = false
end
tools[0..MAX_TOOLS].each do |tool|
process_tool(tool, raw_context, llm, cancel, update_blk, prompt, context)
ongoing_chain &&= tool.chain_next_response?
end
process_tool(tool, raw_context, llm, cancel, update_blk, prompt, context)
tools_ran += 1
ongoing_chain &&= tool.chain_next_response?
else
needs_newlines = true
update_blk.call(partial, cancel)
if partial.is_a?(DiscourseAi::Completions::ToolCall)
Rails.logger.warn("DiscourseAi: Tool not found: #{partial.name}")
else
update_blk.call(partial, cancel)
end
end
end
lib/ai_bot/personas/persona.rb
+7 -14
View File
@@ -199,23 +199,16 @@ module DiscourseAi
prompt
end
def find_tools(partial, bot_user:, llm:, context:)
return [] if !partial.include?("</invoke>")
parsed_function = Nokogiri::HTML5.fragment(partial)
parsed_function
.css("invoke")
.map do |fragment|
tool_instance(fragment, bot_user: bot_user, llm: llm, context: context)
end
.compact
def find_tool(partial, bot_user:, llm:, context:)
return nil if !partial.is_a?(DiscourseAi::Completions::ToolCall)
tool_instance(partial, bot_user: bot_user, llm: llm, context: context)
end
protected
def tool_instance(parsed_function, bot_user:, llm:, context:)
function_id = parsed_function.at("tool_id")&.text
function_name = parsed_function.at("tool_name")&.text
def tool_instance(tool_call, bot_user:, llm:, context:)
function_id = tool_call.id
function_name = tool_call.name
return nil if function_name.nil?
tool_klass = available_tools.find { |c| c.signature.dig(:name) == function_name }
@@ -224,7 +217,7 @@ module DiscourseAi
arguments = {}
tool_klass.signature[:parameters].to_a.each do |param|
name = param[:name]
value = parsed_function.at(name)&.text
value = tool_call.parameters[name.to_sym]
if param[:type] == "array" && value
value =
lib/completions/anthropic_message_processor.rb
+56 -62
View File
@@ -13,6 +13,11 @@ class DiscourseAi::Completions::AnthropicMessageProcessor
def append(json)
@raw_json << json
end
def to_tool_call
parameters = JSON.parse(raw_json, symbolize_names: true)
DiscourseAi::Completions::ToolCall.new(id: id, name: name, parameters: parameters)
end
end
attr_reader :tool_calls, :input_tokens, :output_tokens
@@ -20,80 +25,69 @@ class DiscourseAi::Completions::AnthropicMessageProcessor
def initialize(streaming_mode:)
@streaming_mode = streaming_mode
@tool_calls = []
@current_tool_call = nil
end
def to_xml_tool_calls(function_buffer)
return function_buffer if @tool_calls.blank?
def to_tool_calls
@tool_calls.map { |tool_call| tool_call.to_tool_call }
end
function_buffer = Nokogiri::HTML5.fragment(<<~TEXT)
<function_calls>
</function_calls>
TEXT
@tool_calls.each do |tool_call|
node =
function_buffer.at("function_calls").add_child(
Nokogiri::HTML5::DocumentFragment.parse(
DiscourseAi::Completions::Endpoints::Base.noop_function_call_text + "\n",
),
)
params = JSON.parse(tool_call.raw_json, symbolize_names: true)
xml =
params.map { |name, value| "<#{name}>#{CGI.escapeHTML(value.to_s)}</#{name}>" }.join("\n")
node.at("tool_name").content = tool_call.name
node.at("tool_id").content = tool_call.id
node.at("parameters").children = Nokogiri::HTML5::DocumentFragment.parse(xml) if xml.present?
def process_streamed_message(parsed)
result = nil
if parsed[:type] == "content_block_start" && parsed.dig(:content_block, :type) == "tool_use"
tool_name = parsed.dig(:content_block, :name)
tool_id = parsed.dig(:content_block, :id)
result = @current_tool_call.to_tool_call if @current_tool_call
@current_tool_call = AnthropicToolCall.new(tool_name, tool_id) if tool_name
elsif parsed[:type] == "content_block_start" || parsed[:type] == "content_block_delta"
if @current_tool_call
tool_delta = parsed.dig(:delta, :partial_json).to_s
@current_tool_call.append(tool_delta)
else
result = parsed.dig(:delta, :text).to_s
end
elsif parsed[:type] == "content_block_stop"
if @current_tool_call
result = @current_tool_call.to_tool_call
@current_tool_call = nil
end
elsif parsed[:type] == "message_start"
@input_tokens = parsed.dig(:message, :usage, :input_tokens)
elsif parsed[:type] == "message_delta"
@output_tokens =
parsed.dig(:usage, :output_tokens) || parsed.dig(:delta, :usage, :output_tokens)
elsif parsed[:type] == "message_stop"
# bedrock has this ...
if bedrock_stats = parsed.dig("amazon-bedrock-invocationMetrics".to_sym)
@input_tokens = bedrock_stats[:inputTokenCount] || @input_tokens
@output_tokens = bedrock_stats[:outputTokenCount] || @output_tokens
end
end
function_buffer
result
end
def process_message(payload)
result = ""
parsed = JSON.parse(payload, symbolize_names: true)
parsed = payload
parsed = JSON.parse(payload, symbolize_names: true) if payload.is_a?(String)
if @streaming_mode
if parsed[:type] == "content_block_start" && parsed.dig(:content_block, :type) == "tool_use"
tool_name = parsed.dig(:content_block, :name)
tool_id = parsed.dig(:content_block, :id)
@tool_calls << AnthropicToolCall.new(tool_name, tool_id) if tool_name
elsif parsed[:type] == "content_block_start" || parsed[:type] == "content_block_delta"
if @tool_calls.present?
result = parsed.dig(:delta, :partial_json).to_s
@tool_calls.last.append(result)
else
result = parsed.dig(:delta, :text).to_s
content = parsed.dig(:content)
if content.is_a?(Array)
result =
content.map do |data|
if data[:type] == "tool_use"
call = AnthropicToolCall.new(data[:name], data[:id])
call.append(data[:input].to_json)
call.to_tool_call
else
data[:text]
end
end
elsif parsed[:type] == "message_start"
@input_tokens = parsed.dig(:message, :usage, :input_tokens)
elsif parsed[:type] == "message_delta"
@output_tokens =
parsed.dig(:usage, :output_tokens) || parsed.dig(:delta, :usage, :output_tokens)
elsif parsed[:type] == "message_stop"
# bedrock has this ...
if bedrock_stats = parsed.dig("amazon-bedrock-invocationMetrics".to_sym)
@input_tokens = bedrock_stats[:inputTokenCount] || @input_tokens
@output_tokens = bedrock_stats[:outputTokenCount] || @output_tokens
end
end
else
content = parsed.dig(:content)
if content.is_a?(Array)
tool_call = content.find { |c| c[:type] == "tool_use" }
if tool_call
@tool_calls << AnthropicToolCall.new(tool_call[:name], tool_call[:id])
@tool_calls.last.append(tool_call[:input].to_json)
else
result = parsed.dig(:content, 0, :text).to_s
end
end
@input_tokens = parsed.dig(:usage, :input_tokens)
@output_tokens = parsed.dig(:usage, :output_tokens)
end
@input_tokens = parsed.dig(:usage, :input_tokens)
@output_tokens = parsed.dig(:usage, :output_tokens)
result
end
end
lib/completions/dialects/ollama.rb
+17 -2
View File
@@ -63,8 +63,23 @@ module DiscourseAi
def user_msg(msg)
user_message = { role: "user", content: msg[:content] }
# TODO: Add support for user messages with empbeded user ids
# TODO: Add support for user messages with attachments
encoded_uploads = prompt.encoded_uploads(msg)
if encoded_uploads.present?
images =
encoded_uploads
.map do |upload|
if upload[:mime_type].start_with?("image/")
upload[:base64]
else
nil
end
end
.compact
user_message[:images] = images if images.present?
end
# TODO: Add support for user messages with embedded user ids
user_message
end
lib/completions/endpoints/anthropic.rb
+19 -16
View File
@@ -63,6 +63,10 @@ module DiscourseAi
URI(llm_model.url)
end
def xml_tools_enabled?
!@native_tool_support
end
def prepare_payload(prompt, model_params, dialect)
@native_tool_support = dialect.native_tool_support?
@@ -90,35 +94,34 @@ module DiscourseAi
Net::HTTP::Post.new(model_uri, headers).tap { |r| r.body = payload }
end
def decode_chunk(partial_data)
@decoder ||= JsonStreamDecoder.new
(@decoder << partial_data)
.map { |parsed_json| processor.process_streamed_message(parsed_json) }
.compact
end
def decode(response_data)
processor.process_message(response_data)
end
def processor
@processor ||=
DiscourseAi::Completions::AnthropicMessageProcessor.new(streaming_mode: @streaming_mode)
end
def add_to_function_buffer(function_buffer, partial: nil, payload: nil)
processor.to_xml_tool_calls(function_buffer) if !partial
end
def extract_completion_from(response_raw)
processor.process_message(response_raw)
end
def has_tool?(_response_data)
processor.tool_calls.present?
end
def tool_calls
processor.to_tool_calls
end
def final_log_update(log)
log.request_tokens = processor.input_tokens if processor.input_tokens
log.response_tokens = processor.output_tokens if processor.output_tokens
end
def native_tool_support?
@native_tool_support
end
def partials_from(decoded_chunk)
decoded_chunk.split("\n").map { |line| line.split("data: ", 2)[1] }.compact
end
end
end
end
lib/completions/endpoints/aws_bedrock.rb
+22 -26
View File
@@ -117,7 +117,24 @@ module DiscourseAi
end
end
def decode(chunk)
def decode_chunk(partial_data)
bedrock_decode(partial_data)
.map do |decoded_partial_data|
@raw_response ||= +""
@raw_response << decoded_partial_data
@raw_response << "\n"
parsed_json = JSON.parse(decoded_partial_data, symbolize_names: true)
processor.process_streamed_message(parsed_json)
end
.compact
end
def decode(response_data)
processor.process_message(response_data)
end
def bedrock_decode(chunk)
@decoder ||= Aws::EventStream::Decoder.new
decoded, _done = @decoder.decode_chunk(chunk)
@@ -147,12 +164,13 @@ module DiscourseAi
Aws::EventStream::Errors::MessageChecksumError,
Aws::EventStream::Errors::PreludeChecksumError => e
Rails.logger.error("#{self.class.name}: #{e.message}")
nil
[]
end
def final_log_update(log)
log.request_tokens = processor.input_tokens if processor.input_tokens
log.response_tokens = processor.output_tokens if processor.output_tokens
log.raw_response_payload = @raw_response
end
def processor
@@ -160,30 +178,8 @@ module DiscourseAi
DiscourseAi::Completions::AnthropicMessageProcessor.new(streaming_mode: @streaming_mode)
end
def add_to_function_buffer(function_buffer, partial: nil, payload: nil)
processor.to_xml_tool_calls(function_buffer) if !partial
end
def extract_completion_from(response_raw)
processor.process_message(response_raw)
end
def has_tool?(_response_data)
processor.tool_calls.present?
end
def partials_from(decoded_chunks)
decoded_chunks
end
def native_tool_support?
@native_tool_support
end
def chunk_to_string(chunk)
joined = +chunk.join("\n")
joined << "\n" if joined.length > 0
joined
def xml_tools_enabled?
!@native_tool_support
end
end
end
lib/completions/endpoints/base.rb
+119 -197
View File
File diff suppressed because it is too large Load Diff
lib/completions/endpoints/canned_response.rb
+13 -9
View File
@@ -45,17 +45,21 @@ module DiscourseAi
cancel_fn = lambda { cancelled = true }
# We buffer and return tool invocations in one go.
if is_tool?(response)
yield(response, cancel_fn)
else
response.each_char do |char|
break if cancelled
yield(char, cancel_fn)
as_array = response.is_a?(Array) ? response : [response]
as_array.each do |response|
if is_tool?(response)
yield(response, cancel_fn)
else
response.each_char do |char|
break if cancelled
yield(char, cancel_fn)
end
end
end
else
response
end
response = response.first if response.is_a?(Array) && response.length == 1
response
end
def tokenizer
@@ -65,7 +69,7 @@ module DiscourseAi
private
def is_tool?(response)
Nokogiri::HTML5.fragment(response).at("function_calls").present?
response.is_a?(DiscourseAi::Completions::ToolCall)
end
end
end
lib/completions/endpoints/cohere.rb
+55 -34
View File
@@ -49,6 +49,47 @@ module DiscourseAi
Net::HTTP::Post.new(model_uri, headers).tap { |r| r.body = payload }
end
def decode(response_raw)
rval = []
parsed = JSON.parse(response_raw, symbolize_names: true)
text = parsed[:text]
rval << parsed[:text] if !text.to_s.empty? # also allow " "
# TODO tool calls
update_usage(parsed)
rval
end
def decode_chunk(chunk)
@tool_idx ||= -1
@json_decoder ||= JsonStreamDecoder.new(line_regex: /^\s*({.*})$/)
(@json_decoder << chunk)
.map do |parsed|
update_usage(parsed)
rval = []
rval << parsed[:text] if !parsed[:text].to_s.empty?
if tool_calls = parsed[:tool_calls]
tool_calls&.each do |tool_call|
@tool_idx += 1
tool_name = tool_call[:name]
tool_params = tool_call[:parameters]
tool_id = "tool_#{@tool_idx}"
rval << ToolCall.new(id: tool_id, name: tool_name, parameters: tool_params)
end
end
rval
end
.flatten
.compact
end
def extract_completion_from(response_raw)
parsed = JSON.parse(response_raw, symbolize_names: true)
@@ -77,36 +118,8 @@ module DiscourseAi
end
end
def has_tool?(_ignored)
@has_tool
end
def native_tool_support?
true
end
def add_to_function_buffer(function_buffer, partial: nil, payload: nil)
if partial
tools = JSON.parse(partial)
tools.each do |tool|
name = tool["name"]
parameters = tool["parameters"]
xml_params = parameters.map { |k, v| "<#{k}>#{v}</#{k}>\n" }.join
current_function = function_buffer.at("invoke")
if current_function.nil? || current_function.at("tool_name").content.present?
current_function =
function_buffer.at("function_calls").add_child(
Nokogiri::HTML5::DocumentFragment.parse(noop_function_call_text + "\n"),
)
end
current_function.at("tool_name").content = name == "search_local" ? "search" : name
current_function.at("parameters").children =
Nokogiri::HTML5::DocumentFragment.parse(xml_params)
end
end
function_buffer
def xml_tools_enabled?
false
end
def final_log_update(log)
@@ -114,10 +127,6 @@ module DiscourseAi
log.response_tokens = @output_tokens if @output_tokens
end
def partials_from(decoded_chunk)
decoded_chunk.split("\n").compact
end
def extract_prompt_for_tokenizer(prompt)
text = +""
if prompt[:chat_history]
@@ -131,6 +140,18 @@ module DiscourseAi
text
end
private
def update_usage(parsed)
input_tokens = parsed.dig(:meta, :billed_units, :input_tokens)
input_tokens ||= parsed.dig(:response, :meta, :billed_units, :input_tokens)
@input_tokens = input_tokens if input_tokens.present?
output_tokens = parsed.dig(:meta, :billed_units, :output_tokens)
output_tokens ||= parsed.dig(:response, :meta, :billed_units, :output_tokens)
@output_tokens = output_tokens if output_tokens.present?
end
end
end
end
lib/completions/endpoints/fake.rb
+25 -21
View File
@@ -133,31 +133,35 @@ module DiscourseAi
content = content.shift if content.is_a?(Array)
if block_given?
split_indices = (1...content.length).to_a.sample(self.class.chunk_count - 1).sort
indexes = [0, *split_indices, content.length]
if content.is_a?(DiscourseAi::Completions::ToolCall)
yield(content, -> {})
else
split_indices = (1...content.length).to_a.sample(self.class.chunk_count - 1).sort
indexes = [0, *split_indices, content.length]
original_content = content
content = +""
original_content = content
content = +""
cancel = false
cancel_proc = -> { cancel = true }
cancel = false
cancel_proc = -> { cancel = true }
i = 0
indexes
.each_cons(2)
.map { |start, finish| original_content[start...finish] }
.each do |chunk|
break if cancel
if self.class.delays.present? &&
(delay = self.class.delays[i % self.class.delays.length])
sleep(delay)
i += 1
i = 0
indexes
.each_cons(2)
.map { |start, finish| original_content[start...finish] }
.each do |chunk|
break if cancel
if self.class.delays.present? &&
(delay = self.class.delays[i % self.class.delays.length])
sleep(delay)
i += 1
end
break if cancel
content << chunk
yield(chunk, cancel_proc)
end
break if cancel
content << chunk
yield(chunk, cancel_proc)
end
end
end
content
lib/completions/endpoints/gemini.rb
+76 -40
View File
@@ -103,15 +103,7 @@ module DiscourseAi
end
end
def partials_from(decoded_chunk)
decoded_chunk
end
def chunk_to_string(chunk)
chunk.to_s
end
class Decoder
class GeminiStreamingDecoder
def initialize
@buffer = +""
end
@@ -151,43 +143,87 @@ module DiscourseAi
end
def decode(chunk)
@decoder ||= Decoder.new
@decoder.decode(chunk)
json = JSON.parse(chunk, symbolize_names: true)
idx = -1
json
.dig(:candidates, 0, :content, :parts)
.map do |part|
if part[:functionCall]
idx += 1
ToolCall.new(
id: "tool_#{idx}",
name: part[:functionCall][:name],
parameters: part[:functionCall][:args],
)
else
part = part[:text]
if part != ""
part
else
nil
end
end
end
end
def decode_chunk(chunk)
@tool_index ||= -1
streaming_decoder
.decode(chunk)
.map do |parsed|
update_usage(parsed)
parsed
.dig(:candidates, 0, :content, :parts)
.map do |part|
if part[:text]
part = part[:text]
if part != ""
part
else
nil
end
elsif part[:functionCall]
@tool_index += 1
ToolCall.new(
id: "tool_#{@tool_index}",
name: part[:functionCall][:name],
parameters: part[:functionCall][:args],
)
end
end
end
.flatten
.compact
end
def update_usage(parsed)
usage = parsed.dig(:usageMetadata)
if usage
if prompt_token_count = usage[:promptTokenCount]
@prompt_token_count = prompt_token_count
end
if candidate_token_count = usage[:candidatesTokenCount]
@candidate_token_count = candidate_token_count
end
end
end
def final_log_update(log)
log.request_tokens = @prompt_token_count if @prompt_token_count
log.response_tokens = @candidate_token_count if @candidate_token_count
end
def streaming_decoder
@decoder ||= GeminiStreamingDecoder.new
end
def extract_prompt_for_tokenizer(prompt)
prompt.to_s
end
def has_tool?(_response_data)
@has_function_call
end
def native_tool_support?
true
end
def add_to_function_buffer(function_buffer, payload: nil, partial: nil)
if @streaming_mode
return function_buffer if !partial
else
partial = payload
end
function_buffer.at("tool_name").content = partial[:name] if partial[:name].present?
if partial[:args]
argument_fragments =
partial[:args].reduce(+"") do |memo, (arg_name, value)|
memo << "\n<#{arg_name}>#{CGI.escapeHTML(value.to_s)}</#{arg_name}>"
end
argument_fragments << "\n"
function_buffer.at("parameters").children =
Nokogiri::HTML5::DocumentFragment.parse(argument_fragments)
end
function_buffer
def xml_tools_enabled?
false
end
end
end
lib/completions/endpoints/hugging_face.rb
+22 -14
View File
@@ -59,22 +59,30 @@ module DiscourseAi
Net::HTTP::Post.new(model_uri, headers).tap { |r| r.body = payload }
end
def extract_completion_from(response_raw)
parsed = JSON.parse(response_raw, symbolize_names: true).dig(:choices, 0)
# half a line sent here
return if !parsed
response_h = @streaming_mode ? parsed.dig(:delta) : parsed.dig(:message)
response_h.dig(:content)
def xml_tools_enabled?
true
end
def partials_from(decoded_chunk)
decoded_chunk
.split("\n")
.map do |line|
data = line.split("data:", 2)[1]
data&.squish == "[DONE]" ? nil : data
def decode(response_raw)
parsed = JSON.parse(response_raw, symbolize_names: true)
text = parsed.dig(:choices, 0, :message, :content)
if text.to_s.empty?
[""]
else
[text]
end
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
lib/completions/endpoints/ollama.rb
+20 -68
View File
@@ -37,12 +37,8 @@ module DiscourseAi
URI(llm_model.url)
end
def native_tool_support?
@native_tool_support
end
def has_tool?(_response_data)
@has_function_call
def xml_tools_enabled?
!@native_tool_support
end
def prepare_payload(prompt, model_params, dialect)
@@ -67,74 +63,30 @@ module DiscourseAi
Net::HTTP::Post.new(model_uri, headers).tap { |r| r.body = payload }
end
def partials_from(decoded_chunk)
decoded_chunk.split("\n").compact
def decode_chunk(chunk)
# Native tool calls are not working right in streaming mode, use XML
@json_decoder ||= JsonStreamDecoder.new(line_regex: /^\s*({.*})$/)
(@json_decoder << chunk).map { |parsed| parsed.dig(:message, :content) }.compact
end
def extract_completion_from(response_raw)
def decode(response_raw)
rval = []
parsed = JSON.parse(response_raw, symbolize_names: true)
return if !parsed
content = parsed.dig(:message, :content)
rval << content if !content.to_s.empty?
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"),
)
idx = -1
parsed
.dig(:message, :tool_calls)
&.each do |tool_call|
idx += 1
id = "tool_#{idx}"
name = tool_call.dig(:function, :name)
args = tool_call.dig(:function, :arguments)
rval << ToolCall.new(id: id, name: name, parameters: args)
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
rval
end
end
end
lib/completions/endpoints/open_ai.rb
+16 -80
View File
@@ -93,98 +93,34 @@ module DiscourseAi
end
def final_log_update(log)
log.request_tokens = @prompt_tokens if @prompt_tokens
log.response_tokens = @completion_tokens if @completion_tokens
log.request_tokens = processor.prompt_tokens if processor.prompt_tokens
log.response_tokens = processor.completion_tokens if processor.completion_tokens
end
def extract_completion_from(response_raw)
json = JSON.parse(response_raw, symbolize_names: true)
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
response_h = @streaming_mode ? parsed.dig(:delta) : 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)
def decode(response_raw)
processor.process_message(JSON.parse(response_raw, symbolize_names: true))
end
def partials_from(decoded_chunk)
decoded_chunk
.split("\n")
.map do |line|
data = line.split("data: ", 2)[1]
data == "[DONE]" ? nil : data
end
def decode_chunk(chunk)
@decoder ||= JsonStreamDecoder.new
(@decoder << chunk)
.map { |parsed_json| processor.process_streamed_message(parsed_json) }
.flatten
.compact
end
def has_tool?(_response_data)
@has_function_call
def decode_chunk_finish
@processor.finish
end
def native_tool_support?
true
def xml_tools_enabled?
false
end
def add_to_function_buffer(function_buffer, partial: nil, payload: nil)
if @streaming_mode
return function_buffer if !partial
else
partial = payload
end
private
@args_buffer ||= +""
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
begin
json_args = JSON.parse(@args_buffer, symbolize_names: true)
argument_fragments =
json_args.reduce(+"") do |memo, (arg_name, value)|
memo << "\n<#{arg_name}>#{CGI.escapeHTML(value.to_s)}</#{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
def processor
@processor ||= OpenAiMessageProcessor.new
end
end
end
lib/completions/endpoints/samba_nova.rb
+22 -18
View File
@@ -55,27 +55,31 @@ module DiscourseAi
log.response_tokens = @completion_tokens if @completion_tokens
end
def extract_completion_from(response_raw)
json = JSON.parse(response_raw, symbolize_names: true)
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)
def xml_tools_enabled?
true
end
def partials_from(decoded_chunk)
decoded_chunk
.split("\n")
.map do |line|
data = line.split("data: ", 2)[1]
data == "[DONE]" ? nil : data
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 |json|
text = json.dig(:choices, 0, :delta, :content)
@prompt_tokens ||= json.dig(:usage, :prompt_tokens)
@completion_tokens ||= json.dig(:usage, :completion_tokens)
if !text.to_s.empty?
text
else
nil
end
end
.flatten
.compact
end
end
lib/completions/endpoints/vllm.rb
+36 -15
View File
@@ -42,7 +42,10 @@ module DiscourseAi
def prepare_payload(prompt, model_params, dialect)
payload = default_options.merge(model_params).merge(messages: prompt)
payload[:stream] = true if @streaming_mode
if @streaming_mode
payload[:stream] = true if @streaming_mode
payload[:stream_options] = { include_usage: true }
end
payload
end
@@ -56,24 +59,42 @@ module DiscourseAi
Net::HTTP::Post.new(model_uri, headers).tap { |r| r.body = payload }
end
def partials_from(decoded_chunk)
decoded_chunk
.split("\n")
.map do |line|
data = line.split("data: ", 2)[1]
data == "[DONE]" ? nil : data
end
.compact
def xml_tools_enabled?
true
end
def extract_completion_from(response_raw)
parsed = JSON.parse(response_raw, symbolize_names: true).dig(:choices, 0)
# half a line sent here
return if !parsed
def final_log_update(log)
log.request_tokens = @prompt_tokens if @prompt_tokens
log.response_tokens = @completion_tokens if @completion_tokens
end
response_h = @streaming_mode ? parsed.dig(:delta) : parsed.dig(:message)
def decode(response_raw)
json = JSON.parse(response_raw, symbolize_names: true)
@prompt_tokens = json.dig(:usage, :prompt_tokens)
@completion_tokens = json.dig(:usage, :completion_tokens)
[json.dig(:choices, 0, :message, :content)]
end
response_h.dig(:content)
def decode_chunk(chunk)
@json_decoder ||= JsonStreamDecoder.new
(@json_decoder << chunk)
.map do |parsed|
# vLLM keeps sending usage over and over again
prompt_tokens = parsed.dig(:usage, :prompt_tokens)
completion_tokens = parsed.dig(:usage, :completion_tokens)
@prompt_tokens = prompt_tokens if prompt_tokens
@completion_tokens = completion_tokens if completion_tokens
text = parsed.dig(:choices, 0, :delta, :content)
if text.to_s.empty?
nil
else
text
end
end
.compact
end
end
end
lib/completions/function_call_normalizer.rb
-113
View File
@@ -1,113 +0,0 @@
# frozen_string_literal: true
class DiscourseAi::Completions::FunctionCallNormalizer
attr_reader :done
# blk is the block to call with filtered data
def initialize(blk, cancel)
@blk = blk
@cancel = cancel
@done = false
@in_tool = false
@buffer = +""
@function_buffer = +""
end
def self.normalize(data)
text = +""
cancel = -> {}
blk = ->(partial, _) { text << partial }
normalizer = self.new(blk, cancel)
normalizer << data
[text, normalizer.function_calls]
end
def function_calls
return nil if @function_buffer.blank?
xml = Nokogiri::HTML5.fragment(@function_buffer)
self.class.normalize_function_ids!(xml)
last_invoke = xml.at("invoke:last")
if last_invoke
last_invoke.next_sibling.remove while last_invoke.next_sibling
xml.at("invoke:last").add_next_sibling("\n") if !last_invoke.next_sibling
end
xml.at("function_calls").to_s.dup.force_encoding("UTF-8")
end
def <<(text)
@buffer << text
if !@in_tool
# double check if we are clearly in a tool
search_length = text.length + 20
search_string = @buffer[-search_length..-1] || @buffer
index = search_string.rindex("<function_calls>")
@in_tool = !!index
if @in_tool
@function_buffer = @buffer[index..-1]
text_index = text.rindex("<function_calls>")
@blk.call(text[0..text_index - 1].strip, @cancel) if text_index && text_index > 0
end
else
@function_buffer << text
end
if !@in_tool
if maybe_has_tool?(@buffer)
split_index = text.rindex("<").to_i - 1
if split_index >= 0
@function_buffer = text[split_index + 1..-1] || ""
text = text[0..split_index] || ""
else
@function_buffer << text
text = ""
end
else
if @function_buffer.length > 0
@blk.call(@function_buffer, @cancel)
@function_buffer = +""
end
end
@blk.call(text, @cancel) if text.length > 0
else
if text.include?("</function_calls>")
@done = true
@cancel.call
end
end
end
def self.normalize_function_ids!(function_buffer)
function_buffer
.css("invoke")
.each_with_index do |invoke, index|
if invoke.at("tool_id")
invoke.at("tool_id").content = "tool_#{index}" if invoke.at("tool_id").content.blank?
else
invoke.add_child("<tool_id>tool_#{index}</tool_id>\n") if !invoke.at("tool_id")
end
end
end
private
def maybe_has_tool?(text)
# 16 is the length of function calls
substring = text[-16..-1] || text
split = substring.split("<")
if split.length > 1
match = "<" + split.last
"<function_calls>".start_with?(match)
else
substring.ends_with?("<")
end
end
end
lib/completions/json_stream_decoder.rb
+48
View File
@@ -0,0 +1,48 @@
# frozen_string_literal: true
module DiscourseAi
module Completions
# will work for anthropic and open ai compatible
class JsonStreamDecoder
attr_reader :buffer
LINE_REGEX = /data: ({.*})\s*$/
def initialize(symbolize_keys: true, line_regex: LINE_REGEX)
@symbolize_keys = symbolize_keys
@buffer = +""
@line_regex = line_regex
end
def <<(raw)
@buffer << raw.to_s
rval = []
split = @buffer.scan(/.*\n?/)
split.pop if split.last.blank?
@buffer = +(split.pop.to_s)
split.each do |line|
matches = line.match(@line_regex)
next if !matches
rval << JSON.parse(matches[1], symbolize_names: @symbolize_keys)
end
if @buffer.present?
matches = @buffer.match(@line_regex)
if matches
begin
rval << JSON.parse(matches[1], symbolize_names: @symbolize_keys)
@buffer = +""
rescue JSON::ParserError
# maybe it is a partial line
end
end
end
rval
end
end
end
end
lib/completions/open_ai_message_processor.rb
+103
View File
@@ -0,0 +1,103 @@
# frozen_string_literal: true
module DiscourseAi::Completions
class OpenAiMessageProcessor
attr_reader :prompt_tokens, :completion_tokens
def initialize
@tool = nil
@tool_arguments = +""
@prompt_tokens = nil
@completion_tokens = nil
end
def process_message(json)
result = []
tool_calls = json.dig(:choices, 0, :message, :tool_calls)
message = json.dig(:choices, 0, :message, :content)
result << message if message.present?
if tool_calls.present?
tool_calls.each do |tool_call|
id = tool_call.dig(:id)
name = tool_call.dig(:function, :name)
arguments = tool_call.dig(:function, :arguments)
parameters = arguments.present? ? JSON.parse(arguments, symbolize_names: true) : {}
result << ToolCall.new(id: id, name: name, parameters: parameters)
end
end
update_usage(json)
result
end
def process_streamed_message(json)
rval = nil
tool_calls = json.dig(:choices, 0, :delta, :tool_calls)
content = json.dig(:choices, 0, :delta, :content)
finished_tools = json.dig(:choices, 0, :finish_reason) || tool_calls == []
if tool_calls.present?
id = tool_calls.dig(0, :id)
name = tool_calls.dig(0, :function, :name)
arguments = tool_calls.dig(0, :function, :arguments)
# TODO: multiple tool support may require index
#index = tool_calls[0].dig(:index)
if id.present? && @tool && @tool.id != id
process_arguments
rval = @tool
@tool = nil
end
if id.present? && name.present?
@tool_arguments = +""
@tool = ToolCall.new(id: id, name: name)
end
@tool_arguments << arguments.to_s
elsif finished_tools && @tool
parsed_args = JSON.parse(@tool_arguments, symbolize_names: true)
@tool.parameters = parsed_args
rval = @tool
@tool = nil
elsif content.present?
rval = content
end
update_usage(json)
rval
end
def finish
rval = []
if @tool
process_arguments
rval << @tool
@tool = nil
end
rval
end
private
def process_arguments
if @tool_arguments.present?
parsed_args = JSON.parse(@tool_arguments, symbolize_names: true)
@tool.parameters = parsed_args
@tool_arguments = nil
end
end
def update_usage(json)
@prompt_tokens ||= json.dig(:usage, :prompt_tokens)
@completion_tokens ||= json.dig(:usage, :completion_tokens)
end
end
end
lib/completions/tool_call.rb
+29
View File
@@ -0,0 +1,29 @@
# frozen_string_literal: true
module DiscourseAi
module Completions
class ToolCall
attr_reader :id, :name, :parameters
def initialize(id:, name:, parameters: nil)
@id = id
@name = name
self.parameters = parameters if parameters
@parameters ||= {}
end
def parameters=(parameters)
raise ArgumentError, "parameters must be a hash" unless parameters.is_a?(Hash)
@parameters = parameters.symbolize_keys
end
def ==(other)
id == other.id && name == other.name && parameters == other.parameters
end
def to_s
"#{name} - #{id} (\n#{parameters.map(&:to_s).join("\n")}\n)"
end
end
end
end
lib/completions/xml_tool_processor.rb
+124
View File
@@ -0,0 +1,124 @@
# frozen_string_literal: true
# This class can be used to process a stream of text that may contain XML tool
# calls.
# It will return either text or ToolCall objects.
module DiscourseAi
module Completions
class XmlToolProcessor
def initialize
@buffer = +""
@function_buffer = +""
@should_cancel = false
@in_tool = false
end
def <<(text)
@buffer << text
result = []
if !@in_tool
# double check if we are clearly in a tool
search_length = text.length + 20
search_string = @buffer[-search_length..-1] || @buffer
index = search_string.rindex("<function_calls>")
@in_tool = !!index
if @in_tool
@function_buffer = @buffer[index..-1]
text_index = text.rindex("<function_calls>")
result << text[0..text_index - 1].strip if text_index && text_index > 0
end
else
@function_buffer << text
end
if !@in_tool
if maybe_has_tool?(@buffer)
split_index = text.rindex("<").to_i - 1
if split_index >= 0
@function_buffer = text[split_index + 1..-1] || ""
text = text[0..split_index] || ""
else
@function_buffer << text
text = ""
end
else
if @function_buffer.length > 0
result << @function_buffer
@function_buffer = +""
end
end
result << text if text.length > 0
else
@should_cancel = true if text.include?("</function_calls>")
end
result
end
def finish
return [] if @function_buffer.blank?
xml = Nokogiri::HTML5.fragment(@function_buffer)
normalize_function_ids!(xml)
last_invoke = xml.at("invoke:last")
if last_invoke
last_invoke.next_sibling.remove while last_invoke.next_sibling
xml.at("invoke:last").add_next_sibling("\n") if !last_invoke.next_sibling
end
xml
.css("invoke")
.map do |invoke|
tool_name = invoke.at("tool_name").content.force_encoding("UTF-8")
tool_id = invoke.at("tool_id").content.force_encoding("UTF-8")
parameters = {}
invoke
.at("parameters")
&.children
&.each do |node|
next if node.text?
name = node.name
value = node.content.to_s
parameters[name.to_sym] = value.to_s.force_encoding("UTF-8")
end
ToolCall.new(id: tool_id, name: tool_name, parameters: parameters)
end
end
def should_cancel?
@should_cancel
end
private
def normalize_function_ids!(function_buffer)
function_buffer
.css("invoke")
.each_with_index do |invoke, index|
if invoke.at("tool_id")
invoke.at("tool_id").content = "tool_#{index}" if invoke.at("tool_id").content.blank?
else
invoke.add_child("<tool_id>tool_#{index}</tool_id>\n") if !invoke.at("tool_id")
end
end
end
def maybe_has_tool?(text)
# 16 is the length of function calls
substring = text[-16..-1] || text
split = substring.split("<")
if split.length > 1
match = "<" + split.last
"<function_calls>".start_with?(match)
else
substring.ends_with?("<")
end
end
end
end
end
spec/lib/completions/endpoints/anthropic_spec.rb
+26 -22
View File
@@ -104,7 +104,7 @@ RSpec.describe DiscourseAi::Completions::Endpoints::Anthropic do
data: {"type":"message_stop"}
STRING
result = +""
result = []
body = body.scan(/.*\n/)
EndpointMock.with_chunk_array_support do
stub_request(:post, url).to_return(status: 200, body: body)
@@ -114,18 +114,17 @@ RSpec.describe DiscourseAi::Completions::Endpoints::Anthropic do
end
end
expected = (<<~TEXT).strip
<function_calls>
<invoke>
<tool_name>search</tool_name>
<parameters><search_query>s&lt;a&gt;m sam</search_query>
<category>general</category></parameters>
<tool_id>toolu_01DjrShFRRHp9SnHYRFRc53F</tool_id>
</invoke>
</function_calls>
TEXT
tool_call =
DiscourseAi::Completions::ToolCall.new(
name: "search",
id: "toolu_01DjrShFRRHp9SnHYRFRc53F",
parameters: {
search_query: "s<a>m sam",
category: "general",
},
)
expect(result.strip).to eq(expected)
expect(result).to eq([tool_call])
end
it "can stream a response" do
@@ -191,6 +190,8 @@ RSpec.describe DiscourseAi::Completions::Endpoints::Anthropic do
expect(log.feature_name).to eq("testing")
expect(log.response_tokens).to eq(15)
expect(log.request_tokens).to eq(25)
expect(log.raw_request_payload).to eq(expected_body.to_json)
expect(log.raw_response_payload.strip).to eq(body.strip)
end
it "supports non streaming tool calls" do
@@ -242,17 +243,20 @@ RSpec.describe DiscourseAi::Completions::Endpoints::Anthropic do
result = llm.generate(prompt, user: Discourse.system_user)
expected = <<~TEXT.strip
<function_calls>
<invoke>
<tool_name>calculate</tool_name>
<parameters><expression>2758975 + 21.11</expression></parameters>
<tool_id>toolu_012kBdhG4eHaV68W56p4N94h</tool_id>
</invoke>
</function_calls>
TEXT
tool_call =
DiscourseAi::Completions::ToolCall.new(
name: "calculate",
id: "toolu_012kBdhG4eHaV68W56p4N94h",
parameters: {
expression: "2758975 + 21.11",
},
)
expect(result.strip).to eq(expected)
expect(result).to eq(["Here is the calculation:", tool_call])
log = AiApiAuditLog.order(:id).last
expect(log.request_tokens).to eq(345)
expect(log.response_tokens).to eq(65)
end
it "can send images via a completion prompt" do
spec/lib/completions/endpoints/aws_bedrock_spec.rb
+23 -26
View File
@@ -79,7 +79,7 @@ RSpec.describe DiscourseAi::Completions::Endpoints::AwsBedrock do
}
prompt.tools = [tool]
response = +""
response = []
proxy.generate(prompt, user: user) { |partial| response << partial }
expect(request.headers["Authorization"]).to be_present
@@ -90,21 +90,18 @@ RSpec.describe DiscourseAi::Completions::Endpoints::AwsBedrock do
expect(parsed_body["tools"]).to eq(nil)
expect(parsed_body["stop_sequences"]).to eq(["</function_calls>"])
# note we now have a tool_id cause we were normalized
function_call = <<~XML.strip
hello
expected = [
"hello\n",
DiscourseAi::Completions::ToolCall.new(
id: "tool_0",
name: "google",
parameters: {
query: "sydney weather today",
},
),
]
<function_calls>
<invoke>
<tool_name>google</tool_name>
<parameters><query>sydney weather today</query></parameters>
<tool_id>tool_0</tool_id>
</invoke>
</function_calls>
XML
expect(response.strip).to eq(function_call)
expect(response).to eq(expected)
end
end
@@ -230,23 +227,23 @@ RSpec.describe DiscourseAi::Completions::Endpoints::AwsBedrock do
}
prompt.tools = [tool]
response = +""
response = []
proxy.generate(prompt, user: user) { |partial| response << partial }
expect(request.headers["Authorization"]).to be_present
expect(request.headers["X-Amz-Content-Sha256"]).to be_present
expected_response = (<<~RESPONSE).strip
<function_calls>
<invoke>
<tool_name>google</tool_name>
<parameters><query>sydney weather today</query></parameters>
<tool_id>toolu_bdrk_014CMjxtGmKUtGoEFPgc7PF7</tool_id>
</invoke>
</function_calls>
RESPONSE
expected_response = [
DiscourseAi::Completions::ToolCall.new(
id: "toolu_bdrk_014CMjxtGmKUtGoEFPgc7PF7",
name: "google",
parameters: {
query: "sydney weather today",
},
),
]
expect(response.strip).to eq(expected_response)
expect(response).to eq(expected_response)
expected = {
"max_tokens" => 3000,
spec/lib/completions/endpoints/cohere_spec.rb
+12 -13
View File
@@ -66,7 +66,7 @@ RSpec.describe DiscourseAi::Completions::Endpoints::Cohere do
TEXT
parsed_body = nil
result = +""
result = []
sig = {
name: "google",
@@ -91,21 +91,20 @@ RSpec.describe DiscourseAi::Completions::Endpoints::Cohere do
},
).to_return(status: 200, body: body.split("|"))
result = llm.generate(prompt, user: user) { |partial, cancel| result << partial }
llm.generate(prompt, user: user) { |partial, cancel| result << partial }
end
expected = <<~TEXT
<function_calls>
<invoke>
<tool_name>google</tool_name>
<parameters><query>who is sam saffron</query>
</parameters>
<tool_id>tool_0</tool_id>
</invoke>
</function_calls>
TEXT
text = "I will search for 'who is sam saffron' and relay the information to the user."
tool_call =
DiscourseAi::Completions::ToolCall.new(
id: "tool_0",
name: "google",
parameters: {
query: "who is sam saffron",
},
)
expect(result.strip).to eq(expected.strip)
expect(result).to eq([text, tool_call])
expected = {
model: "command-r-plus",

Some files were not shown because too many files have changed in this diff Show More