feat: support save to notebook (#6)

* feat: support add to notebook

Signed-off-by: tygao <[email protected]>

* fix comments

Signed-off-by: tygao <[email protected]>

* seperate save chat and disable save when notebook no input

Signed-off-by: tygao <[email protected]>

* seperate save chat and disable save when notebook no input

Signed-off-by: tygao <[email protected]>

---------

Signed-off-by: tygao <[email protected]>

common/constants/llm.ts

@@ -6,6 +6,7 @@
 export const API_BASE = '/api/assistant';
 export const DSL_BASE = '/api/dsl';
 export const DSL_SEARCH = '/search';
+export const NOTEBOOK_PREFIX = '/api/observability/notebooks';
 
 export const ASSISTANT_API = {
   SEND_MESSAGE: `${API_BASE}/send_message`,
@@ -23,3 +24,8 @@ export const LLM_INDEX = {
   SESSIONS: '.assistant-sessions',
   VECTOR_STORE: '.llm-vector-store',
 };
+
+export const NOTEBOOK_API = {
+  CREATE_NOTEBOOK: `${NOTEBOOK_PREFIX}/note`,
+  SET_PARAGRAPH: `${NOTEBOOK_PREFIX}/set_paragraphs/`,
+};

common/types/chat_saved_object_attributes.ts

@@ -19,7 +19,7 @@ export interface ISessionFindResponse {
   total: number;
 }
 
-interface IInput {
+export interface IInput {
   type: 'input';
   contentType: 'text';
   content: string;

opensearch_dashboards.json

@@ -9,5 +9,8 @@
     "embeddable",
     "opensearchDashboardsReact",
     "opensearchDashboardsUtils"
+  ],
+  "optionalPlugins":[
+    "observabilityDashboards"
   ]
 }

public/components/notebook/notebook_name_modal.tsx

@@ -0,0 +1,64 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import {
+  EuiButton,
+  EuiButtonEmpty,
+  EuiForm,
+  EuiFormRow,
+  EuiModal,
+  EuiModalBody,
+  EuiModalFooter,
+  EuiModalHeader,
+  EuiFieldText,
+  EuiModalHeaderTitle,
+} from '@elastic/eui';
+import React, { useState, useCallback } from 'react';
+
+interface Props {
+  onClose: () => void;
+  // SaveChat hook depends on context. Runtime modal component can't get context, so saveChat needs to be passed in.
+  saveChat: (name: string) => void;
+}
+
+export const NotebookNameModal = ({ onClose, saveChat }: Props) => {
+  const [name, setName] = useState('');
+  const [loading, setLoading] = useState(false);
+
+  const onSubmit = useCallback(async () => {
+    setLoading(true);
+    await saveChat(name);
+    onClose();
+  }, [name, saveChat, onclose]);
+
+  return (
+    <>
+      <EuiModal onClose={onClose}>
+        <EuiModalHeader>
+          <EuiModalHeaderTitle>Save to notebook</EuiModalHeaderTitle>
+        </EuiModalHeader>
+
+        <EuiModalBody>
+          <EuiFormRow label="Please enter a name for your notebook.">
+            <EuiFieldText value={name} onChange={(e) => setName(e.target.value)} />
+          </EuiFormRow>
+        </EuiModalBody>
+
+        <EuiModalFooter>
+          <EuiButtonEmpty onClick={onClose}>Cancel</EuiButtonEmpty>
+          <EuiButton
+            type="submit"
+            fill
+            isLoading={loading}
+            disabled={name.length < 1}
+            onClick={onSubmit}
+          >
+            Confirm name
+          </EuiButton>
+        </EuiModalFooter>
+      </EuiModal>
+    </>
+  );
+};

public/hooks/use_chat_actions.tsx

@@ -16,6 +16,10 @@ interface SendResponse {
   messages: IMessage[];
 }
 
+interface SetParagraphResponse {
+  objectId: string;
+}
+
 let abortControllerRef: AbortController;
 
 export const useChatActions = (): AssistantActions => {

public/hooks/use_save_chat.tsx

@@ -0,0 +1,99 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import React from 'react';
+import { useCallback } from 'react';
+import { EuiLink } from '@elastic/eui';
+import { NOTEBOOK_API } from '../../common/constants/llm';
+import { useCore } from '../contexts/core_context';
+import { useChatState } from './use_chat_state';
+import { convertMessagesToParagraphs, Paragraphs } from '../utils';
+import { getCoreStart } from '../plugin';
+import { toMountPoint } from '../../../../src/plugins/opensearch_dashboards_react/public';
+
+interface SetParagraphResponse {
+  objectId: string;
+}
+
+export const useSaveChat = () => {
+  const core = useCore();
+  const { chatState } = useChatState();
+
+  const createNotebook = useCallback(
+    async (name: string) => {
+      const id = await core.services.http.post<string>(NOTEBOOK_API.CREATE_NOTEBOOK, {
+        // do not send abort signal to http client to allow LLM call run in background
+        body: JSON.stringify({
+          name,
+        }),
+      });
+      if (!id) {
+        throw new Error('create notebook error');
+      }
+      return id;
+    },
+    [core]
+  );
+
+  const setParagraphs = useCallback(
+    async (id: string, paragraphs: Paragraphs) => {
+      const response = await core.services.http.post<SetParagraphResponse>(
+        NOTEBOOK_API.SET_PARAGRAPH,
+        {
+          // do not send abort signal to http client to allow LLM call run in background
+          body: JSON.stringify({
+            noteId: id,
+            paragraphs,
+          }),
+        }
+      );
+      const { objectId } = response;
+      if (!objectId) {
+        throw new Error('set paragraphs error');
+      }
+      return objectId;
+    },
+    [core]
+  );
+
+  const saveChat = useCallback(
+    async (name: string) => {
+      try {
+        const id = await createNotebook(name);
+        const paragraphs = convertMessagesToParagraphs(chatState.messages);
+        await setParagraphs(id, paragraphs);
+        const notebookLink = `./observability-notebooks#/${id}?view=view_both`;
+
+        getCoreStart().notifications.toasts.addSuccess({
+          text: toMountPoint(
+            <>
+              <p>
+                This conversation was saved as{' '}
+                <EuiLink href={notebookLink} target="_blank">
+                  {name}
+                </EuiLink>
+                .
+              </p>
+            </>
+          ),
+        });
+      } catch (error) {
+        if (error.message === 'Not Found') {
+          getCoreStart().notifications.toasts.addError(error, {
+            title:
+              'This feature depends on the observability plugin, please install it before use.',
+          });
+        } else {
+          getCoreStart().notifications.toasts.addError(error, {
+            title: 'Failed to save to notebook',
+          });
+        }
+      }
+    },
+    [chatState, createNotebook, setParagraphs]
+  );
+
+  return { saveChat };
+};

public/tabs/chat_window_header.tsx

@@ -16,6 +16,10 @@ import React, { useCallback, useState } from 'react';
 import { EditConversationNameModal } from '../components/edit_conversation_name_modal';
 import { useChatContext } from '../contexts/chat_context';
 import { useChatActions } from '../hooks/use_chat_actions';
+import { NotebookNameModal } from '../components/notebook/notebook_name_modal';
+import { useCore } from '../contexts/core_context';
+import { useChatState } from '../hooks/use_chat_state';
+import { useSaveChat } from '../hooks/use_save_chat';
 
 interface ChatWindowHeaderProps {
   flyoutFullScreen: boolean;
@@ -25,8 +29,11 @@ interface ChatWindowHeaderProps {
 export const ChatWindowHeader: React.FC<ChatWindowHeaderProps> = React.memo((props) => {
   const chatContext = useChatContext();
   const { loadChat } = useChatActions();
+  const core = useCore();
   const [isPopoverOpen, setPopover] = useState(false);
   const [isRenameModelOpen, setRenameModelOpen] = useState(false);
+  const { chatState } = useChatState();
+  const { saveChat } = useSaveChat();
 
   const onButtonClick = () => {
     setPopover(!isPopoverOpen);
@@ -101,10 +108,15 @@ export const ChatWindowHeader: React.FC<ChatWindowHeaderProps> = React.memo((pro
     <EuiContextMenuItem
       key="save-as-notebook"
       onClick={() => {
+        const modal = core.overlays.openModal(
+          <NotebookNameModal onClose={() => modal.close()} saveChat={saveChat} />
+        );
         closePopover();
       }}
+      // There is only one message in initial discussion, which will not be stored.
+      disabled={chatState.messages.length <= 1}
     >
-      Save as notebook
+      Save to notebook
     </EuiContextMenuItem>,
   ];
 

public/utils/index.ts

@@ -0,0 +1,6 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+export * from './notebook';

public/utils/notebook.ts

@@ -0,0 +1,117 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { v4 as uuid } from 'uuid';
+import { htmlIdGenerator } from '@elastic/eui';
+import { IMessage } from '../../common/types/chat_saved_object_attributes';
+
+const buildBasicGraph = () => ({
+  id: 'paragraph_' + uuid(),
+  dateCreated: new Date().toISOString(),
+  dateModified: new Date().toISOString(),
+  input: {
+    inputText: '',
+    inputType: '',
+  },
+  output: [{ result: '', outputType: '', execution_time: '0 ms' }],
+});
+
+const ASSISTANT_MESSAGE_PREFIX = 'OpenSearch Assistant: ';
+const USER_MESSAGE_PREFIX = 'User: ';
+
+const createDashboardVizObject = (objectId: string) => {
+  const vizUniqueId = htmlIdGenerator()();
+  // a dashboard container object for new visualization
+  const basicVizObject = {
+    viewMode: 'view',
+    panels: {
+      '1': {
+        gridData: {
+          x: 0,
+          y: 0,
+          w: 50,
+          h: 20,
+          i: '1',
+        },
+        type: 'visualization',
+        explicitInput: {
+          id: '1',
+          savedObjectId: objectId,
+        },
+      },
+    },
+    isFullScreenMode: false,
+    filters: [],
+    useMargins: false,
+    id: vizUniqueId,
+    timeRange: {
+      // We support last 15minutes here to keep consistent with chat bot preview.
+      to: 'now',
+      from: 'now-15m',
+    },
+    title: 'embed_viz_' + vizUniqueId,
+    query: {
+      query: '',
+      language: 'lucene',
+    },
+    refreshConfig: {
+      pause: true,
+      value: 15,
+    },
+  } as const;
+  return basicVizObject;
+};
+
+export const convertMessagesToParagraphs = (messages: IMessage[]) => {
+  return messages.map((message: IMessage) => {
+    const paragraph = buildBasicGraph();
+
+    switch (message.contentType) {
+      // markdown,text and error are all text formatted in notebook.
+      case 'markdown':
+      case 'text':
+      case 'error':
+        const messageText =
+          // markdown and error represents assistant, text represents user.
+          message.contentType === 'text'
+            ? USER_MESSAGE_PREFIX + message.content
+            : ASSISTANT_MESSAGE_PREFIX + message.content;
+
+        Object.assign(paragraph, {
+          input: { inputText: `%md\n${messageText}`, inputType: 'MARKDOWN' },
+          output: [
+            {
+              result: messageText,
+              outputType: 'MARKDOWN',
+              execution_time: '0 ms',
+            },
+          ],
+        });
+        break;
+
+      case 'visualization':
+        const visualizationObjectId = message.content;
+        const inputText = JSON.stringify(createDashboardVizObject(visualizationObjectId));
+        Object.assign(paragraph, {
+          input: { inputText, inputType: 'VISUALIZATION' },
+          output: [
+            {
+              result: '',
+              outputType: 'VISUALIZATION',
+              execution_time: '0 ms',
+            },
+          ],
+        });
+        break;
+
+      // error and ppl_visualization contentType will not be handled currently.
+      default:
+        break;
+    }
+    return paragraph;
+  });
+};
+
+export type Paragraphs = ReturnType<typeof convertMessagesToParagraphs>;