Merge pull request elastic#8 from joemcelroy/add-hook

Add API Service

package.json

@@ -99,6 +99,7 @@
     "@dnd-kit/core": "^3.1.1",
     "@dnd-kit/sortable": "^4.0.0",
     "@dnd-kit/utilities": "^2.0.0",
+    "@elastic/ai-assist": "0.0.8",
     "@elastic/apm-rum": "^5.16.0",
     "@elastic/apm-rum-react": "^2.0.2",
     "@elastic/charts": "63.0.0",
@@ -1694,4 +1695,4 @@
     "zod-to-json-schema": "^3.22.3"
   },
   "packageManager": "yarn@1.22.21"
-}
+}

x-pack/plugins/enterprise_search/server/routes/enterprise_search/ai_playground.ts

@@ -0,0 +1,86 @@
+import { RouteDependencies } from '../../plugin';
+import { elasticsearchErrorHandler } from '../../utils/elasticsearch_error_handler';
+import Assist, { ConversationalChain } from '@elastic/ai-assist';
+import { ChatOpenAI } from '@elastic/ai-assist/models';
+import { Prompt } from '@elastic/ai-assist';
+import { schema } from '@kbn/config-schema';
+import Stream from 'stream';
+
+export function registerAIPlaygroundRoutes({ log, router }: RouteDependencies) {
+  router.post(
+    {
+      path: '/internal/enterprise_search/ai_playground/chat',
+      validate: {
+        body: schema.object({
+          data: schema.any(),
+          messages: schema.any(),
+        }),
+      },
+    },
+    elasticsearchErrorHandler(log, async (context, request, response) => {
+      const { client } = (await context.core).elasticsearch;
+
+      const aiClient = Assist({
+        es_client: client.asCurrentUser,
+      });
+
+      const { messages, data } = await request.body;
+
+      const model = new ChatOpenAI({
+        openAIApiKey: data.api_key,
+      });
+
+      const chain = ConversationalChain({
+        model,
+        rag: {
+          index: data.indices,
+          retriever: (question: string) => {
+            return {
+              text_expansion: {
+                'vector.tokens': {
+                  model_id: '.elser_model_2',
+                  model_text: question,
+                },
+              },
+            };
+          },
+        },
+        prompt: Prompt(data.prompt, {
+          citations: true,
+          context: true,
+          type: 'openai',
+        }),
+      });
+
+      const stream = await chain.stream(aiClient, messages);
+
+      const reader = (stream as ReadableStream).getReader();
+
+      class UIStream extends Stream.Readable {
+        _read() {
+          const that = this;
+
+          function read() {
+            reader.read().then(({ done, value }: { done: boolean; value?: string }) => {
+              if (done) {
+                that.push(null);
+                return;
+              }
+              that.push(value);
+              read();
+            });
+          }
+          read();
+        }
+      }
+
+      return response.custom({
+        body: new UIStream(),
+        statusCode: 200,
+        headers: {
+          'Content-Type': 'text/plain; charset=utf-8',
+        },
+      });
+    })
+  );
+}

x-pack/plugins/enterprise_search/server/routes/enterprise_search/index.ts

@@ -7,6 +7,7 @@
 
 import { RouteDependencies } from '../../plugin';
 
+import { registerAIPlaygroundRoutes } from './ai_playground';
 import { registerDocumentRoute } from './documents';
 import { registerIndexRoutes } from './indices';
 import { registerMappingRoute } from './mapping';
@@ -19,4 +20,5 @@ export const registerEnterpriseSearchRoutes = (dependencies: RouteDependencies)
   registerSearchRoute(dependencies);
   registerDocumentRoute(dependencies);
   registerSearchApplicationsRoutes(dependencies);
+  registerAIPlaygroundRoutes(dependencies);
 };