fix: reimplement ProxyConfigProvider in AI SDK (#344)

Yonom · Jun 27, 2024 · 179e869 · 179e869
1 parent 6de7c9a
commit 179e869
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 17 deletions.
diff --git a/packages/react-ai-sdk/src/rsc/VercelRSCRuntime.tsx b/packages/react-ai-sdk/src/rsc/VercelRSCRuntime.tsx
@@ -11,6 +11,7 @@ import { type StoreApi, type UseBoundStore, create } from "zustand";
 import type { VercelRSCAdapter } from "./VercelRSCAdapter";
 import type { VercelRSCMessage } from "./VercelRSCMessage";
 import { useVercelRSCSync } from "./useVercelRSCSync";
+import { ModelConfigProvider } from "@assistant-ui/react";
 
 const { ProxyConfigProvider, BaseAssistantRuntime } = INTERNAL;
 
@@ -19,6 +20,8 @@ const EMPTY_BRANCHES: readonly string[] = Object.freeze([]);
 export class VercelRSCRuntime<
   T extends WeakKey = VercelRSCMessage,
 > extends BaseAssistantRuntime<VercelRSCThreadRuntime<T>> {
+  private readonly _proxyConfigProvider = new ProxyConfigProvider();
+
   constructor(adapter: VercelRSCAdapter<T>) {
     super(new VercelRSCThreadRuntime(adapter));
   }
@@ -31,9 +34,12 @@ export class VercelRSCRuntime<
     return this.thread.onAdapterUpdated();
   }
 
-  public registerModelConfigProvider() {
-    // no-op
-    return () => {};
+  public getModelConfig() {
+    return this._proxyConfigProvider.getModelConfig();
+  }
+
+  public registerModelConfigProvider(provider: ModelConfigProvider) {
+    return this._proxyConfigProvider.registerModelConfigProvider(provider);
   }
 
   public switchToThread() {
@@ -42,7 +48,6 @@ export class VercelRSCRuntime<
 }
 
 class VercelRSCThreadRuntime<T extends WeakKey = VercelRSCMessage>
-  extends ProxyConfigProvider
   implements ReactThreadRuntime
 {
   private useAdapter: UseBoundStore<StoreApi<{ adapter: VercelRSCAdapter<T> }>>;
@@ -53,8 +58,6 @@ class VercelRSCThreadRuntime<T extends WeakKey = VercelRSCMessage>
   public messages: ThreadMessage[] = [];
 
   constructor(public adapter: VercelRSCAdapter<T>) {
-    super();
-
     this.useAdapter = create(() => ({
       adapter,
     }));

diff --git a/packages/react-ai-sdk/src/ui/VercelAIRuntime.tsx b/packages/react-ai-sdk/src/ui/VercelAIRuntime.tsx
@@ -12,6 +12,7 @@ import type { VercelHelpers } from "./utils/VercelHelpers";
 import { sliceMessagesUntil } from "./utils/sliceMessagesUntil";
 import { useVercelAIComposerSync } from "./utils/useVercelAIComposerSync";
 import { useVercelAIThreadSync } from "./utils/useVercelAIThreadSync";
+import { ModelConfigProvider } from "@assistant-ui/react";
 
 const { ProxyConfigProvider, MessageRepository, BaseAssistantRuntime } =
   INTERNAL;
@@ -21,6 +22,8 @@ const hasUpcomingMessage = (isRunning: boolean, messages: ThreadMessage[]) => {
 };
 
 export class VercelAIRuntime extends BaseAssistantRuntime<VercelAIThreadRuntime> {
+  private readonly _proxyConfigProvider = new ProxyConfigProvider();
+
   constructor(vercel: VercelHelpers) {
     super(new VercelAIThreadRuntime(vercel));
   }
@@ -33,9 +36,12 @@ export class VercelAIRuntime extends BaseAssistantRuntime<VercelAIThreadRuntime>
     return this.thread.onVercelUpdated();
   }
 
-  public registerModelConfigProvider() {
-    // no-op
-    return () => {};
+  public getModelConfig() {
+    return this._proxyConfigProvider.getModelConfig();
+  }
+
+  public registerModelConfigProvider(provider: ModelConfigProvider) {
+    return this._proxyConfigProvider.registerModelConfigProvider(provider);
   }
 
   public switchToThread(threadId: string | null) {
@@ -53,10 +59,7 @@ export class VercelAIRuntime extends BaseAssistantRuntime<VercelAIThreadRuntime>
   }
 }
 
-class VercelAIThreadRuntime
-  extends ProxyConfigProvider
-  implements ReactThreadRuntime
-{
+class VercelAIThreadRuntime implements ReactThreadRuntime {
   private _subscriptions = new Set<() => void>();
   private repository = new MessageRepository();
   private assistantOptimisticId: string | null = null;
@@ -67,7 +70,6 @@ class VercelAIThreadRuntime
   public isRunning = false;
 
   constructor(public vercel: VercelHelpers) {
-    super();
     this.useVercel = create(() => ({
       vercel,
     }));

diff --git a/packages/react-ui/components/ui/tooltip.tsx b/packages/react-ui/components/ui/tooltip.tsx
@@ -5,8 +5,6 @@ import * as TooltipPrimitive from "@radix-ui/react-tooltip";
 
 import { cn } from "@/lib/utils";
 
-const TooltipProvider = TooltipPrimitive.Provider;
-
 const Tooltip = TooltipPrimitive.Root;
 
 const TooltipTrigger = TooltipPrimitive.Trigger;
@@ -27,4 +25,4 @@ const TooltipContent = React.forwardRef<
 ));
 TooltipContent.displayName = TooltipPrimitive.Content.displayName;
 
-export { Tooltip, TooltipTrigger, TooltipContent, TooltipProvider };
+export { Tooltip, TooltipTrigger, TooltipContent };