refactor(modal): wire modal ask-ai flow to the Typesense hook

Author: tharropoulos

packages/docsearch-react/src/DocSearchModal.tsx

@@ -6,7 +6,6 @@ import {
 } from '@algolia/autocomplete-core';
 import type { InitialAskAiMessage, OnAskAiToggle } from '@docsearch/core';
 import { useTheme } from '@docsearch/core/useTheme';
-import type { ChatRequestOptions } from 'ai';
 import type { SearchResponse } from 'algoliasearch/lite';
 import React, { type JSX } from 'react';
 import type { MultiSearchRequestSchema } from 'typesense/lib/Typesense/Types';
@@ -34,7 +33,6 @@ import type {
 import type { AIMessage, AskAiState } from './types/AskiAi';
 import { useAskAi } from './useAskAi';
 import { useSearchClient } from './useSearchClient';
-import { useSuggestedQuestions } from './useSuggestedQuestions';
 import { useTouchEvents } from './useTouchEvents';
 import { useTrapFocus } from './useTrapFocus';
 import { groupBy, identity, noop, removeHighlightTags, isModifierEvent, scrollTo as scrollToUtils } from './utils';
@@ -468,17 +466,8 @@ export function DocSearchModal({
 
   const searchClient = useSearchClient(transformSearchClient, typesenseServerConfig);
 
-  const askAiConfig = typeof askAi === 'object' ? askAi : null;
-  const askAiConfigurationId = typeof askAi === 'string' ? askAi : askAiConfig?.assistantId || null;
-  const askAiSearchParameters = askAiConfig?.searchParameters;
-  const askAiUseStagingEnv = askAiConfig?.useStagingEnv || false;
+  const askAiConfig = askAi ?? null;
   const [askAiState, setAskAiState] = React.useState<AskAiState>('initial');
-  const suggestedQuestions = useSuggestedQuestions({
-    assistantId: askAiConfigurationId,
-    searchClient,
-    suggestedQuestionsEnabled: askAiConfig?.suggestedQuestions,
-  });
-  const agentStudio = askAiConfig?.agentStudio ?? false;
 
   // Format the `indexes` to be used until `indexName` and `searchParameters` props are fully removed.
   const indexes: DocSearchIndex[] = [];
@@ -505,7 +494,7 @@ export function DocSearchModal({
   // storage
   const conversations = React.useRef(
     createStoredConversations<StoredAskAiState>({
-      key: `__DOCSEARCH_ASKAI_CONVERSATIONS__${askAiConfig?.indexName || defaultIndexName}`,
+      key: `__DOCSEARCH_ASKAI_CONVERSATIONS__${defaultIndexName}`,
       limit: 10,
     }),
   ).current;
@@ -524,14 +513,14 @@ export function DocSearchModal({
 
   const [stoppedStream, setStoppedStream] = React.useState(false);
 
-  const { messages, status, setMessages, sendMessage, stopAskAiStreaming, askAiError, sendFeedback } = useAskAi({
-    assistantId: askAiConfigurationId,
-    apiKey: askAiConfig?.apiKey ?? 'testkey',
-    appId: askAiConfig?.appId ?? 'testappid',
-    indexName: askAiConfig?.indexName || defaultIndexName,
-    searchParameters: askAiSearchParameters,
-    useStagingEnv: askAiUseStagingEnv,
-    agentStudio,
+  const { messages, status, setMessages, sendMessage, stopAskAiStreaming, askAiError } = useAskAi({
+    typesenseServerConfig,
+    storageKey: `__DOCSEARCH_ASKAI_CONVERSATIONS__${defaultIndexName}`,
+    collection: askAiConfig?.collection || defaultIndexName,
+    conversationModelId: askAiConfig?.conversationModelId || '',
+    queryBy: askAiConfig?.queryBy || 'embedding',
+    excludeFields: askAiConfig?.excludeFields || 'embedding',
+    searchParameters: askAiConfig?.searchParameters,
   });
 
   const prevStatus = React.useRef(status);
@@ -544,13 +533,14 @@ export function DocSearchModal({
       // if we stopped the stream, store it on the most recent message
       if (stoppedStream && messages.at(-1)) {
         messages.at(-1)!.metadata = {
+          ...messages.at(-1)!.metadata,
           stopped: true,
         };
       }
 
       for (const part of messages[0].parts) {
         if (part.type === 'text') {
-          conversations.add(buildDummyAskAiHit(part.text, messages));
+          conversations.add(buildDummyAskAiHit(part.text, messages, messages.at(-1)?.metadata?.conversationId));
         }
       }
     }
@@ -659,26 +649,9 @@ export function DocSearchModal({
 
       setStoppedStream(false);
 
-      const messageOptions: ChatRequestOptions = {};
-
-      if (suggestedQuestion) {
-        messageOptions.body = {
-          suggestedQuestionId: suggestedQuestion.objectID,
-        };
-      }
-
-      sendMessage(
-        {
-          role: 'user',
-          parts: [
-            {
-              type: 'text',
-              text: query,
-            },
-          ],
-        },
-        messageOptions,
-      );
+      void sendMessage(query, {
+        suggestedQuestionId: suggestedQuestion?.objectID,
+      });
 
       if (dropdownRef.current) {
         // some test environments (like jsdom) don't implement element.scrollTo
@@ -700,13 +673,9 @@ export function DocSearchModal({
   );
 
   // feedback handler
-  const handleFeedbackSubmit = React.useCallback(
-    async (messageId: string, thumbs: 0 | 1): Promise<void> => {
-      if (!askAiConfigurationId) return;
-      await sendFeedback(messageId, thumbs);
-    },
-    [askAiConfigurationId, sendFeedback],
-  );
+  const handleFeedbackSubmit = React.useCallback(async (_messageId: string, _thumbs: 0 | 1): Promise<void> => {
+    return;
+  }, []);
 
   if (!autocompleteRef.current) {
     autocompleteRef.current = createAutocomplete({
@@ -1041,9 +1010,8 @@ export function DocSearchModal({
               hasCollections={hasCollections}
               askAiState={askAiState}
               selectAskAiQuestion={handleSelectAskAiQuestion}
-              suggestedQuestions={suggestedQuestions}
+              suggestedQuestions={[]}
               selectSuggestedQuestion={selectSuggestedQuestion}
-              agentStudio={agentStudio}
               onAskAiToggle={onAskAiToggle}
               onNewConversation={handleNewConversation}
               onItemClick={(item, event) => {

packages/docsearch-react/src/ScreenState.tsx

@@ -1,4 +1,3 @@
-import type { UseChatHelpers } from '@ai-sdk/react';
 import type { AutocompleteApi, AutocompleteState, BaseItem } from '@algolia/autocomplete-core';
 import React from 'react';
 
@@ -18,7 +17,7 @@ import type { StartScreenTranslations } from './StartScreen';
 import { StartScreen } from './StartScreen';
 import type { StoredSearchPlugin } from './stored-searches';
 import type { InternalDocSearchHit, StoredAskAiState, StoredDocSearchHit, SuggestedQuestionHit } from './types';
-import type { AIMessage, AskAiState } from './types/AskiAi';
+import type { AIMessage, AskAiState, AskAiStatus } from './types/AskiAi';
 
 export type ScreenStateTranslations = Partial<{
   errorScreen: ErrorScreenTranslations;
@@ -42,8 +41,8 @@ export interface ScreenStateProps<TItem extends BaseItem>
   inputRef: React.MutableRefObject<HTMLInputElement | null>;
   hitComponent: DocSearchProps['hitComponent'];
   indexName: DocSearchProps['typesenseCollectionName'];
-  messages: UseChatHelpers<AIMessage>['messages'];
-  status: UseChatHelpers<AIMessage>['status'];
+  messages: AIMessage[];
+  status: AskAiStatus;
   askAiError?: Error;
   disableUserPersonalization: boolean;
   resultsFooterComponent: DocSearchProps['resultsFooterComponent'];
@@ -56,7 +55,6 @@ export interface ScreenStateProps<TItem extends BaseItem>
   suggestedQuestions: SuggestedQuestionHit[];
   selectSuggestedQuestion: (question: SuggestedQuestionHit) => void;
   onNewConversation: () => void;
-  agentStudio?: boolean;
 }
 
 export const ScreenState = React.memo(
@@ -83,7 +81,6 @@ export const ScreenState = React.memo(
           status={props.status}
           askAiError={props.askAiError}
           translations={translations?.askAiScreen}
-          agentStudio={props.agentStudio}
         />
       );
     }

packages/docsearch-react/src/SearchBox.tsx

@@ -1,4 +1,3 @@
-import type { UseChatHelpers } from '@ai-sdk/react';
 import type { AutocompleteApi, AutocompleteState } from '@algolia/autocomplete-core';
 import React, { type JSX, type RefObject } from 'react';
 
@@ -16,7 +15,7 @@ import { BackIcon } from './icons/BackIcon';
 import { Menu } from './Menu';
 import { ModalHeading } from './ModalHeading';
 import type { InternalDocSearchHit } from './types';
-import type { AIMessage, AskAiState } from './types/AskiAi';
+import type { AskAiState, AskAiStatus } from './types/AskiAi';
 
 export type SearchBoxTranslations = Partial<{
   clearButtonTitle: string;
@@ -49,7 +48,7 @@ interface SearchBoxProps
   onStopAskAiStreaming: () => Promise<void>;
   placeholder: string;
   isAskAiActive: boolean;
-  askAiStatus: UseChatHelpers<AIMessage>['status'];
+  askAiStatus: AskAiStatus;
   askAiError?: Error;
   isFromSelection: boolean;
   translations?: SearchBoxTranslations;

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

@@ -7,8 +7,20 @@ import '@testing-library/jest-dom/vitest';
 import { DocSearch as DocSearchComponent } from '../DocSearch';
 import type { DocSearchProps } from '../DocSearch';
 
+const typesenseServerConfig = {
+  apiKey: 'test-key',
+  nodes: [{ host: 'localhost', port: 8108, protocol: 'http' as const }],
+};
+
 function DocSearch(props: Partial<DocSearchProps>): JSX.Element {
-  return <DocSearchComponent appId="woo" apiKey="foo" indexName="bar" {...props} />;
+  return (
+    <DocSearchComponent
+      typesenseCollectionName="docs"
+      typesenseServerConfig={typesenseServerConfig}
+      typesenseSearchParameters={{}}
+      {...props}
+    />
+  );
 }
 
 // mock empty response
@@ -264,7 +276,7 @@ describe('api', () => {
 
   describe('ask AI integration', () => {
     it('updates placeholder when ask AI is available', async () => {
-      render(<DocSearch askAi="assistant" />);
+      render(<DocSearch askAi={{ conversationModelId: 'conv-model-1' }} />);
 
       await act(async () => {
         fireEvent.click(await screen.findByText('Search'));
@@ -276,7 +288,7 @@ describe('api', () => {
     it('opens ask AI screen and returns to search', async () => {
       render(
         <DocSearch
-          askAi="assistant"
+          askAi={{ conversationModelId: 'conv-model-1' }}
           transformSearchClient={(searchClient) => ({
             ...searchClient,
             search: noResultSearch,
@@ -300,8 +312,9 @@ describe('api', () => {
 
       expect(document.querySelector('.DocSearch-AskAiScreen')).toBeInTheDocument();
 
-      // could be "Answering..." or "Ask another question..."
-      expect(screen.getByText('Answering...')).toBeInTheDocument();
+      expect(
+        screen.queryByText('Answering...') ?? screen.getByPlaceholderText('Ask another question...'),
+      ).toBeInTheDocument();
     });
   });
 

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

@@ -7,13 +7,26 @@ import '@testing-library/jest-dom/vitest';
 
 import { type DocSearchProps, DocSearch as DocSearchComponent, type DocSearchRef } from '../DocSearch';
 
+const typesenseServerConfig = {
+  apiKey: 'test-key',
+  nodes: [{ host: 'localhost', port: 8108, protocol: 'http' as const }],
+};
+
 type TestDocSearchProps = DocSearchCallbacks & Partial<DocSearchProps> & { refObj?: RefObject<DocSearchRef | null> };
 
 function DocSearch(props: TestDocSearchProps): JSX.Element {
   const internalRef = useRef<DocSearchRef>(null);
   const ref = props.refObj ?? internalRef;
 
-  return <DocSearchComponent ref={ref} appId="woo" apiKey="foo" indexName="bar" {...props} />;
+  return (
+    <DocSearchComponent
+      ref={ref}
+      typesenseCollectionName="docs"
+      typesenseServerConfig={typesenseServerConfig}
+      typesenseSearchParameters={{}}
+      {...props}
+    />
+  );
 }
 
 describe('imperative handle', () => {

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

@@ -7,8 +7,20 @@ import '@testing-library/jest-dom/vitest';
 import { DocSearch as DocSearchComponent } from '../DocSearch';
 import type { DocSearchProps } from '../DocSearch';
 
+const typesenseServerConfig = {
+  apiKey: 'test-key',
+  nodes: [{ host: 'localhost', port: 8108, protocol: 'http' as const }],
+};
+
 function DocSearch(props: Partial<DocSearchProps>): JSX.Element {
-  return <DocSearchComponent appId="woo" apiKey="foo" indexName="bar" {...props} />;
+  return (
+    <DocSearchComponent
+      typesenseCollectionName="docs"
+      typesenseServerConfig={typesenseServerConfig}
+      typesenseSearchParameters={{}}
+      {...props}
+    />
+  );
 }
 
 describe('keyboard shortcuts', () => {