feat: Unified Runtime & Context API (Part 1) (#873)

* feat: Unified Runtime & Context API * add export/import/reset functions * cleanup exports
Yonom · Sep 23, 2024 · 88957ac · 88957ac
1 parent 3187013
commit 88957ac
Show file tree

Hide file tree

Showing 72 changed files with 1,271 additions and 506 deletions.
diff --git a/.changeset/mean-scissors-push.md b/.changeset/mean-scissors-push.md
@@ -0,0 +1,9 @@
+---
+"@assistant-ui/react-playground": patch
+"@assistant-ui/react-hook-form": patch
+"@assistant-ui/react-langgraph": patch
+"@assistant-ui/react-ai-sdk": patch
+"@assistant-ui/react": patch
+---
+
+feat: New unified Runtime API (part 1/n)
diff --git a/apps/docs/components/docs/parameters/context.tsx b/apps/docs/components/docs/parameters/context.tsx
@@ -1,20 +1,14 @@
 import { ParametersTableProps } from "../ParametersTable";
 
-export const AssistantActionsState: ParametersTableProps = {
-  type: "AssistantActionsState",
+export const AssistantRuntimeState: ParametersTableProps = {
+  type: "AssistantRuntime",
   parameters: [
     {
-      name: "switchToThread",
-      type: "(threadId: string | null) => void",
+      name: "switchToNewThread",
+      type: "() => void",
       description: "Switch to a new thread.",
       required: true,
     },
-    {
-      name: "getRuntime",
-      type: "() => AssistantRuntime",
-      description: "Get the current runtime.",
-      required: true,
-    },
   ],
 };
 
@@ -88,8 +82,8 @@ export const ThreadMessagesState: ParametersTableProps = {
   ],
 };
 
-export const ThreadActionsState: ParametersTableProps = {
-  type: "ThreadActionsState",
+export const ThreadRuntimeState: ParametersTableProps = {
+  type: "ThreadRuntime",
   parameters: [
     {
       name: "getBranches",

diff --git a/apps/docs/content/docs/reference/context.mdx b/apps/docs/content/docs/reference/context.mdx
@@ -13,11 +13,11 @@ The context is split into four hierarchies:
 
 import { ParametersTable } from "@/components/docs";
 import {
-  AssistantActionsState,
+  AssistantRuntimeState,
   AssistantToolUIsState,
   ThreadState,
   ThreadMessagesState,
-  ThreadActionsState,
+  ThreadRuntimeState,
   ComposerState,
   EditComposerState,
   ThreadViewportState,
@@ -30,16 +30,15 @@ import {
 
 ## Assistant Context
 
-#### `useAssistantActions`
+#### `useAssistantRuntime`
 
 ```tsx
-import { useAssistantActions } from "@assistant-ui/react";
+import { useAssistantRuntime } from "@assistant-ui/react";
 
-const actions = useAssistantActions();
-const switchToNewThread = useAssistantActions((m) => m.switchToThread);
+const runtime = useAssistantRuntime();
 ```
 
-<ParametersTable {...AssistantActionsState} />
+<ParametersTable {...AssistantRuntimeState} />
 
 ### `useToolUIs`
 
@@ -54,6 +53,16 @@ const getToolUI = useToolUIs((m) => m.getToolUI);
 
 ## Thread Context
 
+### `useThreadRuntime`
+
+```tsx
+import { useThreadRuntime } from "@assistant-ui/react";
+
+const thread = useThreadRuntime();
+```
+
+<ParametersTable {...ThreadRuntimeState} />
+
 ### `useThread`
 
 ```tsx
@@ -76,25 +85,6 @@ const firstMessage = useThreadMessages((m) => m[0]);
 
 <ParametersTable {...ThreadMessagesState} />
 
-### `useThreadActions`
-
-```tsx
-import { useThreadActions } from "@assistant-ui/react";
-
-const actions = useThreadActions();
-const getBranches = useThreadActions((m) => m.getBranches);
-```
-
-<ParametersTable {...ThreadActionsState} />
-
-### `useThreadRuntime`
-
-```tsx
-import { useThreadRuntime } from "@assistant-ui/react";
-
-const runtime = useThreadRuntime();
-```
-
 ### `useThreadComposer`
 
 ```tsx

diff --git a/packages/react-ai-sdk/package.json b/packages/react-ai-sdk/package.json
@@ -27,6 +27,7 @@
     "build": "tsup src/index.ts --format cjs,esm --dts --sourcemap --clean"
   },
   "dependencies": {
+    "@radix-ui/react-use-callback-ref": "^1.1.0",
     "zod": "^3.23.8",
     "zustand": "^4.5.5"
   },

diff --git a/packages/react-ai-sdk/src/ui/utils/useInputSync.tsx b/packages/react-ai-sdk/src/ui/utils/useInputSync.tsx
@@ -1,33 +1,31 @@
-import { useRef, useEffect } from "react";
-import {
-  ExternalStoreRuntime,
-  subscribeToMainThread,
-} from "@assistant-ui/react";
+import { useEffect } from "react";
 import { useAssistant, useChat } from "ai/react";
+import { AssistantRuntime } from "@assistant-ui/react";
+import { useCallbackRef } from "@radix-ui/react-use-callback-ref";
 
 type VercelHelpers =
   | ReturnType<typeof useChat>
   | ReturnType<typeof useAssistant>;
 
 export const useInputSync = (
   helpers: VercelHelpers,
-  runtime: ExternalStoreRuntime,
+  runtime: AssistantRuntime,
 ) => {
   // sync input from vercel to assistant-ui
-  const helpersRef = useRef(helpers);
   useEffect(() => {
-    helpersRef.current = helpers;
-    if (runtime.thread.composer.text !== helpers.input) {
+    if (runtime.thread.composer.getState().text !== helpers.input) {
       runtime.thread.composer.setText(helpers.input);
     }
   }, [helpers, runtime]);
 
   // sync input from assistant-ui to vercel
+  const handleThreadUpdate = useCallbackRef(() => {
+    if (runtime.thread.composer.getState().text !== helpers.input) {
+      helpers.setInput(runtime.thread.composer.getState().text);
+    }
+  });
+
   useEffect(() => {
-    return subscribeToMainThread(runtime, () => {
-      if (runtime.thread.composer.text !== helpersRef.current.input) {
-        helpersRef.current.setInput(runtime.thread.composer.text);
-      }
-    });
-  }, [runtime]);
+    return runtime.thread.subscribe(handleThreadUpdate);
+  }, [runtime, handleThreadUpdate]);
 };
diff --git a/packages/react-hook-form/src/useAssistantForm.tsx b/packages/react-hook-form/src/useAssistantForm.tsx
@@ -3,6 +3,7 @@
 import {
   type ModelConfig,
   type ToolCallContentPartComponent,
+  useAssistantRuntime,
   useAssistantToolUI,
 } from "@assistant-ui/react";
 import { useEffect } from "react";
@@ -14,7 +15,6 @@ import {
 } from "react-hook-form";
 import type { z } from "zod";
 import { formTools } from "./formTools";
-import { useAssistantActionsStore } from "@assistant-ui/react";
 
 export type UseAssistantFormProps<
   TFieldValues extends FieldValues,
@@ -62,7 +62,7 @@ export const useAssistantForm = <
   const form = useForm<TFieldValues, TContext, TTransformedValues>(props);
   const { control, getValues, setValue } = form;
 
-  const assistantActionsStore = useAssistantActionsStore();
+  const assistantRuntime = useAssistantRuntime();
   useEffect(() => {
     const value: ModelConfig = {
       system: `Form State:\n${JSON.stringify(getValues())}`,
@@ -107,10 +107,10 @@ export const useAssistantForm = <
         },
       },
     };
-    return assistantActionsStore.getState().registerModelConfigProvider({
+    return assistantRuntime.registerModelConfigProvider({
       getModelConfig: () => value,
     });
-  }, [control, setValue, getValues, assistantActionsStore]);
+  }, [control, setValue, getValues, assistantRuntime]);
 
   const renderFormFieldTool = props?.assistant?.tools?.set_form_field?.render;
   useAssistantToolUI(

diff --git a/packages/react-langgraph/src/useLangGraphRuntime.ts b/packages/react-langgraph/src/useLangGraphRuntime.ts
@@ -6,7 +6,6 @@ import {
 } from "@assistant-ui/react";
 import { convertLangchainMessages } from "./convertLangchainMessages";
 import { useLangGraphMessages } from "./useLangGraphMessages";
-import { ExternalStoreRuntime } from "@assistant-ui/react";
 
 const getPendingToolCalls = (messages: LangChainMessage[]) => {
   const pendingToolCalls = new Map<string, LangChainToolCall>();
@@ -41,7 +40,7 @@ export const useLangGraphRuntime = ({
   onSwitchToThread?: (
     threadId: string,
   ) => Promise<{ messages: LangChainMessage[] }>;
-}): ExternalStoreRuntime => {
+}) => {
   const { messages, sendMessage, setMessages } = useLangGraphMessages({
     stream,
   });

diff --git a/packages/react-playground/src/components/ui/assistant-ui/assistant-playground.tsx b/packages/react-playground/src/components/ui/assistant-ui/assistant-playground.tsx
@@ -1,7 +1,7 @@
 "use client";
 
 import { ChangeEvent, FC, PropsWithChildren, useState } from "react";
-import { Tool, useAssistantActionsStore } from "@assistant-ui/react";
+import { Tool, useAssistantRuntime } from "@assistant-ui/react";
 import { PayloadEditorButton } from "../../payload-editor-button";
 import { Thread } from "./thread";
 import { Button } from "../button";
@@ -322,9 +322,9 @@ const Sidebar: FC<AssistantPlaygroundProps> = ({
   modelSelector,
   apiKey = true,
 }) => {
-  const assistantActionsStore = useAssistantActionsStore();
+  const assistantRuntime = useAssistantRuntime();
   const handleReset = () => {
-    assistantActionsStore.getState().switchToNewThread();
+    assistantRuntime.switchToNewThread();
   };
 
   return (

diff --git a/packages/react-playground/src/components/ui/assistant-ui/thread.tsx b/packages/react-playground/src/components/ui/assistant-ui/thread.tsx
@@ -10,7 +10,7 @@ import {
   useMessage,
   useMessageStore,
   useThread,
-  useThreadActionsStore,
+  useThreadRuntime,
 } from "@assistant-ui/react";
 import { useState, type FC, type KeyboardEvent, type MouseEvent } from "react";
 import {
@@ -106,7 +106,7 @@ const Composer: FC = () => {
   const isRunning = useThread((t) => t.isRunning);
   const hasText = useComposer((c) => c.text.length > 0);
 
-  const threadActionsStore = useThreadActionsStore();
+  const threadRuntime = useThreadRuntime();
   const composerStore = useComposerStore();
 
   const performAdd = () => {
@@ -118,7 +118,7 @@ const Composer: FC = () => {
 
   const performSubmit = () => {
     performAdd();
-    threadActionsStore.getState().startRun(null);
+    threadRuntime.startRun(null);
   };
 
   const handleAdd = (e: MouseEvent) => {