From 93e12e0dc952ca4bf3eb43dc781889714a8b8e58 Mon Sep 17 00:00:00 2001
From: Saikat Mitra <saikatmitra91@gmail.com>
Date: Thu, 2 May 2024 11:41:01 +0530
Subject: [PATCH] feat: add support for tool calls in assistant output (#199)

---
 .changeset/seven-oranges-decide.md            |  9 +++
 apps/web/components/json-as-tab.tsx           | 12 +++-
 apps/web/components/sample-output-card.tsx    | 60 +++++++++++--------
 apps/web/components/tool-calls-view.tsx       | 53 ++++++++++++++++
 apps/web/components/ui/tabs.tsx               |  2 +-
 packages/ai/src/providers/openai/index.ts     |  7 +--
 .../executors/run/transformers/assistant.ts   |  5 +-
 packages/types/src/ai/index.ts                |  5 +-
 packages/types/src/index.ts                   |  3 +
 9 files changed, 117 insertions(+), 39 deletions(-)
 create mode 100644 .changeset/seven-oranges-decide.md
 create mode 100644 apps/web/components/tool-calls-view.tsx

diff --git a/.changeset/seven-oranges-decide.md b/.changeset/seven-oranges-decide.md
new file mode 100644
index 00000000..74d1d08a
--- /dev/null
+++ b/.changeset/seven-oranges-decide.md
@@ -0,0 +1,9 @@
+---
+"@empiricalrun/types": minor
+"@empiricalrun/core": minor
+"@empiricalrun/ai": minor
+"web": minor
+"@empiricalrun/cli": minor
+---
+
+feat: add support for assistant tool calls
diff --git a/apps/web/components/json-as-tab.tsx b/apps/web/components/json-as-tab.tsx
index 8d0aa059..4e7208e4 100644
--- a/apps/web/components/json-as-tab.tsx
+++ b/apps/web/components/json-as-tab.tsx
@@ -26,6 +26,9 @@ export function JsonAsTab({
   onSampleRemove,
   onEditorContentUpdate,
   onClickRunAll,
+  showExpandOption = true,
+  readonlyContent = false,
+  scrollableContent = true,
 }: {
   storeKey: string;
   data: { [key: string]: any };
@@ -35,6 +38,9 @@ export function JsonAsTab({
   onSampleRemove?: () => void;
   onEditorContentUpdate?: (key: string, value: string) => void;
   onClickRunAll?: () => void;
+  showExpandOption?: boolean;
+  readonlyContent?: boolean;
+  scrollableContent?: boolean;
 }) {
   const tabs = useMemo(
     () => defaultTabs || Object.keys(data),
@@ -70,7 +76,7 @@ export function JsonAsTab({
     <>
       <div className="flex flex-row space-x-2 justify-end">
         <>
-          {activeTabValue && (
+          {activeTabValue && showExpandOption && (
             <Sheet>
               <SheetTrigger asChild>
                 <Button
@@ -169,8 +175,8 @@ export function JsonAsTab({
                       : JSON.stringify(value, null, 2)
                   }
                   language="text"
-                  readOnly={false}
-                  scrollable
+                  readOnly={readonlyContent}
+                  scrollable={scrollableContent}
                   onChange={(value) => onEditorContentUpdate?.(key, value!)}
                 />
               </TabsContent>
diff --git a/apps/web/components/sample-output-card.tsx b/apps/web/components/sample-output-card.tsx
index e12901cd..8a2a05db 100644
--- a/apps/web/components/sample-output-card.tsx
+++ b/apps/web/components/sample-output-card.tsx
@@ -25,6 +25,7 @@ import SampleCompletionError from "./sample-completion-error";
 import { JsonAsTab } from "./json-as-tab";
 import { RunSampleOutputMetric } from "./run-response-metadata";
 import { Scores } from "./scores";
+import { ToolCalls } from "./tool-calls-view";
 
 type Diff = {
   type: string;
@@ -56,12 +57,13 @@ export default function SampleOutputCard({
 
   const showCompareAgainst = useMemo(
     () =>
-      baseSample?.expected?.value ||
-      comparisonSamples?.some(
-        (comparisonSample, index) =>
-          comparisonSample?.output &&
-          comparisonResults?.[index]?.id !== baseResult?.id,
-      ),
+      baseSample?.output.value &&
+      (baseSample?.expected?.value ||
+        comparisonSamples?.some(
+          (comparisonSample, index) =>
+            comparisonSample?.output &&
+            comparisonResults?.[index]?.id !== baseResult?.id,
+        )),
     [
       baseResult?.id,
       baseSample?.expected?.value,
@@ -139,12 +141,12 @@ export default function SampleOutputCard({
           <CardTitle className="flex flex-row space-x-2 items-center">
             <Scores scores={baseSample?.scores || []} />
             <div className="flex flex-row space-x-2 justify-end items-start self-baseline">
-              <DropdownMenu>
-                <DropdownMenuTrigger>
-                  <DotsVerticalIcon />
-                </DropdownMenuTrigger>
-                <DropdownMenuContent align="end">
-                  {showCompareAgainst ? (
+              {showCompareAgainst && (
+                <DropdownMenu>
+                  <DropdownMenuTrigger>
+                    <DotsVerticalIcon />
+                  </DropdownMenuTrigger>
+                  <DropdownMenuContent align="end">
                     <>
                       <DropdownMenuLabel className="text-xs">
                         Compare against
@@ -200,9 +202,9 @@ export default function SampleOutputCard({
                         );
                       })}
                     </>
-                  ) : null}
-                </DropdownMenuContent>
-              </DropdownMenu>
+                  </DropdownMenuContent>
+                </DropdownMenu>
+              )}
             </div>
           </CardTitle>
         )}
@@ -216,9 +218,6 @@ export default function SampleOutputCard({
         ref={containerWrapper}
       >
         <section className="flex flex-col">
-          {showOutput && hasMetadata && (
-            <p className=" text-sm font-medium mb-2">Output</p>
-          )}
           {diffView.enabled && baseSample && (
             <DiffEditor
               original={baseSample?.output.value || ""}
@@ -236,11 +235,16 @@ export default function SampleOutputCard({
             />
           )}
           {showOutput && (
-            <CodeViewer
-              value={baseSample?.output.value || ""}
-              language="json"
-              readOnly
-            />
+            <>
+              {!baseSample?.output.tool_calls && (
+                <CodeViewer
+                  value={baseSample?.output.value || ""}
+                  language="json"
+                  readOnly
+                />
+              )}
+              <ToolCalls toolCalls={baseSample?.output.tool_calls} />
+            </>
           )}
         </section>
         {showOutput && (
@@ -250,7 +254,11 @@ export default function SampleOutputCard({
               value={baseSample?.output?.tokens_used}
               hideSeparator
             />
-            <RunSampleOutputMetric title="Latency" value={latency} />
+            <RunSampleOutputMetric
+              title="Latency"
+              value={latency}
+              hideSeparator={!baseSample?.output?.tokens_used}
+            />
             <RunSampleOutputMetric
               title="Finish reason"
               value={baseSample?.output?.finish_reason}
@@ -259,11 +267,13 @@ export default function SampleOutputCard({
         )}
         {!diffView.enabled && hasMetadata && (
           <section className="flex flex-col h-[200px] mt-2">
-            <p className=" text-sm font-medium mt-2">Metadata</p>
             <section className="relative flex flex-col flex-1">
               <JsonAsTab
                 storeKey={baseResult?.id!}
                 data={baseSample?.output.metadata!}
+                showExpandOption={false}
+                scrollableContent={false}
+                readonlyContent
               />
             </section>
           </section>
diff --git a/apps/web/components/tool-calls-view.tsx b/apps/web/components/tool-calls-view.tsx
new file mode 100644
index 00000000..e73a955a
--- /dev/null
+++ b/apps/web/components/tool-calls-view.tsx
@@ -0,0 +1,53 @@
+import { ChatCompletionMessageToolCall } from "@empiricalrun/types";
+import CodeViewer from "./ui/code-viewer";
+import { Tabs, TabsContent, TabsList, TabsTrigger } from "./ui/tabs";
+
+export function ToolCalls({
+  toolCalls,
+}: {
+  toolCalls: ChatCompletionMessageToolCall[] | undefined;
+}) {
+  if (!toolCalls) {
+    return null;
+  }
+  if (toolCalls.length === 0) {
+    return null;
+  }
+  const tabs = toolCalls.filter((t) => t.type === "function");
+  return (
+    <>
+      <Tabs defaultValue={tabs[0]?.id} className="h-full">
+        <TabsList className=" rounded-sm w-full overflow-x-scroll justify-start no-scrollbar">
+          <p className=" font-semibold text-sm mr-2 text-white">Tool Calls</p>
+          {tabs.map((tab) => (
+            <TabsTrigger
+              key={tab.id}
+              value={tab.id}
+              className="text-xs rounded-sm"
+            >
+              {tab.function.name}
+            </TabsTrigger>
+          ))}
+        </TabsList>
+        {tabs.map((tab) => (
+          <TabsContent
+            key={tab.id}
+            value={tab.id}
+            // 2.25rem as the height of the tabs is h-9 by default. change this if tab height changes
+            className="h-[calc(100%-3rem)]"
+          >
+            <CodeViewer
+              value={JSON.stringify(
+                JSON.parse(tab.function.arguments),
+                null,
+                2,
+              )}
+              language="text"
+              readOnly
+            />
+          </TabsContent>
+        ))}
+      </Tabs>
+    </>
+  );
+}
diff --git a/apps/web/components/ui/tabs.tsx b/apps/web/components/ui/tabs.tsx
index b707e779..79c95c9c 100644
--- a/apps/web/components/ui/tabs.tsx
+++ b/apps/web/components/ui/tabs.tsx
@@ -44,7 +44,7 @@ const TabsContent = React.forwardRef<
   <TabsPrimitive.Content
     ref={ref}
     className={cn(
-      "mt-2 ring-offset-background focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-2",
+      "ring-offset-background focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-2",
       className,
     )}
     {...props}
diff --git a/packages/ai/src/providers/openai/index.ts b/packages/ai/src/providers/openai/index.ts
index 189d37f3..a6144df6 100644
--- a/packages/ai/src/providers/openai/index.ts
+++ b/packages/ai/src/providers/openai/index.ts
@@ -138,12 +138,9 @@ const runAssistant: ICreateAndRunAssistantThread = async (body) => {
               if (eventData.event === "thread.run.requires_action") {
                 const { tool_calls } = eventData.data.required_action
                   ?.submit_tool_outputs || {
-                  tool_calls: [],
+                  tool_calls: undefined,
                 };
-                const toolSummary = tool_calls.map((tc) => {
-                  return `${tc.function.name} with args ${tc.function.arguments}`;
-                });
-                asstRunResp.content = `Attempting to make tool call: ${toolSummary.join(", ") || ""}`;
+                asstRunResp.content = "";
                 asstRunResp.tool_calls = tool_calls;
               }
 
diff --git a/packages/core/src/executors/run/transformers/assistant.ts b/packages/core/src/executors/run/transformers/assistant.ts
index 4baff16c..f765b827 100644
--- a/packages/core/src/executors/run/transformers/assistant.ts
+++ b/packages/core/src/executors/run/transformers/assistant.ts
@@ -56,14 +56,10 @@ export const assistantExecutor: Transformer = async function (
     });
 
     const hasCitation = message.citations && message.citations.length > 0;
-    const hasToolCall = message.tool_calls && message.tool_calls.length > 0;
     let metadata: any = {};
     if (hasCitation) {
       metadata.citations = message.citations;
     }
-    if (hasToolCall) {
-      metadata.tool_calls = message.tool_calls;
-    }
 
     return {
       output: {
@@ -71,6 +67,7 @@ export const assistantExecutor: Transformer = async function (
         metadata,
         tokens_used: message.usage?.total_tokens,
         latency: message.latency,
+        tool_calls: message.tool_calls,
       },
     };
   } catch (e: any) {
diff --git a/packages/types/src/ai/index.ts b/packages/types/src/ai/index.ts
index 17656c81..ceea6708 100644
--- a/packages/types/src/ai/index.ts
+++ b/packages/types/src/ai/index.ts
@@ -14,6 +14,9 @@ export interface ICreateChatCompletion {
   (body: IChatCompletionCreateParams): Promise<IChatCompletion>;
 }
 
+export interface ChatCompletionMessageToolCall
+  extends OpenAI.ChatCompletionMessageToolCall {}
+
 export interface Citation {
   file_id?: string;
   quote?: string;
@@ -23,7 +26,7 @@ export interface Citation {
 export interface IAssistantRunResponse {
   content: string;
   citations: Citation[];
-  tool_calls?: any[];
+  tool_calls?: ChatCompletionMessageToolCall[];
   usage?: OpenAI.CompletionUsage;
   latency?: number;
 }
diff --git a/packages/types/src/index.ts b/packages/types/src/index.ts
index 4af7e547..1e97c646 100644
--- a/packages/types/src/index.ts
+++ b/packages/types/src/index.ts
@@ -1,3 +1,5 @@
+import { ChatCompletionMessageToolCall } from "./ai";
+
 export * from "./ai";
 
 export enum RoleType {
@@ -193,6 +195,7 @@ export type RunOutput = {
   finish_reason?: string;
   tokens_used?: number;
   latency?: number;
+  tool_calls?: ChatCompletionMessageToolCall[];
 };
 
 export type RunSampleOutput = {