From cf65633c54e5e2b66eb1c791a384ec297b8d91d4 Mon Sep 17 00:00:00 2001
From: "opensearch-trigger-bot[bot]"
 <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com>
Date: Thu, 12 Sep 2024 22:17:04 +0800
Subject: [PATCH] Add discovery summary API (#295) (#306)

* Add discovery summary API

Signed-off-by: Liyun Xiu <xiliyun@amazon.com>

* Remove local change

Signed-off-by: Liyun Xiu <xiliyun@amazon.com>

* Update changelog

Signed-off-by: Liyun Xiu <xiliyun@amazon.com>

---------

Signed-off-by: Liyun Xiu <xiliyun@amazon.com>
(cherry picked from commit b72ba97d6c8f35dfde719625982ee005c60f2ee8)
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

# Conflicts:
#	CHANGELOG.md

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
---
 common/constants/llm.ts         |  1 +
 server/plugin.ts                |  5 ++++-
 server/routes/summary_routes.ts | 37 +++++++++++++++++++++++++++++++++
 3 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/common/constants/llm.ts b/common/constants/llm.ts
index 21bb0ab4..aef7e830 100644
--- a/common/constants/llm.ts
+++ b/common/constants/llm.ts
@@ -31,6 +31,7 @@ export const AGENT_API = {
 export const SUMMARY_ASSISTANT_API = {
   SUMMARIZE: `${API_BASE}/summary`,
   INSIGHT: `${API_BASE}/insight`,
+  DATA2SUMMARY: `${API_BASE}/data2summary`,
 };
 
 export const NOTEBOOK_API = {
diff --git a/server/plugin.ts b/server/plugin.ts
index a0636ad2..f6440bb3 100644
--- a/server/plugin.ts
+++ b/server/plugin.ts
@@ -19,7 +19,7 @@ import { registerChatRoutes } from './routes/chat_routes';
 import { registerText2VizRoutes } from './routes/text2viz_routes';
 import { AssistantService } from './services/assistant_service';
 import { registerAgentRoutes } from './routes/agent_routes';
-import { registerSummaryAssistantRoutes } from './routes/summary_routes';
+import { registerSummaryAssistantRoutes, registerData2SummaryRoutes } from './routes/summary_routes';
 import { capabilitiesProvider as visNLQCapabilitiesProvider } from './vis_type_nlq/capabilities_provider';
 import { visNLQSavedObjectType } from './vis_type_nlq/saved_object_type';
 import { capabilitiesProvider } from './capabilities';
@@ -72,6 +72,9 @@ export class AssistantPlugin implements Plugin<AssistantPluginSetup, AssistantPl
     }
 
     core.capabilities.registerProvider(capabilitiesProvider);
+    
+    // Register router for discovery summary
+    registerData2SummaryRoutes(router, assistantServiceSetup);
 
     const registerMessageParser = (messageParser: MessageParser) => {
       const findItem = this.messageParsers.find((item) => item.id === messageParser.id);
diff --git a/server/routes/summary_routes.ts b/server/routes/summary_routes.ts
index dcdc6879..6a099d5f 100644
--- a/server/routes/summary_routes.ts
+++ b/server/routes/summary_routes.ts
@@ -12,6 +12,7 @@ import { AssistantServiceSetup } from '../services/assistant_service';
 
 const SUMMARY_AGENT_CONFIG_ID = 'os_summary';
 const OS_INSIGHT_AGENT_CONFIG_ID = 'os_insight';
+const DATA2SUMMARY_AGENT_CONFIG_ID = 'os_data2summary';
 let osInsightAgentId: string | undefined;
 let userInsightAgentId: string | undefined;
 
@@ -125,3 +126,39 @@ export function registerSummaryAssistantRoutes(
     })
   );
 }
+
+export function registerData2SummaryRoutes(router: IRouter, assistantService: AssistantServiceSetup) {
+  router.post(
+    {
+      path: SUMMARY_ASSISTANT_API.DATA2SUMMARY,
+      validate: {
+        body: schema.object({
+          sample_data: schema.string(),
+          sample_count: schema.maybe(schema.number()),
+          total_count: schema.maybe(schema.number()),
+          question: schema.maybe(schema.string()),
+          ppl: schema.maybe(schema.string()),
+        }),
+        query: schema.object({
+          dataSourceId: schema.maybe(schema.string()),
+        }),
+      },
+    },
+    router.handleLegacyErrors(async (context, req, res) => {
+      const assistantClient = assistantService.getScopedClient(req, context);
+      try {
+        const response = await assistantClient.executeAgentByName(DATA2SUMMARY_AGENT_CONFIG_ID, {
+          sample_data: req.body.sample_data,
+          total_count: req.body.total_count,
+          sample_count: req.body.sample_count,
+          ppl: req.body.ppl,
+          question: req.body.question,
+        });
+        const result = response.body.inference_results[0].output[0].result;
+        return res.ok({ body: result });
+      } catch (e) {
+        return res.internalError();
+      }
+    })
+  );
+}