Adding Tech Preview badge for Reranker

x-pack/packages/ml/trained_models_utils/src/constants/trained_models.ts

@@ -311,7 +311,7 @@ export type InferenceServiceSettings =
 export type InferenceAPIConfigResponse = {
   // Refers to a deployment id
   inference_id: string;
-  task_type: 'sparse_embedding' | 'text_embedding';
+  task_type: 'sparse_embedding' | 'text_embedding' | 'rerank';
   task_settings: {
     model?: string;
   };

x-pack/plugins/search_inference_endpoints/public/components/all_inference_endpoints/render_table_columns/render_endpoint/endpoint_info.test.tsx

@@ -11,8 +11,60 @@ import { EndpointInfo } from './endpoint_info';
 
 describe('RenderEndpoint component tests', () => {
   it('renders the component with inference id', () => {
-    render(<EndpointInfo inferenceId={'cohere-2'} />);
+    const mockProvider = {
+      inference_id: 'cohere-2',
+      service: 'cohere',
+      service_settings: {
+        similarity: 'cosine',
+        dimensions: 384,
+        model_id: 'embed-english-light-v3.0',
+        rate_limit: {
+          requests_per_minute: 10000,
+        },
+        embedding_type: 'byte',
+      },
+      task_settings: {},
+    } as any;
+
+    render(<EndpointInfo inferenceId={'cohere-2'} provider={mockProvider} />);
 
     expect(screen.getByText('cohere-2')).toBeInTheDocument();
   });
+
+  it('renders correctly without model_id in service_settings', () => {
+    const mockProvider = {
+      inference_id: 'azure-openai-1',
+      service: 'azureopenai',
+      service_settings: {
+        resource_name: 'resource-xyz',
+        deployment_id: 'deployment-123',
+        api_version: 'v1',
+      },
+    } as any;
+
+    render(<EndpointInfo inferenceId={'azure-openai-1'} provider={mockProvider} />);
+
+    expect(screen.getByText('azure-openai-1')).toBeInTheDocument();
+  });
+
+  it('renders with tech preview badge when endpoint is reranker type', () => {
+    const mockProvider = {
+      inference_id: 'elastic-rerank',
+      task_type: 'rerank',
+      service: 'elasticsearch',
+      service_settings: {
+        num_allocations: 1,
+        num_threads: 1,
+        model_id: '.rerank-v1',
+      },
+      task_settings: {
+        return_documents: true,
+      },
+    } as any;
+
+    render(<EndpointInfo inferenceId={'elastic-rerank'} provider={mockProvider} />);
+
+    expect(screen.getByText('elastic-rerank')).toBeInTheDocument();
+    expect(screen.getByText('TECH PREVIEW')).toBeInTheDocument();
+  });
 });

x-pack/plugins/search_inference_endpoints/public/components/all_inference_endpoints/render_table_columns/render_endpoint/endpoint_info.tsx

@@ -7,26 +7,49 @@
 
 import { EuiBetaBadge, EuiFlexGroup, EuiFlexItem } from '@elastic/eui';
 import React from 'react';
+import { InferenceAPIConfigResponse } from '@kbn/ml-trained-models-utils';
 import { isEndpointPreconfigured } from '../../../../utils/preconfigured_endpoint_helper';
 import * as i18n from './translations';
+import { isProviderTechPreview } from '../../../../utils/reranker_helper';
 
 export interface EndpointInfoProps {
   inferenceId: string;
+  provider: InferenceAPIConfigResponse;
 }
 
-export const EndpointInfo: React.FC<EndpointInfoProps> = ({ inferenceId }) => (
-  <EuiFlexGroup justifyContent="spaceBetween">
-    <EuiFlexItem grow={false}>
-      <span>
-        <strong>{inferenceId}</strong>
-      </span>
-    </EuiFlexItem>
-    <EuiFlexItem grow={false}>
-      <span>
-        {isEndpointPreconfigured(inferenceId) ? (
-          <EuiBetaBadge label={i18n.PRECONFIGURED_LABEL} size="s" color="hollow" />
-        ) : null}
-      </span>
-    </EuiFlexItem>
-  </EuiFlexGroup>
-);
+export const EndpointInfo: React.FC<EndpointInfoProps> = ({ inferenceId, provider }) => {
+  const isInTechPreview = isProviderTechPreview(provider);
+
+  return (
+    <EuiFlexGroup justifyContent="spaceBetween">
+      <EuiFlexItem grow={false}>
+        <EuiFlexGroup gutterSize="s" alignItems="center">
+          <EuiFlexItem grow={false}>
+            <span>
+              <strong>{inferenceId}</strong>
+            </span>
+          </EuiFlexItem>
+          {isInTechPreview ? (
+            <EuiFlexItem grow={false}>
+              <span>
+                <EuiBetaBadge
+                  label={i18n.TECH_PREVIEW_LABEL}
+                  size="s"
+                  color="hollow"
+                  alignment="middle"
+                />
+              </span>
+            </EuiFlexItem>
+          ) : null}
+        </EuiFlexGroup>
+      </EuiFlexItem>
+      <EuiFlexItem grow={false}>
+        <span>
+          {isEndpointPreconfigured(inferenceId) ? (
+            <EuiBetaBadge label={i18n.PRECONFIGURED_LABEL} size="s" color="hollow" />
+          ) : null}
+        </span>
+      </EuiFlexItem>
+    </EuiFlexGroup>
+  );
+};

x-pack/plugins/search_inference_endpoints/public/components/all_inference_endpoints/render_table_columns/render_endpoint/translations.ts

@@ -13,3 +13,10 @@ export const PRECONFIGURED_LABEL = i18n.translate(
     defaultMessage: 'PRECONFIGURED',
   }
 );
+
+export const TECH_PREVIEW_LABEL = i18n.translate(
+  'xpack.searchInferenceEndpoints.elasticsearch.endpointInfo.techPreview',
+  {
+    defaultMessage: 'TECH PREVIEW',
+  }
+);

x-pack/plugins/search_inference_endpoints/public/components/all_inference_endpoints/tabular_page.test.tsx

@@ -67,6 +67,19 @@ const inferenceEndpoints = [
     },
     task_settings: {},
   },
+  {
+    inference_id: 'elastic-rerank',
+    task_type: 'rerank',
+    service: 'elasticsearch',
+    service_settings: {
+      num_allocations: 1,
+      num_threads: 1,
+      model_id: '.rerank-v1',
+    },
+    task_settings: {
+      return_documents: true,
+    },
+  },
 ] as InferenceAPIConfigResponse[];
 
 jest.mock('../../hooks/use_delete_endpoint', () => ({
@@ -82,9 +95,10 @@ describe('When the tabular page is loaded', () => {
     const rows = screen.getAllByRole('row');
     expect(rows[1]).toHaveTextContent('.elser-2-elasticsearch');
     expect(rows[2]).toHaveTextContent('.multilingual-e5-small-elasticsearch');
-    expect(rows[3]).toHaveTextContent('local-model');
-    expect(rows[4]).toHaveTextContent('my-elser-model-05');
-    expect(rows[5]).toHaveTextContent('third-party-model');
+    expect(rows[3]).toHaveTextContent('elastic-rerank');
+    expect(rows[4]).toHaveTextContent('local-model');
+    expect(rows[5]).toHaveTextContent('my-elser-model-05');
+    expect(rows[6]).toHaveTextContent('third-party-model');
   });
 
   it('should display all service and model ids in the table', () => {
@@ -98,13 +112,16 @@ describe('When the tabular page is loaded', () => {
     expect(rows[2]).toHaveTextContent('.multilingual-e5-small');
 
     expect(rows[3]).toHaveTextContent('Elasticsearch');
-    expect(rows[3]).toHaveTextContent('.own_model');
+    expect(rows[3]).toHaveTextContent('.rerank-v1');
 
     expect(rows[4]).toHaveTextContent('Elasticsearch');
-    expect(rows[4]).toHaveTextContent('.elser_model_2');
+    expect(rows[4]).toHaveTextContent('.own_model');
 
-    expect(rows[5]).toHaveTextContent('OpenAI');
-    expect(rows[5]).toHaveTextContent('.own_model');
+    expect(rows[5]).toHaveTextContent('Elasticsearch');
+    expect(rows[5]).toHaveTextContent('.elser_model_2');
+
+    expect(rows[6]).toHaveTextContent('OpenAI');
+    expect(rows[6]).toHaveTextContent('.own_model');
   });
 
   it('should only disable delete action for preconfigured endpoints', () => {
@@ -131,4 +148,18 @@ describe('When the tabular page is loaded', () => {
     expect(rows[4]).not.toHaveTextContent(preconfigured);
     expect(rows[5]).not.toHaveTextContent(preconfigured);
   });
+
+  it('should show tech preview badge only for reranker-v1 model', () => {
+    render(<TabularPage inferenceEndpoints={inferenceEndpoints} />);
+
+    const techPreview = 'TECH PREVIEW';
+
+    const rows = screen.getAllByRole('row');
+    expect(rows[1]).not.toHaveTextContent(techPreview);
+    expect(rows[2]).not.toHaveTextContent(techPreview);
+    expect(rows[3]).toHaveTextContent(techPreview);
+    expect(rows[4]).not.toHaveTextContent(techPreview);
+    expect(rows[5]).not.toHaveTextContent(techPreview);
+    expect(rows[6]).not.toHaveTextContent(techPreview);
+  });
 });

x-pack/plugins/search_inference_endpoints/public/components/all_inference_endpoints/tabular_page.tsx

@@ -53,9 +53,9 @@ export const TabularPage: React.FC<TabularPageProps> = ({ inferenceEndpoints })
       field: 'endpoint',
       name: i18n.ENDPOINT,
       'data-test-subj': 'endpointCell',
-      render: (endpoint: string) => {
+      render: (endpoint: string, additionalInfo: InferenceEndpointUI) => {
         if (endpoint) {
-          return <EndpointInfo inferenceId={endpoint} />;
+          return <EndpointInfo inferenceId={endpoint} provider={additionalInfo.provider} />;
         }
 
         return null;

x-pack/plugins/search_inference_endpoints/public/utils/reranker_helper.test.ts

@@ -0,0 +1,54 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { isProviderTechPreview } from './reranker_helper';
+
+describe('Reranker Tech preview badge', () => {
+  const mockProvider = {
+    inference_id: 'elastic-rerank',
+    task_type: 'rerank',
+    service: 'elasticsearch',
+    service_settings: {
+      num_allocations: 1,
+      num_threads: 1,
+      model_id: '.rerank-v1',
+    },
+    task_settings: {
+      return_documents: true,
+    },
+  } as any;
+
+  it('return true for reranker', () => {
+    expect(isProviderTechPreview(mockProvider)).toEqual(true);
+  });
+
+  it('return false for other provider', () => {
+    const otherProviderServiceSettings = {
+      ...mockProvider.service_settings,
+      model_id: '.elser_model_2',
+    };
+    const otherProvider = {
+      ...mockProvider,
+      task_type: 'sparse_embedding',
+      service_settings: otherProviderServiceSettings,
+    } as any;
+    expect(isProviderTechPreview(otherProvider)).toEqual(false);
+  });
+
+  it('return false for other provider without model_id', () => {
+    const mockThirdPartyProvider = {
+      inference_id: 'azure-openai-1',
+      service: 'azureopenai',
+      service_settings: {
+        resource_name: 'resource-xyz',
+        deployment_id: 'deployment-123',
+        api_version: 'v1',
+      },
+    } as any;
+    expect(isProviderTechPreview(mockThirdPartyProvider)).toEqual(false);
+  });
+});

x-pack/plugins/search_inference_endpoints/public/utils/reranker_helper.ts

@@ -0,0 +1,21 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { InferenceAPIConfigResponse } from '@kbn/ml-trained-models-utils';
+export const isProviderTechPreview = (provider: InferenceAPIConfigResponse) => {
+  if (hasModelId(provider)) {
+    return provider.task_type === 'rerank' && provider.service_settings?.model_id?.startsWith('.');
+  }
+
+  return false;
+};
+
+function hasModelId(
+  service: InferenceAPIConfigResponse
+): service is Extract<InferenceAPIConfigResponse, { service_settings: { model_id: string } }> {
+  return 'model_id' in service.service_settings;
+}