From 133d6b975689401f16e470763233e261f8d8a301 Mon Sep 17 00:00:00 2001
From: Coen Warmer <coen.warmer@gmail.com>
Date: Fri, 11 Aug 2023 22:21:29 +0200
Subject: [PATCH 1/5] Add autotitling of conversations

---
 .../action_menu_item/action_menu_item.tsx     |  3 +
 .../components/chat/chat_timeline.stories.tsx | 21 ++++++-
 .../components/chat/conversation_list.tsx     |  1 +
 .../public/hooks/use_conversation.ts          | 37 +++++++++--
 .../conversations/conversation_view.tsx       | 37 ++++++-----
 .../server/routes/chat/route.ts               |  3 +-
 .../server/routes/conversations/route.ts      | 62 ++++++++++++++++++-
 .../server/service/client/index.ts            | 17 ++---
 .../server/service/types.ts                   |  7 ++-
 9 files changed, 149 insertions(+), 39 deletions(-)

diff --git a/x-pack/plugins/observability_ai_assistant/public/components/action_menu_item/action_menu_item.tsx b/x-pack/plugins/observability_ai_assistant/public/components/action_menu_item/action_menu_item.tsx
index b7d843506a550..1e9a847348e5b 100644
--- a/x-pack/plugins/observability_ai_assistant/public/components/action_menu_item/action_menu_item.tsx
+++ b/x-pack/plugins/observability_ai_assistant/public/components/action_menu_item/action_menu_item.tsx
@@ -10,6 +10,7 @@ import React, { useState } from 'react';
 import { ObservabilityAIAssistantChatServiceProvider } from '../../context/observability_ai_assistant_chat_service_provider';
 import { useAbortableAsync } from '../../hooks/use_abortable_async';
 import { useConversation } from '../../hooks/use_conversation';
+import { useGenAIConnectors } from '../../hooks/use_genai_connectors';
 import { useObservabilityAIAssistant } from '../../hooks/use_observability_ai_assistant';
 import { EMPTY_CONVERSATION_TITLE } from '../../i18n';
 import { AssistantAvatar } from '../assistant_avatar';
@@ -17,6 +18,7 @@ import { ChatFlyout } from '../chat/chat_flyout';
 
 export function ObservabilityAIAssistantActionMenuItem() {
   const service = useObservabilityAIAssistant();
+  const connectors = useGenAIConnectors();
 
   const [isOpen, setIsOpen] = useState(false);
 
@@ -34,6 +36,7 @@ export function ObservabilityAIAssistantActionMenuItem() {
 
   const { conversation, displayedMessages, setDisplayedMessages, save } = useConversation({
     conversationId,
+    connectorId: connectors.selectedConnector,
   });
 
   if (!service.isEnabled()) {
diff --git a/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_timeline.stories.tsx b/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_timeline.stories.tsx
index 94bcb27d9810b..abbbb6806fd38 100644
--- a/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_timeline.stories.tsx
+++ b/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_timeline.stories.tsx
@@ -77,12 +77,22 @@ const defaultProps: ComponentProps<typeof Component> = {
         arguments: '{ "foo": "bar" }',
         trigger: MessageRole.Assistant,
       },
-      canEdit: true,
+      actions: {
+        canEdit: false,
+        canCopy: true,
+        canGiveFeedback: true,
+        canRegenerate: true,
+      },
     }),
     buildFunctionChatItem({
       content: '{ "message": "The arguments are wrong" }',
       error: new Error(),
-      canRegenerate: false,
+      actions: {
+        canRegenerate: false,
+        canEdit: true,
+        canGiveFeedback: false,
+        canCopy: true,
+      },
     }),
     buildAssistantChatItem({
       content: '',
@@ -92,7 +102,12 @@ const defaultProps: ComponentProps<typeof Component> = {
         arguments: '{ "bar": "foo" }',
         trigger: MessageRole.Assistant,
       },
-      canEdit: true,
+      actions: {
+        canEdit: true,
+        canCopy: true,
+        canGiveFeedback: true,
+        canRegenerate: true,
+      },
     }),
     buildFunctionChatItem({
       content: '',
diff --git a/x-pack/plugins/observability_ai_assistant/public/components/chat/conversation_list.tsx b/x-pack/plugins/observability_ai_assistant/public/components/chat/conversation_list.tsx
index 05408f62ea6ba..3c260dfbaf7de 100644
--- a/x-pack/plugins/observability_ai_assistant/public/components/chat/conversation_list.tsx
+++ b/x-pack/plugins/observability_ai_assistant/public/components/chat/conversation_list.tsx
@@ -101,6 +101,7 @@ export function ConversationList({
                       isActive={conversation.id === selected}
                       isDisabled={loading}
                       href={conversation.href}
+                      wrapText
                       extraAction={
                         conversation.id
                           ? {
diff --git a/x-pack/plugins/observability_ai_assistant/public/hooks/use_conversation.ts b/x-pack/plugins/observability_ai_assistant/public/hooks/use_conversation.ts
index b0793c59e1f6b..74aa98c5351ea 100644
--- a/x-pack/plugins/observability_ai_assistant/public/hooks/use_conversation.ts
+++ b/x-pack/plugins/observability_ai_assistant/public/hooks/use_conversation.ts
@@ -7,7 +7,7 @@
 import { i18n } from '@kbn/i18n';
 import { merge, omit } from 'lodash';
 import { Dispatch, SetStateAction, useState } from 'react';
-import type { Conversation, Message } from '../../common';
+import { type Conversation, type Message } from '../../common';
 import type { ConversationCreateRequest } from '../../common/types';
 import { ObservabilityAIAssistantChatService } from '../types';
 import { useAbortableAsync, type AbortableAsyncState } from './use_abortable_async';
@@ -18,14 +18,16 @@ import { createNewConversation } from './use_timeline';
 export function useConversation({
   conversationId,
   chatService,
+  connectorId,
 }: {
   conversationId?: string;
   chatService?: ObservabilityAIAssistantChatService;
+  connectorId: string | undefined;
 }): {
   conversation: AbortableAsyncState<ConversationCreateRequest | Conversation | undefined>;
   displayedMessages: Message[];
   setDisplayedMessages: Dispatch<SetStateAction<Message[]>>;
-  save: (messages: Message[]) => Promise<Conversation>;
+  save: (messages: Message[], handleRefreshConversations?: () => void) => Promise<Conversation>;
 } {
   const service = useObservabilityAIAssistant();
 
@@ -67,11 +69,12 @@ export function useConversation({
     conversation,
     displayedMessages,
     setDisplayedMessages,
-    save: (messages: Message[]) => {
+    save: (messages: Message[], handleRefreshConversations?: () => void) => {
       const conversationObject = conversation.value!;
+
       return conversationId
         ? service
-            .callApi(`POST /internal/observability_ai_assistant/conversation/{conversationId}`, {
+            .callApi(`PUT /internal/observability_ai_assistant/conversation/{conversationId}`, {
               signal: null,
               params: {
                 path: {
@@ -100,7 +103,7 @@ export function useConversation({
               throw err;
             })
         : service
-            .callApi(`PUT /internal/observability_ai_assistant/conversation`, {
+            .callApi(`POST /internal/observability_ai_assistant/conversation`, {
               signal: null,
               params: {
                 body: {
@@ -108,6 +111,30 @@ export function useConversation({
                 },
               },
             })
+            .then((nextConversation) => {
+              if (connectorId) {
+                service
+                  .callApi(
+                    `PUT /internal/observability_ai_assistant/conversation/{conversationId}/auto_title`,
+                    {
+                      signal: null,
+                      params: {
+                        path: {
+                          conversationId: nextConversation.conversation.id,
+                        },
+                        body: {
+                          connectorId,
+                        },
+                      },
+                    }
+                  )
+                  .then(() => {
+                    handleRefreshConversations?.();
+                    return conversation.refresh();
+                  });
+              }
+              return nextConversation;
+            })
             .catch((err) => {
               notifications.toasts.addError(err, {
                 title: i18n.translate('xpack.observabilityAiAssistant.errorCreatingConversation', {
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 3573a819da842..d4c8e94193f40 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
@@ -33,14 +33,14 @@ const chatBodyContainerClassNameWithError = css`
 `;
 
 export function ConversationView() {
-  const connectors = useGenAIConnectors();
-
-  const knowledgeBase = useKnowledgeBase();
-
   const currentUser = useCurrentUser();
 
   const service = useObservabilityAIAssistant();
 
+  const connectors = useGenAIConnectors();
+
+  const knowledgeBase = useKnowledgeBase();
+
   const observabilityAIAssistantRouter = useObservabilityAIAssistantRouter();
 
   const { path } = useObservabilityAIAssistantParams('/conversations/*');
@@ -75,6 +75,7 @@ export function ConversationView() {
   const { conversation, displayedMessages, setDisplayedMessages, save } = useConversation({
     conversationId,
     chatService: chatService.value,
+    connectorId: connectors.selectedConnector,
   });
 
   const conversations = useAbortableAsync(
@@ -111,11 +112,15 @@ export function ConversationView() {
     );
   }
 
+  function handleRefreshConversations() {
+    conversations.refresh();
+  }
+
   return (
     <>
       {confirmDeleteElement}
       <EuiFlexGroup direction="row" className={containerClassName}>
-        <EuiFlexItem grow={false}>
+        <EuiFlexItem grow={false} style={{ minWidth: 300, maxWidth: 300 }}>
           <ConversationList
             selected={conversationId ?? ''}
             loading={conversations.loading || isUpdatingList}
@@ -210,7 +215,7 @@ export function ConversationView() {
               })}
             </EuiCallOut>
           ) : null}
-          {chatService.loading || conversation.loading ? (
+          {chatService.loading ? (
             <EuiFlexGroup direction="column" alignItems="center" gutterSize="l">
               <EuiFlexItem grow={false}>
                 <EuiSpacer size="xl" />
@@ -218,27 +223,27 @@ export function ConversationView() {
               </EuiFlexItem>
             </EuiFlexGroup>
           ) : null}
-          {!conversation.error && conversation.value && chatService.value ? (
+          {conversation.value && chatService.value && !conversation.error ? (
             <ObservabilityAIAssistantChatServiceProvider value={chatService.value}>
               <ChatBody
-                currentUser={currentUser}
                 connectors={connectors}
-                knowledgeBase={knowledgeBase}
-                title={conversation.value.conversation.title}
                 connectorsManagementHref={getConnectorsManagementHref(http)}
+                currentUser={currentUser}
+                knowledgeBase={knowledgeBase}
                 messages={displayedMessages}
+                title={conversation.value.conversation.title}
+                onChatUpdate={(messages) => {
+                  setDisplayedMessages(messages);
+                }}
                 onChatComplete={(messages) => {
-                  save(messages)
+                  save(messages, handleRefreshConversations)
                     .then((nextConversation) => {
                       conversations.refresh();
-                      if (!conversationId) {
+                      if (!conversationId && nextConversation?.conversation?.id) {
                         navigateToConversation(nextConversation.conversation.id);
                       }
                     })
-                    .catch(() => {});
-                }}
-                onChatUpdate={(messages) => {
-                  setDisplayedMessages(messages);
+                    .catch((e) => {});
                 }}
               />
             </ObservabilityAIAssistantChatServiceProvider>
diff --git a/x-pack/plugins/observability_ai_assistant/server/routes/chat/route.ts b/x-pack/plugins/observability_ai_assistant/server/routes/chat/route.ts
index c8cc5f0af45fd..152bc5fec1ac1 100644
--- a/x-pack/plugins/observability_ai_assistant/server/routes/chat/route.ts
+++ b/x-pack/plugins/observability_ai_assistant/server/routes/chat/route.ts
@@ -5,7 +5,6 @@
  * 2.0.
  */
 import * as t from 'io-ts';
-import type { IncomingMessage } from 'http';
 import { notImplemented } from '@hapi/boom';
 import { createObservabilityAIAssistantServerRoute } from '../create_observability_ai_assistant_server_route';
 import { messageRt } from '../runtime_types';
@@ -28,7 +27,7 @@ const chatRoute = createObservabilityAIAssistantServerRoute({
       ),
     }),
   }),
-  handler: async (resources): Promise<IncomingMessage> => {
+  handler: async (resources) => {
     const { request, params, service } = resources;
 
     const client = await service.getClient({ request });
diff --git a/x-pack/plugins/observability_ai_assistant/server/routes/conversations/route.ts b/x-pack/plugins/observability_ai_assistant/server/routes/conversations/route.ts
index 625594cc74578..3c42c9a70b4e2 100644
--- a/x-pack/plugins/observability_ai_assistant/server/routes/conversations/route.ts
+++ b/x-pack/plugins/observability_ai_assistant/server/routes/conversations/route.ts
@@ -7,7 +7,7 @@
 import { notImplemented } from '@hapi/boom';
 import * as t from 'io-ts';
 import { merge } from 'lodash';
-import { Conversation } from '../../../common/types';
+import { Conversation, MessageRole } from '../../../common/types';
 import { createObservabilityAIAssistantServerRoute } from '../create_observability_ai_assistant_server_route';
 import { conversationCreateRt, conversationUpdateRt } from '../runtime_types';
 
@@ -58,7 +58,7 @@ const findConversationsRoute = createObservabilityAIAssistantServerRoute({
 });
 
 const createConversationRoute = createObservabilityAIAssistantServerRoute({
-  endpoint: 'PUT /internal/observability_ai_assistant/conversation',
+  endpoint: 'POST /internal/observability_ai_assistant/conversation',
   params: t.type({
     body: t.type({
       conversation: conversationCreateRt,
@@ -81,7 +81,7 @@ const createConversationRoute = createObservabilityAIAssistantServerRoute({
 });
 
 const updateConversationRoute = createObservabilityAIAssistantServerRoute({
-  endpoint: 'POST /internal/observability_ai_assistant/conversation/{conversationId}',
+  endpoint: 'PUT /internal/observability_ai_assistant/conversation/{conversationId}',
   params: t.type({
     path: t.type({
       conversationId: t.string,
@@ -108,6 +108,61 @@ const updateConversationRoute = createObservabilityAIAssistantServerRoute({
   },
 });
 
+const updateConversationTitleBasedOnMessages = createObservabilityAIAssistantServerRoute({
+  endpoint: 'PUT /internal/observability_ai_assistant/conversation/{conversationId}/auto_title',
+  params: t.type({
+    path: t.type({
+      conversationId: t.string,
+    }),
+    body: t.type({
+      connectorId: t.string,
+    }),
+  }),
+  options: {
+    tags: ['access:ai_assistant'],
+  },
+  handler: async (resources): Promise<Conversation> => {
+    const { service, request, params } = resources;
+
+    const client = await service.getClient({ request });
+
+    if (!client) {
+      throw notImplemented();
+    }
+
+    const conversation = await client.get(params.path.conversationId);
+
+    const response = await client.chat({
+      messages: [
+        {
+          '@timestamp': new Date().toISOString(),
+          message: {
+            role: MessageRole.Assistant,
+            content: conversation.messages.slice(1).reduce((acc, curr) => {
+              return `${acc} ${curr.message.role}: ${curr.message.content}`;
+            }, 'You are a helpful assistant for Elastic Observability. Assume the following message is the start of a conversation between you and a user; give this conversation a title based on this content: '),
+          },
+        },
+      ],
+      connectorId: params.body.connectorId,
+      stream: false,
+    });
+
+    if ('object' in response && response.object === 'chat.completion') {
+      const title =
+        response.choices[0].message?.content?.slice(1, -1) ||
+        `Conversation on ${conversation['@timestamp']}`;
+
+      return client.update({
+        ...conversation,
+        conversation: { ...conversation.conversation, title },
+      });
+    }
+
+    return Promise.resolve(conversation);
+  },
+});
+
 const deleteConversationRoute = createObservabilityAIAssistantServerRoute({
   endpoint: 'DELETE /internal/observability_ai_assistant/conversation/{conversationId}',
   params: t.type({
@@ -136,5 +191,6 @@ export const conversationRoutes = {
   ...findConversationsRoute,
   ...createConversationRoute,
   ...updateConversationRoute,
+  ...updateConversationTitleBasedOnMessages,
   ...deleteConversationRoute,
 };
diff --git a/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts b/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts
index e55144a089b8a..29cd817965e90 100644
--- a/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts
+++ b/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts
@@ -17,6 +17,7 @@ import type {
   ChatCompletionFunctions,
   ChatCompletionRequestMessage,
   CreateChatCompletionRequest,
+  CreateChatCompletionResponse,
 } from 'openai';
 import pRetry from 'p-retry';
 import { v4 } from 'uuid';
@@ -138,15 +139,17 @@ export class ObservabilityAIAssistantClient implements IObservabilityAIAssistant
     });
   };
 
-  chat = async ({
+  chat = async <TStream extends boolean | undefined = true>({
     messages,
     connectorId,
     functions,
+    stream = true,
   }: {
     messages: Message[];
     connectorId: string;
-    functions: Array<Pick<FunctionDefinition['options'], 'name' | 'description' | 'parameters'>>;
-  }): Promise<IncomingMessage> => {
+    functions?: Array<Pick<FunctionDefinition['options'], 'name' | 'description' | 'parameters'>>;
+    stream?: TStream;
+  }): Promise<TStream extends false ? CreateChatCompletionResponse : IncomingMessage> => {
     const messagesForOpenAI: ChatCompletionRequestMessage[] = compact(
       messages
         .filter((message) => message.message.content || message.message.function_call?.name)
@@ -165,7 +168,7 @@ export class ObservabilityAIAssistantClient implements IObservabilityAIAssistant
         })
     );
 
-    const functionsForOpenAI: ChatCompletionFunctions[] = functions;
+    const functionsForOpenAI: ChatCompletionFunctions[] | undefined = functions;
 
     const request: Omit<CreateChatCompletionRequest, 'model'> & { model?: string } = {
       messages: messagesForOpenAI,
@@ -177,10 +180,10 @@ export class ObservabilityAIAssistantClient implements IObservabilityAIAssistant
     const executeResult = await this.dependencies.actionsClient.execute({
       actionId: connectorId,
       params: {
-        subAction: 'stream',
+        subAction: stream ? 'stream' : 'run',
         subActionParams: {
           body: JSON.stringify(request),
-          stream: true,
+          ...(stream ? { stream: true } : {}),
         },
       },
     });
@@ -189,7 +192,7 @@ export class ObservabilityAIAssistantClient implements IObservabilityAIAssistant
       throw internal(`${executeResult?.message} - ${executeResult?.serviceMessage}`);
     }
 
-    return executeResult.data as IncomingMessage;
+    return executeResult.data as any;
   };
 
   find = async (options?: { query?: string }): Promise<{ conversations: Conversation[] }> => {
diff --git a/x-pack/plugins/observability_ai_assistant/server/service/types.ts b/x-pack/plugins/observability_ai_assistant/server/service/types.ts
index cc6c4c744f969..9d2326df00ebc 100644
--- a/x-pack/plugins/observability_ai_assistant/server/service/types.ts
+++ b/x-pack/plugins/observability_ai_assistant/server/service/types.ts
@@ -5,7 +5,6 @@
  * 2.0.
  */
 
-import { IncomingMessage } from 'http';
 import { KibanaRequest } from '@kbn/core/server';
 import type {
   Conversation,
@@ -15,13 +14,15 @@ import type {
   KnowledgeBaseEntry,
   Message,
 } from '../../common/types';
+import { ObservabilityAIAssistantClient } from './client';
 
 export interface IObservabilityAIAssistantClient {
   chat: (options: {
     messages: Message[];
     connectorId: string;
-    functions: Array<Pick<FunctionDefinition['options'], 'name' | 'description' | 'parameters'>>;
-  }) => Promise<IncomingMessage>;
+    functions?: Array<Pick<FunctionDefinition['options'], 'name' | 'description' | 'parameters'>>;
+    stream?: boolean | undefined;
+  }) => ReturnType<ObservabilityAIAssistantClient['chat']>;
   get: (conversationId: string) => Promise<Conversation>;
   find: (options?: { query?: string }) => Promise<{ conversations: Conversation[] }>;
   create: (conversation: ConversationCreateRequest) => Promise<Conversation>;

From c45fd1d80ec094d5e0ec111e5cab56d3573f9d92 Mon Sep 17 00:00:00 2001
From: Coen Warmer <coen.warmer@gmail.com>
Date: Sat, 12 Aug 2023 02:28:12 +0200
Subject: [PATCH 2/5] Add filtering to function list

---
 .../components/chat/conversation_list.tsx     |   6 +-
 .../components/chat/function_list_popover.tsx | 123 +++++++++++++-----
 .../conversations/conversation_view.tsx       |   2 -
 3 files changed, 93 insertions(+), 38 deletions(-)

diff --git a/x-pack/plugins/observability_ai_assistant/public/components/chat/conversation_list.tsx b/x-pack/plugins/observability_ai_assistant/public/components/chat/conversation_list.tsx
index 3c260dfbaf7de..23ac0cecebcbc 100644
--- a/x-pack/plugins/observability_ai_assistant/public/components/chat/conversation_list.tsx
+++ b/x-pack/plugins/observability_ai_assistant/public/components/chat/conversation_list.tsx
@@ -28,6 +28,10 @@ const titleClassName = css`
   text-transform: uppercase;
 `;
 
+const newChatButtonWrapperClassName = css`
+  padding-bottom: 5px;
+`;
+
 export function ConversationList({
   selected,
   onClickNewChat,
@@ -136,7 +140,7 @@ export function ConversationList({
         <EuiFlexItem grow={false}>
           <EuiPanel paddingSize="s" hasBorder={false} hasShadow={false}>
             <EuiFlexGroup alignItems="center">
-              <EuiFlexItem grow>
+              <EuiFlexItem grow className={newChatButtonWrapperClassName}>
                 <NewChatButton onClick={onClickNewChat} />
               </EuiFlexItem>
             </EuiFlexGroup>
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
index 64e4233ab1d7b..f77c715a3c7e0 100644
--- 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
@@ -8,16 +8,22 @@
 import React, { useEffect, useState } from 'react';
 import {
   EuiButtonEmpty,
-  EuiContextMenu,
-  EuiContextMenuPanel,
+  EuiHighlight,
   EuiPopover,
+  EuiSelectable,
+  EuiSelectableOption,
   EuiSpacer,
   EuiText,
 } from '@elastic/eui';
+import type { EuiSelectableOptionCheckedType } from '@elastic/eui/src/components/selectable/selectable_option';
 import { i18n } from '@kbn/i18n';
-import { FunctionDefinition } from '../../../common/types';
 import { useObservabilityAIAssistantChatService } from '../../hooks/use_observability_ai_assistant_chat_service';
 
+interface FunctionListOption {
+  label: string;
+  searchableLabel: string;
+}
+
 export function FunctionListPopover({
   selectedFunctionName,
   onSelectFunction,
@@ -25,7 +31,11 @@ export function FunctionListPopover({
   selectedFunctionName?: string;
   onSelectFunction: (func: string) => void;
 }) {
-  const chatService = useObservabilityAIAssistantChatService();
+  const { getFunctions } = useObservabilityAIAssistantChatService();
+
+  const [functionOptions, setFunctionOptions] = useState<
+    Array<EuiSelectableOption<FunctionListOption>>
+  >([]);
 
   const [isFunctionListOpen, setIsFunctionListOpen] = useState(false);
 
@@ -33,9 +43,9 @@ export function FunctionListPopover({
     setIsFunctionListOpen(!isFunctionListOpen);
   };
 
-  const handleSelectFunction = (func: FunctionDefinition) => {
+  const handleSelectFunction = (func: EuiSelectableOption<FunctionListOption>) => {
     setIsFunctionListOpen(false);
-    onSelectFunction(func.options.name);
+    onSelectFunction(func.label);
   };
 
   useEffect(() => {
@@ -52,6 +62,42 @@ export function FunctionListPopover({
     };
   }, []);
 
+  useEffect(() => {
+    const options = getFunctions().map((func) => ({
+      label: func.options.name,
+      searchableLabel: func.options.descriptionForUser,
+      checked:
+        func.options.name === selectedFunctionName
+          ? ('on' as EuiSelectableOptionCheckedType)
+          : ('off' as EuiSelectableOptionCheckedType),
+    }));
+
+    setFunctionOptions(options);
+  }, [getFunctions, selectedFunctionName]);
+
+  const renderCountryOption = (
+    option: EuiSelectableOption<FunctionListOption>,
+    searchValue: string
+  ) => {
+    return (
+      <>
+        <EuiText size="s">
+          <p>
+            <strong>
+              <EuiHighlight search={searchValue}>{option.label}</EuiHighlight>{' '}
+            </strong>
+          </p>
+        </EuiText>
+        <EuiSpacer size="xs" />
+        <EuiText size="s">
+          <p style={{ textOverflow: 'ellipsis', overflow: 'hidden' }}>
+            <EuiHighlight search={searchValue}>{option.searchableLabel || ''}</EuiHighlight>
+          </p>
+        </EuiText>
+      </>
+    );
+  };
+
   return (
     <EuiPopover
       anchorPosition="downLeft"
@@ -64,43 +110,50 @@ export function FunctionListPopover({
         >
           {selectedFunctionName
             ? selectedFunctionName
-            : i18n.translate('xpack.observabilityAiAssistant.prompt.callFunction', {
+            : i18n.translate('xpack.observabilityAiAssistant.prompt.functionList.callFunction', {
                 defaultMessage: 'Call function',
               })}
         </EuiButtonEmpty>
       }
       closePopover={handleClickFunctionList}
-      css={{ maxWidth: 400 }}
       panelPaddingSize="none"
       isOpen={isFunctionListOpen}
     >
-      <EuiContextMenuPanel size="s">
-        <EuiContextMenu
-          initialPanelId={0}
-          panels={[
-            {
-              id: 0,
-              width: 500,
-              items: chatService.getFunctions().map((func) => ({
-                name: (
-                  <>
-                    <EuiText size="s">
-                      <p>
-                        <strong>{func.options.name}</strong>
-                      </p>
-                    </EuiText>
-                    <EuiSpacer size="xs" />
-                    <EuiText size="s">
-                      <p>{func.options.descriptionForUser}</p>
-                    </EuiText>
-                  </>
-                ),
-                onClick: () => handleSelectFunction(func),
-              })),
-            },
-          ]}
-        />
-      </EuiContextMenuPanel>
+      <EuiSelectable
+        aria-label={i18n.translate(
+          'xpack.observabilityAiAssistant.prompt.functionList.functionList',
+          {
+            defaultMessage: 'Function list',
+          }
+        )}
+        listProps={{
+          isVirtualized: false,
+          showIcons: false,
+        }}
+        options={functionOptions}
+        renderOption={renderCountryOption}
+        searchable
+        searchProps={{
+          'data-test-subj': 'selectableSearchHere',
+          placeholder: i18n.translate('xpack.observabilityAiAssistant.prompt.functionList.filter', {
+            defaultMessage: 'Filter',
+          }),
+        }}
+        singleSelection
+        onChange={(functions) => {
+          const selectedFunction = functions.filter((fn) => fn.checked !== 'off');
+          if (selectedFunction && selectedFunction.length === 1) {
+            handleSelectFunction({ ...selectedFunction[0], checked: 'on' });
+          }
+        }}
+      >
+        {(list, search) => (
+          <div style={{ overflow: 'hidden' }}>
+            {search}
+            <div style={{ width: 500, height: 350, overflowY: 'scroll' }}>{list}</div>
+          </div>
+        )}
+      </EuiSelectable>
     </EuiPopover>
   );
 }
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 d4c8e94193f40..ded84630d12cc 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
@@ -191,7 +191,6 @@ export function ConversationView() {
                 });
             }}
           />
-          <EuiSpacer size="m" />
         </EuiFlexItem>
         <EuiFlexItem
           grow
@@ -248,7 +247,6 @@ export function ConversationView() {
               />
             </ObservabilityAIAssistantChatServiceProvider>
           ) : null}
-          <EuiSpacer size="m" />
         </EuiFlexItem>
       </EuiFlexGroup>
     </>

From 3920c0fe2aa5dd3e4ecca74b587124633b2b81c8 Mon Sep 17 00:00:00 2001
From: Coen Warmer <coen.warmer@gmail.com>
Date: Mon, 14 Aug 2023 10:54:19 +0200
Subject: [PATCH 3/5] Add focus to functionlist on opening

---
 .../public/components/chat/function_list_popover.tsx | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

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
index f77c715a3c7e0..ca81564042b1e 100644
--- 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
@@ -5,7 +5,7 @@
  * 2.0.
  */
 
-import React, { useEffect, useState } from 'react';
+import React, { useEffect, useRef, useState } from 'react';
 import {
   EuiButtonEmpty,
   EuiHighlight,
@@ -32,6 +32,7 @@ export function FunctionListPopover({
   onSelectFunction: (func: string) => void;
 }) {
   const { getFunctions } = useObservabilityAIAssistantChatService();
+  const filterRef = useRef<HTMLInputElement | null>(null);
 
   const [functionOptions, setFunctionOptions] = useState<
     Array<EuiSelectableOption<FunctionListOption>>
@@ -62,6 +63,12 @@ export function FunctionListPopover({
     };
   }, []);
 
+  useEffect(() => {
+    if (isFunctionListOpen && filterRef.current) {
+      filterRef.current.focus();
+    }
+  }, [isFunctionListOpen]);
+
   useEffect(() => {
     const options = getFunctions().map((func) => ({
       label: func.options.name,
@@ -134,7 +141,8 @@ export function FunctionListPopover({
         renderOption={renderCountryOption}
         searchable
         searchProps={{
-          'data-test-subj': 'selectableSearchHere',
+          'data-test-subj': 'searchFiltersList',
+          inputRef: (node) => (filterRef.current = node),
           placeholder: i18n.translate('xpack.observabilityAiAssistant.prompt.functionList.filter', {
             defaultMessage: 'Filter',
           }),

From b41d0dc5461d0287aa303b43b0dad2ee6bab27a1 Mon Sep 17 00:00:00 2001
From: Coen Warmer <coen.warmer@gmail.com>
Date: Mon, 14 Aug 2023 11:02:07 +0200
Subject: [PATCH 4/5] Set Conversationlist width to 250

---
 .../public/routes/conversations/conversation_view.tsx           | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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 ded84630d12cc..298b389aed598 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
@@ -120,7 +120,7 @@ export function ConversationView() {
     <>
       {confirmDeleteElement}
       <EuiFlexGroup direction="row" className={containerClassName}>
-        <EuiFlexItem grow={false} style={{ minWidth: 300, maxWidth: 300 }}>
+        <EuiFlexItem grow={false} style={{ width: 250 }}>
           <ConversationList
             selected={conversationId ?? ''}
             loading={conversations.loading || isUpdatingList}

From ceba8dad888609563d5b7659bbf7ec74d1aab409 Mon Sep 17 00:00:00 2001
From: Coen Warmer <coen.warmer@gmail.com>
Date: Mon, 14 Aug 2023 11:02:24 +0200
Subject: [PATCH 5/5] Remove ObsAiAssistant interface

---
 .../server/service/client/index.ts            |  7 ++--
 .../server/service/index.ts                   |  3 +-
 .../server/service/types.ts                   | 33 +------------------
 3 files changed, 4 insertions(+), 39 deletions(-)

diff --git a/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts b/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts
index 29cd817965e90..42beb515edd82 100644
--- a/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts
+++ b/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts
@@ -30,10 +30,7 @@ import {
   type FunctionDefinition,
   type Message,
 } from '../../../common/types';
-import type {
-  IObservabilityAIAssistantClient,
-  ObservabilityAIAssistantResourceNames,
-} from '../types';
+import type { ObservabilityAIAssistantResourceNames } from '../types';
 
 const ELSER_MODEL_ID = '.elser_model_1';
 
@@ -41,7 +38,7 @@ function throwKnowledgeBaseNotReady(body: any) {
   throw serverUnavailable(`Knowledge base is not ready yet`, body);
 }
 
-export class ObservabilityAIAssistantClient implements IObservabilityAIAssistantClient {
+export class ObservabilityAIAssistantClient {
   constructor(
     private readonly dependencies: {
       actionsClient: PublicMethodsOf<ActionsClient>;
diff --git a/x-pack/plugins/observability_ai_assistant/server/service/index.ts b/x-pack/plugins/observability_ai_assistant/server/service/index.ts
index 5a76bd9125797..643294b2051f5 100644
--- a/x-pack/plugins/observability_ai_assistant/server/service/index.ts
+++ b/x-pack/plugins/observability_ai_assistant/server/service/index.ts
@@ -16,7 +16,6 @@ import { ObservabilityAIAssistantClient } from './client';
 import { conversationComponentTemplate } from './conversation_component_template';
 import { kbComponentTemplate } from './kb_component_template';
 import type {
-  IObservabilityAIAssistantClient,
   IObservabilityAIAssistantService,
   ObservabilityAIAssistantResourceNames,
 } from './types';
@@ -185,7 +184,7 @@ export class ObservabilityAIAssistantService implements IObservabilityAIAssistan
     request,
   }: {
     request: KibanaRequest;
-  }): Promise<IObservabilityAIAssistantClient> {
+  }): Promise<ObservabilityAIAssistantClient> {
     const [_, [coreStart, plugins]] = await Promise.all([
       this.init(),
       this.core.getStartServices() as Promise<
diff --git a/x-pack/plugins/observability_ai_assistant/server/service/types.ts b/x-pack/plugins/observability_ai_assistant/server/service/types.ts
index 9d2326df00ebc..8e74bb5c7abff 100644
--- a/x-pack/plugins/observability_ai_assistant/server/service/types.ts
+++ b/x-pack/plugins/observability_ai_assistant/server/service/types.ts
@@ -6,43 +6,12 @@
  */
 
 import { KibanaRequest } from '@kbn/core/server';
-import type {
-  Conversation,
-  ConversationCreateRequest,
-  ConversationUpdateRequest,
-  FunctionDefinition,
-  KnowledgeBaseEntry,
-  Message,
-} from '../../common/types';
 import { ObservabilityAIAssistantClient } from './client';
 
-export interface IObservabilityAIAssistantClient {
-  chat: (options: {
-    messages: Message[];
-    connectorId: string;
-    functions?: Array<Pick<FunctionDefinition['options'], 'name' | 'description' | 'parameters'>>;
-    stream?: boolean | undefined;
-  }) => ReturnType<ObservabilityAIAssistantClient['chat']>;
-  get: (conversationId: string) => Promise<Conversation>;
-  find: (options?: { query?: string }) => Promise<{ conversations: Conversation[] }>;
-  create: (conversation: ConversationCreateRequest) => Promise<Conversation>;
-  update: (conversation: ConversationUpdateRequest) => Promise<Conversation>;
-  delete: (conversationId: string) => Promise<void>;
-  recall: (query: string) => Promise<{ entries: KnowledgeBaseEntry[] }>;
-  summarise: (options: { entry: Omit<KnowledgeBaseEntry, '@timestamp'> }) => Promise<void>;
-  getKnowledgeBaseStatus: () => Promise<{
-    ready: boolean;
-    error?: any;
-    deployment_state?: string;
-    allocation_state?: string;
-  }>;
-  setupKnowledgeBase: () => Promise<void>;
-}
-
 export interface IObservabilityAIAssistantService {
   getClient: (options: {
     request: KibanaRequest;
-  }) => Promise<IObservabilityAIAssistantClient | undefined>;
+  }) => Promise<ObservabilityAIAssistantClient | undefined>;
 }
 
 export interface ObservabilityAIAssistantResourceNames {