forked from opensearch-project/dashboards-observability
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
support searching agent by name (opensearch-project#1359)
Signed-off-by: Joshua Li <[email protected]>
- Loading branch information
1 parent
38957cd
commit 7b4d148
Showing
4 changed files
with
255 additions
and
60 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
125 changes: 125 additions & 0 deletions
125
server/routes/query_assist/utils/__tests__/agents.test.ts
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,125 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| import { CoreRouteHandlerContext } from '../../../../../../../src/core/server/core_route_handler_context'; | ||
| import { coreMock, httpServerMock } from '../../../../../../../src/core/server/mocks'; | ||
| import { agentIdMap, requestWithRetryAgentSearch, searchAgentIdByName } from '../agents'; | ||
|
|
||
| describe('Agents helper functions', () => { | ||
| const coreContext = new CoreRouteHandlerContext( | ||
| coreMock.createInternalStart(), | ||
| httpServerMock.createOpenSearchDashboardsRequest() | ||
| ); | ||
| const client = coreContext.opensearch.client.asCurrentUser; | ||
| const mockedTransport = client.transport.request as jest.Mock; | ||
|
|
||
| afterEach(() => { | ||
| jest.clearAllMocks(); | ||
| }); | ||
|
|
||
| it('searches agent id by name', async () => { | ||
| mockedTransport.mockResolvedValueOnce({ | ||
| body: { hits: { total: { value: 1 }, hits: [{ _id: 'agentId' }] } }, | ||
| }); | ||
| const id = await searchAgentIdByName(client, 'test agent'); | ||
| expect(id).toEqual('agentId'); | ||
| expect(mockedTransport.mock.calls[0]).toMatchInlineSnapshot(` | ||
| Array [ | ||
| Object { | ||
| "body": Object { | ||
| "query": Object { | ||
| "term": Object { | ||
| "name.keyword": "test agent", | ||
| }, | ||
| }, | ||
| "sort": Object { | ||
| "created_time": "desc", | ||
| }, | ||
| }, | ||
| "method": "GET", | ||
| "path": "/_plugins/_ml/agents/_search", | ||
| }, | ||
| ] | ||
| `); | ||
| }); | ||
|
|
||
| it('handles not found errors', async () => { | ||
| mockedTransport.mockResolvedValueOnce({ body: { hits: { total: 0 } } }); | ||
| await expect( | ||
| searchAgentIdByName(client, 'test agent') | ||
| ).rejects.toThrowErrorMatchingInlineSnapshot( | ||
| `"search agent 'test agent' failed, reason: Error: cannot find any agent by name: test agent"` | ||
| ); | ||
| }); | ||
|
|
||
| it('handles search errors', async () => { | ||
| mockedTransport.mockRejectedValueOnce('request failed'); | ||
| await expect( | ||
| searchAgentIdByName(client, 'test agent') | ||
| ).rejects.toThrowErrorMatchingInlineSnapshot( | ||
| `"search agent 'test agent' failed, reason: request failed"` | ||
| ); | ||
| }); | ||
|
|
||
| it('requests with valid agent id', async () => { | ||
| agentIdMap['test agent'] = 'test-id'; | ||
| mockedTransport.mockResolvedValueOnce({ | ||
| body: { inference_results: [{ output: [{ result: 'test response' }] }] }, | ||
| }); | ||
| const response = await requestWithRetryAgentSearch({ | ||
| client, | ||
| agentName: 'test agent', | ||
| shouldRetryAgentSearch: true, | ||
| body: { parameters: { param1: 'value1' } }, | ||
| }); | ||
| expect(mockedTransport).toBeCalledWith( | ||
| expect.objectContaining({ | ||
| path: '/_plugins/_ml/agents/test-id/_execute', | ||
| }), | ||
| expect.anything() | ||
| ); | ||
| expect(response.body.inference_results[0].output[0].result).toEqual('test response'); | ||
| }); | ||
|
|
||
| it('searches for agent id if id is undefined', async () => { | ||
| mockedTransport | ||
| .mockResolvedValueOnce({ body: { hits: { total: { value: 1 }, hits: [{ _id: 'new-id' }] } } }) | ||
| .mockResolvedValueOnce({ | ||
| body: { inference_results: [{ output: [{ result: 'test response' }] }] }, | ||
| }); | ||
| const response = await requestWithRetryAgentSearch({ | ||
| client, | ||
| agentName: 'new agent', | ||
| shouldRetryAgentSearch: true, | ||
| body: { parameters: { param1: 'value1' } }, | ||
| }); | ||
| expect(mockedTransport).toBeCalledWith( | ||
| expect.objectContaining({ path: '/_plugins/_ml/agents/new-id/_execute' }), | ||
| expect.anything() | ||
| ); | ||
| expect(response.body.inference_results[0].output[0].result).toEqual('test response'); | ||
| }); | ||
|
|
||
| it('searches for agent id if id is not found', async () => { | ||
| agentIdMap['test agent'] = 'non-exist-agent'; | ||
| mockedTransport | ||
| .mockRejectedValueOnce({ statusCode: 404, body: {}, headers: {} }) | ||
| .mockResolvedValueOnce({ body: { hits: { total: { value: 1 }, hits: [{ _id: 'new-id' }] } } }) | ||
| .mockResolvedValueOnce({ | ||
| body: { inference_results: [{ output: [{ result: 'test response' }] }] }, | ||
| }); | ||
| const response = await requestWithRetryAgentSearch({ | ||
| client, | ||
| agentName: 'test agent', | ||
| shouldRetryAgentSearch: true, | ||
| body: { parameters: { param1: 'value1' } }, | ||
| }); | ||
| expect(mockedTransport).toBeCalledWith( | ||
| expect.objectContaining({ path: '/_plugins/_ml/agents/new-id/_execute' }), | ||
| expect.anything() | ||
| ); | ||
| expect(response.body.inference_results[0].output[0].result).toEqual('test response'); | ||
| }); | ||
| }); |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,96 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| import { ApiResponse } from '@opensearch-project/opensearch/.'; | ||
| import { SearchResponse, SearchTotalHits } from '@opensearch-project/opensearch/api/types'; | ||
| import { RequestBody } from '@opensearch-project/opensearch/lib/Transport'; | ||
| import { OpenSearchClient } from '../../../../../../src/core/server'; | ||
| import { isResponseError } from '../../../../../../src/core/server/opensearch/client/errors'; | ||
| import { ML_COMMONS_API_PREFIX } from '../../../../common/constants/query_assist'; | ||
|
|
||
| const AGENT_REQUEST_OPTIONS = { | ||
| /** | ||
| * It is time-consuming for LLM to generate final answer | ||
| * Give it a large timeout window | ||
| */ | ||
| requestTimeout: 5 * 60 * 1000, | ||
| /** | ||
| * Do not retry | ||
| */ | ||
| maxRetries: 0, | ||
| }; | ||
|
|
||
| type AgentResponse = ApiResponse<{ | ||
| inference_results: Array<{ | ||
| output: Array<{ name: string; result?: string }>; | ||
| }>; | ||
| }>; | ||
|
|
||
| export const agentIdMap: Record<string, string> = {}; | ||
|
|
||
| export const searchAgentIdByName = async ( | ||
| opensearchClient: OpenSearchClient, | ||
| name: string | ||
| ): Promise<string> => { | ||
| try { | ||
| const response = (await opensearchClient.transport.request({ | ||
| method: 'GET', | ||
| path: `${ML_COMMONS_API_PREFIX}/agents/_search`, | ||
| body: { | ||
| query: { | ||
| term: { | ||
| 'name.keyword': name, | ||
| }, | ||
| }, | ||
| sort: { | ||
| created_time: 'desc', | ||
| }, | ||
| }, | ||
| })) as ApiResponse<SearchResponse>; | ||
|
|
||
| if ( | ||
| !response || | ||
| (typeof response.body.hits.total === 'number' && response.body.hits.total === 0) || | ||
| (response.body.hits.total as SearchTotalHits).value === 0 | ||
| ) { | ||
| throw new Error('cannot find any agent by name: ' + name); | ||
| } | ||
| const id = response.body.hits.hits[0]._id; | ||
| return id; | ||
| } catch (error) { | ||
| const errorMessage = JSON.stringify(error.meta?.body) || error; | ||
| throw new Error(`search agent '${name}' failed, reason: ` + errorMessage); | ||
| } | ||
| }; | ||
|
|
||
| export const requestWithRetryAgentSearch = async (options: { | ||
| client: OpenSearchClient; | ||
| agentName: string; | ||
| shouldRetryAgentSearch?: boolean; | ||
| body: RequestBody; | ||
| }): Promise<AgentResponse> => { | ||
| const { client, agentName, shouldRetryAgentSearch = true, body } = options; | ||
| let retry = shouldRetryAgentSearch; | ||
| if (!agentIdMap[agentName]) { | ||
| agentIdMap[agentName] = await searchAgentIdByName(client, agentName); | ||
| retry = false; | ||
| } | ||
| return client.transport | ||
| .request( | ||
| { | ||
| method: 'POST', | ||
| path: `${ML_COMMONS_API_PREFIX}/agents/${agentIdMap[agentName]}/_execute`, | ||
| body, | ||
| }, | ||
| AGENT_REQUEST_OPTIONS | ||
| ) | ||
| .catch(async (error) => { | ||
| if (retry && isResponseError(error) && error.statusCode === 404) { | ||
| agentIdMap[agentName] = await searchAgentIdByName(client, agentName); | ||
| return requestWithRetryAgentSearch({ ...options, shouldRetryAgentSearch: false }); | ||
| } | ||
| return Promise.reject(error); | ||
| }) as Promise<AgentResponse>; | ||
| }; |