Refactor ChatTemplateFormatter to support structured message content, including text, reasoning, tool calls, and tool results. Update tests to reflect new content format and ensure compatibility with unknown types.

Author: michalwarda

packages/torque/src/formatter.test.ts

@@ -66,7 +66,7 @@ describe("ChatTemplateFormatter", () => {
     expect(result.messages).toHaveLength(1);
     expect(result.messages[0]).toEqual({
       role: "user",
-      content: "Hello",
+      content: [{ type: "text", text: "Hello" }],
     });
   });
 
@@ -93,15 +93,13 @@ describe("ChatTemplateFormatter", () => {
     expect(result.messages).toHaveLength(1);
     expect(result.messages[0]).toEqual({
       role: "assistant",
-      content: "Thinking...",
-      tool_calls: [
+      content: [
+        { type: "text", text: "Thinking..." },
         {
+          type: "tool_call",
           id: "call_1",
-          type: "function",
-          function: {
-            name: "calc",
-            arguments: { a: 1 },
-          },
+          name: "calc",
+          arguments: { a: 1 },
         },
       ],
     });
@@ -125,12 +123,14 @@ describe("ChatTemplateFormatter", () => {
     expect(result.messages).toHaveLength(1);
     expect(result.messages[0]).toEqual({
       role: "assistant",
-      content: "Checking weather...",
-      reasoning: "I should check the weather.",
+      content: [
+        { type: "reasoning", text: "I should check the weather." },
+        { type: "text", text: "Checking weather..." },
+      ],
     });
   });
 
-  it("should flatten tool result messages", () => {
+  it("should transform tool result messages", () => {
     const row = createMockRow({
       messages: [
         {
@@ -155,18 +155,49 @@ describe("ChatTemplateFormatter", () => {
     });
 
     const result = formatter.format(row);
-    expect(result.messages).toHaveLength(2);
+    expect(result.messages).toHaveLength(1);
     expect(result.messages[0]).toEqual({
       role: "tool",
-      tool_call_id: "call_1",
-      name: "calc",
-      content: "2",
+      content: [
+        {
+          type: "tool_result",
+          tool_call_id: "call_1",
+          name: "calc",
+          content: { type: "text", value: "2" },
+        },
+        {
+          type: "tool_result",
+          tool_call_id: "call_2",
+          name: "calc",
+          content: { type: "text", value: "4" },
+        },
+      ],
     });
-    expect(result.messages[1]).toEqual({
-      role: "tool",
-      tool_call_id: "call_2",
-      name: "calc",
-      content: "4",
+  });
+
+  it("should passthrough unknown types", () => {
+    const videoObject = { id: "vid_1" };
+    const row = createMockRow({
+      messages: [
+        {
+          role: "user",
+          content: [
+            { type: "video", video: videoObject } as any,
+            { type: "text", text: "What do you see?" },
+          ],
+          generationId: "1",
+        },
+      ],
+    });
+
+    const result = formatter.format(row);
+    expect(result.messages).toHaveLength(1);
+    expect(result.messages[0]).toEqual({
+      role: "user",
+      content: [
+        { type: "video", video: videoObject },
+        { type: "text", text: "What do you see?" },
+      ],
     });
   });
 });

packages/torque/src/formatter.ts

@@ -34,83 +34,51 @@ export class ChatTemplateFormatter implements IDatasetFormatter {
       },
     }));
 
-    const messages = row.messages.flatMap((msg) => {
-      if (msg.role === "tool") {
-        // Flatten tool results
-        if (Array.isArray(msg.content)) {
-          return msg.content
-            .map((part: any) => {
-              if (part.type === "tool-result") {
-                return {
-                  role: "tool",
-                  tool_call_id: part.toolCallId,
-                  name: part.toolName,
-                  content: JSON.stringify(part.output),
-                };
-              }
-              return null;
-            })
-            .filter(Boolean);
-        }
-        return [];
-      }
-
-      if (msg.role === "assistant") {
-        const toolCalls: any[] = [];
-        let contentString = "";
-        let reasoningString = "";
+    const messages = row.messages.map((msg) => {
+      let contentParts: any[] = [];
 
-        if (Array.isArray(msg.content)) {
-          for (const part of msg.content) {
-            if (part.type === "tool-call") {
-              toolCalls.push({
-                id: part.toolCallId,
-                type: "function",
-                function: {
-                  name: part.toolName,
-                  arguments: part.input,
-                },
-              });
-            } else if (part.type === "text") {
-              contentString += part.text;
-            } else if (part.type === "reasoning") {
-              // ai sdk types might vary but usually it's text or reasoning
-              reasoningString +=
-                (part as any).text || (part as any).reasoning || "";
-            }
+      if (typeof msg.content === "string") {
+        contentParts = [{ type: "text", text: msg.content }];
+      } else if (Array.isArray(msg.content)) {
+        contentParts = msg.content.map((part: any) => {
+          if (part.type === "text") {
+            return { type: "text", text: part.text };
           }
-        } else if (typeof msg.content === "string") {
-          contentString = msg.content;
-        }
-
-        const newMsg: any = {
-          role: "assistant",
-          content: contentString || null,
-        };
-        if (toolCalls.length > 0) {
-          newMsg.tool_calls = toolCalls;
-        }
-        if (reasoningString) {
-          newMsg.reasoning = reasoningString;
-        }
-        return [newMsg];
-      }
-
-      // User / System
-      let content = msg.content;
-      if (Array.isArray(content)) {
-        // Ensure content parts are compatible
-        // OpenAI accepts array of text/image parts.
-        // We'll assume they are compatible or simplify if needed.
-        // For now, pass through.
+          if (part.type === "image") {
+            // Pass through image parts
+            return { ...part };
+          }
+          if (part.type === "reasoning") {
+            return {
+              type: "reasoning",
+              text: part.text || part.reasoning || "",
+            };
+          }
+          if (part.type === "tool-call") {
+            return {
+              type: "tool_call",
+              id: part.toolCallId,
+              name: part.toolName,
+              arguments: part.input, // Assuming input is the arguments object
+            };
+          }
+          if (part.type === "tool-result") {
+            return {
+              type: "tool_result",
+              tool_call_id: part.toolCallId,
+              name: part.toolName,
+              content: part.output || part.result,
+            };
+          }
+          // Pass through any other parts (e.g. video)
+          return part;
+        });
       }
 
-      return [
-        {
-          role: msg.role,
-          content,
-        },
-      ];
+      return {
+        role: msg.role,
+        content: contentParts,
+      };
     });
 
     return { tools, messages };