[WIP]detect serverless data source

changelogs/fragments/10840.yml

@@ -0,0 +1,2 @@
+fix:
+- Detect serverless data source ([#10840](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/10840))

src/plugins/query_enhancements/server/routes/query_assist/agents.test.ts

@@ -9,6 +9,7 @@ import { loggerMock } from '@osd/logging/target/mocks';
 import { RequestHandlerContext } from 'src/core/server';
 import { coreMock } from '../../../../../core/server/mocks';
 import { getAgentIdByConfig, requestAgentByConfig } from './agents';
+import { DataSourceEngineType } from '../../../../data_source/common/data_sources';
 
 describe('Agents helper functions', () => {
   const coreContext = coreMock.createRequestHandlerContext();
@@ -133,4 +134,28 @@ describe('Agents helper functions', () => {
     );
     expect(response.body.inference_results[0].output[0].result).toEqual('test response');
   });
+
+  it('should able to detect serverless data source and treat config name as agent id', async () => {
+    context.core.savedObjects.client.get = jest.fn().mockResolvedValue({
+      attributes: {
+        dataSourceEngineType: DataSourceEngineType.OpenSearchServerless,
+      },
+    });
+
+    await requestAgentByConfig({
+      context,
+      configName: 'new_agent',
+      body: { parameters: { param1: 'value1' } },
+      dataSourceId: 'test-datasource-id',
+    });
+
+    mockedTransport.mockResolvedValueOnce({
+      body: { inference_results: [{ output: [{ result: 'test response' }] }] },
+    });
+
+    expect(mockedTransport).toBeCalledWith(
+      expect.objectContaining({ path: '/_plugins/_ml/agents/new_agent/_execute' }),
+      expect.anything()
+    );
+  });
 });

src/plugins/query_enhancements/server/routes/query_assist/agents.ts

@@ -7,6 +7,10 @@ import { ApiResponse } from '@opensearch-project/opensearch';
 import { RequestBody, TransportRequestPromise } from '@opensearch-project/opensearch/lib/Transport';
 import { RequestHandlerContext } from 'src/core/server';
 import { URI } from '../../../common';
+import {
+  DataSourceEngineType,
+  DataSourceAttributes,
+} from '../../../../data_source/common/data_sources';
 
 const AGENT_REQUEST_OPTIONS = {
   /**
@@ -51,6 +55,21 @@ export const getAgentIdByConfig = async (
     throw new Error(`Get agent '${configName}' failed, reason: ` + errorMessage);
   }
 };
+const detectServerlessDatasource = async (
+  context: RequestHandlerContext,
+  dsId: string | undefined
+) => {
+  if (dsId) {
+    const savedObject = await context.core.savedObjects.client.get<DataSourceAttributes>(
+      'data-source',
+      dsId
+    );
+
+    return (
+      savedObject?.attributes?.dataSourceEngineType === DataSourceEngineType.OpenSearchServerless
+    );
+  }
+};
 
 export const requestAgentByConfig = async (options: {
   context: RequestHandlerContext;
@@ -59,12 +78,16 @@ export const requestAgentByConfig = async (options: {
   dataSourceId?: string;
 }): Promise<AgentResponse> => {
   const { context, configName, body, dataSourceId } = options;
+  const isServerlessDatasource = await detectServerlessDatasource(context, dataSourceId);
   const client =
     // @ts-expect-error TS2339 TODO(ts-error): fixme
     context.query_assist.dataSourceEnabled && dataSourceId
       ? await context.dataSource.opensearch.getClient(dataSourceId)
       : context.core.opensearch.client.asCurrentUser;
-  const agentId = await getAgentIdByConfig(client, configName);
+
+  const agentId = isServerlessDatasource
+    ? configName
+    : await getAgentIdByConfig(client, configName);
   return client.transport.request(
     {
       method: 'POST',

src/plugins/query_enhancements/server/routes/query_assist/routes.test.ts

@@ -13,6 +13,7 @@ import { API } from '../../../common';
 import { getAgentIdByConfig, requestAgentByConfig } from './agents';
 import { getUnselectedTimeFields, parseTimeRangeXML } from './ppl/time_parser_utils';
 import { HttpService } from 'opensearch-dashboards/server/http';
+import { DataSourceEngineType } from '../../../../data_source/common/data_sources';
 
 // Mock the dependencies
 jest.mock('./agents');
@@ -31,7 +32,7 @@ describe('Query Assist Routes', () => {
   let testServer: HttpService;
   let opensearchClient: ReturnType<typeof opensearchClientMock.createInternalClient>;
 
-  const testSetup = async () => {
+  const testSetup = async (saveObjectClient: { get: any } = { get: jest.fn() }) => {
     const { server, httpSetup, handlerContext } = await setupServer();
     const router = httpSetup.createRouter('');
 
@@ -46,6 +47,11 @@ describe('Query Assist Routes', () => {
             asCurrentUser: opensearchClient,
           },
         },
+        savedObjects: {
+          client: {
+            saveObjectClient,
+          },
+        },
       },
       query_assist: {
         dataSourceEnabled: true,
@@ -108,7 +114,6 @@ describe('Query Assist Routes', () => {
     testServer = server;
     return httpSetup;
   };
-
   afterEach(async () => {
     if (testServer) {
       await testServer.stop();
@@ -131,6 +136,24 @@ describe('Query Assist Routes', () => {
       });
     });
 
+    it('should able to detect serverless data source gracefully', async () => {
+      const mockSavedObjectsClient = {
+        get: jest.fn().mockResolvedValue({
+          attributes: {
+            dataSourceEngineType: DataSourceEngineType.OpenSearchServerless,
+          },
+        }),
+      };
+
+      const httpSetup = await testSetup(mockSavedObjectsClient);
+      const result = await supertest(httpSetup.server.listener)
+        .get(API.QUERY_ASSIST.LANGUAGES)
+        .expect(200);
+
+      expect(result.body).toEqual({
+        configuredLanguages: ['PPL', 'SQL'],
+      });
+    });
     it('should handle agent configuration errors gracefully', async () => {
       // Mock agent configuration failure
       mockGetAgentIdByConfig.mockRejectedValue(new Error('Agent not found'));

src/plugins/query_enhancements/server/routes/query_assist/routes.ts

@@ -11,6 +11,10 @@ import { API, ERROR_DETAILS } from '../../../common';
 import { getAgentIdByConfig, requestAgentByConfig } from './agents';
 import { createResponseBody } from './createResponse';
 import { parseTimeRangeXML, getUnselectedTimeFields } from './ppl/time_parser_utils';
+import {
+  DataSourceAttributes,
+  DataSourceEngineType,
+} from '../../../../data_source/common/data_sources';
 
 export function registerQueryAssistRoutes(router: IRouter) {
   router.get(
@@ -25,6 +29,18 @@ export function registerQueryAssistRoutes(router: IRouter) {
     async (context, request, response) => {
       // @ts-expect-error TS2339 TODO(ts-error): fixme
       const config = await context.query_assist.configPromise;
+
+      let isServerlessDatasource = false;
+      if (request.query.dataSourceId) {
+        const savedObject = await context.core.savedObjects.client.get<DataSourceAttributes>(
+          'data-source',
+          request.query.dataSourceId
+        );
+        isServerlessDatasource =
+          savedObject?.attributes?.dataSourceEngineType ===
+          DataSourceEngineType.OpenSearchServerless;
+      }
+
       const client =
         // @ts-expect-error TS2339 TODO(ts-error): fixme
         context.query_assist.dataSourceEnabled && request.query.dataSourceId
@@ -36,9 +52,10 @@ export function registerQueryAssistRoutes(router: IRouter) {
           // @ts-expect-error TS7006 TODO(ts-error): fixme
           config.queryAssist.supportedLanguages.map((languageConfig) =>
             // if the call does not throw any error, then the agent is properly configured
-            getAgentIdByConfig(client, languageConfig.agentConfig).then(() =>
-              configuredLanguages.push(languageConfig.language)
-            )
+            (isServerlessDatasource
+              ? Promise.resolve(languageConfig.agentConfig)
+              : getAgentIdByConfig(client, languageConfig.agentConfig)
+            ).then(() => configuredLanguages.push(languageConfig.language))
           )
         );
         return response.ok({ body: { configuredLanguages } });