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 b53bc97.

* fix: types

Author: 8bittitan

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);
+};