feat(askai): Allow Agent Studio specific search params (#2842)

* feat(askai): Allow Agent Studio specific search params

* Simplify DocSearchProps types definitions

* Revert "Simplify DocSearchProps types definitions"

This reverts commit b53bc97688.

* fix: types

examples/demo-react/src/examples/agent-studio.tsx

@@ -21,9 +21,8 @@ export function AgentStudioExample(): JSX.Element {
           apiKey="a00716d83c64f6c61905c078b7d5ab66"
           askAi={{
             assistantId: 'ccdec697-e3fe-465b-a1c3-657e7bf18aef',
-            indexName: 'docsearch-markdown',
+            agentStudio: true,
           }}
-          agentStudio={true}
         />
       </DocSearch>
 

packages/docsearch-react/src/DocSearch.tsx

@@ -31,9 +31,11 @@ export type AskAiSearchParameters = {
   filters?: string;
   attributesToRetrieve?: string[];
   restrictSearchableAttributes?: string[];
-  distinct?: boolean;
+  distinct?: boolean | number | string;
 };
 
+export type AgentStudioSearchParameters = Record<string, Omit<AskAiSearchParameters, 'facetFilters'>>;
+
 export type DocSearchAskAi = {
   /**
    * The index name to use for the ask AI feature. Your assistant will search this index for relevant documents.
@@ -54,10 +56,6 @@ export type DocSearchAskAi = {
    * The assistant ID to use for the ask AI feature.
    */
   assistantId: string;
-  /**
-   * The search parameters to use for the ask AI feature.
-   */
-  searchParameters?: AskAiSearchParameters;
   /**
    * Enables displaying suggested questions on Ask AI's new conversation screen.
    *
@@ -66,7 +64,43 @@ export type DocSearchAskAi = {
   suggestedQuestions?: boolean;
   // HACK: This is a hack for testing staging, remove before releasing
   useStagingEnv?: boolean;
-};
+} & (
+  | {
+      /**
+       * **Experimental:** Whether to use Agent Studio as the chat backend.
+       *
+       * This is an experimental feature and its API may change without notice in future releases.
+       * Use with caution in production environments.
+       *
+       * @default false
+       */
+      agentStudio?: never;
+      /**
+       * The search parameters to use for the ask AI feature.
+       *
+       * **NOTE**: If using `agentStudio = true`, the `searchParameters` object is
+       * keyed by the index name.
+       */
+      searchParameters?: AskAiSearchParameters;
+    }
+  | {
+      agentStudio: false;
+      searchParameters?: AskAiSearchParameters;
+    }
+  | {
+      agentStudio: true;
+      /**
+       * The search parameters to use for the ask AI feature.
+       * Keyed by the index name.
+       *
+       * @example
+       * {
+       *   "INDEX_NAME": { distinct: false }
+       * }
+       */
+      searchParameters?: AgentStudioSearchParameters;
+    }
+);
 
 export interface DocSearchIndex {
   name: string;
@@ -105,15 +139,6 @@ export interface DocSearchProps {
    * Useful to route Ask AI into a different UI (e.g. `@docsearch/sidepanel-js`) without flicker.
    */
   interceptAskAiEvent?: (initialMessage: InitialAskAiMessage) => boolean | void;
-  /**
-   * **Experimental:** Whether to use Agent Studio as the chat backend.
-   *
-   * This is an experimental feature and its API may change without notice in future releases.
-   * Use with caution in production environments.
-   *
-   * @default false
-   */
-  agentStudio?: boolean;
   /**
    * Theme overrides applied to the modal and related components.
    */

packages/docsearch-react/src/DocSearchModal.tsx

@@ -310,7 +310,6 @@ export function DocSearchModal({
   indexName,
   searchParameters,
   isHybridModeSupported = false,
-  agentStudio = false,
   ...props
 }: DocSearchModalProps): JSX.Element {
   const { footer: footerTranslations, searchBox: searchBoxTranslations, ...screenStateTranslations } = translations;
@@ -360,6 +359,7 @@ export function DocSearchModal({
     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[] = [];

packages/docsearch-react/src/Sidepanel.tsx

@@ -4,12 +4,49 @@ import type { JSX } from 'react';
 import React from 'react';
 import { createPortal } from 'react-dom';
 
-import type { AskAiSearchParameters } from './DocSearch';
-import type { SidepanelButtonProps, SidepanelProps } from './Sidepanel/index';
+import type { AgentStudioSearchParameters, AskAiSearchParameters } from './DocSearch';
+import type { SidepanelButtonProps, SidepanelProps as SidepanelPanelProps } from './Sidepanel/index';
 import { SidepanelButton, Sidepanel } from './Sidepanel/index';
 
 export type { DocSearchRef, DocSearchCallbacks } from '@docsearch/core';
 
+export type SidepanelSearchParameters =
+  | {
+      /**
+       * **Experimental:** Whether to use Agent Studio as the chat backend.
+       *
+       * This is an experimental feature and its API may change without notice in future releases.
+       * Use with caution in production environments.
+       *
+       * @default false
+       */
+      agentStudio?: never;
+      /**
+       * The search parameters to use for the ask AI feature.
+       *
+       * **NOTE**: If using `agentStudio = true`, the `searchParameters` object is
+       * keyed by the index name.
+       */
+      searchParameters?: AskAiSearchParameters;
+    }
+  | {
+      agentStudio: false;
+      searchParameters?: AskAiSearchParameters;
+    }
+  | {
+      agentStudio: true;
+      /**
+       * The search parameters to use for the ask AI feature.
+       * Keyed by the index name.
+       *
+       * @example
+       * {
+       *   "INDEX_NAME": { distinct: false }
+       * }
+       */
+      searchParameters?: AgentStudioSearchParameters;
+    };
+
 export type DocSearchSidepanelProps = DocSearchCallbacks & {
   /**
    * The assistant ID to use for the ask AI feature.
@@ -27,10 +64,6 @@ export type DocSearchSidepanelProps = DocSearchCallbacks & {
    * The index name to use for the ask AI feature. Your assistant will search this index for relevant documents.
    */
   indexName: string;
-  /**
-   * The search parameters to use for the ask AI feature.
-   */
-  searchParameters?: AskAiSearchParameters;
   /**
    * Configuration for keyboard shortcuts. Allows enabling/disabling specific shortcuts.
    *
@@ -50,29 +83,13 @@ export type DocSearchSidepanelProps = DocSearchCallbacks & {
   /**
    * Props specific to the Sidepanel panel.
    */
-  panel?: Omit<SidepanelProps, 'keyboardShortcuts'>;
-  /**
-   * **Experimental:** Whether to use Agent Studio as the chat backend.
-   *
-   * This is an experimental feature and its API may change without notice in future releases.
-   * Use with caution in production environments.
-   *
-   * @default false
-   */
-  agentStudio?: boolean;
+  panel?: Omit<SidepanelPanelProps, 'keyboardShortcuts'>;
 };
 
+type SidepanelProps = DocSearchSidepanelProps & SidepanelSearchParameters;
+
 function DocSearchSidepanelComponent(
-  {
-    keyboardShortcuts,
-    theme,
-    onReady,
-    onOpen,
-    onClose,
-    onSidepanelOpen,
-    onSidepanelClose,
-    ...props
-  }: DocSearchSidepanelProps,
+  { keyboardShortcuts, theme, onReady, onOpen, onClose, onSidepanelOpen, onSidepanelClose, ...props }: SidepanelProps,
   ref: React.ForwardedRef<DocSearchRef>,
 ): JSX.Element {
   return (

packages/docsearch-react/src/Sidepanel/ConversationScreen.tsx

@@ -57,6 +57,10 @@ export type ConversationScreenTranslations = Partial<
      * Message displayed after feedback action.
      **/
     thanksForFeedbackText: string;
+    /**
+     * Error title shown if there is an error while chatting.
+     */
+    errorTitleText;
   }
 >;
 
@@ -98,6 +102,7 @@ const ConversationExchange = React.forwardRef<HTMLDivElement, ConversationnExcha
       toolCallResultText = 'Searched for',
       copyButtonText = 'Copy',
       copyButtonCopiedText = 'Copied!',
+      errorTitleText = 'Chat error',
     } = translations;
 
     const assistantContent = useMemo(() => getMessageContent(assistantMessage), [assistantMessage]);
@@ -110,7 +115,8 @@ const ConversationExchange = React.forwardRef<HTMLDivElement, ConversationnExcha
     const urlsToDisplay = React.useMemo(() => extractLinksFromMessage(assistantMessage), [assistantMessage]);
 
     const wasStopped = userMessage.metadata?.stopped || assistantMessage?.metadata?.stopped;
-    const isThinking = !assistantParts.some((part) => part.type !== 'step-start');
+    const isThinking =
+      ['submitted', 'streaming'].includes(status) && !assistantParts.some((part) => part.type !== 'step-start');
     const showActions =
       !wasStopped && (!isLastExchange || (isLastExchange && status === 'ready' && Boolean(assistantMessage)));
 
@@ -125,12 +131,15 @@ const ConversationExchange = React.forwardRef<HTMLDivElement, ConversationnExcha
               {status === 'error' && streamError && isLastExchange && (
                 <div className="DocSearch-AskAiScreen-MessageContent DocSearch-AskAiScreen-Error">
                   <AlertIcon />
-                  <MemoizedMarkdown
-                    content={streamError.message}
-                    copyButtonText=""
-                    copyButtonCopiedText=""
-                    isStreaming={false}
-                  />
+                  <div className="DocSearch-AskAiScreen-Error-Content">
+                    <h4 className="DocSearch-AskAiScreen-Error-Title">{errorTitleText}</h4>
+                    <MemoizedMarkdown
+                      content={streamError.message}
+                      copyButtonText=""
+                      copyButtonCopiedText=""
+                      isStreaming={false}
+                    />
+                  </div>
                 </div>
               )}
 

packages/docsearch-react/src/Sidepanel/Sidepanel.tsx

@@ -3,7 +3,7 @@ import React, { useCallback } from 'react';
 import type { JSX } from 'react';
 
 import { AlgoliaLogo, type AlgoliaLogoTranslations } from '../AlgoliaLogo';
-import type { DocSearchSidepanelProps } from '../Sidepanel';
+import type { DocSearchSidepanelProps, SidepanelSearchParameters } from '../Sidepanel';
 import type { StoredAskAiState, SuggestedQuestionHit } from '../types';
 import { useAskAi } from '../useAskAi';
 import { useSearchClient } from '../useSearchClient';
@@ -122,7 +122,8 @@ export type SidepanelProps = {
 };
 
 type Props = Omit<DocSearchSidepanelProps, 'button' | 'panel'> &
-  SidepanelProps & {
+  SidepanelProps &
+  SidepanelSearchParameters & {
     isOpen?: boolean;
     onOpen: () => void;
     onClose: () => void;

packages/docsearch-react/src/askai.ts

@@ -102,3 +102,35 @@ export const postFeedback = async ({
     headers,
   });
 };
+
+interface AgentStudioValidationError extends Error {
+  name: 'ValidationError';
+  detail?: Array<{ type: string; loc: string[]; msg: string }>;
+}
+
+// Parse Agent Studio errors as they are returned as JSON rather than Markdown/text
+export const getAgentStudioErrorMessage = (error: Error): Error => {
+  let errorMessage = error.message;
+
+  try {
+    const parsedError = JSON.parse(error.message) as Error;
+
+    // Check for known errors that we know how to parse
+    if (parsedError.name === 'ValidationError') {
+      const validationError = parsedError as AgentStudioValidationError;
+
+      if (validationError.detail && validationError.detail.length > 0) {
+        const { msg, loc } = validationError.detail[0];
+        const field = loc.at(-1);
+
+        errorMessage = `${msg}: ${field}`;
+      }
+    } else {
+      errorMessage = parsedError.message;
+    }
+  } catch {
+    // We don't care about this catch, we default to the error.message above
+  }
+
+  return new Error(errorMessage);
+};

packages/docsearch-react/src/useAskAi.ts

@@ -3,14 +3,14 @@ import { useChat } from '@ai-sdk/react';
 import { DefaultChatTransport, lastAssistantMessageIsCompleteWithToolCalls } from 'ai';
 import { useCallback, useMemo, useRef } from 'react';
 
-import { getValidToken, postFeedback } from './askai';
+import { getAgentStudioErrorMessage, getValidToken, postFeedback } from './askai';
 import type { Exchange } from './AskAiScreen';
 import { ASK_AI_API_URL, BETA_ASK_AI_API_URL } from './constants';
 import type { StoredSearchPlugin } from './stored-searches';
 import { createStoredConversations } from './stored-searches';
 import type { AIMessage } from './types/AskiAi';
 
-import type { AskAiSearchParameters, StoredAskAiState } from '.';
+import type { AgentStudioSearchParameters, AskAiSearchParameters, StoredAskAiState } from '.';
 
 type UseChat = UseChatHelpers<AIMessage>;
 
@@ -19,10 +19,19 @@ type UseAskAiParams = {
   apiKey: string;
   appId: string;
   indexName: string;
-  searchParameters?: AskAiSearchParameters;
   useStagingEnv?: boolean;
-  agentStudio?: boolean;
-};
+  searchParameters?: AskAiSearchParameters;
+  agentStudio: boolean;
+} & (
+  | {
+      agentStudio: false;
+      searchParameters?: AskAiSearchParameters;
+    }
+  | {
+      agentStudio: true;
+      searchParameters?: AgentStudioSearchParameters;
+    }
+);
 
 type UseAskAiReturn = {
   messages: AIMessage[];
@@ -39,17 +48,23 @@ type UseAskAiReturn = {
 
 type UseAskAi = (params: UseAskAiParams) => UseAskAiReturn;
 
+type AgentStudioTransportParams = Pick<UseAskAiParams, 'apiKey' | 'appId' | 'assistantId'> & {
+  searchParameters?: AgentStudioSearchParameters;
+};
+
 const getAgentStudioTransport = ({
   appId,
   apiKey,
   assistantId,
-}: Pick<UseAskAiParams, 'apiKey' | 'appId' | 'assistantId'>): DefaultChatTransport<AIMessage> => {
+  searchParameters,
+}: AgentStudioTransportParams): DefaultChatTransport<AIMessage> => {
   return new DefaultChatTransport({
     api: `https://${appId}.algolia.net/agent-studio/1/agents/${assistantId}/completions?stream=true&compatibilityMode=ai-sdk-5`,
     headers: {
       'x-algolia-application-id': appId,
       'x-algolia-api-key': apiKey,
     },
+    body: searchParameters ? { algolia: { searchParameters } } : {},
   });
 };
 
@@ -90,35 +105,28 @@ const getAskAiTransport = ({
   });
 };
 
-export const useAskAi: UseAskAi = ({
-  assistantId,
-  apiKey,
-  appId,
-  indexName,
-  searchParameters,
-  useStagingEnv = false,
-  agentStudio = false,
-}) => {
+export const useAskAi: UseAskAi = ({ assistantId, apiKey, appId, indexName, useStagingEnv = false, ...params }) => {
   const abortControllerRef = useRef(new AbortController());
 
   const askAiTransport = useMemo(
     () =>
-      agentStudio
+      params.agentStudio
         ? getAgentStudioTransport({
             apiKey,
             appId,
             assistantId: assistantId ?? '',
+            searchParameters: params.searchParameters,
           })
         : getAskAiTransport({
             assistantId: assistantId ?? '',
             apiKey,
             appId,
             indexName,
-            searchParameters,
+            searchParameters: params.searchParameters,
             abortController: abortControllerRef.current,
             useStagingEnv,
           }),
-    [apiKey, appId, assistantId, indexName, searchParameters, agentStudio, useStagingEnv],
+    [apiKey, appId, assistantId, indexName, useStagingEnv, params],
   );
 
   const { messages, sendMessage, status, setMessages, error, stop } = useChat({
@@ -176,12 +184,20 @@ export const useAskAi: UseAskAi = ({
 
   const isStreaming = status === 'streaming' || status === 'submitted';
 
+  const askAiError = useMemo((): Error | undefined => {
+    if (!error) return undefined;
+
+    if (!params.agentStudio) return error;
+
+    return getAgentStudioErrorMessage(error);
+  }, [error, params.agentStudio]);
+
   return {
     messages,
     sendMessage,
     status,
     setMessages,
-    askAiError: error,
+    askAiError,
     stopAskAiStreaming: onStopStreaming,
     isStreaming,
     exchanges,

packages/docsearch-sidepanel/src/Sidepanel.tsx

@@ -1,11 +1,16 @@
 import { useDocSearch } from '@docsearch/core';
-import { Sidepanel as SidepanelComp, type DocSearchSidepanelProps } from '@docsearch/react/sidepanel';
+import {
+  Sidepanel as SidepanelComp,
+  type DocSearchSidepanelProps,
+  type SidepanelSearchParameters,
+} from '@docsearch/react/sidepanel';
 import React from 'react';
 import type { JSX } from 'react';
 import { createPortal } from 'react-dom';
 
 export type SidepanelProps = DocSearchSidepanelProps['panel'] &
-  Omit<DocSearchSidepanelProps, 'button' | 'panel' | 'theme'>;
+  Omit<DocSearchSidepanelProps, 'button' | 'panel' | 'theme'> &
+  SidepanelSearchParameters;
 
 export function Sidepanel({ portalContainer, ...props }: SidepanelProps): JSX.Element {
   const { docsearchState, setDocsearchState, keyboardShortcuts, registerView, initialAskAiMessage } = useDocSearch();