feat: ThreadConfig.tools (#755)

Yonom · Sep 7, 2024 · b4a826c · b4a826c
1 parent 0a4b8d7
commit b4a826c
Show file tree

Hide file tree

Showing 5 changed files with 42 additions and 10 deletions.
diff --git a/packages/react/src/model-config/makeAssistantTool.tsx b/packages/react/src/model-config/makeAssistantTool.tsx
@@ -1,15 +1,21 @@
 "use client";
+import { FC } from "react";
 import { type AssistantToolProps, useAssistantTool } from "./useAssistantTool";
 
+export type AssistantTool = FC & {
+  unstable_tool: AssistantToolProps<any, any>;
+};
+
 export const makeAssistantTool = <
   TArgs extends Record<string, unknown>,
   TResult,
 >(
   tool: AssistantToolProps<TArgs, TResult>,
 ) => {
-  const Tool = () => {
+  const Tool: AssistantTool = () => {
     useAssistantTool(tool);
     return null;
   };
+  Tool.unstable_tool = tool;
   return Tool;
 };
diff --git a/packages/react/src/model-config/makeAssistantToolUI.tsx b/packages/react/src/model-config/makeAssistantToolUI.tsx
@@ -1,18 +1,24 @@
 "use client";
+import { FC } from "react";
 import {
   type AssistantToolUIProps,
   useAssistantToolUI,
 } from "./useAssistantToolUI";
 
+export type AssistantToolUI = FC & {
+  unstable_tool: AssistantToolUIProps<any, any>;
+};
+
 export const makeAssistantToolUI = <
   TArgs extends Record<string, unknown>,
   TResult,
 >(
   tool: AssistantToolUIProps<TArgs, TResult>,
 ) => {
-  const ToolUI = () => {
+  const ToolUI: AssistantToolUI = () => {
     useAssistantToolUI(tool);
     return null;
   };
+  ToolUI.unstable_tool = tool;
   return ToolUI;
 };
diff --git a/packages/react/src/primitives/message/MessageContent.tsx b/packages/react/src/primitives/message/MessageContent.tsx
@@ -30,12 +30,14 @@ export type MessagePrimitiveContentProps = {
         Text?: TextContentPartComponent | undefined;
         Image?: ImageContentPartComponent | undefined;
         UI?: UIContentPartComponent | undefined;
-        tools?: {
-          by_name?:
-            | Record<string, ToolCallContentPartComponent | undefined>
-            | undefined;
-          Fallback?: ComponentType<ToolCallContentPartProps> | undefined;
-        };
+        tools?:
+          | {
+              by_name?:
+                | Record<string, ToolCallContentPartComponent | undefined>
+                | undefined;
+              Fallback?: ComponentType<ToolCallContentPartProps> | undefined;
+            }
+          | undefined;
       }
     | undefined;
 };

diff --git a/packages/react/src/ui/assistant-message.tsx b/packages/react/src/ui/assistant-message.tsx
@@ -44,13 +44,26 @@ const AssistantMessageContent = forwardRef<
   HTMLDivElement,
   AssistantMessageContentProps
 >(({ components: componentsProp, ...rest }, ref) => {
-  const { assistantMessage: { components = {} } = {} } = useThreadConfig();
+  const { tools, assistantMessage: { components = {} } = {} } =
+    useThreadConfig();
+
   return (
     <AssistantMessageContentWrapper {...rest} ref={ref}>
       <MessagePrimitive.Content
         components={{
           ...componentsProp,
           Text: componentsProp?.Text ?? components.Text ?? ContentPart.Text,
+          tools: {
+            by_name: !tools
+              ? undefined
+              : Object.fromEntries(
+                  tools.map((t) => [
+                    t.unstable_tool.toolName,
+                    t.unstable_tool.render,
+                  ]),
+                ),
+            Fallback: components.ToolFallback,
+          },
         }}
       />
     </AssistantMessageContentWrapper>

diff --git a/packages/react/src/ui/thread-config.tsx b/packages/react/src/ui/thread-config.tsx
@@ -1,6 +1,7 @@
 "use client";
 
 import {
+  ComponentType,
   FC,
   PropsWithChildren,
   ReactNode,
@@ -9,9 +10,10 @@ import {
 } from "react";
 
 import { AvatarProps } from "./base/avatar";
-import { TextContentPartComponent } from "../types";
+import { TextContentPartComponent, ToolCallContentPartProps } from "../types";
 import { AssistantRuntime } from "../runtimes";
 import { AssistantRuntimeProvider, useAssistantContext } from "../context";
+import { AssistantToolUI } from "../model-config";
 
 export type SuggestionConfig = {
   text?: ReactNode;
@@ -34,6 +36,7 @@ export type AssistantMessageConfig = {
   components?:
     | {
         Text?: TextContentPartComponent | undefined;
+        ToolFallback?: ComponentType<ToolCallContentPartProps> | undefined;
       }
     | undefined;
 };
@@ -133,6 +136,8 @@ export type ThreadConfig = {
   branchPicker?: BranchPickerConfig;
 
   strings?: StringsConfig;
+
+  tools?: AssistantToolUI[]; // TODO add AssistantTool support
 };
 
 export const useThreadConfig = (): Omit<ThreadConfig, "runtime"> => {