feat(askai): add Agent Studio memory support (#2888)

Author: 8bittitan

packages/docsearch-react/src/AskAiScreen.tsx

@@ -3,7 +3,7 @@ import React, { type JSX, useMemo, useState, useEffect } from 'react';
 
 import { AggregatedSearchBlock } from './AggregatedSearchBlock';
 import type { AskAiScreenStateProps } from './AskAiScreenState';
-import { ToolCall } from './components/ui/ToolCall';
+import { ToolCall, type ToolCallTranslations } from './components/ToolCall';
 import { AlertIcon, LoadingIcon } from './icons';
 import { MemoizedMarkdown } from './MemoizedMarkdown';
 import type { StoredSearchPlugin } from './stored-searches';
@@ -12,56 +12,88 @@ import { type AIMessage, type ToolCalls } from './types/AskiAi';
 import { extractLinksFromMessage, getMessageContent, isThreadDepthError, isAIToolPart } from './utils/ai';
 import { groupConsecutiveToolResults } from './utils/groupConsecutiveToolResults';
 
-export type AskAiScreenTranslations = Partial<{
-  // Misc texts
-  disclaimerText: string;
-  relatedSourcesText: string;
-  thinkingText: string;
-  copyButtonText: string;
-  copyButtonCopiedText: string;
-  // Feedback buttons
-  copyButtonTitle: string;
-  likeButtonTitle: string;
-  dislikeButtonTitle: string;
-  thanksForFeedbackText: string;
-  // Tool call texts
-  preToolCallText: string;
-  duringToolCallText: string;
-  afterToolCallText: string;
-  /**
-   * Build the full jsx element for the aggregated search block.
-   * If provided, completely overrides the default english renderer.
-   */
-  aggregatedToolCallNode?: (queries: string[], onSearchQueryClick: (query: string) => void) => React.ReactNode;
-
-  /**
-   * Generate the list connective parts only (backwards compatibility).
-   * Receives full list of queries and should return translation parts for before/after/separators.
-   * Example: (qs) => ({ before: 'searched for ', separator: ', ', lastSeparator: ' and ', after: '' }).
-   */
-  aggregatedToolCallText?: (queries: string[]) => {
-    before?: string;
-    separator?: string;
-    lastSeparator?: string;
-    after?: string;
+export type AskAiScreenTranslations = Partial<
+  // Inherit the shared tool-call translations, but expose the search-related
+  // keys under AskAiScreen's own public names (see mapping below).
+  Omit<ToolCallTranslations, 'searchingText' | 'toolCallResultText'> & {
+    // Misc texts
+    disclaimerText: string;
+    relatedSourcesText: string;
+    thinkingText: string;
+    copyButtonText: string;
+    copyButtonCopiedText: string;
+    // Feedback buttons
+    copyButtonTitle: string;
+    likeButtonTitle: string;
+    dislikeButtonTitle: string;
+    thanksForFeedbackText: string;
+    // Tool call texts
+    /**
+     * Text shown while assistant is performing search tool call.
+     * Maps to `ToolCallTranslations.searchingText`.
+     */
+    duringToolCallText: string;
+    /**
+     * Text shown while assistant is finished performing tool call.
+     * Maps to `ToolCallTranslations.toolCallResultText`.
+     */
+    afterToolCallText: string;
+    /**
+     * Build the full jsx element for the aggregated search block.
+     * If provided, completely overrides the default english renderer.
+     */
+    aggregatedToolCallNode?: (queries: string[], onSearchQueryClick: (query: string) => void) => React.ReactNode;
+    /**
+     * Generate the list connective parts only (backwards compatibility).
+     * Receives full list of queries and should return translation parts for before/after/separators.
+     * Example: (qs) => ({ before: 'searched for ', separator: ', ', lastSeparator: ' and ', after: '' }).
+     */
+    aggregatedToolCallText?: (queries: string[]) => {
+      before?: string;
+      separator?: string;
+      lastSeparator?: string;
+      after?: string;
+    };
+    /**
+     * Message that's shown when user has stopped the streaming of a message.
+     */
+    stoppedStreamingText: string;
+    /**
+     * Error title shown if there is an error while chatting.
+     */
+    errorTitleText: string;
+    /**
+     * Message shown when thread depth limit is exceeded (AI-217 error).
+     */
+    threadDepthExceededMessage: string;
+    /**
+     * Button text for starting a new conversation after thread depth error.
+     */
+    startNewConversationButtonText: string;
+  }
+>;
+
+/**
+ * Maps AskAiScreen's public translation keys to the shared `ToolCallTranslations`
+ * shape consumed by the `ToolCall` component, applying default English values.
+ */
+function toToolCallTranslations(translations: AskAiScreenTranslations): ToolCallTranslations {
+  const {
+    preToolCallText = 'Searching...',
+    duringToolCallText = 'Searching...',
+    afterToolCallText = 'Searched for',
+    savedMemoryToolResultText = 'Saved to memory',
+    memoryToolResultText = 'Used memory to enhance results',
+  } = translations;
+
+  return {
+    preToolCallText,
+    searchingText: duringToolCallText,
+    toolCallResultText: afterToolCallText,
+    savedMemoryToolResultText,
+    memoryToolResultText,
   };
-  /**
-   * Message that's shown when user has stopped the streaming of a message.
-   */
-  stoppedStreamingText: string;
-  /**
-   * Error title shown if there is an error while chatting.
-   */
-  errorTitleText: string;
-  /**
-   * Message shown when thread depth limit is exceeded (AI-217 error).
-   */
-  threadDepthExceededMessage: string;
-  /**
-   * Button text for starting a new conversation after thread depth error.
-   */
-  startNewConversationButtonText: string;
-}>;
+}
 
 type AskAiScreenProps = Omit<AskAiScreenStateProps<InternalDocSearchHit>, 'translations'> & {
   messages: AIMessage[];
@@ -71,6 +103,7 @@ type AskAiScreenProps = Omit<AskAiScreenStateProps<InternalDocSearchHit>, 'trans
   translations?: AskAiScreenTranslations;
   onNewConversation: () => void;
   agentStudio?: boolean;
+  memoryEnabled?: boolean;
 };
 
 interface AskAiScreenHeaderProps {
@@ -98,6 +131,7 @@ interface AskAiExchangeCardProps {
   conversations: StoredSearchPlugin<StoredAskAiState>;
   onFeedback?: (messageId: string, thumbs: 0 | 1) => Promise<void>;
   agentStudio?: boolean;
+  memoryEnabled?: boolean;
 }
 
 function AskAiExchangeCard({
@@ -111,16 +145,13 @@ function AskAiExchangeCard({
   conversations,
   onFeedback,
   agentStudio,
+  memoryEnabled,
 }: AskAiExchangeCardProps): JSX.Element {
   const { userMessage, assistantMessage } = exchange;
 
-  const {
-    stoppedStreamingText = 'You stopped this response',
-    errorTitleText = 'Chat error',
-    preToolCallText = 'Searching...',
-    afterToolCallText = 'Searched for',
-    duringToolCallText = 'Searching...',
-  } = translations;
+  const { stoppedStreamingText = 'You stopped this response', errorTitleText = 'Chat error' } = translations;
+
+  const toolCallTranslations = useMemo(() => toToolCallTranslations(translations), [translations]);
 
   const isThreadDepth = isThreadDepthError(askAiError);
 
@@ -191,13 +222,10 @@ function AskAiExchangeCard({
                 return (
                   <ToolCall
                     key={index}
-                    translations={{
-                      preToolCallText,
-                      searchingText: duringToolCallText,
-                      toolCallResultText: afterToolCallText,
-                    }}
+                    translations={toolCallTranslations}
                     part={part}
                     tools={tools}
+                    memoryEnabled={memoryEnabled}
                     onSearchQueryClick={onSearchQueryClick}
                   />
                 );
@@ -377,7 +405,7 @@ export function AskAiScreen({ translations = {}, ...props }: AskAiScreenProps):
     startNewConversationButtonText = 'Start a new conversation',
   } = translations;
 
-  const { messages, tools, askAiError, status, agentStudio } = props;
+  const { messages, tools, askAiError, status, agentStudio, memoryEnabled } = props;
 
   // Check if there's a thread depth error
   const hasThreadDepthError = useMemo(() => {
@@ -454,6 +482,7 @@ export function AskAiScreen({ translations = {}, ...props }: AskAiScreenProps):
                 tools={tools}
                 conversations={props.conversations}
                 agentStudio={agentStudio}
+                memoryEnabled={memoryEnabled}
                 onSearchQueryClick={handleSearchQueryClick}
                 onFeedback={props.onFeedback}
               />

packages/docsearch-react/src/AskAiScreenState.tsx

@@ -58,6 +58,7 @@ export interface AskAiScreenStateProps<TItem extends BaseItem>
   selectSuggestedQuestion: (question: SuggestedQuestionHit) => void;
   onNewConversation: () => void;
   agentStudio?: boolean;
+  memoryEnabled?: boolean;
 }
 
 export const AskAiScreenState = React.memo(
@@ -86,6 +87,7 @@ export const AskAiScreenState = React.memo(
           askAiError={props.askAiError}
           translations={translations?.askAiScreen}
           agentStudio={props.agentStudio}
+          memoryEnabled={props.memoryEnabled}
         />
       );
     }

packages/docsearch-react/src/DocSearch.tsx

@@ -234,6 +234,21 @@ export interface DocSearchProps {
   keyboardShortcuts?: DocSearchModalShortcuts;
 }
 
+export interface Memory {
+  /**
+   * Determines whether or not to display the memory based tool calls.
+   *
+   * @default false
+   */
+  enabled?: boolean;
+  /**
+   * The JWT used by the agent to know which user's memory to read.
+   *
+   * @see https://www.algolia.com/doc/guides/algolia-ai/agent-studio/how-to/user-authentication
+   */
+  userToken?: string;
+}
+
 export interface DocSearchAIProps extends DocSearchProps {
   /**
    * Configuration or assistant id to enable ask ai mode. Pass a string assistant id or a full config object.
@@ -254,6 +269,10 @@ export interface DocSearchAIProps extends DocSearchProps {
    * render but will not affect correctness.
    **/
   tools?: ToolCalls;
+  /**
+   * Configuration for the Agent Studio memory feature.
+   */
+  memory?: Memory;
 }
 
 function DocSearchComponent(props: DocSearchProps, ref: React.ForwardedRef<DocSearchRef>): JSX.Element {

packages/docsearch-react/src/DocSearchAskAiModal.tsx

@@ -117,6 +117,7 @@ export function DocSearchAskAiModal({
     suggestedQuestionsEnabled: askAiConfig?.suggestedQuestions,
   });
   const agentStudio = askAiConfig?.agentStudio ?? false;
+  const memoryEnabled = props.memory?.enabled ?? false;
 
   const indexes = normalizeDocSearchIndexes({
     indexName,
@@ -143,6 +144,7 @@ export function DocSearchAskAiModal({
       useStagingEnv: askAiUseStagingEnv,
       agentStudio,
       tools,
+      memory: props.memory,
     });
 
   const prevStatus = React.useRef(status);
@@ -464,6 +466,7 @@ export function DocSearchAskAiModal({
           suggestedQuestions={suggestedQuestions}
           selectSuggestedQuestion={selectSuggestedQuestion}
           agentStudio={agentStudio}
+          memoryEnabled={memoryEnabled}
           onAskAiToggle={onAskAiToggle}
           onNewConversation={handleNewConversation}
           onItemClick={(item, event) => {

packages/docsearch-react/src/Sidepanel.tsx

@@ -4,7 +4,7 @@ import type { JSX } from 'react';
 import React from 'react';
 import { createPortal } from 'react-dom';
 
-import type { AgentStudioSearchParameters, AskAiSearchParameters } from './DocSearch';
+import type { AgentStudioSearchParameters, AskAiSearchParameters, Memory } from './DocSearch';
 import type { SidepanelButtonProps, SidepanelProps as SidepanelPanelProps } from './Sidepanel/index';
 import { SidepanelButton, Sidepanel } from './Sidepanel/index';
 import type { ToolCalls } from './types/AskiAi';
@@ -93,6 +93,10 @@ export type DocSearchSidepanelProps = DocSearchCallbacks & {
    * render but will not affect correctness.
    **/
   tools?: ToolCalls;
+  /**
+   * Configuration for the Agent Studio memory feature.
+   */
+  memory?: Memory;
 };
 
 type SidepanelProps = DocSearchSidepanelProps & SidepanelSearchParameters;

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

@@ -3,7 +3,7 @@ import type { JSX } from 'react';
 import React, { memo, useMemo } from 'react';
 
 import { AskAiSourcesPanel, type Exchange } from '../AskAiScreen';
-import { ToolCall, type ToolCallTranslations } from '../components/ui/ToolCall';
+import { ToolCall, type ToolCallTranslations } from '../components/ToolCall';
 import { AlertIcon, LoadingIcon } from '../icons';
 import { MemoizedMarkdown } from '../MemoizedMarkdown';
 import type { StoredSearchPlugin } from '../stored-searches';
@@ -72,6 +72,7 @@ export type ConversationScreenProps = {
   handleFeedback?: (messageId: string, thumbs: 0 | 1) => Promise<void>;
   streamError?: Error;
   agentStudio?: boolean;
+  memoryEnabled?: boolean;
   tools?: ToolCalls;
 };
 
@@ -84,12 +85,24 @@ type ConversationnExchangeProps = {
   onFeedback?: ConversationScreenProps['handleFeedback'];
   streamError?: ConversationScreenProps['streamError'];
   agentStudio?: boolean;
+  memoryEnabled?: boolean;
   tools: ToolCalls;
 };
 
 const ConversationExchange = React.forwardRef<HTMLDivElement, ConversationnExchangeProps>(
   (
-    { exchange, translations = {}, isLastExchange, conversations, onFeedback, status, streamError, agentStudio, tools },
+    {
+      exchange,
+      translations = {},
+      isLastExchange,
+      conversations,
+      onFeedback,
+      status,
+      streamError,
+      agentStudio,
+      memoryEnabled,
+      tools,
+    },
     conversationRef,
   ): JSX.Element => {
     const { userMessage, assistantMessage } = exchange;
@@ -105,6 +118,8 @@ const ConversationExchange = React.forwardRef<HTMLDivElement, ConversationnExcha
       copyButtonText = 'Copy',
       copyButtonCopiedText = 'Copied!',
       errorTitleText = 'Chat error',
+      savedMemoryToolResultText = 'Saved to memory',
+      memoryToolResultText = 'Used memory to enhance results',
     } = translations;
 
     const assistantContent = useMemo(() => getMessageContent(assistantMessage), [assistantMessage]);
@@ -172,8 +187,11 @@ const ConversationExchange = React.forwardRef<HTMLDivElement, ConversationnExcha
                         preToolCallText,
                         searchingText,
                         toolCallResultText,
+                        savedMemoryToolResultText,
+                        memoryToolResultText,
                       }}
                       tools={tools}
+                      memoryEnabled={memoryEnabled}
                     />
                   );
                 }

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

@@ -153,6 +153,7 @@ function SidepanelInner(
     useStagingEnv = false,
     agentStudio = false,
     tools = EMPTY_TOOLS,
+    memory,
   }: Props,
   ref: React.ForwardedRef<SidepanelRef>,
 ): JSX.Element {
@@ -197,6 +198,7 @@ function SidepanelInner(
     useStagingEnv,
     agentStudio,
     tools,
+    memory,
   });
 
   const suggestedQuestions = useSuggestedQuestions({
@@ -400,6 +402,7 @@ function SidepanelInner(
               translations={translations.conversationScreen}
               streamError={askAiError}
               agentStudio={agentStudio}
+              memoryEnabled={memory?.enabled ?? false}
               tools={tools}
             />
           )}