[OpenAI] Add OpenAI-API-compatible function calling support (mlc-ai#321)

Author: rickzx

examples/openai-api/package.json

@@ -15,6 +15,6 @@
         "url": "^0.11.3"
     },
     "dependencies": {
-        "@mlc-ai/web-llm": "^0.2.28"
+        "@mlc-ai/web-llm": "file:../.."
     }
 }

examples/openai-api/src/openai_api.ts

@@ -137,7 +137,62 @@ async function mainStateful() {
   console.log(await chat.runtimeStatsText());
 }
 
+async function mainFunctionCalling() {
+  const chat: webllm.ChatInterface = new webllm.ChatModule();
+
+  chat.setInitProgressCallback((report: webllm.InitProgressReport) => {
+    setLabel("init-label", report.text);
+  });
+
+  const myAppConfig: webllm.AppConfig = {
+    model_list: [
+      {
+        "model_url": "https://huggingface.co/mlc-ai/gorilla-openfunctions-v2-q4f16_1-MLC/resolve/main/",
+        "local_id": "gorilla-openfunctions-v2-q4f16_1",
+        "model_lib_url": "https://raw.githubusercontent.com/mlc-ai/binary-mlc-llm-libs/main/gorilla-openfunctions-v2/gorilla-openfunctions-v2-q4f16_1.wasm",
+      },
+    ]
+  }
+  const selectedModel = "gorilla-openfunctions-v2-q4f16_1"
+  await chat.reload(selectedModel, undefined, myAppConfig);
+
+  const tools: Array<webllm.ChatCompletionTool> = [
+    {
+      type: "function",
+      function: {
+        name: "get_current_weather",
+        description: "Get the current weather in a given location",
+        parameters: {
+          "type": "object",
+          "properties": {
+            "location": {
+              "type": "string",
+              "description": "The city and state, e.g. San Francisco, CA",
+            },
+            "unit": { "type": "string", "enum": ["celsius", "fahrenheit"] },
+          },
+          "required": ["location"],
+        },
+      },
+    }
+  ]
+
+  const request: webllm.ChatCompletionRequest = {
+    stream: false,
+    messages: [
+      { "role": "user", "content": "What is the current weather in celsius in Pittsburgh and Tokyo?" },
+    ],
+    tool_choice: 'auto',
+    tools: tools,
+  };
+
+  const reply0 = await chat.chatCompletion(request);
+  console.log(reply0.choices[0].message.content);
+
+  console.log(await chat.runtimeStatsText());
+}
+
 // Run one of the functions
-mainNonStreaming();
+// mainNonStreaming();
 // mainStreaming();
-// mainStateful();
+mainFunctionCalling();

package-lock.json

@@ -1,7 +1,7 @@
 import { nodeResolve } from '@rollup/plugin-node-resolve';
 import ignore from "rollup-plugin-ignore";
 import commonjs from '@rollup/plugin-commonjs';
-import typescript from 'rollup-plugin-typescript2'
+import typescript from 'rollup-plugin-typescript2';
 
 export default {
     input: 'src/index.ts',

rollup.config.js

@@ -7,7 +7,8 @@ import {
   prebuiltAppConfig,
   GenerationConfig,
   postInitAndCheckGenerationConfigValues,
-  ModelRecord
+  ModelRecord,
+  Role
 } from "./config";
 import { LLMChatPipeline } from "./llm_chat"
 import {
@@ -20,6 +21,7 @@ import {
   ChatCompletionRequestStreaming,
   ChatCompletionRequestBase,
   CompletionUsage,
+  ChatCompletionUserMessageParam,
 } from "./openai_api_protocols/index";
 import * as ChatCompletionAPI from "./openai_api_protocols/index";
 import {
@@ -316,6 +318,11 @@ export class ChatModule implements ChatInterface {
       top_logprobs: request.top_logprobs,
     }
 
+    const error_msg = this.checkFunctionCallUsage(request);
+    if (error_msg) {
+      throw new Error(error_msg);
+    }
+
     // 1. If request is streaming, return an AsyncIterable (an iterable version of `generate()`)
     if (request.stream) {
       return this.chatCompletionAsyncChunkGenerator(request, genConfig);
@@ -506,24 +513,65 @@ export class ChatModule implements ChatInterface {
           throw new Error("Last messages should be a string from the `user`.");
         }
         this.getPipeline().appendConversationMessage(
-          message.name ? message.name : roles[0],
+          Role.User,
           message.content,
+          message.name
         );
       } else if (message.role === "assistant") {
         if (typeof message.content !== "string") {
-          // TODO(Charlie): Remove when we support function calling
           throw new Error("Assistant message should have string content.");
         }
         this.getPipeline().appendConversationMessage(
-          message.name ? message.name : roles[1],
+          Role.Assistant,
           message.content,
+          message.name
         );
       } else {
         throw new Error("Unsupported role: " + message.role);
       }
     }
   }
 
+  private checkFunctionCallUsage(request: ChatCompletionRequest): string | null {
+    if (request.tools == undefined || 
+        (typeof request.tool_choice == "string" && request.tool_choice == "none")) {
+        this.getPipeline().overrideFunctionCalling(false, "");
+        return null;
+    }
+
+    if (typeof request.tool_choice == "string" && request.tool_choice !== "auto") {
+      return `Invalid tool choice value: ${request.tool_choice}`;
+    }
+
+    if (typeof request.tool_choice !== "string" && request.tool_choice?.type) {
+      return "Only 'function' tool choice is supported";
+    }
+
+    const singleFunctionToCall = typeof request.tool_choice !== "string" && request.tool_choice?.function?.name;
+
+    if (singleFunctionToCall) {
+      for (const f of request.tools) {
+        if (singleFunctionToCall == f.function.name) {
+          this.getPipeline().overrideFunctionCalling(true, JSON.stringify([f.function]));
+          return null;
+        }
+      }
+
+      return `The tool choice function ${singleFunctionToCall} is not found in the tools list`;
+    }
+
+    let function_list = [];
+    for (const f of request.tools) {
+      if (f.type !== "function") {
+        return "Only 'function' tool type is supported";
+      }
+
+      function_list.push(f.function);
+    }
+    this.getPipeline().overrideFunctionCalling(true, JSON.stringify(function_list));
+    return null;
+  }
+
   /**
    * Run a prefill step with a given input.
    * @param input The input prompt, or `messages` in OpenAI-like APIs.
@@ -533,15 +581,18 @@ export class ChatModule implements ChatInterface {
     genConfig?: GenerationConfig
   ) {
     let input_str: string;
+    let input_role_str : string | undefined;
     if (typeof input === "string") {
       input_str = input;
     } else {
       // Process ChatCompletionMessageParam
       // We treat the last message as our usual input
       this.updateConversationWithChatCompletionMessages(input);
-      input_str = input[input.length - 1].content as string;
+      const last_msg = input[input.length - 1] as ChatCompletionUserMessageParam;
+      input_str = last_msg.content as string;
+      input_role_str = last_msg.name ? last_msg.name : undefined;
     }
-    return this.getPipeline().prefillStep(input_str, genConfig);
+    return this.getPipeline().prefillStep(input_str, input_role_str, genConfig);
   }
 
   /**

src/chat_module.ts

@@ -4,7 +4,9 @@
  */
 export interface ConvTemplateConfig {
   system: string;
-  roles: Array<string>;
+  roles: Record<Role, string>;
+  role_templates?: Partial<Record<Role, string>>;
+  function_calling_template?: string;
   seps: Array<string>;
   separator_style: string;
   offset: number;
@@ -13,6 +15,26 @@ export interface ConvTemplateConfig {
   stop_tokens: Array<number>;
 }
 
+export enum Role {
+  User,
+  Assistant
+}
+
+/**
+ * Place holders that can be used in role templates.
+ * For example, a role template of
+ * `<<question>> ${MessagePlaceholders.USER} <<function>> ${MessagePlaceholders.FUNCTION}`
+ * will insert the user message to ${MessagePlaceholders.USER}
+ * and insert the function message to ${MessagePlaceholders.FUNCTION}
+ * at run time.
+ */
+export enum MessagePlaceholders {
+  User = "{user_message}",
+  Assitant = "{assistant_message}",
+  Tool = "{tool_message}",
+  Function = "{function_string}"
+}
+
 /**
  * Config of one chat model, a data structure representing `mlc-chat-config.json`.
  * This only corresponds to the chat-related fields and `tokenizer_files` of `mlc-chat-config.json`.