code abstraction for tool calling

samolego · samolego · commit 623691885c60 · 2025-05-08T14:13:45.000+02:00
diff --git a/tools/server/webui/src/components/ChatMessage.tsx b/tools/server/webui/src/components/ChatMessage.tsx
@@ -9,6 +9,8 @@ interface SplitMessage {
   content: PendingMessage['content'];
   thought?: string;
   isThinking?: boolean;
+  toolOutput?: string;
+  toolTitle?: string;
 }
 
 export default function ChatMessage({
@@ -48,32 +50,71 @@ export default function ChatMessage({
   const nextSibling = siblingLeafNodeIds[siblingCurrIdx + 1];
   const prevSibling = siblingLeafNodeIds[siblingCurrIdx - 1];
 
-  // for reasoning model, we split the message into content and thought
-  // TODO: implement this as remark/rehype plugin in the future
-  const { content, thought, isThinking }: SplitMessage = useMemo(() => {
-    if (msg.content === null || msg.role !== 'assistant') {
-      return { content: msg.content };
-    }
-    let actualContent = '';
-    let thought = '';
-    let isThinking = false;
-    let thinkSplit = msg.content.split('<think>', 2);
-    actualContent += thinkSplit[0];
-    while (thinkSplit[1] !== undefined) {
-      // <think> tag found
-      thinkSplit = thinkSplit[1].split('</think>', 2);
-      thought += thinkSplit[0];
-      isThinking = true;
-      if (thinkSplit[1] !== undefined) {
-        // </think> closing tag found
-        isThinking = false;
-        thinkSplit = thinkSplit[1].split('<think>', 2);
-        actualContent += thinkSplit[0];
+  // for reasoning model, we split the message into content, thought, and tool output
+  const { content, thought, isThinking, toolOutput, toolTitle }: SplitMessage =
+    useMemo(() => {
+      if (msg.content === null || msg.role !== 'assistant') {
+        return { content: msg.content };
       }
-    }
-    return { content: actualContent, thought, isThinking };
-  }, [msg]);
+      let currentContent = msg.content;
+      let extractedThought: string | undefined = undefined;
+      let isCurrentlyThinking = false;
+      let extractedToolOutput: string | undefined = undefined;
+      let extractedToolTitle: string | undefined = 'Tool Output';
 
+      // Process <think> tags
+      const thinkParts = currentContent.split('<think>');
+      currentContent = thinkParts[0];
+      if (thinkParts.length > 1) {
+        isCurrentlyThinking = true;
+        const tempThoughtArray: string[] = [];
+        for (let i = 1; i < thinkParts.length; i++) {
+          const thinkSegment = thinkParts[i].split('</think>');
+          tempThoughtArray.push(thinkSegment[0]);
+          if (thinkSegment.length > 1) {
+            isCurrentlyThinking = false; // Closing tag found
+            currentContent += thinkSegment[1];
+          }
+        }
+        extractedThought = tempThoughtArray.join('\n');
+      }
+
+      // Process <tool> tags (after thoughts are processed)
+      const toolParts = currentContent.split('<tool>');
+      console.log(toolParts);
+      currentContent = toolParts[0];
+      if (toolParts.length > 1) {
+        const tempToolOutputArray: string[] = [];
+        for (let i = 1; i < toolParts.length; i++) {
+          const toolSegment = toolParts[i].split('</tool>');
+          const toolContent = toolSegment[0].trim();
+
+          const firstLineEnd = toolContent.indexOf('\n');
+          if (firstLineEnd !== -1) {
+            extractedToolTitle = toolContent.substring(0, firstLineEnd);
+            tempToolOutputArray.push(
+              toolContent.substring(firstLineEnd + 1).trim()
+            );
+          } else {
+            // If no newline, extractedToolTitle keeps its default; toolContent is pushed as is.
+            tempToolOutputArray.push(toolContent);
+          }
+
+          if (toolSegment.length > 1) {
+            currentContent += toolSegment[1];
+          }
+        }
+        extractedToolOutput = tempToolOutputArray.join('\n\n');
+      }
+
+      return {
+        content: currentContent.trim(),
+        thought: extractedThought,
+        isThinking: isCurrentlyThinking,
+        toolOutput: extractedToolOutput,
+        toolTitle: extractedToolTitle,
+      };
+    }, [msg]);
   if (!viewingChat) return null;
 
   return (
@@ -192,6 +233,24 @@ export default function ChatMessage({
                       content={content}
                       isGenerating={isPending}
                     />
+
+                    {toolOutput && (
+                      <details
+                        className="collapse bg-base-200 collapse-arrow mb-4"
+                        open={true} // todo: make this configurable like showThoughtInProgress
+                      >
+                        <summary className="collapse-title">
+                          <b>{toolTitle || 'Tool Output'}</b>
+                        </summary>
+                        <div className="collapse-content">
+                          <MarkdownDisplay
+                            content={toolOutput}
+                            // Tool output is not "generating" in the same way
+                            isGenerating={false}
+                          />
+                        </div>
+                      </details>
+                    )}
                   </div>
                 </>
               )}
diff --git a/tools/server/webui/src/utils/app.context.tsx b/tools/server/webui/src/utils/app.context.tsx
@@ -16,7 +16,7 @@ import {
 } from './misc';
 import { BASE_URL, CONFIG_DEFAULT, isDev } from '../Config';
 import { matchPath, useLocation, useNavigate } from 'react-router';
-import { JS_TOOL_CALL_SPEC } from './js_tool_call';
+import { AVAILABLE_TOOLS } from './tool_calling/available_tools';
 
 interface AppContextValue {
   // conversations and messages
@@ -211,7 +211,11 @@ export const AppContextProvider = ({
         dry_penalty_last_n: config.dry_penalty_last_n,
         max_tokens: config.max_tokens,
         timings_per_token: !!config.showTokensPerSecond,
-        tools: config.jsInterpreterToolUse ? [JS_TOOL_CALL_SPEC] : undefined,
+        tools: config.jsInterpreterToolUse
+          ? Array.from(AVAILABLE_TOOLS, ([_name, tool], _index) => tool).filter(
+              (tool) => tool.enabled()
+            )
+          : undefined,
         ...(config.custom.length ? JSON.parse(config.custom) : {}),
       };
 
@@ -278,45 +282,45 @@ export const AppContextProvider = ({
             messageFromAPI.tool_calls &&
             messageFromAPI.tool_calls.length > 0
           ) {
-            console.log(messageFromAPI.tool_calls[0]);
+            console.log(messageFromAPI.tool_calls);
             for (let i = 0; i < messageFromAPI.tool_calls.length; i++) {
               console.log('Tool use #' + i);
               const tc = messageFromAPI.tool_calls[i] as ToolCall;
               console.log(tc);
 
               if (tc) {
-                if (
-                  tc.function.name === 'javascript_interpreter' &&
-                  config.jsInterpreterToolUse
-                ) {
-                  // Execute code provided
-                  const args = JSON.parse(tc.function.arguments);
-                  console.log('Arguments for tool call:');
-                  console.log(args);
-                  const result = eval(args.code);
-                  console.log(result);
-
-                  newContent += `<tool_result>${result}</tool_result>`;
+                // Set up call id
+                tc.call_id ??= `call_${i}`;
+
+                // Process tool call
+                const toolResult = AVAILABLE_TOOLS.get(
+                  tc.function.name
+                )?.processCall(tc);
+
+                if (toolResult) {
+                  newContent += `<tool>Tool use: ${tc.function.name}\n\n`;
+                  newContent += toolResult.output;
+                  newContent += '</tool>';
+                } else {
+                  newContent += `<tool>Tool use: ${tc.function.name}\n\nError: invalid tool call!\n</tool>`;
                 }
               }
             }
 
-            const toolCallsInfo = messageFromAPI.tool_calls
+            /*const toolCallsInfo = messageFromAPI.tool_calls
               .map(
                 (
                   tc: any // Use 'any' for tc temporarily if type is not imported/defined here
                 ) =>
-                  `Tool Call Invoked: ${tc.function.name}\nArguments: ${tc.function.arguments}`
+                  `Tool invoked: ${tc.function.name}\nArguments: ${tc.function.arguments}`
               )
               .join('\n\n');
 
             if (newContent.length > 0) {
               newContent += `\n\n${toolCallsInfo}`;
             } else {
               newContent = toolCallsInfo;
-            }
-            // TODO: Ideally, store structured tool_calls in pendingMsg if its type supports it.
-            // pendingMsg.tool_calls = messageFromAPI.tool_calls;
+              }*/
           }
 
           pendingMsg = {
@@ -338,6 +342,13 @@ export const AppContextProvider = ({
           }
           setPending(convId, pendingMsg);
           onChunk(); // Update UI to show the processed message
+
+          // if message ended due to "finish_reason": "tool_calls"
+          // resend it to assistant to process the result.
+          if (choice.finish_reason === 'tool_calls') {
+            console.log('Ended due to tool call. Interpreting results ...');
+            generateMessage(convId, pendingId, onChunk);
+          }
         } else {
           console.error(
             'API response missing choices or message:',
diff --git a/tools/server/webui/src/utils/js_tool_call.ts b/tools/server/webui/src/utils/js_tool_call.ts
diff --git a/tools/server/webui/src/utils/tool_calling/available_tools.ts b/tools/server/webui/src/utils/tool_calling/available_tools.ts
@@ -0,0 +1,53 @@
+import { ToolCall, ToolCallOutput, ToolCallSpec } from '../types';
+
+/**
+ * Map of available tools for function calling.
+ * Note that these tools are not necessarily enabled by the user.
+ */
+export const AVAILABLE_TOOLS = new Map<string, AgentTool>();
+
+export abstract class AgentTool {
+  id: string;
+  isEnabled: () => boolean;
+
+  constructor(id: string, enabled: () => boolean) {
+    this.id = id;
+    this.isEnabled = enabled;
+    AVAILABLE_TOOLS.set(id, this);
+  }
+
+  /**
+   * "Public" wrapper for the tool call processing logic.
+   * @param call The tool call object from the API response.
+   * @returns The tool call output or undefined if the tool is not enabled.
+   */
+  public processCall(call: ToolCall): ToolCallOutput | undefined {
+    if (this.enabled()) {
+      return this._process(call);
+    }
+
+    return undefined;
+  }
+
+  /**
+   * Whether calling this tool is enabled.
+   * User can toggle the status from the settings panel.
+   * @returns enabled status.
+   */
+  public enabled(): boolean {
+    return this.isEnabled();
+  }
+
+  /**
+   * Specifications for the tool call.
+   * https://github.com/ggml-org/llama.cpp/blob/master/docs/function-calling.md
+   * https://platform.openai.com/docs/guides/function-calling?api-mode=responses#defining-functions
+   */
+  public abstract specs(): ToolCallSpec;
+
+  /**
+   * The actual tool call processing logic.
+   * @param call: The tool call object from the API response.
+   */
+  protected abstract _process(call: ToolCall): ToolCallOutput;
+}
diff --git a/tools/server/webui/src/utils/tool_calling/js_tool_call.ts b/tools/server/webui/src/utils/tool_calling/js_tool_call.ts
@@ -0,0 +1,84 @@
+import StorageUtils from '../storage';
+import { ToolCall, ToolCallOutput, ToolCallSpec } from '../types';
+import { AgentTool } from './available_tools';
+
+class JSReplAgentTool extends AgentTool {
+  private static readonly id = 'javascript_interpreter';
+  private fakeLogger: FakeConsoleLog;
+
+  constructor() {
+    super(
+      JSReplAgentTool.id,
+      () => StorageUtils.getConfig().jsInterpreterToolUse
+    );
+    this.fakeLogger = new FakeConsoleLog();
+  }
+
+  _process(tc: ToolCall): ToolCallOutput {
+    const args = JSON.parse(tc.function.arguments);
+    console.log('Arguments for tool call:');
+    console.log(args);
+
+    // Redirect console.log which agent will use to
+    // the fake logger so that later we can get the content
+    const originalConsoleLog = console.log;
+    console.log = this.fakeLogger.log;
+
+    let result = '';
+    try {
+      // Evaluate the provided agent code
+      result = eval(args.code);
+    } catch (err) {
+      result = String(err);
+    } finally {
+      // Ensure original console.log is restored even if eval throws
+      console.log = originalConsoleLog;
+    }
+
+    result = this.fakeLogger.content + result;
+
+    this.fakeLogger.clear();
+
+    return { call_id: tc.call_id, output: result } as ToolCallOutput;
+  }
+
+  public specs(): ToolCallSpec {
+    return {
+      type: 'function',
+      function: {
+        name: this.id,
+        description:
+          'Executes JavaScript code in the browser console. The code should be self-contained valid javascript. You can use console.log(variable) to print out intermediate values..',
+        parameters: {
+          type: 'object',
+          properties: {
+            code: {
+              type: 'string',
+              description: 'Valid JavaScript code to execute.',
+            },
+          },
+          required: ['code'],
+        },
+      },
+    };
+  }
+}
+export const jsAgentTool = new JSReplAgentTool();
+
+class FakeConsoleLog {
+  private _content: string = '';
+
+  public get content(): string {
+    return this._content;
+  }
+
+  // Use an arrow function for log to correctly bind 'this'
+  public log = (...args: any[]): void => {
+    // Convert arguments to strings and join them.
+    this._content += args.map((arg) => String(arg)).join(' ') + '\n';
+  };
+
+  public clear = (): void => {
+    this._content = '';
+  };
+}
diff --git a/tools/server/webui/src/utils/types.ts b/tools/server/webui/src/utils/types.ts
