-
Notifications
You must be signed in to change notification settings - Fork 416
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
fix: Vercel useAssistant BranchPicker duplicates bug (#356)
- Loading branch information
Showing
7 changed files
with
240 additions
and
68 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,6 @@ | ||
| --- | ||
| "@assistant-ui/react-ai-sdk": patch | ||
| "@assistant-ui/react": patch | ||
| --- | ||
|
|
||
| fix: Vercel useAssistant BranchPicker duplicates bug |
51 changes: 51 additions & 0 deletions
51
packages/react-ai-sdk/src/ui/use-assistant/VercelUseAssistantRuntime.tsx
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,51 @@ | ||
| import { type ThreadMessage, INTERNAL } from "@assistant-ui/react"; | ||
| import { ModelConfigProvider } from "@assistant-ui/react"; | ||
| import { UseAssistantHelpers } from "@ai-sdk/react"; | ||
| import { VercelUseAssistantThreadRuntime } from "./VercelUseAssistantThreadRuntime"; | ||
|
|
||
| const { ProxyConfigProvider, BaseAssistantRuntime } = INTERNAL; | ||
|
|
||
| export const hasUpcomingMessage = ( | ||
| isRunning: boolean, | ||
| messages: ThreadMessage[], | ||
| ) => { | ||
| return isRunning && messages[messages.length - 1]?.role !== "assistant"; | ||
| }; | ||
|
|
||
| export class VercelUseAssistantRuntime extends BaseAssistantRuntime<VercelUseAssistantThreadRuntime> { | ||
| private readonly _proxyConfigProvider = new ProxyConfigProvider(); | ||
|
|
||
| constructor(vercel: UseAssistantHelpers) { | ||
| super(new VercelUseAssistantThreadRuntime(vercel)); | ||
| } | ||
|
|
||
| public set vercel(vercel: UseAssistantHelpers) { | ||
| this.thread.vercel = vercel; | ||
| } | ||
|
|
||
| public onVercelUpdated() { | ||
| return this.thread.onVercelUpdated(); | ||
| } | ||
|
|
||
| public getModelConfig() { | ||
| return this._proxyConfigProvider.getModelConfig(); | ||
| } | ||
|
|
||
| public registerModelConfigProvider(provider: ModelConfigProvider) { | ||
| return this._proxyConfigProvider.registerModelConfigProvider(provider); | ||
| } | ||
|
|
||
| public switchToThread(threadId: string | null) { | ||
| if (threadId) { | ||
| throw new Error("VercelAIRuntime does not yet support switching threads"); | ||
| } | ||
|
|
||
| // clear the vercel state (otherwise, it will be captured by the MessageRepository) | ||
| this.thread.vercel.messages = []; | ||
| this.thread.vercel.input = ""; | ||
| this.thread.vercel.setMessages([]); | ||
| this.thread.vercel.setInput(""); | ||
|
|
||
| this.thread = new VercelUseAssistantThreadRuntime(this.thread.vercel); | ||
| } | ||
| } |
110 changes: 110 additions & 0 deletions
110
packages/react-ai-sdk/src/ui/use-assistant/VercelUseAssistantThreadRuntime.tsx
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,110 @@ | ||
| import { | ||
| type ReactThreadRuntime, | ||
| type Unsubscribe, | ||
| type AppendMessage, | ||
| type ThreadMessage, | ||
| } from "@assistant-ui/react"; | ||
| import { type StoreApi, type UseBoundStore, create } from "zustand"; | ||
| import { useVercelAIComposerSync } from "../utils/useVercelAIComposerSync"; | ||
| import { useVercelAIThreadSync } from "../utils/useVercelAIThreadSync"; | ||
| import { UseAssistantHelpers } from "@ai-sdk/react"; | ||
| import { hasUpcomingMessage } from "./VercelUseAssistantRuntime"; | ||
|
|
||
| const EMPTY_BRANCHES: readonly string[] = Object.freeze([]); | ||
|
|
||
| export class VercelUseAssistantThreadRuntime implements ReactThreadRuntime { | ||
| private _subscriptions = new Set<() => void>(); | ||
|
|
||
| private useVercel: UseBoundStore<StoreApi<{ vercel: UseAssistantHelpers }>>; | ||
|
|
||
| public messages: readonly ThreadMessage[] = []; | ||
| public isRunning = false; | ||
|
|
||
| constructor(public vercel: UseAssistantHelpers) { | ||
| this.useVercel = create(() => ({ | ||
| vercel, | ||
| })); | ||
| } | ||
|
|
||
| public getBranches(): readonly string[] { | ||
| return EMPTY_BRANCHES; | ||
| } | ||
|
|
||
| public switchToBranch(): void { | ||
| throw new Error( | ||
| "VercelUseAssistantRuntime does not support switching branches.", | ||
| ); | ||
| } | ||
|
|
||
| public async append(message: AppendMessage): Promise<void> { | ||
| // add user message | ||
| if (message.content.length !== 1 || message.content[0]?.type !== "text") | ||
| throw new Error("VercelUseAssistantRuntime only supports text content."); | ||
|
|
||
| if (message.parentId !== (this.messages.at(-1)?.id ?? null)) | ||
| throw new Error( | ||
| "VercelUseAssistantRuntime does not support editing messages.", | ||
| ); | ||
|
|
||
| await this.vercel.append({ | ||
| role: "user", | ||
| content: message.content[0].text, | ||
| }); | ||
| } | ||
|
|
||
| public async startRun(): Promise<void> { | ||
| throw new Error("VercelUseAssistantRuntime does not support reloading."); | ||
| } | ||
|
|
||
| public cancelRun(): void { | ||
| const previousMessage = this.vercel.messages.at(-1); | ||
|
|
||
| this.vercel.stop(); | ||
| if (previousMessage?.role === "user") { | ||
| this.vercel.setInput(previousMessage.content); | ||
| } | ||
| } | ||
|
|
||
| public subscribe(callback: () => void): Unsubscribe { | ||
| this._subscriptions.add(callback); | ||
| return () => this._subscriptions.delete(callback); | ||
| } | ||
|
|
||
| public onVercelUpdated() { | ||
| if (this.useVercel.getState().vercel !== this.vercel) { | ||
| this.useVercel.setState({ vercel: this.vercel }); | ||
| } | ||
| } | ||
|
|
||
| private updateData = (isRunning: boolean, vm: ThreadMessage[]) => { | ||
| if (hasUpcomingMessage(isRunning, vm)) { | ||
| vm.push({ | ||
| id: "__optimistic__result", | ||
| createdAt: new Date(), | ||
| status: "in_progress", | ||
| role: "assistant", | ||
| content: [{ type: "text", text: "" }], | ||
| }); | ||
| } | ||
|
|
||
| this.messages = vm; | ||
| this.isRunning = isRunning; | ||
|
|
||
| for (const callback of this._subscriptions) callback(); | ||
| }; | ||
|
|
||
| unstable_synchronizer = () => { | ||
| const { vercel } = this.useVercel(); | ||
|
|
||
| useVercelAIThreadSync(vercel, this.updateData); | ||
| useVercelAIComposerSync(vercel); | ||
|
|
||
| return null; | ||
| }; | ||
|
|
||
| addToolResult() { | ||
| throw new Error( | ||
| "VercelUseAssistantRuntime does not support adding tool results.", | ||
| ); | ||
| } | ||
| } |
6 changes: 4 additions & 2 deletions
6
packages/react-ai-sdk/src/ui/use-assistant/useVercelUseAssistantRuntime.tsx
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
46 changes: 46 additions & 0 deletions
46
packages/react-ai-sdk/src/ui/use-chat/VercelUseChatRuntime.tsx
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,46 @@ | ||
| import { INTERNAL } from "@assistant-ui/react"; | ||
| import { ModelConfigProvider } from "@assistant-ui/react"; | ||
| import { useChat } from "@ai-sdk/react"; | ||
| import { VercelUseChatThreadRuntime } from "./VercelUseChatThreadRuntime"; | ||
|
|
||
| const { ProxyConfigProvider, BaseAssistantRuntime } = INTERNAL; | ||
|
|
||
| export class VercelUseChatRuntime extends BaseAssistantRuntime<VercelUseChatThreadRuntime> { | ||
| private readonly _proxyConfigProvider = new ProxyConfigProvider(); | ||
|
|
||
| constructor(vercel: ReturnType<typeof useChat>) { | ||
| super(new VercelUseChatThreadRuntime(vercel)); | ||
| } | ||
|
|
||
| public set vercel(vercel: ReturnType<typeof useChat>) { | ||
| this.thread.vercel = vercel; | ||
| } | ||
|
|
||
| public onVercelUpdated() { | ||
| return this.thread.onVercelUpdated(); | ||
| } | ||
|
|
||
| public getModelConfig() { | ||
| return this._proxyConfigProvider.getModelConfig(); | ||
| } | ||
|
|
||
| public registerModelConfigProvider(provider: ModelConfigProvider) { | ||
| return this._proxyConfigProvider.registerModelConfigProvider(provider); | ||
| } | ||
|
|
||
| public switchToThread(threadId: string | null) { | ||
| if (threadId) { | ||
| throw new Error( | ||
| "VercelAIRuntime does not yet support switching threads.", | ||
| ); | ||
| } | ||
|
|
||
| // clear the vercel state (otherwise, it will be captured by the MessageRepository) | ||
| this.thread.vercel.messages = []; | ||
| this.thread.vercel.input = ""; | ||
| this.thread.vercel.setMessages([]); | ||
| this.thread.vercel.setInput(""); | ||
|
|
||
| this.thread = new VercelUseChatThreadRuntime(this.thread.vercel); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
4 changes: 2 additions & 2 deletions
4
packages/react-ai-sdk/src/ui/use-chat/useVercelUseChatRuntime.tsx
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters