refactor: streamline chat API imports and enhance DocsAssistant with settings provider

Author: Crokily

app/api/chat/route.ts

@@ -1,5 +1,5 @@
-import { openai, createOpenAI } from "@ai-sdk/openai";
-import { google, createGoogleGenerativeAI } from "@ai-sdk/google";
+import { createOpenAI } from "@ai-sdk/openai";
+import { createGoogleGenerativeAI } from "@ai-sdk/google";
 import { streamText, UIMessage, convertToModelMessages } from "ai";
 
 // Allow streaming responses up to 30 seconds
@@ -15,7 +15,7 @@ export async function POST(req: Request) {
   }: {
     messages: UIMessage[];
     system?: string; // System message forwarded from AssistantChatTransport
-    tools?: any; // Frontend tools forwarded from AssistantChatTransport
+    tools?: unknown; // Frontend tools forwarded from AssistantChatTransport
     pageContext?: {
       title?: string;
       description?: string;
@@ -67,13 +67,13 @@ export async function POST(req: Request) {
       const customGoogle = createGoogleGenerativeAI({
         apiKey: apiKey,
       });
-      model = customGoogle("models/gemini-1.5-pro-latest");
+      model = customGoogle("models/gemini-2.0-flash");
     } else {
       // Default to OpenAI
       const customOpenAI = createOpenAI({
         apiKey: apiKey,
       });
-      model = customOpenAI("gpt-4o-mini");
+      model = customOpenAI("gpt-4.1-nano");
     }
 
     const result = streamText({

app/components/DocsAssistant.tsx

@@ -5,7 +5,10 @@ import { useAISDKRuntime } from "@assistant-ui/react-ai-sdk";
 import { useChat } from "@ai-sdk/react";
 import { DefaultChatTransport } from "ai";
 import { AssistantModal } from "@/app/components/assistant-ui/assistant-modal";
-import { useAssistantSettings } from "@/app/hooks/useAssistantSettings";
+import {
+  AssistantSettingsProvider,
+  useAssistantSettings,
+} from "@/app/hooks/useAssistantSettings";
 
 interface PageContext {
   title?: string;
@@ -19,13 +22,19 @@ interface DocsAssistantProps {
 }
 
 export function DocsAssistant({ pageContext }: DocsAssistantProps) {
+  return (
+    <AssistantSettingsProvider>
+      <DocsAssistantInner pageContext={pageContext} />
+    </AssistantSettingsProvider>
+  );
+}
+
+function DocsAssistantInner({ pageContext }: DocsAssistantProps) {
   const { provider, openaiApiKey, geminiApiKey } = useAssistantSettings();
 
-  // Use DefaultChatTransport with request-level body configuration
   const chat = useChat({
     transport: new DefaultChatTransport({
       api: "/api/chat",
-      // Use function to ensure dynamic values are captured at request time
       body: () => {
         const apiKey = provider === "openai" ? openaiApiKey : geminiApiKey;
         return {

app/hooks/useAssistantSettings.ts

@@ -0,0 +1,135 @@
+"use client";
+
+import { createContext, useContext, useEffect, useMemo, useState } from "react";
+import type { ReactNode } from "react";
+
+type Provider = "openai" | "gemini";
+
+interface AssistantSettingsState {
+  provider: Provider;
+  openaiApiKey: string;
+  geminiApiKey: string;
+}
+
+interface AssistantSettingsContextValue extends AssistantSettingsState {
+  setProvider: (provider: Provider) => void;
+  setOpenaiApiKey: (key: string) => void;
+  setGeminiApiKey: (key: string) => void;
+}
+
+const SETTINGS_KEY = "assistant-settings-storage";
+
+const defaultSettings: AssistantSettingsState = {
+  provider: "openai",
+  openaiApiKey: "",
+  geminiApiKey: "",
+};
+
+const AssistantSettingsContext = createContext<
+  AssistantSettingsContextValue | undefined
+>(undefined);
+
+const parseStoredSettings = (raw: string | null): AssistantSettingsState => {
+  if (!raw) {
+    return { ...defaultSettings };
+  }
+
+  try {
+    const parsed = JSON.parse(raw) as Partial<AssistantSettingsState>;
+    return {
+      provider: parsed.provider === "gemini" ? "gemini" : "openai",
+      openaiApiKey:
+        typeof parsed.openaiApiKey === "string" ? parsed.openaiApiKey : "",
+      geminiApiKey:
+        typeof parsed.geminiApiKey === "string" ? parsed.geminiApiKey : "",
+    };
+  } catch (error) {
+    console.error(
+      "Failed to parse assistant settings from localStorage",
+      error,
+    );
+    return { ...defaultSettings };
+  }
+};
+
+const readStoredSettings = (): AssistantSettingsState => {
+  if (typeof window === "undefined") {
+    return { ...defaultSettings };
+  }
+
+  const raw = window.localStorage.getItem(SETTINGS_KEY);
+  return parseStoredSettings(raw);
+};
+
+export const AssistantSettingsProvider = ({
+  children,
+}: {
+  children: ReactNode;
+}) => {
+  const [settings, setSettings] = useState<AssistantSettingsState>(() =>
+    readStoredSettings(),
+  );
+
+  useEffect(() => {
+    if (typeof window === "undefined") {
+      return;
+    }
+
+    try {
+      window.localStorage.setItem(SETTINGS_KEY, JSON.stringify(settings));
+    } catch (error) {
+      console.error("Failed to save assistant settings to localStorage", error);
+    }
+  }, [settings]);
+
+  useEffect(() => {
+    if (typeof window === "undefined") {
+      return;
+    }
+
+    const handleStorage = (event: StorageEvent) => {
+      if (event.key !== SETTINGS_KEY) {
+        return;
+      }
+
+      setSettings(parseStoredSettings(event.newValue));
+    };
+
+    window.addEventListener("storage", handleStorage);
+    return () => window.removeEventListener("storage", handleStorage);
+  }, []);
+
+  const value = useMemo(
+    (): AssistantSettingsContextValue => ({
+      ...settings,
+      setProvider: (provider: Provider) => {
+        setSettings((prev) => ({ ...prev, provider }));
+      },
+      setOpenaiApiKey: (key: string) => {
+        setSettings((prev) => ({ ...prev, openaiApiKey: key }));
+      },
+      setGeminiApiKey: (key: string) => {
+        setSettings((prev) => ({ ...prev, geminiApiKey: key }));
+      },
+    }),
+    [settings],
+  );
+
+  return (
+    <AssistantSettingsContext.Provider value={value}>
+      {children}
+    </AssistantSettingsContext.Provider>
+  );
+};
+
+export const useAssistantSettings = () => {
+  const context = useContext(AssistantSettingsContext);
+
+  if (!context) {
+    throw new Error(
+      "useAssistantSettings must be used within an AssistantSettingsProvider",
+    );
+  }
+
+  return context;
+};