Expose a general function for agent execution

CHANGELOG.md

@@ -4,6 +4,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 
 ### Unreleased
 - fix: make sure $schema always added to LLM generated vega json object([252](https://github.com/opensearch-project/dashboards-assistant/pull/252))
+- feat: expose a general function for agent execution([268](https://github.com/opensearch-project/dashboards-assistant/pull/268))
 
 ### 📈 Features/Enhancements
 

common/constants/llm.ts

@@ -24,6 +24,10 @@ export const TEXT2VIZ_API = {
   TEXT2VEGA: `${API_BASE}/text2vega`,
 };
 
+export const AGENT_API = {
+  EXECUTE: `${API_BASE}/agents/{agentId}/_execute`,
+};
+
 export const NOTEBOOK_API = {
   CREATE_NOTEBOOK: `${NOTEBOOK_PREFIX}/note`,
   SET_PARAGRAPH: `${NOTEBOOK_PREFIX}/set_paragraphs/`,

public/plugin.tsx

@@ -44,6 +44,7 @@ import { ConfigSchema } from '../common/types/config';
 import { DataSourceService } from './services/data_source_service';
 import { ASSISTANT_API, DEFAULT_USER_NAME } from '../common/constants/llm';
 import { IncontextInsightProps } from './components/incontext_insight';
+import { AssistantService } from './services/assistant_service';
 
 export const [getCoreStart, setCoreStart] = createGetterSetter<CoreStart>('CoreStart');
 
@@ -72,6 +73,7 @@ export class AssistantPlugin
   incontextInsightRegistry: IncontextInsightRegistry | undefined;
   private dataSourceService: DataSourceService;
   private resetChatSubscription: Subscription | undefined;
+  private assistantService = new AssistantService();
 
   constructor(initializerContext: PluginInitializerContext) {
     this.config = initializerContext.config.get<ConfigSchema>();
@@ -82,6 +84,7 @@ export class AssistantPlugin
     core: CoreSetup<AssistantPluginStartDependencies>,
     setupDeps: AssistantPluginSetupDependencies
   ): AssistantSetup {
+    this.assistantService.setup();
     this.incontextInsightRegistry = new IncontextInsightRegistry();
     this.incontextInsightRegistry?.setIsEnabled(this.config.incontextInsight.enabled);
     setIncontextInsightRegistry(this.incontextInsightRegistry);
@@ -213,18 +216,21 @@ export class AssistantPlugin
   }
 
   public start(core: CoreStart): AssistantStart {
+    const assistantServiceStart = this.assistantService.start(core.http);
     setCoreStart(core);
     setChrome(core.chrome);
     setNotifications(core.notifications);
     setConfigSchema(this.config);
 
     return {
       dataSource: this.dataSourceService.start(),
+      assistantClient: assistantServiceStart.client,
     };
   }
 
   public stop() {
     this.dataSourceService.stop();
+    this.assistantService.stop();
     this.resetChatSubscription?.unsubscribe();
   }
 }

public/services/assistant_client.ts

@@ -0,0 +1,25 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { API_BASE } from '../../common/constants/llm';
+import { HttpSetup } from '../../../../src/core/public';
+
+interface Options {
+  dataSourceId?: string;
+}
+
+export class AssistantClient {
+  constructor(private http: HttpSetup) {}
+
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  executeAgent = (agentName: string, parameters: Record<string, any>, options?: Options) => {
+    return this.http.fetch({
+      method: 'POST',
+      path: `${API_BASE}/agents/${agentName}/_execute`,
+      body: JSON.stringify(parameters),
+      query: { dataSourceId: options?.dataSourceId },
+    });
+  };
+}

public/services/assistant_service.ts

@@ -0,0 +1,26 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { HttpSetup } from '../../../../src/core/public';
+import { AssistantClient } from './assistant_client';
+
+export interface AssistantServiceStart {
+  client: AssistantClient;
+}
+
+export class AssistantService {
+  constructor() {}
+
+  setup() {}
+
+  start(http: HttpSetup): AssistantServiceStart {
+    const assistantClient = new AssistantClient(http);
+    return {
+      client: assistantClient,
+    };
+  }
+
+  stop() {}
+}

public/types.ts

@@ -16,6 +16,7 @@ import {
 } from '../../../src/plugins/visualizations/public';
 import { DataPublicPluginSetup, DataPublicPluginStart } from '../../../src/plugins/data/public';
 import { AppMountParameters, CoreStart } from '../../../src/core/public';
+import { AssistantClient } from './services/assistant_client';
 
 export interface RenderProps {
   props: MessageContentProps;
@@ -67,6 +68,7 @@ export interface AssistantSetup {
 
 export interface AssistantStart {
   dataSource: DataSourceServiceContract;
+  assistantClient: AssistantClient;
 }
 
 export type StartServices = CoreStart &

server/plugin.ts

@@ -17,10 +17,13 @@ import { BasicInputOutputParser } from './parsers/basic_input_output_parser';
 import { VisualizationCardParser } from './parsers/visualization_card_parser';
 import { registerChatRoutes } from './routes/chat_routes';
 import { registerText2VizRoutes } from './routes/text2viz_routes';
+import { AssistantService } from './services/assistant_service';
+import { registerAgentRoutes } from './routes/agent_routes';
 
 export class AssistantPlugin implements Plugin<AssistantPluginSetup, AssistantPluginStart> {
   private readonly logger: Logger;
   private messageParsers: MessageParser[] = [];
+  private assistantService = new AssistantService();
 
   constructor(private readonly initializerContext: PluginInitializerContext) {
     this.logger = initializerContext.logger.get();
@@ -33,6 +36,8 @@ export class AssistantPlugin implements Plugin<AssistantPluginSetup, AssistantPl
       .pipe(first())
       .toPromise();
 
+    const assistantServiceSetup = this.assistantService.setup();
+
     const router = core.http.createRouter();
 
     core.http.registerRouteHandlerContext('assistant_plugin', () => {
@@ -42,6 +47,8 @@ export class AssistantPlugin implements Plugin<AssistantPluginSetup, AssistantPl
       };
     });
 
+    registerAgentRoutes(router, assistantServiceSetup);
+
     // Register server side APIs
     registerChatRoutes(router, {
       messageParsers: this.messageParsers,
@@ -50,7 +57,7 @@ export class AssistantPlugin implements Plugin<AssistantPluginSetup, AssistantPl
 
     // Register router for text to visualization
     if (config.next.enabled) {
-      registerText2VizRoutes(router);
+      registerText2VizRoutes(router, assistantServiceSetup);
     }
 
     core.capabilities.registerProvider(() => ({
@@ -72,6 +79,7 @@ export class AssistantPlugin implements Plugin<AssistantPluginSetup, AssistantPl
     registerMessageParser(VisualizationCardParser);
 
     return {
+      assistantService: assistantServiceSetup,
       registerMessageParser,
       removeMessageParser: (parserId: MessageParser['id']) => {
         const findIndex = this.messageParsers.findIndex((item) => item.id === parserId);
@@ -86,8 +94,11 @@ export class AssistantPlugin implements Plugin<AssistantPluginSetup, AssistantPl
 
   public start(core: CoreStart) {
     this.logger.debug('Assistant: Started');
+    this.assistantService.start();
     return {};
   }
 
-  public stop() {}
+  public stop() {
+    this.assistantService.stop();
+  }
 }

server/routes/agent_routes.ts

@@ -0,0 +1,35 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { schema } from '@osd/config-schema';
+import { IRouter } from '../../../../src/core/server';
+import { AGENT_API } from '../../common/constants/llm';
+import { AssistantServiceSetup } from '../services/assistant_service';
+
+export function registerAgentRoutes(router: IRouter, assistantService: AssistantServiceSetup) {
+  router.post(
+    {
+      path: AGENT_API.EXECUTE,
+      validate: {
+        body: schema.any(),
+        query: schema.object({
+          dataSourceId: schema.maybe(schema.string()),
+        }),
+        params: schema.object({
+          agentId: schema.string(),
+        }),
+      },
+    },
+    router.handleLegacyErrors(async (context, req, res) => {
+      try {
+        const assistantClient = assistantService.getScopedClient(req, context);
+        const response = await assistantClient.executeAgent(req.params.agentId, req.body);
+        return res.ok({ body: response });
+      } catch (e) {
+        return res.internalError();
+      }
+    })
+  );
+}

server/routes/text2viz_routes.ts

@@ -6,14 +6,12 @@
 import { schema } from '@osd/config-schema';
 import { IRouter } from '../../../../src/core/server';
 import { TEXT2VIZ_API } from '../../common/constants/llm';
-import { getOpenSearchClientTransport } from '../utils/get_opensearch_client_transport';
-import { ML_COMMONS_BASE_API } from '../utils/constants';
-import { getAgent } from './get_agent';
+import { AssistantServiceSetup } from '../services/assistant_service';
 
 const TEXT2VEGA_AGENT_CONFIG_ID = 'text2vega';
 const TEXT2PPL_AGENT_CONFIG_ID = 'text2ppl';
 
-export function registerText2VizRoutes(router: IRouter) {
+export function registerText2VizRoutes(router: IRouter, assistantService: AssistantServiceSetup) {
   router.post(
     {
       path: TEXT2VIZ_API.TEXT2VEGA,
@@ -30,25 +28,15 @@ export function registerText2VizRoutes(router: IRouter) {
       },
     },
     router.handleLegacyErrors(async (context, req, res) => {
-      const client = await getOpenSearchClientTransport({
-        context,
-        dataSourceId: req.query.dataSourceId,
-      });
-      const agentId = await getAgent(TEXT2VEGA_AGENT_CONFIG_ID, client);
-      const response = await client.request({
-        method: 'POST',
-        path: `${ML_COMMONS_BASE_API}/agents/${agentId}/_execute`,
-        body: {
-          parameters: {
-            input: req.body.input,
-            ppl: req.body.ppl,
-            dataSchema: req.body.dataSchema,
-            sampleData: req.body.sampleData,
-          },
-        },
-      });
-
+      const assistantClient = assistantService.getScopedClient(req, context);
       try {
+        const response = await assistantClient.executeAgent(TEXT2VEGA_AGENT_CONFIG_ID, {
+          input: req.body.input,
+          ppl: req.body.ppl,
+          dataSchema: req.body.dataSchema,
+          sampleData: req.body.sampleData,
+        });
+
         // let result = response.body.inference_results[0].output[0].dataAsMap;
         let result = JSON.parse(response.body.inference_results[0].output[0].result);
         // sometimes llm returns {response: <schema>} instead of <schema>
@@ -83,22 +71,13 @@ export function registerText2VizRoutes(router: IRouter) {
       },
     },
     router.handleLegacyErrors(async (context, req, res) => {
-      const client = await getOpenSearchClientTransport({
-        context,
-        dataSourceId: req.query.dataSourceId,
-      });
-      const agentId = await getAgent(TEXT2PPL_AGENT_CONFIG_ID, client);
-      const response = await client.request({
-        method: 'POST',
-        path: `${ML_COMMONS_BASE_API}/agents/${agentId}/_execute`,
-        body: {
-          parameters: {
-            question: req.body.question,
-            index: req.body.index,
-          },
-        },
-      });
+      const assistantClient = assistantService.getScopedClient(req, context);
       try {
+        const response = await assistantClient.executeAgent(TEXT2PPL_AGENT_CONFIG_ID, {
+          question: req.body.question,
+          index: req.body.index,
+        });
+
         const result = JSON.parse(response.body.inference_results[0].output[0].result);
         return res.ok({ body: result });
       } catch (e) {

server/services/assistant_client.ts

@@ -0,0 +1,63 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { ApiResponse } from '@opensearch-project/opensearch';
+
+import {
+  OpenSearchClient,
+  OpenSearchDashboardsRequest,
+  RequestHandlerContext,
+} from '../../../../src/core/server';
+import { ML_COMMONS_BASE_API } from '../utils/constants';
+import { getAgent } from '../routes/get_agent';
+
+interface AgentExecuteResponse {
+  inference_results: Array<{
+    output: Array<{ result: string }>;
+  }>;
+}
+
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+const isDataSourceQuery = (query: any): query is { dataSourceId: string } => {
+  if ('dataSourceId' in query && query.dataSourceId) {
+    return true;
+  }
+  return false;
+};
+
+export class AssistantClient {
+  constructor(
+    private request: OpenSearchDashboardsRequest,
+    private context: RequestHandlerContext & {
+      dataSource?: {
+        opensearch: {
+          getClient: (dataSourceId: string) => Promise<OpenSearchClient>;
+        };
+      };
+    }
+  ) {}
+
+  executeAgent = async (
+    agentName: string,
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    parameters: Record<string, any>
+  ): Promise<ApiResponse<AgentExecuteResponse>> => {
+    let client = this.context.core.opensearch.client.asCurrentUser;
+    if (isDataSourceQuery(this.request.query) && this.context.dataSource) {
+      client = await this.context.dataSource.opensearch.getClient(this.request.query.dataSourceId);
+    }
+
+    const agentId = await getAgent(agentName, client.transport);
+    const response = await client.transport.request({
+      method: 'POST',
+      path: `${ML_COMMONS_BASE_API}/agents/${agentId}/_execute`,
+      body: {
+        parameters,
+      },
+    });
+
+    return response as ApiResponse<AgentExecuteResponse>;
+  };
+}