Skip to content

Commit

Permalink
FIX: cut completion short after function call is found (#182)
Browse files Browse the repository at this point in the history
Previous to this change we would keep completing and throw away
result
  • Loading branch information
SamSaffron committed Sep 5, 2023
1 parent 86591cf commit 38af2ca
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 0 deletions.
14 changes: 14 additions & 0 deletions lib/modules/ai_bot/bot.rb
Expand Up @@ -8,6 +8,7 @@ def initialize
@functions = []
@current_function = nil
@found = false
@cancel_completion = false
end

def found?
Expand All @@ -18,6 +19,14 @@ def found!
@found = true
end

def cancel_completion?
@cancel_completion
end

def cancel_completion!
@cancel_completion = true
end

def add_function(name)
@current_function = { name: name, arguments: +"" }
@functions << @current_function
Expand Down Expand Up @@ -126,6 +135,8 @@ def reply_to(
functions: functions,
done: false,
)

cancel&.call if functions.cancel_completion?
end

reply << current_delta if !functions.found?
Expand Down Expand Up @@ -315,6 +326,9 @@ def get_delta(partial, context)

def populate_functions(partial:, reply:, functions:, done:)
if !done
if functions.found?
functions.cancel_completion! if reply.split("\n")[-1].match?(/^\s*[^!]+/)
end
functions.found! if reply.match?(/^!/i)
else
reply
Expand Down
30 changes: 30 additions & 0 deletions spec/lib/modules/ai_bot/anthropic_bot_spec.rb
Expand Up @@ -21,6 +21,36 @@ def bot_user
end

describe "parsing a reply prompt" do
it "can correctly predict that a completion needs to be cancelled" do
SiteSetting.ai_bot_enabled_chat_commands = "read|search"
functions = DiscourseAi::AiBot::Bot::FunctionCalls.new

# note anthropic API has a silly leading space, we need to make sure we can handle that
prompt = +<<~REPLY.strip
hello world
!search(search_query: "hello world", random_stuff: 77)
!search(search_query: "hello world 2", random_stuff: 77
REPLY

bot.populate_functions(partial: nil, reply: prompt, functions: functions, done: false)

expect(functions.found?).to eq(true)
expect(functions.cancel_completion?).to eq(false)

prompt << ")\n"

bot.populate_functions(partial: nil, reply: prompt, functions: functions, done: false)

expect(functions.found?).to eq(true)
expect(functions.cancel_completion?).to eq(false)

prompt << "a test test"

bot.populate_functions(partial: nil, reply: prompt, functions: functions, done: false)

expect(functions.cancel_completion?).to eq(true)
end

it "can correctly detect commands from a prompt" do
SiteSetting.ai_bot_enabled_chat_commands = "read|search"
functions = DiscourseAi::AiBot::Bot::FunctionCalls.new
Expand Down

0 comments on commit 38af2ca

Please sign in to comment.