From 38af2ca63ec2d8a0a6f8415455e0e3a0765e05ba Mon Sep 17 00:00:00 2001
From: Sam <sam.saffron@gmail.com>
Date: Tue, 5 Sep 2023 10:37:58 +1000
Subject: [PATCH] FIX: cut completion short after function call is found (#182)

Previous to this change we would keep completing and throw away
result
---
 lib/modules/ai_bot/bot.rb                     | 14 +++++++++
 spec/lib/modules/ai_bot/anthropic_bot_spec.rb | 30 +++++++++++++++++++
 2 files changed, 44 insertions(+)

diff --git a/lib/modules/ai_bot/bot.rb b/lib/modules/ai_bot/bot.rb
index cab1bb6e..b62cc702 100644
--- a/lib/modules/ai_bot/bot.rb
+++ b/lib/modules/ai_bot/bot.rb
@@ -8,6 +8,7 @@ module DiscourseAi
           @functions = []
           @current_function = nil
           @found = false
+          @cancel_completion = false
         end
 
         def found?
@@ -18,6 +19,14 @@ module DiscourseAi
           @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
@@ -126,6 +135,8 @@ module DiscourseAi
               functions: functions,
               done: false,
             )
+
+            cancel&.call if functions.cancel_completion?
           end
 
           reply << current_delta if !functions.found?
@@ -315,6 +326,9 @@ module DiscourseAi
 
       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
diff --git a/spec/lib/modules/ai_bot/anthropic_bot_spec.rb b/spec/lib/modules/ai_bot/anthropic_bot_spec.rb
index 075d5ff6..73d4838e 100644
--- a/spec/lib/modules/ai_bot/anthropic_bot_spec.rb
+++ b/spec/lib/modules/ai_bot/anthropic_bot_spec.rb
@@ -21,6 +21,36 @@ module ::DiscourseAi
       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