handle thinking in tool implementation

Author: SamSaffron

lib/ai_bot/bot.rb

@@ -112,6 +112,7 @@ def reply(context, &update_blk)
 
           allow_partial_tool_calls = persona.allow_partial_tool_calls?
           existing_tools = Set.new
+          current_thinking = []
 
           result =
             llm.generate(
@@ -149,7 +150,17 @@ def reply(context, &update_blk)
                   needs_newlines = false
                 end
 
-                process_tool(tool, raw_context, llm, cancel, update_blk, prompt, context)
+                process_tool(
+                  tool: tool,
+                  raw_context: raw_context,
+                  llm: llm,
+                  cancel: cancel,
+                  update_blk: update_blk,
+                  prompt: prompt,
+                  context: context,
+                  current_thinking: current_thinking,
+                )
+
                 tools_ran += 1
                 ongoing_chain &&= tool.chain_next_response?
 
@@ -167,6 +178,7 @@ def reply(context, &update_blk)
                     if !partial.partial?
                       # this will be dealt with later
                       raw_context << partial
+                      current_thinking << partial
                     end
                   else
                     update_blk.call(partial, cancel)
@@ -222,7 +234,16 @@ def embed_thinking(raw_context)
         embedded_thinking
       end
 
-      def process_tool(tool, raw_context, llm, cancel, update_blk, prompt, context)
+      def process_tool(
+        tool:,
+        raw_context:,
+        llm:,
+        cancel:,
+        update_blk:,
+        prompt:,
+        context:,
+        current_thinking:
+      )
         tool_call_id = tool.tool_call_id
         invocation_result_json = invoke_tool(tool, llm, cancel, context, &update_blk).to_json
 
@@ -233,6 +254,17 @@ def process_tool(tool, raw_context, llm, cancel, update_blk, prompt, context)
           name: tool.name,
         }
 
+        if current_thinking.present?
+          current_thinking.each do |thinking|
+            if thinking.redacted
+              tool_call_message[:redacted_thinking_signature] = thinking.signature
+            else
+              tool_call_message[:thinking] = thinking.message
+              tool_call_message[:thinking_signature] = thinking.signature
+            end
+          end
+        end
+
         tool_message = {
           type: :tool,
           id: tool_call_id,

lib/completions/dialects/claude_tools.rb

@@ -45,15 +45,35 @@ def instructions
 
         def from_raw_tool_call(raw_message)
           call_details = JSON.parse(raw_message[:content], symbolize_names: true)
+          result = []
+
+          if raw_message[:thinking] || raw_message[:redacted_thinking_signature]
+            if raw_message[:thinking]
+              result << {
+                type: "thinking",
+                thinking: raw_message[:thinking],
+                signature: raw_message[:thinking_signature],
+              }
+            end
+
+            if raw_message[:redacted_thinking_signature]
+              result << {
+                type: "redacted_thinking",
+                data: raw_message[:redacted_thinking_signature],
+              }
+            end
+          end
+
           tool_call_id = raw_message[:id]
-          [
-            {
-              type: "tool_use",
-              id: tool_call_id,
-              name: raw_message[:name],
-              input: call_details[:arguments],
-            },
-          ]
+
+          result << {
+            type: "tool_use",
+            id: tool_call_id,
+            name: raw_message[:name],
+            input: call_details[:arguments],
+          }
+
+          result
         end
 
         def from_raw_tool(raw_message)