[Inference Endpoints View] Adds Amazon Bedrock to Deletion, search an…

…d filtering of inference endpoints (elastic#188670) ## Summary Adds Amazon Bedrock support to the [Inference Endpoints management UI](elastic#186206) (`relevance/inference_endpoints`) management list view. ### Checklist - [x] [Unit or functional tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html) were updated or added to match the most common scenarios --------- Co-authored-by: kibanamachine <[email protected]> Co-authored-by: Elastic Machine <[email protected]>
lcawl · Jul 26, 2024 · 25960d2 · 25960d2
1 parent 841e95c
commit 25960d2
Show file tree

Hide file tree

Showing 6 changed files with 61 additions and 0 deletions.
diff --git a/x-pack/packages/ml/trained_models_utils/src/constants/trained_models.ts b/x-pack/packages/ml/trained_models_utils/src/constants/trained_models.ts
@@ -259,6 +259,16 @@ export type InferenceServiceSettings =
         api_key: string;
         url: string;
       };
+    }
+  | {
+      service: 'amazonbedrock';
+      service_settings: {
+        access_key: string;
+        secret_key: string;
+        region: string;
+        provider: string;
+        model: string;
+      };
     };
 
 export type InferenceAPIConfigResponse = {

diff --git a/...ns/search_inference_endpoints/public/assets/images/providers/amazon_bedrock.svg b/...ns/search_inference_endpoints/public/assets/images/providers/amazon_bedrock.svg
diff --git a/...nents/all_inference_endpoints/render_table_columns/render_endpoint/endpoint_info.test.tsx b/...nents/all_inference_endpoints/render_table_columns/render_endpoint/endpoint_info.test.tsx
@@ -265,6 +265,26 @@ describe('RenderEndpoint component tests', () => {
     });
   });
 
+  describe('with amazonbedrock service', () => {
+    const mockEndpoint = {
+      model_id: 'amazon-bedrock-1',
+      service: 'amazonbedrock',
+      service_settings: {
+        region: 'us-west-1',
+        provider: 'AMAZONTITAN',
+        model: 'model-bedrock-xyz',
+      },
+    } as any;
+
+    it('renders the component with endpoint details', () => {
+      render(<EndpointInfo endpoint={mockEndpoint} />);
+
+      expect(screen.getByText('amazon-bedrock-1')).toBeInTheDocument();
+      expect(screen.getByText('model-bedrock-xyz')).toBeInTheDocument();
+      expect(screen.getByText('region: us-west-1, provider: amazontitan')).toBeInTheDocument();
+    });
+  });
+
   describe('for MIT licensed models', () => {
     const mockEndpointWithMitLicensedModel = {
       model_id: 'model-123',

diff --git a/...components/all_inference_endpoints/render_table_columns/render_endpoint/endpoint_info.tsx b/...components/all_inference_endpoints/render_table_columns/render_endpoint/endpoint_info.tsx
@@ -94,6 +94,8 @@ function endpointModelAtrributes(endpoint: InferenceAPIConfigResponse) {
       return mistralAttributes(endpoint);
     case ServiceProviderKeys.googleaistudio:
       return googleAIStudioAttributes(endpoint);
+    case ServiceProviderKeys.amazonbedrock:
+      return amazonBedrockAttributes(endpoint);
     default:
       return null;
   }
@@ -175,6 +177,18 @@ function mistralAttributes(endpoint: InferenceAPIConfigResponse) {
     .join(', ');
 }
 
+function amazonBedrockAttributes(endpoint: InferenceAPIConfigResponse) {
+  const serviceSettings = endpoint.service_settings;
+
+  const region = 'region' in serviceSettings ? serviceSettings.region : undefined;
+  const provider =
+    'provider' in serviceSettings ? serviceSettings.provider.toLocaleLowerCase() : undefined;
+
+  return [region && `region: ${region}`, provider && `provider: ${provider}`]
+    .filter(Boolean)
+    .join(', ');
+}
+
 function googleAIStudioAttributes(endpoint: InferenceAPIConfigResponse) {
   const serviceSettings = endpoint.service_settings;
 

diff --git a/...all_inference_endpoints/render_table_columns/render_service_provider/service_provider.tsx b/...all_inference_endpoints/render_table_columns/render_service_provider/service_provider.tsx
@@ -15,6 +15,7 @@ import azureAIStudioIcon from '../../../../assets/images/providers/azure_ai_stud
 import azureOpenAIIcon from '../../../../assets/images/providers/azure_open_ai.svg';
 import googleAIStudioIcon from '../../../../assets/images/providers/google_ai_studio.svg';
 import mistralIcon from '../../../../assets/images/providers/mistral.svg';
+import amazonBedrockIcon from '../../../../assets/images/providers/amazon_bedrock.svg';
 import { ServiceProviderKeys } from '../../types';
 
 interface ServiceProviderProps {
@@ -27,6 +28,10 @@ interface ServiceProviderRecord {
 }
 
 export const SERVICE_PROVIDERS: Record<ServiceProviderKeys, ServiceProviderRecord> = {
+  [ServiceProviderKeys.amazonbedrock]: {
+    icon: amazonBedrockIcon,
+    name: 'Amazon Bedrock',
+  },
   [ServiceProviderKeys.azureaistudio]: {
     icon: azureAIStudioIcon,
     name: 'Azure AI Studio',

diff --git a/x-pack/plugins/search_inference_endpoints/public/components/all_inference_endpoints/types.ts b/x-pack/plugins/search_inference_endpoints/public/components/all_inference_endpoints/types.ts
@@ -9,6 +9,7 @@ import { InferenceAPIConfigResponse } from '@kbn/ml-trained-models-utils';
 export const INFERENCE_ENDPOINTS_TABLE_PER_PAGE_VALUES = [10, 25, 50, 100];
 
 export enum ServiceProviderKeys {
+  amazonbedrock = 'amazonbedrock',
   azureopenai = 'azureopenai',
   azureaistudio = 'azureaistudio',
   cohere = 'cohere',