[Observability AI Assistant] Add isStream param to chat/complete endpoint to support non-streaming responses (#240819)

Closes #239439

This adds a parameter `isStream` to both internal and public
`chat/complete` endpoint.

## Solution

The internal ~and public~ `chat/complete` APIs should return a single
JSON response instead of a stream of events when passing
`isStream:false`

### Internal API
```sh
curl -X POST "http://localhost:5601/internal/observability_ai_assistant/chat/complete" \
  -u elastic:changeme \
  -H "kbn-version: 9.3.0" \
  -H "Content-Type: application/json" \
  -H "x-elastic-internal-origin: kibana" \
  -d '{
    "messages": [
      {
        "@timestamp": "'$(date -u +"%Y-%m-%dT%H:%M:%S.000Z")'",
        "message": {
          "role": "user",
          "content": "This is a test. Respond with Ack."
        }
      }
    ],
    "connectorId": "<your-connector-id>",
    "scopes": ["observability"],
    "screenContexts": [],
    "persist": false,
    "isStream": false
  }' | jq
```

### Public API

<details>

<summary>Disabled for the initial rollout</summary>

```sh
curl -X POST "http://localhost:5601/api/observability_ai_assistant/chat/complete" \
  -u elastic:changeme \
  -H "kbn-xsrf: true" \
  -H "Content-Type: application/json" \
  -d '{
    "messages": [
      {
        "@timestamp": "'$(date -u +"%Y-%m-%dT%H:%M:%S.000Z")'",
        "message": {
          "role": "user",
          "content": "This is a test. Respond with Ack."
        }
      }
    ],
    "connectorId": "<your-connector-id>",
    "persist": false,
    "isStream": false
  }' | jq
```

</details>


### Response

```json
{
  "conversationId": "a408ec61-9401-4c7d-9ae6-349fb4844684",
  "data": "Why do programmers prefer dark mode?\n\nBecause light attracts bugs! 🐛",
  "connectorId": "bedrock-claude-40"
}
```

---------

Co-authored-by: kibanamachine <[email protected]>

Author: sorenlouv

x-pack/platform/plugins/shared/observability_ai_assistant/server/routes/chat/route.ts

@@ -7,10 +7,12 @@
 import { notImplemented } from '@hapi/boom';
 import { toBooleanRt } from '@kbn/io-ts-utils';
 import * as t from 'io-ts';
+import type { Observable } from 'rxjs';
 import { from, map } from 'rxjs';
 import { v4 } from 'uuid';
 import type { Readable } from 'stream';
 import type { AssistantScope } from '@kbn/ai-assistant-common';
+import { last } from 'lodash';
 import { aiAssistantSimulatedFunctionCalling } from '../..';
 import { createFunctionResponseMessage } from '../../../common/utils/create_function_response_message';
 import { flushBuffer } from '../../service/util/flush_buffer';
@@ -28,6 +30,11 @@ import {
   screenContextRt,
 } from '../runtime_types';
 import type { ObservabilityAIAssistantRouteHandlerResources } from '../types';
+import type {
+  BufferFlushEvent,
+  StreamingChatResponseEventWithoutError,
+} from '../../../common/conversation_complete';
+import type { ConversationCreateRequest } from '../../../common/types';
 
 const chatCompleteBaseRt = (apiType: 'public' | 'internal') =>
   t.type({
@@ -38,6 +45,7 @@ const chatCompleteBaseRt = (apiType: 'public' | 'internal') =>
         persist: toBooleanRt,
       }),
       t.partial({
+        isStream: t.union([t.undefined, toBooleanRt]), // accept undefined in order to default to true
         conversationId: t.string,
         title: t.string,
         disableFunctions: toBooleanRt,
@@ -245,7 +253,10 @@ async function chatComplete(
   resources: ObservabilityAIAssistantRouteHandlerResources & {
     params: t.TypeOf<typeof chatCompleteInternalRt>;
   }
-) {
+): Promise<{
+  response$: Observable<StreamingChatResponseEventWithoutError | BufferFlushEvent>;
+  getConversation: () => Promise<ConversationCreateRequest>;
+}> {
   const { params, service } = resources;
 
   const {
@@ -286,7 +297,7 @@ async function chatComplete(
         : userInstructionOrString
   );
 
-  const response$ = client.complete({
+  const { response$, getConversation } = client.complete({
     messages,
     connectorId,
     conversationId,
@@ -299,7 +310,8 @@ async function chatComplete(
     disableFunctions,
   });
 
-  return response$.pipe(flushBuffer(isCloudEnabled));
+  const responseWithFlushBuffer$ = response$.pipe(flushBuffer(isCloudEnabled));
+  return { response$: responseWithFlushBuffer$, getConversation };
 }
 
 const chatCompleteRoute = createObservabilityAIAssistantServerRoute({
@@ -310,8 +322,22 @@ const chatCompleteRoute = createObservabilityAIAssistantServerRoute({
     },
   },
   params: chatCompleteInternalRt,
-  handler: async (resources): Promise<Readable> => {
-    return observableIntoStream(await chatComplete(resources));
+  handler: async (resources) => {
+    const { params } = resources;
+    const { response$, getConversation } = await chatComplete(resources);
+    const { isStream = true, connectorId } = params.body;
+
+    if (isStream === false) {
+      const response = await getConversation();
+
+      return {
+        conversationId: response.conversation.id,
+        data: last(response.messages)?.message.content,
+        connectorId,
+      };
+    }
+
+    return observableIntoStream(response$);
   },
 });
 
@@ -326,18 +352,15 @@ const publicChatCompleteRoute = createObservabilityAIAssistantServerRoute({
   options: {
     tags: ['observability-ai-assistant'],
   },
-  handler: async (resources): Promise<Readable> => {
+  handler: async (resources) => {
     const { params, logger } = resources;
+    const { actions, ...bodyParams } = params.body;
 
-    const {
-      body: { actions, ...restOfBody },
-    } = params;
-
-    const response$ = await chatComplete({
+    const { response$ } = await chatComplete({
       ...resources,
       params: {
         body: {
-          ...restOfBody,
+          ...bodyParams,
           scopes: ['observability'],
           screenContexts: [
             {

x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.test.ts

@@ -273,15 +273,15 @@ describe('Observability AI Assistant client', () => {
           });
         });
 
-      stream = observableIntoStream(
-        client.complete({
-          connectorId: 'foo',
-          messages: [user('How many alerts do I have?')],
-          functionClient: functionClientMock,
-          signal: new AbortController().signal,
-          persist: true,
-        })
-      );
+      const { response$ } = client.complete({
+        connectorId: 'foo',
+        messages: [user('How many alerts do I have?')],
+        functionClient: functionClientMock,
+        signal: new AbortController().signal,
+        persist: true,
+      });
+
+      stream = observableIntoStream(response$);
     });
 
     describe('when streaming the response from the LLM', () => {
@@ -557,16 +557,16 @@ describe('Observability AI Assistant client', () => {
         return {} as any;
       });
 
-      stream = observableIntoStream(
-        await client.complete({
-          connectorId: 'foo',
-          messages: [user('How many alerts do I have?')],
-          functionClient: functionClientMock,
-          signal: new AbortController().signal,
-          conversationId: 'my-conversation-id',
-          persist: true,
-        })
-      );
+      const { response$ } = client.complete({
+        connectorId: 'foo',
+        messages: [user('How many alerts do I have?')],
+        functionClient: functionClientMock,
+        signal: new AbortController().signal,
+        conversationId: 'my-conversation-id',
+        persist: true,
+      });
+
+      stream = observableIntoStream(response$);
 
       dataHandler = jest.fn();
 
@@ -649,16 +649,16 @@ describe('Observability AI Assistant client', () => {
         });
       });
 
-      stream = observableIntoStream(
-        await client.complete({
-          connectorId: 'foo',
-          messages: [user('How many alerts do I have?')],
-          functionClient: functionClientMock,
-          signal: new AbortController().signal,
-          title: 'My predefined title',
-          persist: true,
-        })
-      );
+      const { response$ } = client.complete({
+        connectorId: 'foo',
+        messages: [user('How many alerts do I have?')],
+        functionClient: functionClientMock,
+        signal: new AbortController().signal,
+        title: 'My predefined title',
+        persist: true,
+      });
+
+      stream = observableIntoStream(response$);
 
       dataHandler = jest.fn();
 
@@ -739,16 +739,16 @@ describe('Observability AI Assistant client', () => {
         });
       });
 
-      stream = observableIntoStream(
-        await client.complete({
-          connectorId: 'foo',
-          messages: [user('How many alerts do I have?')],
-          functionClient: functionClientMock,
-          signal: new AbortController().signal,
-          title: 'My predefined title',
-          persist: true,
-        })
-      );
+      const { response$ } = client.complete({
+        connectorId: 'foo',
+        messages: [user('How many alerts do I have?')],
+        functionClient: functionClientMock,
+        signal: new AbortController().signal,
+        title: 'My predefined title',
+        persist: true,
+      });
+
+      stream = observableIntoStream(response$);
 
       dataHandler = jest.fn();
 
@@ -1178,15 +1178,15 @@ describe('Observability AI Assistant client', () => {
         };
       });
 
-      stream = observableIntoStream(
-        await client.complete({
-          connectorId: 'foo',
-          messages: [user('How many alerts do I have?')],
-          functionClient: functionClientMock,
-          signal: new AbortController().signal,
-          persist: false,
-        })
-      );
+      const { response$ } = client.complete({
+        connectorId: 'foo',
+        messages: [user('How many alerts do I have?')],
+        functionClient: functionClientMock,
+        signal: new AbortController().signal,
+        persist: false,
+      });
+
+      stream = observableIntoStream(response$);
 
       dataHandler = jest.fn();
 
@@ -1304,16 +1304,16 @@ describe('Observability AI Assistant client', () => {
         };
       });
 
-      stream = observableIntoStream(
-        client.complete({
-          connectorId: 'foo',
-          messages: [user('How many alerts do I have?')],
-          functionClient: functionClientMock,
-          signal: new AbortController().signal,
-          title: 'My predefined title',
-          persist: true,
-        })
-      );
+      const { response$ } = client.complete({
+        connectorId: 'foo',
+        messages: [user('How many alerts do I have?')],
+        functionClient: functionClientMock,
+        signal: new AbortController().signal,
+        title: 'My predefined title',
+        persist: true,
+      });
+
+      stream = observableIntoStream(response$);
 
       dataHandler = jest.fn();
 
@@ -1392,15 +1392,15 @@ describe('Observability AI Assistant client', () => {
         };
       });
 
-      const stream = observableIntoStream(
-        await client.complete({
-          connectorId: 'foo',
-          messages: [user('How many alerts do I have?')],
-          functionClient: functionClientMock,
-          signal: new AbortController().signal,
-          persist: false,
-        })
-      );
+      const { response$ } = client.complete({
+        connectorId: 'foo',
+        messages: [user('How many alerts do I have?')],
+        functionClient: functionClientMock,
+        signal: new AbortController().signal,
+        persist: false,
+      });
+
+      const stream = observableIntoStream(response$);
 
       dataHandler = jest.fn();
 
@@ -1480,16 +1480,16 @@ describe('Observability AI Assistant client', () => {
         });
       });
 
-      stream = observableIntoStream(
-        await client.complete({
-          connectorId: 'foo',
-          messages: [user('How many alerts do I have?')],
-          functionClient: functionClientMock,
-          signal: new AbortController().signal,
-          title: 'My predefined title',
-          persist: true,
-        })
-      );
+      const { response$ } = client.complete({
+        connectorId: 'foo',
+        messages: [user('How many alerts do I have?')],
+        functionClient: functionClientMock,
+        signal: new AbortController().signal,
+        title: 'My predefined title',
+        persist: true,
+      });
+
+      stream = observableIntoStream(response$);
 
       dataHandler = jest.fn();
 
@@ -1553,7 +1553,7 @@ describe('Observability AI Assistant client', () => {
         title: 'My predefined title',
         persist: false,
       })
-      .subscribe(() => {}); // To trigger call to chat
+      .response$.subscribe(() => {}); // To trigger call to chat
     await nextTick();
 
     expect(chatSpy.mock.calls[0][1].systemMessage).toEqual(EXPECTED_STORED_SYSTEM_MESSAGE);
@@ -1571,7 +1571,7 @@ describe('Observability AI Assistant client', () => {
         });
       });
 
-      const complete$ = await client.complete({
+      const { response$ } = client.complete({
         connectorId: 'foo',
         messages: [user('Can you call the my_action function?')],
         functionClient: new ChatFunctionClient([
@@ -1601,7 +1601,7 @@ describe('Observability AI Assistant client', () => {
       const messages: Message[] = [];
 
       completePromise = new Promise<Message[]>((resolve, reject) => {
-        complete$.subscribe({
+        response$.subscribe({
           next: (event) => {
             if (event.type === StreamingChatResponseEventType.MessageAdd) {
               messages.push(event.message);
@@ -1748,8 +1748,8 @@ describe('Observability AI Assistant client', () => {
       // client = createClient(namespace);
       client = createClient();
       (client as any).dependencies.namespace = namespace;
-      (
-        await client.complete({
+      client
+        .complete({
           functionClient: functionClientMock,
           connectorId: 'foo',
           messages: [],
@@ -1758,7 +1758,7 @@ describe('Observability AI Assistant client', () => {
           kibanaPublicUrl: 'http://localhost:5601',
           title: 'Generated title',
         })
-      ).subscribe({});
+        .response$.subscribe({});
 
       await nextTick();
     };