From 280f886a2b0eb07ae17bcae7cd5767f352915b86 Mon Sep 17 00:00:00 2001
From: Coen Warmer <coen.warmer@gmail.com>
Date: Wed, 2 Aug 2023 16:13:20 +0200
Subject: [PATCH] Add JSON schema validation for Prompt Editor

---
 .../public/components/chat/chat_body.tsx      |   8 +-
 .../public/components/chat/chat_header.tsx    |   2 +-
 .../chat/chat_prompt_editor.stories.tsx       |  13 ++
 .../components/chat/chat_prompt_editor.tsx    | 202 ++++++++++++------
 .../chat/function_list_popover.stories.tsx    |  31 +++
 .../components/chat/function_list_popover.tsx |  92 ++++++++
 .../__storybook_mocks__/use_functions.ts      |  60 ------
 .../public/hooks/use_functions.ts             |  22 --
 .../public/hooks/use_json_editor_model.ts     |  53 +++++
 .../conversations/conversation_view.tsx       |   4 +-
 .../public/utils/builders.ts                  |  45 +++-
 .../public/utils/storybook_decorator.tsx      |  37 +++-
 12 files changed, 413 insertions(+), 156 deletions(-)
 create mode 100644 x-pack/plugins/observability_ai_assistant/public/components/chat/function_list_popover.stories.tsx
 create mode 100644 x-pack/plugins/observability_ai_assistant/public/components/chat/function_list_popover.tsx
 delete mode 100644 x-pack/plugins/observability_ai_assistant/public/hooks/__storybook_mocks__/use_functions.ts
 delete mode 100644 x-pack/plugins/observability_ai_assistant/public/hooks/use_functions.ts
 create mode 100644 x-pack/plugins/observability_ai_assistant/public/hooks/use_json_editor_model.ts

diff --git a/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_body.tsx b/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_body.tsx
index 1bac96a9ddf6099..3ccd993b3205dc8 100644
--- a/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_body.tsx
+++ b/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_body.tsx
@@ -7,6 +7,7 @@
 
 import { EuiFlexGroup, EuiFlexItem, EuiHorizontalRule, EuiPanel, useEuiTheme } from '@elastic/eui';
 import { css } from '@emotion/css';
+import { i18n } from '@kbn/i18n';
 import type { AuthenticatedUser } from '@kbn/security-plugin/common';
 import React from 'react';
 import { type ConversationCreateRequest } from '../../../common/types';
@@ -68,7 +69,12 @@ export function ChatBody({
       <EuiFlexItem grow={false}>
         <EuiPanel hasBorder={false} hasShadow={false} paddingSize="m">
           <ChatHeader
-            title={initialConversation?.conversation.title ?? ''}
+            title={
+              initialConversation?.conversation.title ??
+              i18n.translate('xpack.observabilityAiAssistant.chatHeader.newConversation', {
+                defaultMessage: 'New conversation',
+              })
+            }
             connectors={connectors}
           />
         </EuiPanel>
diff --git a/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_header.tsx b/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_header.tsx
index 883c169057adbae..24f61d7f53d2c7f 100644
--- a/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_header.tsx
+++ b/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_header.tsx
@@ -4,8 +4,8 @@
  * 2.0; you may not use this file except in compliance with the Elastic License
  * 2.0.
  */
-import { EuiFlexGroup, EuiFlexItem, EuiTitle } from '@elastic/eui';
 import React from 'react';
+import { EuiFlexGroup, EuiFlexItem, EuiTitle } from '@elastic/eui';
 import { UseGenAIConnectorsResult } from '../../hooks/use_genai_connectors';
 import { AssistantAvatar } from '../assistant_avatar';
 import { ConnectorSelectorBase } from '../connector_selector/connector_selector_base';
diff --git a/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_prompt_editor.stories.tsx b/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_prompt_editor.stories.tsx
index 7506d259d76f0d3..eeb01da9b7d04a8 100644
--- a/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_prompt_editor.stories.tsx
+++ b/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_prompt_editor.stories.tsx
@@ -8,11 +8,24 @@
 import React from 'react';
 import { ComponentStory } from '@storybook/react';
 import { ChatPromptEditor as Component, ChatPromptEditorProps } from './chat_prompt_editor';
+import { KibanaReactStorybookDecorator } from '../../utils/storybook_decorator';
 
+/*
+  JSON Schema validation in the ChatPromptEditor compponent does not work
+  when rendering the component from within Storybook.
+  
+*/
 export default {
   component: Component,
   title: 'app/Molecules/ChatPromptEditor',
   argTypes: {},
+  parameters: {
+    backgrounds: {
+      default: 'white',
+      values: [{ name: 'white', value: '#fff' }],
+    },
+  },
+  decorators: [KibanaReactStorybookDecorator],
 };
 
 const Template: ComponentStory<typeof Component> = (props: ChatPromptEditorProps) => {
diff --git a/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_prompt_editor.tsx b/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_prompt_editor.tsx
index a23c06561fffc85..93317e9dfeb0f56 100644
--- a/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_prompt_editor.tsx
+++ b/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_prompt_editor.tsx
@@ -5,19 +5,23 @@
  * 2.0.
  */
 
-import React, { useState } from 'react';
+import React, { useEffect, useState } from 'react';
 import {
   EuiButtonIcon,
-  EuiContextMenuItem,
-  EuiContextMenuPanel,
+  EuiButtonEmpty,
   EuiFieldText,
   EuiFlexGroup,
   EuiFlexItem,
-  EuiPopover,
+  EuiSpacer,
+  EuiPanel,
 } from '@elastic/eui';
+import { CodeEditor } from '@kbn/kibana-react-plugin/public';
 import { i18n } from '@kbn/i18n';
-import { useFunctions, type Func } from '../../hooks/use_functions';
+import { useObservabilityAIAssistant } from '../../hooks/use_observability_ai_assistant';
+import { useJsonEditorModel } from '../../hooks/use_json_editor_model';
 import { type Message, MessageRole } from '../../../common';
+import type { FunctionDefinition } from '../../../common/types';
+import { FunctionListPopover } from './function_list_popover';
 
 export interface ChatPromptEditorProps {
   disabled: boolean;
@@ -26,80 +30,158 @@ export interface ChatPromptEditorProps {
 }
 
 export function ChatPromptEditor({ onSubmit, disabled, loading }: ChatPromptEditorProps) {
-  const functions = useFunctions();
+  const { getFunctions } = useObservabilityAIAssistant();
+  const functions = getFunctions();
 
   const [prompt, setPrompt] = useState('');
-  const [isFunctionListOpen, setIsFunctionListOpen] = useState(false);
+  const [functionPayload, setFunctionPayload] = useState<string | undefined>('');
+  const [selectedFunction, setSelectedFunction] = useState<FunctionDefinition | undefined>();
+
+  const { model, initialJsonString } = useJsonEditorModel(selectedFunction);
+
+  useEffect(() => {
+    setFunctionPayload(initialJsonString);
+  }, [initialJsonString, selectedFunction]);
 
   const handleChange = (event: React.ChangeEvent<HTMLInputElement>) => {
     setPrompt(event.currentTarget.value);
   };
 
-  const handleSubmit = () => {
-    const currentPrompt = prompt;
-    setPrompt('');
-    onSubmit({
-      '@timestamp': new Date().toISOString(),
-      message: { role: MessageRole.User, content: currentPrompt },
-    })
-      .then(() => {
-        setPrompt('');
-      })
-      .catch(() => {
-        setPrompt(currentPrompt);
-      });
+  const handleChangeFunctionPayload = (params: string) => {
+    setFunctionPayload(params);
   };
 
-  const handleClickFunctionList = () => {
-    setIsFunctionListOpen(!isFunctionListOpen);
+  const handleClearSelection = () => {
+    setSelectedFunction(undefined);
+    setFunctionPayload('');
   };
 
-  const handleSelectFunction = (func: Func) => {
-    setIsFunctionListOpen(false);
+  const handleSubmit = async () => {
+    const currentPrompt = prompt;
+    const currentPayload = functionPayload;
+
+    setPrompt('');
+    setFunctionPayload(undefined);
+
+    try {
+      if (selectedFunction) {
+        await onSubmit({
+          '@timestamp': new Date().toISOString(),
+          message: {
+            role: MessageRole.Function,
+            function_call: {
+              name: selectedFunction.options.name,
+              trigger: MessageRole.User,
+              arguments: currentPayload,
+            },
+          },
+        });
+      } else {
+        await onSubmit({
+          '@timestamp': new Date().toISOString(),
+          message: { role: MessageRole.User, content: currentPrompt },
+        });
+        setPrompt('');
+      }
+    } catch (_) {
+      setPrompt(currentPrompt);
+    }
   };
 
   return (
-    <EuiFlexGroup gutterSize="m">
-      <EuiFlexItem grow={false}>
-        <EuiPopover
-          anchorPosition="downLeft"
-          button={
-            <EuiButtonIcon
-              display="base"
-              iconType="function"
-              size="m"
-              onClick={handleClickFunctionList}
-            />
-          }
-          closePopover={handleClickFunctionList}
-          panelPaddingSize="s"
-          isOpen={isFunctionListOpen}
-        >
-          <EuiContextMenuPanel
-            size="s"
-            items={functions.map((func) => (
-              <EuiContextMenuItem key={func.id} onClick={() => handleSelectFunction(func)}>
-                {func.function_name}
-              </EuiContextMenuItem>
-            ))}
-          />
-        </EuiPopover>
-      </EuiFlexItem>
+    <EuiFlexGroup gutterSize="s" responsive={false}>
       <EuiFlexItem grow>
-        <EuiFieldText
-          fullWidth
-          value={prompt}
-          placeholder={i18n.translate('xpack.observabilityAiAssistant.prompt.placeholder', {
-            defaultMessage: 'Press ‘space’ or ‘$’ for function recommendations',
-          })}
-          onChange={handleChange}
-          onSubmit={handleSubmit}
-        />
+        <EuiFlexGroup direction="column" gutterSize="s">
+          <EuiFlexItem>
+            <EuiFlexGroup responsive={false}>
+              <EuiFlexItem grow>
+                <FunctionListPopover
+                  functions={functions}
+                  selectedFunction={selectedFunction}
+                  onSelectFunction={setSelectedFunction}
+                />
+              </EuiFlexItem>
+              <EuiFlexItem grow={false}>
+                {selectedFunction ? (
+                  <EuiButtonEmpty
+                    iconType="cross"
+                    iconSide="right"
+                    size="xs"
+                    onClick={handleClearSelection}
+                  >
+                    {i18n.translate('xpack.observabilityAiAssistant.prompt.emptySelection', {
+                      defaultMessage: 'Empty selection',
+                    })}
+                  </EuiButtonEmpty>
+                ) : null}
+              </EuiFlexItem>
+            </EuiFlexGroup>
+          </EuiFlexItem>
+          <EuiFlexItem>
+            {selectedFunction ? (
+              <EuiPanel borderRadius="none" color="subdued" hasShadow={false} paddingSize="xs">
+                <CodeEditor
+                  aria-label="payloadEditor"
+                  fullWidth
+                  height="120px"
+                  languageId="json"
+                  value={functionPayload || ''}
+                  onChange={handleChangeFunctionPayload}
+                  isCopyable
+                  languageConfiguration={{
+                    autoClosingPairs: [
+                      {
+                        open: '{',
+                        close: '}',
+                      },
+                    ],
+                  }}
+                  options={{
+                    automaticLayout: true,
+                    autoIndent: 'full',
+                    autoClosingQuotes: 'always',
+                    accessibilitySupport: 'off',
+                    acceptSuggestionOnEnter: 'on',
+                    contextmenu: true,
+                    fontSize: 12,
+                    formatOnType: true,
+                    formatOnPaste: true,
+                    inlineHints: { enabled: true },
+                    lineNumbers: 'on',
+                    minimap: { enabled: false },
+                    model,
+                    overviewRulerBorder: false,
+                    quickSuggestions: true,
+                    scrollbar: { alwaysConsumeMouseWheel: false },
+                    scrollBeyondLastLine: false,
+                    suggestOnTriggerCharacters: true,
+                    tabSize: 2,
+                    wordWrap: 'on',
+                    wrappingIndent: 'indent',
+                  }}
+                  transparentBackground
+                />
+              </EuiPanel>
+            ) : (
+              <EuiFieldText
+                fullWidth
+                value={prompt}
+                placeholder={i18n.translate('xpack.observabilityAiAssistant.prompt.placeholder', {
+                  defaultMessage: 'Press ‘$’ for function recommendations',
+                })}
+                onChange={handleChange}
+                onSubmit={handleSubmit}
+              />
+            )}
+          </EuiFlexItem>
+        </EuiFlexGroup>
       </EuiFlexItem>
       <EuiFlexItem grow={false}>
+        <EuiSpacer size="xl" />
         <EuiButtonIcon
+          aria-label="Submit"
           isLoading={loading}
-          disabled={!prompt || loading || disabled}
+          disabled={selectedFunction ? !functionPayload : !prompt || loading || disabled}
           display={prompt ? 'fill' : 'base'}
           iconType="kqlFunction"
           size="m"
diff --git a/x-pack/plugins/observability_ai_assistant/public/components/chat/function_list_popover.stories.tsx b/x-pack/plugins/observability_ai_assistant/public/components/chat/function_list_popover.stories.tsx
new file mode 100644
index 000000000000000..d4163fabaaa9aeb
--- /dev/null
+++ b/x-pack/plugins/observability_ai_assistant/public/components/chat/function_list_popover.stories.tsx
@@ -0,0 +1,31 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { ComponentStory } from '@storybook/react';
+import React from 'react';
+import { KibanaReactStorybookDecorator } from '../../utils/storybook_decorator';
+import { FunctionListPopover as Component } from './function_list_popover';
+
+export default {
+  component: Component,
+  title: 'app/Organisms/FunctionListPopover',
+  decorators: [KibanaReactStorybookDecorator],
+};
+
+type FunctionListPopover = React.ComponentProps<typeof Component>;
+
+const Template: ComponentStory<typeof Component> = (props: FunctionListPopover) => {
+  return <Component {...props} />;
+};
+
+const defaultProps: FunctionListPopover = {
+  functions: [],
+  onSelectFunction: () => {},
+};
+
+export const ConversationList = Template.bind({});
+ConversationList.args = defaultProps;
diff --git a/x-pack/plugins/observability_ai_assistant/public/components/chat/function_list_popover.tsx b/x-pack/plugins/observability_ai_assistant/public/components/chat/function_list_popover.tsx
new file mode 100644
index 000000000000000..3a8ee36f0083b41
--- /dev/null
+++ b/x-pack/plugins/observability_ai_assistant/public/components/chat/function_list_popover.tsx
@@ -0,0 +1,92 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import React, { useEffect, useState } from 'react';
+import {
+  EuiButtonEmpty,
+  EuiContextMenuItem,
+  EuiContextMenuPanel,
+  EuiPopover,
+  EuiSpacer,
+  EuiText,
+} from '@elastic/eui';
+import { i18n } from '@kbn/i18n';
+import { FunctionDefinition } from '../../../common/types';
+
+export function FunctionListPopover({
+  functions,
+  selectedFunction,
+  onSelectFunction,
+}: {
+  functions: FunctionDefinition[];
+  selectedFunction?: FunctionDefinition;
+  onSelectFunction: (func: FunctionDefinition) => void;
+}) {
+  const [isFunctionListOpen, setIsFunctionListOpen] = useState(false);
+
+  const handleClickFunctionList = () => {
+    setIsFunctionListOpen(!isFunctionListOpen);
+  };
+
+  const handleSelectFunction = (func: FunctionDefinition) => {
+    setIsFunctionListOpen(false);
+    onSelectFunction(func);
+  };
+
+  useEffect(() => {
+    const keyboardListener = (event: KeyboardEvent) => {
+      if (event.code === 'Digit4') {
+        setIsFunctionListOpen(true);
+      }
+    };
+
+    window.addEventListener('keyup', keyboardListener);
+
+    return () => {
+      window.removeEventListener('keyup', keyboardListener);
+    };
+  }, []);
+
+  return (
+    <EuiPopover
+      anchorPosition="downLeft"
+      button={
+        <EuiButtonEmpty
+          iconType="arrowRight"
+          iconSide="right"
+          size="xs"
+          onClick={handleClickFunctionList}
+        >
+          {selectedFunction
+            ? selectedFunction.options.name
+            : i18n.translate('xpack.observabilityAiAssistant.prompt.callFunction', {
+                defaultMessage: 'Call function',
+              })}
+        </EuiButtonEmpty>
+      }
+      closePopover={handleClickFunctionList}
+      panelPaddingSize="none"
+      isOpen={isFunctionListOpen}
+    >
+      <EuiContextMenuPanel size="s">
+        {functions.map((func) => (
+          <EuiContextMenuItem key={func.options.name} onClick={() => handleSelectFunction(func)}>
+            <EuiText size="s">
+              <p>
+                <strong>{func.options.name}</strong>
+              </p>
+            </EuiText>
+            <EuiSpacer size="xs" />
+            <EuiText size="s">
+              <p>{func.options.description}</p>
+            </EuiText>
+          </EuiContextMenuItem>
+        ))}
+      </EuiContextMenuPanel>
+    </EuiPopover>
+  );
+}
diff --git a/x-pack/plugins/observability_ai_assistant/public/hooks/__storybook_mocks__/use_functions.ts b/x-pack/plugins/observability_ai_assistant/public/hooks/__storybook_mocks__/use_functions.ts
deleted file mode 100644
index 06e13df331ea600..000000000000000
--- a/x-pack/plugins/observability_ai_assistant/public/hooks/__storybook_mocks__/use_functions.ts
+++ /dev/null
@@ -1,60 +0,0 @@
-/*
- * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
- * or more contributor license agreements. Licensed under the Elastic License
- * 2.0; you may not use this file except in compliance with the Elastic License
- * 2.0.
- */
-
-interface FuncParameters {
-  name: string;
-  type: 'string' | 'int' | 'boolean' | 'object';
-  description: string;
-}
-
-export interface Func {
-  id: string;
-  function_name: string;
-  parameters: FuncParameters[];
-}
-
-export function useFunctions(): Func[] {
-  return [
-    {
-      id: '1',
-      function_name: 'get_service_summary',
-      parameters: [
-        {
-          name: 'service',
-          type: 'string',
-          description: 'The service',
-        },
-      ],
-    },
-    {
-      id: '2',
-      function_name: 'get_apm_chart',
-      parameters: [],
-    },
-    {
-      id: '3',
-      function_name: 'get_dependencies',
-      parameters: [
-        {
-          name: 'service',
-          type: 'string',
-          description: 'The service',
-        },
-      ],
-    },
-    {
-      id: '4',
-      function_name: 'get_correlation_values',
-      parameters: [],
-    },
-    {
-      id: '4',
-      function_name: 'get_error',
-      parameters: [],
-    },
-  ];
-}
diff --git a/x-pack/plugins/observability_ai_assistant/public/hooks/use_functions.ts b/x-pack/plugins/observability_ai_assistant/public/hooks/use_functions.ts
deleted file mode 100644
index fbb47c31d6a0862..000000000000000
--- a/x-pack/plugins/observability_ai_assistant/public/hooks/use_functions.ts
+++ /dev/null
@@ -1,22 +0,0 @@
-/*
- * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
- * or more contributor license agreements. Licensed under the Elastic License
- * 2.0; you may not use this file except in compliance with the Elastic License
- * 2.0.
- */
-
-interface FuncParameters {
-  name: string;
-  type: 'string' | 'int' | 'boolean' | 'object';
-  description: string;
-}
-
-export interface Func {
-  id: string;
-  function_name: string;
-  parameters: FuncParameters[];
-}
-
-export function useFunctions(): Func[] {
-  return [];
-}
diff --git a/x-pack/plugins/observability_ai_assistant/public/hooks/use_json_editor_model.ts b/x-pack/plugins/observability_ai_assistant/public/hooks/use_json_editor_model.ts
new file mode 100644
index 000000000000000..f04cdf68e9c3e14
--- /dev/null
+++ b/x-pack/plugins/observability_ai_assistant/public/hooks/use_json_editor_model.ts
@@ -0,0 +1,53 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+import { useMemo } from 'react';
+import { monaco } from '@kbn/monaco';
+import { FunctionDefinition } from '../../common/types';
+
+const { editor, languages, Uri } = monaco;
+
+const SCHEMA_URI = 'http://elastic.co/foo.json';
+const modelUri = Uri.parse(SCHEMA_URI);
+
+export const useJsonEditorModel = (functionDefinition?: FunctionDefinition) => {
+  return useMemo(() => {
+    if (!functionDefinition) {
+      return {};
+    }
+
+    const schema = { ...functionDefinition.options.parameters };
+
+    const initialJsonString = functionDefinition.options.parameters.properties
+      ? Object.keys(functionDefinition.options.parameters.properties).reduce(
+          (acc, curr, index, arr) => {
+            const val = `${acc}  "${curr}": "",\n`;
+            return index === arr.length - 1 ? `${val}}` : val;
+          },
+          '{\n'
+        )
+      : '';
+
+    languages.json.jsonDefaults.setDiagnosticsOptions({
+      validate: true,
+      schemas: [
+        {
+          uri: SCHEMA_URI,
+          fileMatch: [String(modelUri)],
+          schema,
+        },
+      ],
+    });
+
+    let model = editor.getModel(modelUri);
+
+    if (model === null) {
+      model = editor.createModel(initialJsonString, 'json', modelUri);
+    }
+
+    return { model, initialJsonString };
+  }, [functionDefinition]);
+};
diff --git a/x-pack/plugins/observability_ai_assistant/public/routes/conversations/conversation_view.tsx b/x-pack/plugins/observability_ai_assistant/public/routes/conversations/conversation_view.tsx
index 47c82a171af6fb3..7d6a7cf8b65d55a 100644
--- a/x-pack/plugins/observability_ai_assistant/public/routes/conversations/conversation_view.tsx
+++ b/x-pack/plugins/observability_ai_assistant/public/routes/conversations/conversation_view.tsx
@@ -4,7 +4,7 @@
  * 2.0; you may not use this file except in compliance with the Elastic License
  * 2.0.
  */
-import { EuiFlexGroup, EuiFlexItem } from '@elastic/eui';
+import { EuiFlexGroup, EuiFlexItem, EuiSpacer } from '@elastic/eui';
 import { css } from '@emotion/css';
 import React from 'react';
 import { ChatBody } from '../../components/chat/chat_body';
@@ -26,7 +26,6 @@ export function ConversationView() {
   const currentUser = useCurrentUser();
 
   const service = useObservabilityAIAssistant();
-
   return (
     <EuiFlexGroup direction="row" className={containerClassName}>
       <EuiFlexItem grow={false} />
@@ -37,6 +36,7 @@ export function ConversationView() {
           initialConversation={undefined}
           service={service}
         />
+        <EuiSpacer size="m" />
       </EuiFlexItem>
     </EuiFlexGroup>
   );
diff --git a/x-pack/plugins/observability_ai_assistant/public/utils/builders.ts b/x-pack/plugins/observability_ai_assistant/public/utils/builders.ts
index 819734b8fdfab2e..7ecd909630e5c6f 100644
--- a/x-pack/plugins/observability_ai_assistant/public/utils/builders.ts
+++ b/x-pack/plugins/observability_ai_assistant/public/utils/builders.ts
@@ -6,7 +6,7 @@
  */
 
 import { uniqueId } from 'lodash';
-import { MessageRole, Conversation } from '../../common/types';
+import { MessageRole, Conversation, FunctionDefinition } from '../../common/types';
 import { ChatTimelineItem } from '../components/chat/chat_timeline';
 
 type ChatItemBuildProps = Partial<ChatTimelineItem> & Pick<ChatTimelineItem, 'role'>;
@@ -96,3 +96,46 @@ export function buildConversation(params?: Partial<Conversation>) {
     ...params,
   };
 }
+
+export function buildFunction(): FunctionDefinition {
+  return {
+    options: {
+      name: 'elasticsearch',
+      contexts: ['core'],
+      description: 'Call Elasticsearch APIs on behalf of the user',
+      parameters: {
+        type: 'object',
+        properties: {
+          method: {
+            type: 'string',
+            description: 'The HTTP method of the Elasticsearch endpoint',
+            enum: ['GET', 'PUT', 'POST', 'DELETE', 'PATCH'] as const,
+          },
+          path: {
+            type: 'string',
+            description: 'The path of the Elasticsearch endpoint, including query parameters',
+          },
+        },
+        required: ['method' as const, 'path' as const],
+      },
+    },
+    respond: async (options: { arguments: any }, signal: AbortSignal) => ({}),
+  };
+}
+
+export const buildFunctionElasticsearch = buildFunction;
+
+export function buildFunctionServiceSummary(): FunctionDefinition {
+  return {
+    options: {
+      name: 'get_service_summary',
+      contexts: ['core'],
+      description:
+        'Gets a summary of a single service, including: the language, service version, deployments, infrastructure, alerting, etc. ',
+      parameters: {
+        type: 'object',
+      },
+    },
+    respond: async (options: { arguments: any }, signal: AbortSignal) => ({}),
+  };
+}
diff --git a/x-pack/plugins/observability_ai_assistant/public/utils/storybook_decorator.tsx b/x-pack/plugins/observability_ai_assistant/public/utils/storybook_decorator.tsx
index 5725ce3473fd2d4..3673bd98794823e 100644
--- a/x-pack/plugins/observability_ai_assistant/public/utils/storybook_decorator.tsx
+++ b/x-pack/plugins/observability_ai_assistant/public/utils/storybook_decorator.tsx
@@ -5,22 +5,41 @@
  * 2.0.
  */
 import React, { ComponentType } from 'react';
+import { Observable } from 'rxjs';
 import { KibanaContextProvider } from '@kbn/kibana-react-plugin/public';
+import type { Serializable } from '@kbn/utility-types';
+import type { AuthenticatedUser } from '@kbn/security-plugin/common';
 import { ObservabilityAIAssistantProvider } from '../context/observability_ai_assistant_provider';
+import { ObservabilityAIAssistantAPIClient } from '../api';
+import type { Message } from '../../common';
+import type { ObservabilityAIAssistantService, PendingMessage } from '../types';
+import { buildFunctionElasticsearch, buildFunctionServiceSummary } from './builders';
 
-const service = {
+const service: ObservabilityAIAssistantService = {
   isEnabled: () => true,
-  chat: async (options: {
-    messages: [];
-    connectorId: string;
-    // signal: new AbortSignal();
-  }) => {},
-  // callApi: ObservabilityAIAssistantAPIClient;
-  getCurrentUser: async () => {},
+  chat: (options: { messages: Message[]; connectorId: string }) => new Observable<PendingMessage>(),
+  callApi: {} as ObservabilityAIAssistantAPIClient,
+  getCurrentUser: async (): Promise<AuthenticatedUser> => ({
+    username: 'user',
+    roles: [],
+    enabled: true,
+    authentication_realm: { name: 'foo', type: '' },
+    lookup_realm: { name: 'foo', type: '' },
+    authentication_provider: { name: '', type: '' },
+    authentication_type: '',
+    elastic_cloud_user: false,
+  }),
+  getContexts: () => [],
+  getFunctions: () => [buildFunctionElasticsearch(), buildFunctionServiceSummary()],
+  executeFunction: async (
+    name: string,
+    args: string | undefined,
+    signal: AbortSignal
+  ): Promise<{ content?: Serializable; data?: Serializable }> => ({}),
+  renderFunction: (name: string, response: {}) => <div>Hello! {name}</div>,
 };
 
 export function KibanaReactStorybookDecorator(Story: ComponentType) {
-  console.log('hello?');
   return (
     <KibanaContextProvider
       services={{