feat(askai): ai-sdk v5 (#2746)

* feat(askai): Initial ai-sdk v5 work

* Add sdk version header, code cleanup

* feat(askai): Get ai sdk v5 working

* fix: tests bundlesize

* fix: not all streaming states were showing

* fix: Show reasoning state, fix CSS issues, extract link sources for multiple text parts

* Remove important for removing text decoration

bundlesize.config.json

@@ -6,11 +6,11 @@
     },
     {
       "path": "packages/docsearch-react/dist/umd/index.js",
-      "maxSize": "75 kB"
+      "maxSize": "108 kB"
     },
     {
       "path": "packages/docsearch-js/dist/umd/index.js",
-      "maxSize": "90 kB"
+      "maxSize": "121 kB"
     }
   ]
 }

packages/docsearch-css/src/modal.css

@@ -782,7 +782,6 @@ assistive tech users */
   display: flex;
   flex: 1 1 auto;
   font-weight: 400;
-  line-height: 1.2em;
   overflow-x: hidden;
   position: relative;
   text-overflow: ellipsis;
@@ -800,6 +799,10 @@ assistive tech users */
   text-overflow: ellipsis;
 }
 
+.DocSearch-Hit-AskAIButton-title mark {
+  text-decoration: none;
+}
+
 @keyframes fade-in {
   0% {
     opacity: 0;
@@ -938,7 +941,14 @@ assistive tech users */
   font-weight: 400;
 }
 
-.DocSearch-AskAiScreen-Error svg {
+.DocSearch-AskAiScreen-MessageContent {
+  display: flex;
+  flex-direction: column;
+  row-gap: 1em;
+}
+
+.DocSearch-AskAiScreen-Error svg,
+.DocSearch-AskAiScreen-MessageContent-Tool svg {
   width: 16px;
   height: 16px;
   flex-shrink: 0;
@@ -1219,10 +1229,13 @@ assistive tech users */
   color: var(--docsearch-muted-color);
 }
 
+.DocSearch-AskAiScreen-MessageContent-Reasoning svg {
+  color: var(--docsearch-icon-color);
+}
+
 .DocSearch-AskAiScreen-MessageContent-Tool {
   display: flex;
-  padding: 1em 0;
+  align-items: baseline;
-  align-items: center;
   width: 100%;
   color: var(--docsearch-muted-color);
 }

packages/docsearch-react/package.json

@@ -36,9 +36,10 @@
     "watch": "nodemon --watch src --ext ts,tsx,js,jsx,json --ignore dist/ --ignore node_modules/ --verbose --delay 250ms --exec \"yarn on:change\""
   },
   "dependencies": {
-    "@ai-sdk/react": "^1.2.12",
+    "@ai-sdk/react": "^2.0.30",
     "@algolia/autocomplete-core": "1.19.2",
     "@docsearch/css": "4.0.0-beta.8",
+    "ai": "^5.0.30",
     "algoliasearch": "^5.28.0",
     "marked": "^15.0.12"
   },

packages/docsearch-react/src/AskAiScreen.tsx

@@ -1,5 +1,5 @@
 import type { UseChatHelpers } from '@ai-sdk/react';
-import React, { type JSX, useState, useEffect, useMemo } from 'react';
+import React, { type JSX, useMemo, useState, useEffect } from 'react';
 
 import { AggregatedSearchBlock } from './AggregatedSearchBlock';
 import { AlertIcon, LoadingIcon, SearchIcon } from './icons';
@@ -7,7 +7,8 @@ import { MemoizedMarkdown } from './MemoizedMarkdown';
 import type { ScreenStateProps } from './ScreenState';
 import type { StoredSearchPlugin } from './stored-searches';
 import type { InternalDocSearchHit, StoredAskAiState } from './types';
-import { extractLinksFromText } from './utils/ai';
+import type { AIMessage } from './types/AskiAi';
+import { extractLinksFromMessage, getMessageContent } from './utils/ai';
 import { groupConsecutiveToolResults } from './utils/groupConsecutiveToolResults';
 
 export type AskAiScreenTranslations = Partial<{
@@ -46,8 +47,8 @@ export type AskAiScreenTranslations = Partial<{
 }>;
 
 type AskAiScreenProps = Omit<ScreenStateProps<InternalDocSearchHit>, 'translations'> & {
-  messages: UseChatHelpers['messages'];
+  messages: AIMessage[];
-  status: UseChatHelpers['status'];
+  status: UseChatHelpers<AIMessage>['status'];
   askAiStreamError: Error | null;
   askAiFetchError: Error | undefined;
   translations?: AskAiScreenTranslations;
@@ -59,8 +60,8 @@ interface AskAiScreenHeaderProps {
 
 interface Exchange {
   id: string;
-  userMessage: UseChatHelpers['messages'][number];
+  userMessage: AIMessage;
-  assistantMessage: UseChatHelpers['messages'][number] | null;
+  assistantMessage: AIMessage | null;
 }
 
 function AskAiScreenHeader({ disclaimerText }: AskAiScreenHeaderProps): JSX.Element {
@@ -71,7 +72,7 @@ interface AskAiExchangeCardProps {
   exchange: Exchange;
   askAiStreamError: Error | null;
   isLastExchange: boolean;
-  loadingStatus: UseChatHelpers['status'];
+  loadingStatus: UseChatHelpers<AIMessage>['status'];
   onSearchQueryClick: (query: string) => void;
   translations: AskAiScreenTranslations;
   conversations: StoredSearchPlugin<StoredAskAiState>;
@@ -92,20 +93,25 @@ function AskAiExchangeCard({
 
   const showActions = !isLastExchange || (isLastExchange && loadingStatus === 'ready' && Boolean(assistantMessage));
 
-  const urlsToDisplay = React.useMemo(() => extractLinksFromText(assistantMessage?.content || ''), [assistantMessage]);
+  const assistantContent = useMemo(() => getMessageContent(assistantMessage), [assistantMessage]);
+  const userContent = useMemo(() => getMessageContent(userMessage), [userMessage]);
+
+  const urlsToDisplay = React.useMemo(() => extractLinksFromMessage(assistantMessage), [assistantMessage]);
 
   const displayParts = React.useMemo(() => {
-    if (!Array.isArray(assistantMessage?.parts)) {
-      return assistantMessage?.content ? [assistantMessage?.content] : [];
-    }
     return groupConsecutiveToolResults(assistantMessage?.parts || []);
   }, [assistantMessage]);
 
+  const isThinking =
+    ['submitted', 'streaming'].includes(loadingStatus) &&
+    isLastExchange &&
+    !displayParts.some((part) => part.type !== 'step-start');
+
   return (
     <div className="DocSearch-AskAiScreen-Response-Container">
       <div className="DocSearch-AskAiScreen-Response">
         <div className="DocSearch-AskAiScreen-Message DocSearch-AskAiScreen-Message--user">
-          <p className="DocSearch-AskAiScreen-Query">{userMessage.content}</p>
+          <p className="DocSearch-AskAiScreen-Query">{userContent?.text ?? ''}</p>
         </div>
         <div className="DocSearch-AskAiScreen-Message DocSearch-AskAiScreen-Message--assistant">
           <div className="DocSearch-AskAiScreen-MessageContent">
@@ -120,127 +126,113 @@ function AskAiExchangeCard({
                 />
               </div>
             )}
-            {loadingStatus === 'submitted' && isLastExchange && (
+            {isThinking && (
               <div className="DocSearch-AskAiScreen-MessageContent-Reasoning">
                 <span className="shimmer">{translations.thinkingText || 'Thinking...'}</span>
               </div>
             )}
-            {Array.isArray(displayParts)
+            {displayParts.map((part, idx) => {
-              ? displayParts.map((part, idx) => {
+              const index = idx;
-                  const index = idx;
 
-                  if (typeof part === 'string') {
+              if (typeof part === 'string') {
-                    return (
+                return (
-                      <MemoizedMarkdown
+                  <MemoizedMarkdown
-                        key={index}
+                    key={index}
-                        content={part}
+                    content={part}
-                        copyButtonText={translations.copyButtonText || 'Copy'}
+                    copyButtonText={translations.copyButtonText || 'Copy'}
-                        copyButtonCopiedText={translations.copyButtonCopiedText || 'Copied!'}
+                    copyButtonCopiedText={translations.copyButtonCopiedText || 'Copied!'}
-                        isStreaming={loadingStatus === 'streaming'}
+                    isStreaming={loadingStatus === 'streaming'}
-                      />
+                  />
-                    );
+                );
-                  }
+              }
 
-                  // aggregated tool call rendering
+              if (part.type === 'aggregated-tool-call') {
-                  if (part && (part as any).type === 'aggregated-tool-call') {
+                return (
-                    return (
+                  <AggregatedSearchBlock
-                      <AggregatedSearchBlock
+                    key={index}
-                        key={index}
+                    queries={part.queries}
-                        queries={(part as any).queries}
+                    translations={translations}
-                        translations={translations}
+                    onSearchQueryClick={onSearchQueryClick}
-                        onSearchQueryClick={onSearchQueryClick}
+                  />
-                      />
+                );
-                    );
+              }
-                  }
 
-                  if (part.type === 'reasoning' && assistantMessage?.parts?.length === 1) {
+              if (part.type === 'reasoning' && part.state === 'streaming') {
+                return (
+                  <div key={index} className="DocSearch-AskAiScreen-MessageContent-Reasoning shimmer">
+                    <LoadingIcon className="DocSearch-AskAiScreen-SmallerLoadingIcon" />
+                    <span className="shimmer">Reasoning...</span>
+                  </div>
+                );
+              }
+
+              if (part.type === 'text') {
+                return (
+                  <MemoizedMarkdown
+                    key={index}
+                    content={part.text}
+                    copyButtonText={translations.copyButtonText || 'Copy'}
+                    copyButtonCopiedText={translations.copyButtonCopiedText || 'Copied!'}
+                    isStreaming={part.state === 'streaming'}
+                  />
+                );
+              }
+              if (part.type === 'tool-searchIndex') {
+                switch (part.state) {
+                  case 'input-streaming':
                     return (
-                      <div key={index} className="DocSearch-AskAiScreen-MessageContent-Reasoning shimmer">
+                      <div key={index} className="DocSearch-AskAiScreen-MessageContent-Tool Tool--PartialCall shimmer">
-                        <span className="shimmer">Reasoning...</span>
+                        <LoadingIcon className="DocSearch-AskAiScreen-SmallerLoadingIcon" />
+                        <span>{translations.preToolCallText || 'Searching...'}</span>
                       </div>
                     );
-                  }
+                  case 'input-available':
-
-                  if (part.type === 'text') {
                     return (
-                      <MemoizedMarkdown
+                      <div key={index} className="DocSearch-AskAiScreen-MessageContent-Tool Tool--Call shimmer">
-                        key={index}
+                        <LoadingIcon className="DocSearch-AskAiScreen-SmallerLoadingIcon" />
-                        content={part.text}
+                        <span>
-                        copyButtonText={translations.copyButtonText || 'Copy'}
+                          {`${translations.duringToolCallText || 'Searching for '} "${part.input || ''}" ...`}
-                        copyButtonCopiedText={translations.copyButtonCopiedText || 'Copied!'}
+                        </span>
-                        isStreaming={loadingStatus === 'streaming'}
+                      </div>
-                      />
                     );
-                  }
+                  case 'output-available':
-                  if (part.type === 'tool-invocation') {
-                    const { toolInvocation } = part;
-                    if (toolInvocation.toolName === 'searchIndex') {
-                      switch (toolInvocation.state) {
-                        case 'partial-call':
-                          return (
-                            <div
-                              key={index}
-                              className="DocSearch-AskAiScreen-MessageContent-Tool Tool--PartialCall shimmer"
-                            >
-                              <LoadingIcon className="DocSearch-AskAiScreen-SmallerLoadingIcon" />
-                              <span>{translations.preToolCallText || 'Searching...'}</span>
-                            </div>
-                          );
-                        case 'call':
-                          return (
-                            <div key={index} className="DocSearch-AskAiScreen-MessageContent-Tool Tool--Call shimmer">
-                              <LoadingIcon className="DocSearch-AskAiScreen-SmallerLoadingIcon" />
-                              <span>
-                                {`${translations.duringToolCallText || 'Searching for '} "${toolInvocation.args?.query || ''}" ...`}
-                              </span>
-                            </div>
-                          );
-                        case 'result':
-                          return (
-                            <div key={index} className="DocSearch-AskAiScreen-MessageContent-Tool Tool--Result">
-                              <SearchIcon size={18} />
-                              <span>
-                                {`${translations.afterToolCallText || 'Searched for'}`}{' '}
-                                <span
-                                  role="button"
-                                  tabIndex={0}
-                                  className="DocSearch-AskAiScreen-MessageContent-Tool-Query"
-                                  onKeyDown={(e) => {
-                                    if (e.key === 'Enter' || e.key === ' ') {
-                                      e.preventDefault();
-                                      onSearchQueryClick(toolInvocation.args?.query || '');
-                                    }
-                                  }}
-                                  onClick={() => onSearchQueryClick(toolInvocation.args?.query || '')}
-                                >
-                                  {' '}
-                                  &quot;{toolInvocation.args?.query || ''}&quot;
-                                </span>
-                              </span>
-                            </div>
-                          );
-                        default:
-                          return null;
-                      }
-                    }
-                    // fallback for unknown tool, should never happen in theory. :shrug:
                     return (
-                      <span key={index} className="text-sm italic shimmer">
+                      <div key={index} className="DocSearch-AskAiScreen-MessageContent-Tool Tool--Result">
-                        {translations.thinkingText || 'Thinking...'}
+                        <SearchIcon />
-                      </span>
+                        <span>
+                          {`${translations.afterToolCallText || 'Searched for'}`}{' '}
+                          <span
+                            role="button"
+                            tabIndex={0}
+                            className="DocSearch-AskAiScreen-MessageContent-Tool-Query"
+                            onKeyDown={(e) => {
+                              if (e.key === 'Enter' || e.key === ' ') {
+                                e.preventDefault();
+                                onSearchQueryClick(part.output.query || '');
+                              }
+                            }}
+                            onClick={() => onSearchQueryClick(part.output.query || '')}
+                          >
+                            {' '}
+                            &quot;{part.output.query || ''}&quot;
+                          </span>
+                        </span>
+                      </div>
                     );
-                  }
+                  default:
-                  // fallback for unknown part type
+                    break;
-                  return null;
+                }
-                })
+              }
-              : assistantMessage?.content}
+              // fallback for unknown part type
+              return null;
+            })}
           </div>
         </div>
         <div className="DocSearch-AskAiScreen-Answer-Footer">
           <AskAiScreenFooterActions
             id={userMessage?.id || exchange.id}
             showActions={showActions}
-            latestAssistantMessageContent={assistantMessage?.content || null}
+            latestAssistantMessageContent={assistantContent?.text || null}
             translations={translations}
             conversations={conversations}
             onFeedback={onFeedback}

packages/docsearch-react/src/DocSearchModal.tsx

@@ -1,4 +1,3 @@
-import type { Message } from '@ai-sdk/react';
 import { useChat } from '@ai-sdk/react';
 import {
   type AutocompleteSource,
@@ -6,6 +5,7 @@ import {
   createAutocomplete,
   type AutocompleteState,
 } from '@algolia/autocomplete-core';
+import { DefaultChatTransport, lastAssistantMessageIsCompleteWithToolCalls } from 'ai';
 import type { SearchResponse } from 'algoliasearch/lite';
 import React, { type JSX } from 'react';
 
@@ -21,6 +21,7 @@ import type { SearchBoxTranslations } from './SearchBox';
 import { SearchBox } from './SearchBox';
 import { createStoredConversations, createStoredSearches } from './stored-searches';
 import type { DocSearchHit, DocSearchState, InternalDocSearchHit, StoredAskAiState, StoredDocSearchHit } from './types';
+import type { AIMessage } from './types/AskiAi';
 import { useSearchClient } from './useSearchClient';
 import { useTheme } from './useTheme';
 import { useTouchEvents } from './useTouchEvents';
@@ -351,35 +352,39 @@ export function DocSearchModal({
 
   const {
     messages,
-    append,
+    sendMessage,
     status,
     setMessages,
     error: askAiFetchError,
-  } = useChat({
+  } = useChat<AIMessage>({
-    api: ASK_AI_API_URL,
+    sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithToolCalls,
-    sendExtraMessageFields: true,
+    transport: new DefaultChatTransport({
-    fetch: async (input, init) => {
+      api: ASK_AI_API_URL,
-      if (!USE_ASK_AI_TOKEN) {
+      headers: async (): Promise<Record<string, string>> => {
-        return fetch(input, init);
+        if (!askAiConfigurationId) {
-      }
+          throw new Error('Ask AI assistant ID is required');
+        }
 
-      if (!askAiConfigurationId) {
+        let token: string | null = null;
-        throw new Error('Ask AI assistant ID is required');
-      }
-      const token = await getValidToken({ assistantId: askAiConfigurationId });
-      const headers = new Headers(init.headers);
-      headers.set('authorization', `TOKEN ${token}`);
 
-      return fetch(input, { ...init, headers });
+        if (USE_ASK_AI_TOKEN) {
-    },
+          token = await getValidToken({
-    headers: {
+            assistantId: askAiConfigurationId,
-      'Content-Type': 'application/json',
+          });
-      'X-Algolia-API-Key': askAiConfig?.apiKey || apiKey,
+        }
-      'X-Algolia-Application-Id': askAiConfig?.appId || appId,
+
-      'X-Algolia-Index-Name': askAiConfig?.indexName || defaultIndexName,
+        return {
-      'X-Algolia-Assistant-Id': askAiConfigurationId || '',
+          ...(token ? { authorization: `TOKEN ${token}` } : {}),
-    },
+          // 'Content-Type': 'application/json',
-    body: askAiSearchParameters ? { searchParameters: askAiSearchParameters } : {},
+          'X-Algolia-API-Key': askAiConfig?.apiKey || apiKey,
+          'X-Algolia-Application-Id': askAiConfig?.appId || appId,
+          'X-Algolia-Index-Name': askAiConfig?.indexName || defaultIndexName,
+          'X-Algolia-Assistant-Id': askAiConfigurationId || '',
+          'X-AI-SDK-Version': 'v5',
+        };
+      },
+      body: askAiSearchParameters ? { searchParameters: askAiSearchParameters } : {},
+    }),
     onError(streamError) {
       setAskAiStreamError(streamError);
     },
@@ -392,7 +397,11 @@ export function DocSearchModal({
     }
     // if we just transitioned from "streaming" → "ready", persist
     if (prevStatus.current === 'streaming' && status === 'ready') {
-      conversations.add(buildDummyAskAiHit(messages[0].content, messages));
+      for (const part of messages[0].parts) {
+        if (part.type === 'text') {
+          conversations.add(buildDummyAskAiHit(part.text, messages));
+        }
+      }
     }
     prevStatus.current = status;
   }, [status, messages, conversations, disableUserPersonalization]);
@@ -465,9 +474,14 @@ export function DocSearchModal({
   const handleAskAiToggle = React.useCallback(
     (toggle: boolean, query: string) => {
       onAskAiToggle(toggle);
-      append({
+      sendMessage({
         role: 'user',
-        content: query,
+        parts: [
+          {
+            type: 'text',
+            text: query,
+          },
+        ],
       });
 
       if (dropdownRef.current) {
@@ -486,7 +500,7 @@ export function DocSearchModal({
         autocompleteRef.current.setQuery('');
       }
     },
-    [onAskAiToggle, append],
+    [onAskAiToggle, sendMessage],
   );
 
   // feedback handler
@@ -532,7 +546,7 @@ export function DocSearchModal({
             canHandleAskAi,
           });
 
-          const recentConversationSource: Array<AutocompleteSource<InternalDocSearchHit & { messages?: Message[] }>> =
+          const recentConversationSource: Array<AutocompleteSource<InternalDocSearchHit & { messages?: AIMessage[] }>> =
             canHandleAskAi
               ? [
                   {

packages/docsearch-react/src/ScreenState.tsx

@@ -15,6 +15,7 @@ import type { StartScreenTranslations } from './StartScreen';
 import { StartScreen } from './StartScreen';
 import type { StoredSearchPlugin } from './stored-searches';
 import type { InternalDocSearchHit, StoredAskAiState, StoredDocSearchHit } from './types';
+import type { AIMessage } from './types/AskiAi';
 
 export type ScreenStateTranslations = Partial<{
   errorScreen: ErrorScreenTranslations;
@@ -37,8 +38,8 @@ export interface ScreenStateProps<TItem extends BaseItem>
   inputRef: React.MutableRefObject<HTMLInputElement | null>;
   hitComponent: DocSearchProps['hitComponent'];
   indexName: DocSearchProps['indexName'];
-  messages: UseChatHelpers['messages'];
+  messages: UseChatHelpers<AIMessage>['messages'];
-  status: UseChatHelpers['status'];
+  status: UseChatHelpers<AIMessage>['status'];
   askAiStreamError: Error | null;
   askAiFetchError: Error | undefined;
   disableUserPersonalization: boolean;

packages/docsearch-react/src/SearchBox.tsx

@@ -6,6 +6,7 @@ import { MAX_QUERY_SIZE } from './constants';
 import { LoadingIcon, CloseIcon, SearchIcon, SparklesIcon } from './icons';
 import { BackIcon } from './icons/BackIcon';
 import type { InternalDocSearchHit } from './types';
+import type { AIMessage } from './types/AskiAi';
 
 export type SearchBoxTranslations = Partial<{
   clearButtonTitle: string;
@@ -32,7 +33,7 @@ interface SearchBoxProps
   onAskAgain: (query: string) => void;
   placeholder: string;
   isAskAiActive: boolean;
-  askAiStatus: UseChatHelpers['status'];
+  askAiStatus: UseChatHelpers<AIMessage>['status'];
   isFromSelection: boolean;
   translations?: SearchBoxTranslations;
 }

packages/docsearch-react/src/__tests__/askai.test.tsx

@@ -1,4 +1,5 @@
 import { render } from '@testing-library/react';
+import type { UIMessage } from 'ai';
 import React from 'react';
 import { describe, it, expect } from 'vitest';
 import '@testing-library/jest-dom/vitest';
@@ -22,7 +23,18 @@ const baseProps = {
 
 describe('AskAiScreen', () => {
   it('displays stream errors in the latest exchange', () => {
-    const messages = [{ id: '1', role: 'user', content: 'hello' } as any];
+    const messages: UIMessage[] = [
+      {
+        id: '1',
+        role: 'user',
+        parts: [
+          {
+            type: 'text',
+            text: 'hello',
+          },
+        ],
+      },
+    ];
 
     const { getByText } = render(
       <AskAiScreen

packages/docsearch-react/src/__tests__/utils.test.ts

@@ -1,6 +1,7 @@
 import { describe, it, expect, beforeEach } from 'vitest';
 
-import { extractLinksFromText } from '../utils/ai';
+import type { AIMessage } from '../types/AskiAi';
+import { extractLinksFromMessage } from '../utils/ai';
 import {
   createObjectStorage,
   createStorage,
@@ -12,12 +13,32 @@ import {
 describe('utils', () => {
   describe('extractLinksFromText', () => {
     it('returns an empty array when no links are present', () => {
-      expect(extractLinksFromText('hello world')).toEqual([]);
+      const message: AIMessage = {
+        id: '123',
+        role: 'assistant',
+        parts: [
+          {
+            type: 'text',
+            text: 'hello world',
+          },
+        ],
+      };
+      expect(extractLinksFromMessage(message)).toEqual([]);
     });
 
     it('extracts markdown and bare URLs', () => {
       const text = 'See [DocSearch](https://docsearch.algolia.com) and https://example.com/docs.';
-      expect(extractLinksFromText(text)).toEqual([
+      const message: AIMessage = {
+        id: '123',
+        role: 'assistant',
+        parts: [
+          {
+            type: 'text',
+            text,
+          },
+        ],
+      };
+      expect(extractLinksFromMessage(message)).toEqual([
         { url: 'https://docsearch.algolia.com', title: 'DocSearch' },
         { url: 'https://example.com/docs' },
       ]);
@@ -25,25 +46,45 @@ describe('utils', () => {
 
     it('deduplicates repeated links and trims punctuation', () => {
       const text = 'Check https://algolia.com, https://algolia.com!';
-      expect(extractLinksFromText(text)).toEqual([{ url: 'https://algolia.com' }]);
+      const message: AIMessage = {
+        id: '123',
+        role: 'assistant',
+        parts: [
+          {
+            type: 'text',
+            text,
+          },
+        ],
+      };
+      expect(extractLinksFromMessage(message)).toEqual([{ url: 'https://algolia.com' }]);
     });
 
     it('does not return links from within code snippets', () => {
       const text = `
-See [Example Docs](https://example.com/docs)
+  See [Example Docs](https://example.com/docs)
 
-This is also ignored \`https://ignored.com\`
+  This is also ignored \`https://ignored.com\`
 
-\`\`\`js
+  \`\`\`js
-  const DOCS_LINK = 'https://algolia.com/doc'
+    const DOCS_LINK = 'https://algolia.com/doc'
-\`\`\`
+  \`\`\`
 
-https://docsearch.algolia.com
+  https://docsearch.algolia.com
 
-https://docsearch.algolia.com/configuration?version=beta
+  https://docsearch.algolia.com/configuration?version=beta
-`;
+  `;
+      const message: AIMessage = {
+        id: '123',
+        role: 'assistant',
+        parts: [
+          {
+            type: 'text',
+            text,
+          },
+        ],
+      };
 
-      const output = extractLinksFromText(text);
+      const output = extractLinksFromMessage(message);
 
       expect(output).toEqual([
         { url: 'https://example.com/docs', title: 'Example Docs' },
@@ -51,6 +92,36 @@ https://docsearch.algolia.com/configuration?version=beta
         { url: 'https://docsearch.algolia.com/configuration?version=beta' },
       ]);
     });
+
+    it('ignores parts that arent text', () => {
+      const text = 'See [DocSearch](https://docsearch.algolia.com)';
+      const message: AIMessage = {
+        id: '123',
+        role: 'assistant',
+        parts: [
+          {
+            type: 'reasoning',
+            state: 'done',
+            text: 'Need to search using the [Algolia](https://algolia.com) searchIndex tool',
+          },
+          {
+            type: 'tool-searchIndex',
+            input: 'What is DocSearch',
+            state: 'output-available',
+            output: {
+              query: 'DocSearch',
+            },
+            toolCallId: 'searchIndex-testing',
+          },
+          {
+            type: 'text',
+            text,
+          },
+        ],
+      };
+
+      expect(extractLinksFromMessage(message)).toEqual([{ url: 'https://docsearch.algolia.com', title: 'DocSearch' }]);
+    });
   });
 
   describe('createObjectStorage', () => {
@@ -103,7 +174,9 @@ https://docsearch.algolia.com/configuration?version=beta
       const storage = createStorage<{ data: string }>(testKey);
 
       // Create a large dataset that might cause quota issues
-      const largeArray = Array.from({ length: 1000 }, (_, i) => ({ data: `test-data-${i}`.repeat(100) }));
+      const largeArray = Array.from({ length: 1000 }, (_, i) => ({
+        data: `test-data-${i}`.repeat(100),
+      }));
 
       // This should not throw an error even if quota is exceeded
       expect(() => {

packages/docsearch-react/src/stored-searches.ts

@@ -56,12 +56,10 @@ export function createStoredConversations<TItem extends StoredAskAiState>({
 
   return {
     add(item: TItem): void {
-      const { objectID, messages } = item;
+      const { objectID, query } = item;
 
       // check if this query is already saved
-      const isQueryAlreadySaved = items.findIndex(
+      const isQueryAlreadySaved = items.findIndex((x) => x.objectID === objectID || x.query === query);
-        (x) => x.objectID === objectID || x.messages?.[0]?.content === messages?.[0]?.content,
-      );
 
       if (isQueryAlreadySaved > -1) {
         items[isQueryAlreadySaved] = item;

packages/docsearch-react/src/types/AskiAi.ts

@@ -0,0 +1,24 @@
+import type { UIMessage } from '@ai-sdk/react';
+import type { UIDataTypes, UIMessagePart } from 'ai';
+
+export interface SearchIndexTool {
+  input: string;
+  output: {
+    query: string;
+  };
+}
+
+export type AIMessage = UIMessage<
+  unknown,
+  UIDataTypes,
+  {
+    searchIndex: SearchIndexTool;
+  }
+>;
+
+export type AIMessagePart = UIMessagePart<
+  UIDataTypes,
+  {
+    searchIndex: SearchIndexTool;
+  }
+>;

packages/docsearch-react/src/types/StoredDocSearchHit.ts

@@ -1,10 +1,9 @@
-import type { Message } from '@ai-sdk/react';
+import type { AIMessage } from './AskiAi';
-
 import type { DocSearchHit } from './DocSearchHit';
 
 export type StoredDocSearchHit = Omit<DocSearchHit, '_highlightResult' | '_snippetResult'>;
 
-export type StoredAskAiMessage = Message & {
+export type StoredAskAiMessage = AIMessage & {
   /** Optional user feedback on this assistant message. */
   feedback?: 'dislike' | 'like';
 };

packages/docsearch-react/src/utils/ai.ts

@@ -1,6 +1,7 @@
-import type { Message } from '@ai-sdk/react';
+import type { TextUIPart } from 'ai';
 
 import type { StoredAskAiState } from '../types';
+import type { AIMessage } from '../types/AskiAi';
 
 type ExtractedLink = {
   url: string;
@@ -8,68 +9,89 @@ type ExtractedLink = {
 };
 
 // utility to extract links (markdown and bare urls) from a string
-export function extractLinksFromText(text: string): ExtractedLink[] {
+export function extractLinksFromMessage(message: AIMessage | null): ExtractedLink[] {
-  const markdownLinkRegex = /\[([^\]]*)\]\(([^)]+)\)/g;
-  const plainLinkRegex = /(?<!\]\()https?:\/\/[^\s<>"{}|\\^`[\]]+/g;
   const links: ExtractedLink[] = [];
   // Used to dedupe multiple urls
   const seen = new Set<string>();
 
-  // Strip out all code blocks e.g. ```
+  if (!message) {
-  const textWithoutCodeBlocks = text.replace(/```[\s\S]*?```/g, '');
+    return [];
-
-  // Strip out all inline code blocks e.g. `
-  const cleanText = textWithoutCodeBlocks.replace(/`[^`]*`/g, '');
-
-  // Get all markdown based links e.g. []()
-  const markdownMatches = cleanText.matchAll(markdownLinkRegex);
-
-  // Parses the title and url from the found links
-  for (const match of markdownMatches) {
-    const title = match[1].trim();
-    const url = match[2];
-
-    if (!seen.has(url)) {
-      seen.add(url);
-      links.push({ url, title: title || undefined });
-    }
   }
 
-  // Get all "plain" links e.g. https://algolia.com/doc
+  message.parts.forEach((part) => {
-  const plainUrls = cleanText.matchAll(plainLinkRegex);
+    if (part.type !== 'text') {
-
+      return;
-  for (const match of plainUrls) {
-    // Strip any extra punctuation
-    const cleanUrl = match[0].replace(/[.,;:!?]+$/, '');
-
-    if (!seen.has(cleanUrl)) {
-      seen.add(cleanUrl);
-      links.push({ url: cleanUrl });
     }
-  }
+
+    if (part.text.length === 0) {
+      return;
+    }
+
+    const markdownLinkRegex = /\[([^\]]*)\]\(([^)]+)\)/g;
+    const plainLinkRegex = /(?<!\]\()https?:\/\/[^\s<>"{}|\\^`[\]]+/g;
+
+    // Strip out all code blocks e.g. ```
+    const textWithoutCodeBlocks = part.text.replace(/```[\s\S]*?```/g, '');
+
+    // Strip out all inline code blocks e.g. `
+    const cleanText = textWithoutCodeBlocks.replace(/`[^`]*`/g, '');
+
+    // Get all markdown based links e.g. []()
+    const markdownMatches = cleanText.matchAll(markdownLinkRegex);
+
+    // Parses the title and url from the found links
+    for (const match of markdownMatches) {
+      const title = match[1].trim();
+      const url = match[2];
+
+      if (!seen.has(url)) {
+        seen.add(url);
+        links.push({ url, title: title || undefined });
+      }
+    }
+
+    // Get all "plain" links e.g. https://algolia.com/doc
+    const plainUrls = cleanText.matchAll(plainLinkRegex);
+
+    for (const match of plainUrls) {
+      // Strip any extra punctuation
+      const cleanUrl = match[0].replace(/[.,;:!?]+$/, '');
+
+      if (!seen.has(cleanUrl)) {
+        seen.add(cleanUrl);
+        links.push({ url: cleanUrl });
+      }
+    }
+  });
 
   return links;
 }
 
-export const buildDummyAskAiHit = (query: string, messages: Message[]): StoredAskAiState => ({
+export const buildDummyAskAiHit = (query: string, messages: AIMessage[]): StoredAskAiState => {
-  query,
+  const textPart = messages[0].parts.find((part) => part.type === 'text');
-  objectID: messages[0].content,
-  messages,
-  type: 'askAI',
-  anchor: 'stored',
 
-  // dummy content to make it a valid hit
+  return {
-  // this is useful to show it among other hits
+    query,
-  content: null,
+    objectID: textPart?.text ?? '',
-  hierarchy: {
+    messages,
-    lvl0: 'askAI',
+    type: 'askAI',
-    lvl1: messages[0].content, // use first message as hit name
+    anchor: 'stored',
-    lvl2: null,
+    // dummy content to make it a valid hit
-    lvl3: null,
+    // this is useful to show it among other hits
-    lvl4: null,
+    content: null,
-    lvl5: null,
+    hierarchy: {
-    lvl6: null,
+      lvl0: 'askAI',
-  },
+      lvl1: textPart?.text ?? '', // use first message as hit name
-  url: '',
+      lvl2: null,
-  url_without_anchor: '',
+      lvl3: null,
-});
+      lvl4: null,
+      lvl5: null,
+      lvl6: null,
+    },
+    url: '',
+    url_without_anchor: '',
+  };
+};
+
+export const getMessageContent = (message: AIMessage | null): TextUIPart | undefined =>
+  message?.parts.find((part) => part.type === 'text');

packages/docsearch-react/src/utils/groupConsecutiveToolResults.ts

@@ -1,3 +1,5 @@
+import type { AIMessagePart } from '../types/AskiAi';
+
 export interface AggregatedToolCallPart {
   type: 'aggregated-tool-call';
   queries: string[];
@@ -7,30 +9,23 @@ export interface AggregatedToolCallPart {
  * Groups consecutive `searchIndex` tool invocation result parts together.
  * Empty or falsy queries are ignored.
  */
-export function groupConsecutiveToolResults<T extends { type: string }>(parts: T[]): Array<AggregatedToolCallPart | T> {
+export function groupConsecutiveToolResults(parts: AIMessagePart[]): Array<AggregatedToolCallPart | AIMessagePart> {
-  const aggregatedParts: Array<AggregatedToolCallPart | T> = [];
+  const aggregatedParts: Array<AggregatedToolCallPart | AIMessagePart> = [];
 
   for (let i = 0; i < parts.length; i++) {
-    const part: any = parts[i];
+    const part = parts[i];
 
-    if (
+    if (part.type === 'tool-searchIndex' && part.state === 'output-available') {
-      part?.type === 'tool-invocation' &&
-      part.toolInvocation?.toolName === 'searchIndex' &&
-      part.toolInvocation?.state === 'result'
-    ) {
       // build list of consecutive result queries
       const queries: string[] = [];
       let j = i;
       while (j < parts.length) {
-        const candidate: any = parts[j];
+        const candidate = parts[j];
-        if (
+        if (candidate.type === 'tool-searchIndex' && candidate.state === 'output-available') {
-          candidate?.type === 'tool-invocation' &&
+          const q = (candidate.output?.query ?? '').trim();
-          candidate.toolInvocation?.toolName === 'searchIndex' &&
+
-          candidate.toolInvocation?.state === 'result'
-        ) {
-          const q = (candidate.toolInvocation?.args?.query || '').trim();
           // eslint-disable-next-line max-depth
-          if (q) {
+          if (q && q.length > 0) {
             queries.push(q);
           }
           j++;

yarn.lock

@@ -12,6 +12,18 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@ai-sdk/gateway@npm:1.0.15":
+  version: 1.0.15
+  resolution: "@ai-sdk/gateway@npm:1.0.15"
+  dependencies:
+    "@ai-sdk/provider": "npm:2.0.0"
+    "@ai-sdk/provider-utils": "npm:3.0.7"
+  peerDependencies:
+    zod: ^3.25.76 || ^4
+  checksum: 10c0/cdd09f119d6618f00c363a27f51dc466a8a64f57f01bcdd127030a804825bd143b0fef2dbdb7802530865d474f4b9d55855670fecd7f2e6c615a5d9ac9fd6e3b
+  languageName: node
+  linkType: hard
+
 "@ai-sdk/provider-utils@npm:2.2.8":
   version: 2.2.8
   resolution: "@ai-sdk/provider-utils@npm:2.2.8"
@@ -25,6 +37,19 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@ai-sdk/provider-utils@npm:3.0.7":
+  version: 3.0.7
+  resolution: "@ai-sdk/provider-utils@npm:3.0.7"
+  dependencies:
+    "@ai-sdk/provider": "npm:2.0.0"
+    "@standard-schema/spec": "npm:^1.0.0"
+    eventsource-parser: "npm:^3.0.5"
+  peerDependencies:
+    zod: ^3.25.76 || ^4
+  checksum: 10c0/7e709289f9e514a6ba56a9b19764eb124ea1bd36d4b3b3e455a1c05353674c152839a4d3cd061af7a4cc36106bd15859a2346e54d4ed0a861feec3b2c4c21513
+  languageName: node
+  linkType: hard
+
 "@ai-sdk/provider@npm:1.1.3":
   version: 1.1.3
   resolution: "@ai-sdk/provider@npm:1.1.3"
@@ -34,6 +59,15 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@ai-sdk/provider@npm:2.0.0":
+  version: 2.0.0
+  resolution: "@ai-sdk/provider@npm:2.0.0"
+  dependencies:
+    json-schema: "npm:^0.4.0"
+  checksum: 10c0/e50e520016c9fc0a8b5009cadd47dae2f1c81ec05c1792b9e312d7d15479f024ca8039525813a33425c884e3449019fed21043b1bfabd6a2626152ca9a388199
+  languageName: node
+  linkType: hard
+
 "@ai-sdk/react@npm:^1.2.12":
   version: 1.2.12
   resolution: "@ai-sdk/react@npm:1.2.12"
@@ -52,6 +86,24 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@ai-sdk/react@npm:^2.0.30":
+  version: 2.0.30
+  resolution: "@ai-sdk/react@npm:2.0.30"
+  dependencies:
+    "@ai-sdk/provider-utils": "npm:3.0.7"
+    ai: "npm:5.0.30"
+    swr: "npm:^2.2.5"
+    throttleit: "npm:2.1.0"
+  peerDependencies:
+    react: ^18 || ^19 || ^19.0.0-rc
+    zod: ^3.25.76 || ^4
+  peerDependenciesMeta:
+    zod:
+      optional: true
+  checksum: 10c0/e03f85a2840cfe39edf2d53df69bde6a866e2304c3de2016428fbdbd85c05d053d6c610a951f449b020737f03fced4201570c9530f60df56c86ee4263d304f2f
+  languageName: node
+  linkType: hard
+
 "@ai-sdk/ui-utils@npm:1.2.11":
   version: 1.2.11
   resolution: "@ai-sdk/ui-utils@npm:1.2.11"
@@ -2422,12 +2474,13 @@ __metadata:
   version: 0.0.0-use.local
   resolution: "@docsearch/react@workspace:packages/docsearch-react"
   dependencies:
-    "@ai-sdk/react": "npm:^1.2.12"
+    "@ai-sdk/react": "npm:^2.0.30"
     "@algolia/autocomplete-core": "npm:1.19.2"
     "@docsearch/css": "npm:4.0.0-beta.8"
     "@rollup/plugin-replace": "npm:6.0.2"
     "@testing-library/jest-dom": "npm:6.6.3"
     "@testing-library/react": "npm:16.2.0"
+    ai: "npm:^5.0.30"
     algoliasearch: "npm:^5.28.0"
     marked: "npm:^15.0.12"
     nodemon: "npm:^3.1.0"
@@ -4511,6 +4564,13 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@opentelemetry/api@npm:1.9.0":
+  version: 1.9.0
+  resolution: "@opentelemetry/api@npm:1.9.0"
+  checksum: 10c0/9aae2fe6e8a3a3eeb6c1fdef78e1939cf05a0f37f8a4fae4d6bf2e09eb1e06f966ece85805626e01ba5fab48072b94f19b835449e58b6d26720ee19a58298add
+  languageName: node
+  linkType: hard
+
 "@pkgjs/parseargs@npm:^0.11.0":
   version: 0.11.0
   resolution: "@pkgjs/parseargs@npm:0.11.0"
@@ -6287,6 +6347,13 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@standard-schema/spec@npm:^1.0.0":
+  version: 1.0.0
+  resolution: "@standard-schema/spec@npm:1.0.0"
+  checksum: 10c0/a1ab9a8bdc09b5b47aa8365d0e0ec40cc2df6437be02853696a0e377321653b0d3ac6f079a8c67d5ddbe9821025584b1fb71d9cc041a6666a96f1fadf2ece15f
+  languageName: node
+  linkType: hard
+
 "@stylistic/eslint-plugin@npm:2.13.0":
   version: 2.13.0
   resolution: "@stylistic/eslint-plugin@npm:2.13.0"
@@ -7901,6 +7968,20 @@ __metadata:
   languageName: node
   linkType: hard
 
+"ai@npm:5.0.30, ai@npm:^5.0.30":
+  version: 5.0.30
+  resolution: "ai@npm:5.0.30"
+  dependencies:
+    "@ai-sdk/gateway": "npm:1.0.15"
+    "@ai-sdk/provider": "npm:2.0.0"
+    "@ai-sdk/provider-utils": "npm:3.0.7"
+    "@opentelemetry/api": "npm:1.9.0"
+  peerDependencies:
+    zod: ^3.25.76 || ^4
+  checksum: 10c0/2d6b52e28aa49ba5a177fa03855abf1a13b7d97172f91fd0027d20b8629b2411cfe89ae9a881892064bca0afc15bb8e6b3d14488c28e07593d10d871436890af
+  languageName: node
+  linkType: hard
+
 "ajv-formats@npm:^2.1.1":
   version: 2.1.1
   resolution: "ajv-formats@npm:2.1.1"
@@ -12498,6 +12579,13 @@ __metadata:
   languageName: node
   linkType: hard
 
+"eventsource-parser@npm:^3.0.5":
+  version: 3.0.6
+  resolution: "eventsource-parser@npm:3.0.6"
+  checksum: 10c0/70b8ccec7dac767ef2eca43f355e0979e70415701691382a042a2df8d6a68da6c2fca35363669821f3da876d29c02abe9b232964637c1b6635c940df05ada78a
+  languageName: node
+  linkType: hard
+
 "exec-buffer@npm:^3.0.0, exec-buffer@npm:^3.2.0":
   version: 3.2.0
   resolution: "exec-buffer@npm:3.2.0"