[Observability AI Assistant]: Function registry (#19)

package.json

@@ -879,6 +879,7 @@
     "js-search": "^1.4.3",
     "js-sha256": "^0.9.0",
     "js-yaml": "^3.14.1",
+    "json-schema-to-ts": "^2.9.1",
     "json-stable-stringify": "^1.0.1",
     "json-stringify-pretty-compact": "1.2.0",
     "json-stringify-safe": "5.0.1",

x-pack/plugins/actions/server/sub_action_framework/sub_action_connector.ts

@@ -11,6 +11,8 @@ import { Logger } from '@kbn/logging';
 import axios, { AxiosInstance, AxiosResponse, AxiosError, AxiosRequestHeaders } from 'axios';
 import { SavedObjectsClientContract } from '@kbn/core-saved-objects-api-server';
 import { ElasticsearchClient } from '@kbn/core-elasticsearch-server';
+import { finished } from 'stream/promises';
+import { IncomingMessage } from 'http';
 import { assertURL } from './helpers/validators';
 import { ActionsConfigurationUtilities } from '../actions_config';
 import { SubAction, SubActionRequestParams } from './types';
@@ -140,6 +142,19 @@ export abstract class SubActionConnector<Config, Secrets> {
           `Request to external service failed. Connector Id: ${this.connector.id}. Connector type: ${this.connector.type}. Method: ${error.config.method}. URL: ${error.config.url}`
         );
 
+        let responseBody = '';
+
+        // The error response body may also be a stream, e.g. for the GenAI connector
+        if (error.response?.config.responseType === 'stream' && error.response?.data) {
+          const incomingMessage = error.response.data as IncomingMessage;
+
+          incomingMessage.on('data', (chunk) => {
+            responseBody += chunk.toString();
+          });
+          await finished(incomingMessage);
+          error.response.data = JSON.parse(responseBody);
+        }
+
         const errorMessage = `Status code: ${
           error.status ?? error.response?.status
         }. Message: ${this.getResponseErrorMessage(error)}`;

x-pack/plugins/observability_ai_assistant/common/types.ts

@@ -5,7 +5,10 @@
  * 2.0.
  */
 
-import { Serializable } from '@kbn/utility-types';
+import type { Serializable } from '@kbn/utility-types';
+import type { FromSchema } from 'json-schema-to-ts';
+import type { JSONSchema } from 'json-schema-to-ts';
+import React from 'react';
 
 export enum MessageRole {
   System = 'system',
@@ -24,10 +27,10 @@ export interface Message {
     role: MessageRole;
     function_call?: {
       name: string;
-      args?: Serializable;
+      arguments?: string;
       trigger: MessageRole.Assistant | MessageRole.User | MessageRole.Elastic;
     };
-    data?: Serializable;
+    data?: string;
   };
 }
 
@@ -54,3 +57,51 @@ export type ConversationRequestBase = Omit<Conversation, 'user' | 'conversation'
 
 export type ConversationCreateRequest = ConversationRequestBase;
 export type ConversationUpdateRequest = ConversationRequestBase & { conversation: { id: string } };
+
+type CompatibleJSONSchema = Exclude<JSONSchema, boolean>;
+
+export interface ContextDefinition {
+  name: string;
+  description: string;
+}
+
+interface FunctionResponse {
+  content?: Serializable;
+  data?: Serializable;
+}
+
+interface FunctionOptions<TParameters extends CompatibleJSONSchema = CompatibleJSONSchema> {
+  name: string;
+  description: string;
+  parameters: TParameters;
+  contexts: string[];
+}
+
+type RespondFunction<
+  TParameters extends CompatibleJSONSchema,
+  TResponse extends FunctionResponse
+> = (options: { arguments: FromSchema<TParameters> }, signal: AbortSignal) => Promise<TResponse>;
+
+type RenderFunction<TResponse extends FunctionResponse> = (options: {
+  response: TResponse;
+}) => React.ReactNode;
+
+export interface FunctionDefinition {
+  options: FunctionOptions;
+  respond: (options: { arguments: any }, signal: AbortSignal) => Promise<FunctionResponse>;
+  render?: RenderFunction<any>;
+}
+
+export type RegisterContextDefinition = (options: ContextDefinition) => void;
+
+export type RegisterFunctionDefinition = <
+  TParameters extends CompatibleJSONSchema,
+  TResponse extends FunctionResponse
+>(
+  options: FunctionOptions<TParameters>,
+  respond: RespondFunction<TParameters, TResponse>,
+  render?: RenderFunction<TResponse>
+) => void;
+
+export type ContextRegistry = Map<string, ContextDefinition>;
+export type FunctionRegistry = Map<string, FunctionDefinition>;

x-pack/plugins/observability_ai_assistant/public/components/chat/chat_body.stories.tsx

@@ -7,6 +7,8 @@
 
 import { ComponentStory } from '@storybook/react';
 import React from 'react';
+import { Observable } from 'rxjs';
+import { ObservabilityAIAssistantService } from '../../types';
 import { ChatBody as Component } from './chat_body';
 
 export default {
@@ -54,13 +56,11 @@ const defaultProps: ChatBodyProps = {
   currentUser: {
     username: 'elastic',
   },
-  chat: {
-    loading: false,
-    abort: () => {},
-    generate: async () => {
-      return {} as any;
+  service: {
+    chat: () => {
+      return new Observable();
     },
-  },
+  } as unknown as ObservabilityAIAssistantService,
 };
 
 export const ChatBody = Template.bind({});

x-pack/plugins/observability_ai_assistant/public/components/chat/chat_body.tsx

@@ -10,9 +10,9 @@ import { css } from '@emotion/css';
 import type { AuthenticatedUser } from '@kbn/security-plugin/common';
 import React from 'react';
 import { type ConversationCreateRequest } from '../../../common/types';
-import type { UseChatResult } from '../../hooks/use_chat';
 import type { UseGenAIConnectorsResult } from '../../hooks/use_genai_connectors';
 import { useTimeline } from '../../hooks/use_timeline';
+import { ObservabilityAIAssistantService } from '../../types';
 import { HideExpandConversationListButton } from '../buttons/hide_expand_conversation_list_button';
 import { ChatHeader } from './chat_header';
 import { ChatPromptEditor } from './chat_prompt_editor';
@@ -30,14 +30,14 @@ export function ChatBody({
   initialConversation,
   connectors,
   currentUser,
-  chat,
+  service,
   isConversationListExpanded,
   onToggleExpandConversationList,
 }: {
   initialConversation?: ConversationCreateRequest;
   connectors: UseGenAIConnectorsResult;
   currentUser?: Pick<AuthenticatedUser, 'full_name' | 'username'>;
-  chat: UseChatResult;
+  service: ObservabilityAIAssistantService;
   isConversationListExpanded?: boolean;
   onToggleExpandConversationList?: () => void;
 }) {
@@ -47,7 +47,7 @@ export function ChatBody({
     initialConversation,
     connectors,
     currentUser,
-    chat,
+    service,
   });
 
   return (
@@ -93,7 +93,7 @@ export function ChatBody({
       <EuiFlexItem grow={false}>
         <EuiPanel hasBorder={false} hasShadow={false} paddingSize="m">
           <ChatPromptEditor
-            loading={chat.loading}
+            loading={false}
             disabled={!connectors.selectedConnector}
             onSubmit={timeline.onSubmit}
           />

x-pack/plugins/observability_ai_assistant/public/components/chat/chat_flyout.tsx

@@ -7,9 +7,9 @@
 import React, { useState } from 'react';
 import { EuiFlexGroup, EuiFlexItem, EuiFlyout, useEuiTheme } from '@elastic/eui';
 import type { ConversationCreateRequest } from '../../../common/types';
-import { useChat } from '../../hooks/use_chat';
 import { useCurrentUser } from '../../hooks/use_current_user';
 import { useGenAIConnectors } from '../../hooks/use_genai_connectors';
+import { useObservabilityAIAssistant } from '../../hooks/use_observability_ai_assistant';
 import { ChatBody } from './chat_body';
 import { ConversationList } from './conversation_list';
 
@@ -22,18 +22,20 @@ export function ChatFlyout({
   isOpen: boolean;
   onClose: () => void;
 }) {
-  const currentUser = useCurrentUser();
   const connectors = useGenAIConnectors();
-  const { euiTheme } = useEuiTheme();
 
-  const chat = useChat();
+  const currentUser = useCurrentUser();
+
+  const { euiTheme } = useEuiTheme();
 
   const [isConversationListExpanded, setIsConversationListExpanded] = useState(false);
 
   const handleClickConversation = (id: string) => {};
   const handleClickNewChat = () => {};
   const handleClickSettings = () => {};
 
+  const service = useObservabilityAIAssistant();
+
   return isOpen ? (
     <EuiFlyout onClose={onClose}>
       <EuiFlexGroup responsive={false} gutterSize="none">
@@ -51,7 +53,7 @@ export function ChatFlyout({
         ) : null}
         <EuiFlexItem>
           <ChatBody
-            chat={chat}
+            service={service}
             connectors={connectors}
             initialConversation={initialConversation}
             currentUser={currentUser}

x-pack/plugins/observability_ai_assistant/public/components/chat/chat_item.tsx

@@ -14,6 +14,7 @@ import {
   EuiFlexItem,
   EuiPopover,
 } from '@elastic/eui';
+import { css } from '@emotion/css';
 import { i18n } from '@kbn/i18n';
 import React, { useState } from 'react';
 import { MessageRole } from '../../../common/types';
@@ -40,6 +41,16 @@ export interface ChatItemProps extends ChatTimelineItem {
   onStopGeneratingClick: () => void;
 }
 
+const euiCommentClassName = css`
+  .euiCommentEvent__headerEvent {
+    flex-grow: 1;
+  }
+
+  > div:last-child {
+    overflow: hidden;
+  }
+`;
+
 export function ChatItem({
   title,
   content,
@@ -143,15 +154,16 @@ export function ChatItem({
           title={title}
         />
       }
-      timelineAvatar={<ChatItemAvatar currentUser={currentUser} role={role} />}
+      className={euiCommentClassName}
+      timelineAvatar={
+        <ChatItemAvatar loading={loading && !content} currentUser={currentUser} role={role} />
+      }
       username={getRoleTranslation(role)}
     >
-      {content !== undefined || error || loading ? (
+      {content || error || loading || controls ? (
         <MessagePanel
           body={
-            content !== undefined || loading ? (
-              <MessageText content={content || ''} loading={loading} />
-            ) : null
+            content || loading ? <MessageText content={content || ''} loading={loading} /> : null
           }
           error={error}
           controls={controls}

x-pack/plugins/observability_ai_assistant/public/components/chat/chat_item_avatar.tsx

@@ -15,16 +15,19 @@ import { MessageRole } from '../../../common/types';
 interface ChatAvatarProps {
   currentUser?: Pick<AuthenticatedUser, 'full_name' | 'username'> | undefined;
   role: MessageRole;
+  loading: boolean;
 }
 
-export function ChatItemAvatar({ currentUser, role }: ChatAvatarProps) {
+export function ChatItemAvatar({ currentUser, role, loading }: ChatAvatarProps) {
+  const isLoading = loading || !currentUser;
+
+  if (isLoading) {
+    return <EuiLoadingSpinner size="xl" />;
+  }
+
   switch (role) {
     case MessageRole.User:
-      return currentUser ? (
-        <UserAvatar user={currentUser} size="m" data-test-subj="userMenuAvatar" />
-      ) : (
-        <EuiLoadingSpinner size="xl" />
-      );
+      return <UserAvatar user={currentUser} size="m" data-test-subj="userMenuAvatar" />;
 
     case MessageRole.Assistant:
     case MessageRole.Elastic:

x-pack/plugins/observability_ai_assistant/public/components/chat/chat_prompt_editor.tsx

@@ -16,15 +16,13 @@ import {
   EuiPopover,
 } from '@elastic/eui';
 import { i18n } from '@kbn/i18n';
-import { useFunctions, Func } from '../../hooks/use_functions';
+import { useFunctions, type Func } from '../../hooks/use_functions';
+import { type Message, MessageRole } from '../../../common';
 
 export interface ChatPromptEditorProps {
   disabled: boolean;
   loading: boolean;
-  onSubmit: (message: {
-    content?: string;
-    function_call?: { name: string; args?: string };
-  }) => Promise<void>;
+  onSubmit: (message: Message) => Promise<void>;
 }
 
 export function ChatPromptEditor({ onSubmit, disabled, loading }: ChatPromptEditorProps) {
@@ -40,7 +38,10 @@ export function ChatPromptEditor({ onSubmit, disabled, loading }: ChatPromptEdit
   const handleSubmit = () => {
     const currentPrompt = prompt;
     setPrompt('');
-    onSubmit({ content: currentPrompt })
+    onSubmit({
+      '@timestamp': new Date().toISOString(),
+      message: { role: MessageRole.User, content: currentPrompt },
+    })
       .then(() => {
         setPrompt('');
       })

x-pack/plugins/observability_ai_assistant/public/components/chat/chat_timeline.tsx

@@ -5,23 +5,17 @@
  * 2.0.
  */
 
-import React from 'react';
 import { EuiCommentList } from '@elastic/eui';
 import type { AuthenticatedUser } from '@kbn/security-plugin/common';
-import { MessageRole } from '../../../common/types';
+import React from 'react';
+import type { Message } from '../../../common';
 import type { Feedback } from '../feedback_buttons';
 import { ChatItem } from './chat_item';
 
-export interface ChatTimelineItem {
+export interface ChatTimelineItem
+  extends Pick<Message['message'], 'role' | 'content' | 'function_call'> {
   id: string;
   title: string;
-  role: MessageRole;
-  content?: string;
-  function_call?: {
-    name: string;
-    args?: string;
-    trigger?: MessageRole;
-  };
   loading: boolean;
   error?: any;
   canEdit: boolean;