check all required agents for t2viz

Signed-off-by: Yulong Ruan <[email protected]>
opensearch-project · Sep 20, 2024 · 2ac8500 · 2ac8500
1 parent fcc2ca1
commit 2ac8500
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 8 deletions.
diff --git a/public/components/visualization/source_selector.tsx b/public/components/visualization/source_selector.tsx
@@ -14,7 +14,11 @@ import {
   DataSourceOption,
 } from '../../../../../src/plugins/data/public';
 import { StartServices } from '../../types';
-import { TEXT2VEGA_RULE_BASED_AGENT_CONFIG_ID } from '../../../common/constants/llm';
+import {
+  TEXT2PPL_AGENT_CONFIG_ID,
+  TEXT2VEGA_RULE_BASED_AGENT_CONFIG_ID,
+  TEXT2VEGA_WITH_INSTRUCTIONS_AGENT_CONFIG_ID,
+} from '../../../common/constants/llm';
 import { getAssistantService } from '../../services';
 
 const DEFAULT_DATA_SOURCE_TYPE = 'DEFAULT_INDEX_PATTERNS';
@@ -110,12 +114,16 @@ export const SourceSelector = ({
 
       const assistantService = getAssistantService();
       /**
-       * Check each data source to see if text to vega agent is configured or not
+       * Check each data source to see if text to vega agents are configured or not
        * If not configured, disable the corresponding index pattern from the selection list
        */
       Object.keys(dataSourceIdToIndexPatternIds).forEach(async (key) => {
         const res = await assistantService.client.agentConfigExists(
-          TEXT2VEGA_RULE_BASED_AGENT_CONFIG_ID,
+          [
+            TEXT2VEGA_RULE_BASED_AGENT_CONFIG_ID,
+            TEXT2VEGA_WITH_INSTRUCTIONS_AGENT_CONFIG_ID,
+            TEXT2PPL_AGENT_CONFIG_ID,
+          ],
           {
             dataSourceId: key !== 'DEFAULT' ? key : undefined,
           }

diff --git a/public/services/assistant_client.ts b/public/services/assistant_client.ts
@@ -38,9 +38,10 @@ export class AssistantClient {
   };
 
   /**
-   * Return if the given agent config name has agent id configured
+   * Check if the given agent config name has agent id configured
+   * Return false if any of the given config name has no agent id configured
    */
-  agentConfigExists = (agentConfigName: string, options?: Options) => {
+  agentConfigExists = (agentConfigName: string | string[], options?: Options) => {
     return this.http.fetch<{ exists: boolean }>({
       method: 'GET',
       path: AGENT_API.CONFIG_EXISTS,

diff --git a/server/routes/agent_routes.ts b/server/routes/agent_routes.ts
@@ -51,16 +51,20 @@ export function registerAgentRoutes(router: IRouter, assistantService: Assistant
         query: schema.oneOf([
           schema.object({
             dataSourceId: schema.maybe(schema.string()),
-            agentConfigName: schema.string(),
+            agentConfigName: schema.oneOf([schema.string(), schema.arrayOf(schema.string())]),
           }),
         ]),
       },
     },
     router.handleLegacyErrors(async (context, req, res) => {
       try {
         const assistantClient = assistantService.getScopedClient(req, context);
-        const agentId = await assistantClient.getAgentIdByConfigName(req.query.agentConfigName);
-        return res.ok({ body: { exists: Boolean(agentId) } });
+        const promises = Array<string>()
+          .concat(req.query.agentConfigName)
+          .map((configName) => assistantClient.getAgentIdByConfigName(configName));
+        const results = await Promise.all(promises);
+        const exists = results.every((r) => Boolean(r));
+        return res.ok({ body: { exists } });
       } catch (e) {
         return res.ok({ body: { exists: false } });
       }