Responses api - ensure proper storage of tool messages (#851)

* improve logic for converting input to db messages

* improved and added tests for tool message logic

Author: ASteinheiser

packages/mongodb-chatbot-server/src/routes/responses/createResponse.test.ts

@@ -5,6 +5,7 @@ import type {
   Conversation,
   ConversationsService,
   SomeMessage,
+  ToolMessage,
 } from "mongodb-rag-core";
 import { type AppConfig } from "../../app";
 import {
@@ -386,15 +387,66 @@ describe("POST /responses", () => {
 
     it("Should store function_call messages when `store: true`", async () => {
       const store = true;
-      const functionCallType = "function_call";
+      const functionCallName = "my_function";
+      const functionCallArguments = `{"query": "value"}`;
       const functionCallOutputType = "function_call_output";
+      const functionCallOutput = `{"result": "success"}`;
       const requestBody: Partial<CreateResponseRequest["body"]> = {
         store,
         input: [
           {
-            type: functionCallType,
+            type: "function_call",
             call_id: "call123",
-            name: "my_function",
+            name: functionCallName,
+            arguments: functionCallArguments,
+            status: "in_progress",
+          },
+          {
+            type: functionCallOutputType,
+            call_id: "call123",
+            output: functionCallOutput,
+            status: "completed",
+          },
+          {
+            role: "user",
+            content: "What is MongoDB?",
+          },
+        ],
+      };
+      const stream = await makeClientAndRequest(requestBody);
+
+      const results = await expectValidResponses({ requestBody, stream });
+
+      const updatedConversation = await conversations.findByMessageId({
+        messageId: getMessageIdFromResults(results),
+      });
+      if (!updatedConversation) {
+        return expect(updatedConversation).not.toBeNull();
+      }
+      const messages = updatedConversation.messages as Array<ToolMessage>;
+
+      expect(updatedConversation.storeMessageContent).toEqual(store);
+
+      expect(messages[0].role).toEqual("tool");
+      expect(messages[0].name).toEqual(functionCallName);
+      expect(messages[0].content).toEqual(functionCallArguments);
+
+      expect(messages[1].role).toEqual("tool");
+      expect(messages[1].name).toEqual(functionCallOutputType);
+      expect(messages[1].content).toEqual(functionCallOutput);
+    });
+
+    it("Should not store function_call message content when `store: false`", async () => {
+      const store = false;
+      const functionCallName = "my_function";
+      const functionCallOutputType = "function_call_output";
+      const requestBody: Partial<CreateResponseRequest["body"]> = {
+        store,
+        input: [
+          {
+            type: "function_call",
+            call_id: "call123",
+            name: functionCallName,
             arguments: `{"query": "value"}`,
             status: "in_progress",
           },
@@ -420,16 +472,17 @@ describe("POST /responses", () => {
       if (!updatedConversation) {
         return expect(updatedConversation).not.toBeNull();
       }
+      const messages = updatedConversation.messages as Array<ToolMessage>;
 
       expect(updatedConversation.storeMessageContent).toEqual(store);
 
-      expect(updatedConversation.messages[0].role).toEqual("system");
-      expect(updatedConversation.messages[0].content).toEqual(functionCallType);
+      expect(messages[0].role).toEqual("tool");
+      expect(messages[0].name).toEqual(functionCallName);
+      expect(messages[0].content).toEqual("");
 
-      expect(updatedConversation.messages[1].role).toEqual("system");
-      expect(updatedConversation.messages[1].content).toEqual(
-        functionCallOutputType
-      );
+      expect(messages[1].role).toEqual("tool");
+      expect(messages[1].name).toEqual(functionCallOutputType);
+      expect(messages[1].content).toEqual("");
     });
   });
 

packages/mongodb-chatbot-server/src/routes/responses/createResponse.ts

@@ -511,16 +511,17 @@ const convertInputToDBMessages = (
   }
 
   return input.map((message) => {
-    // set default role and content for function tool calls and outputs
-    let role: MessagesParam[number]["role"] = "system";
-    let content = message.type ?? "";
-    // set role and content for all other messages
-    if (message.type === "message") {
-      role = message.role;
-      content = formatUserMessageContent(message.content);
+    if (isInputMessage(message)) {
+      const role = message.role;
+      const content = formatUserMessageContent(message.content);
+      return formatMessage({ role, content }, store, metadata);
     }
-
-    return formatMessage({ role, content }, store, metadata);
+    // handle function tool calls and outputs
+    const role = "tool";
+    const name = message.type === "function_call" ? message.name : message.type;
+    const content =
+      message.type === "function_call" ? message.arguments : message.output;
+    return formatMessage({ role, name, content }, store, metadata);
   });
 };