feat: Support for rendering tools in Assistant UI (#267)

inferablehq · Dec 10, 2024 · c62e6ce · c62e6ce
1 parent b4c3add
commit c62e6ce
Show file tree

Hide file tree

Showing 3 changed files with 103 additions and 25 deletions.
diff --git a/assistant-ui/README.md b/assistant-ui/README.md
@@ -34,7 +34,7 @@ pnpm add @inferable/assistant-ui
 
 ```typescript
 import { useInferableRuntime } from '@inferable/assistant-ui';
-import { Thread } from "@assistant-ui/react";
+import { AssistantRuntimeProvider, Thread } from "@assistant-ui/react";
 
 const { runtime, run } = useInferableRuntime({
   clusterId: '<YOUR_CLUSTER_ID>',
@@ -46,7 +46,9 @@ const { runtime, run } = useInferableRuntime({
 
 return (
   <div className="h-full">
-    <Thread runtime={runtime}/>
+    <AssistantRuntimeProvider runtime={runtime}>
+      <Thread/>
+    </AssistantRuntimeProvider>
   </div>
 );
 ```
@@ -68,6 +70,49 @@ You can handle errors by providing an `onError` callback:
   })
 ```
 
+### Rendering function UI
+
+You can provide assistant-ui with [custom UI components](https://www.assistant-ui.com/docs/guides/ToolUI) for rendering Inferable function calls / results.
+
+#### Example
+
+```typescript
+// Fallback UI
+const FallbackToolUI = ({args, result, toolName}) =>
+  <div className="center">
+    <h1>Tool: {toolName}</h1>
+    <h2>Input:</h2>
+    <pre className="whitespace-pre-wrap">{JSON.stringify(args, null, 2)}</pre>
+    <h2>Output:</h2>
+    {result && <pre className="whitespace-pre-wrap">{JSON.stringify(result, null, 2)}</pre>}
+    {!result && <p>No output</p>}
+  </div>
+
+// Custom UI example
+const SearchToolUI = makeAssistantToolUI<any, any>({
+  toolName: "default_webSearch",
+  render: ({ args }) => {
+    return <p>webSearch({args.query})</p>;
+  },
+});
+
+return (
+  <div className="h-full">
+    <AssistantRuntimeProvider runtime={runtime}>
+      <Thread
+        tools={[
+          WebSearchToolUI
+        ]},
+        assistantMessage={{
+          components: {
+            ToolFallback: FallbackToolUI
+          },
+      }} />
+    </AssistantRuntimeProvider>
+  </div>
+);
+```
+
 
 ## Local Development
 

diff --git a/assistant-ui/demo/TestPage.jsx b/assistant-ui/demo/TestPage.jsx
@@ -1,6 +1,15 @@
 import { useInferableRuntime } from '../src'
-import { Thread } from "@assistant-ui/react";
-import toast from "react-hot-toast";
+import { AssistantRuntimeProvider, Thread } from "@assistant-ui/react";
+
+const FallbackToolUI = ({args, result, toolName}) =>
+  <div className="center">
+    <h1>Tool: {toolName}</h1>
+    <h2>Input:</h2>
+    <pre className="whitespace-pre-wrap">{JSON.stringify(args, null, 2)}</pre>
+    <h2>Output:</h2>
+    {result && <pre className="whitespace-pre-wrap">{JSON.stringify(result, null, 2)}</pre>}
+    {!result && <p>No output</p>}
+  </div>
 
 const TestPage = () => {
   const existingRunId = localStorage.getItem("runID")
@@ -24,7 +33,13 @@ const TestPage = () => {
 
   return (
     <div className="h-full">
-      <Thread runtime={runtime}/>
+      <AssistantRuntimeProvider runtime={runtime}>
+        <Thread assistantMessage={{
+          components: {
+            ToolFallback: FallbackToolUI
+          },
+        }} />
+      </AssistantRuntimeProvider>
     </div>
   );
 };

diff --git a/assistant-ui/src/inferable-provider-runtime.ts b/assistant-ui/src/inferable-provider-runtime.ts
@@ -81,15 +81,15 @@ export function useInferableRuntime({
     runtime: useExternalStoreRuntime({
       isRunning,
       messages,
-      convertMessage,
+      convertMessage: (message) => convertMessage(message, messages),
       onNew,
     }),
     run,
   }
 
 }
 
-const convertMessage = (message: any): ThreadMessageLike => {
+const convertMessage = (message: any, allMessages: any): ThreadMessageLike => {
   switch (message.type) {
     case "human": {
       const parsedData = genericMessageDataSchema.parse(message.data);
@@ -114,39 +114,57 @@ const convertMessage = (message: any): ThreadMessageLike => {
         });
       }
 
-      return {
-        id: message.id,
-        role: "assistant",
-        content: content
+      if (parsedData.invocations) {
+
+        parsedData.invocations.forEach((invocation) => {
+
+          // Attempt to find corresponding `invocation-result` message
+          let result = null;
+          allMessages.forEach((message: any) => {
+            if ('type' in message && message.type !== "invocation-result") {
+              return false
+            }
+
+            const parsedResult = resultDataSchema.parse(message.data);
+
+            if (parsedResult.id === invocation.id) {
+              result = parsedResult.result;
+              return true;
+            }
+          });
+
+          content.push({
+            type: "tool-call",
+            toolName: invocation.toolName,
+            args: invocation.input,
+            toolCallId: invocation.id,
+            result
+          });
+        })
       }
-    }
-    case "invocation-result": {
-      const parsedData = resultDataSchema.parse(message.data);
 
-      // TODO: Search chat history for the corresponding invocation mesasge (With args)
+      if (content.length === 0) {
+        return {
+          id: message.id,
+          role: "system",
+          content: "MESSAGE HAS NO CONTENT"
+        }
+      }
 
       return {
         id: message.id,
         role: "assistant",
-        content: [{
-          type: "tool-call",
-          toolName: "inferable",
-          args: {},
-          toolCallId: parsedData.id,
-          result: parsedData.result,
-        }]
+        content: content
       }
-
     }
-
   }
 
   return {
     id: message.id,
     role: "system",
     content: [{
       type: "text",
-      text: ""
+      text: `UNKNON MESSAGE TYPE: ${message.type}`
     }],
   };
 }