From ba01d34b60735de75ba122e7041bac9a4b7a1380 Mon Sep 17 00:00:00 2001
From: Lin Wang <wonglam@amazon.com>
Date: Mon, 3 Jul 2023 15:18:36 +0800
Subject: [PATCH] feat: exclude remote model for admin UI (#225)

Signed-off-by: Lin Wang <wonglam@amazon.com>
---
 public/apis/model.ts                          |  1 +
 .../components/monitoring/use_monitoring.ts   |  1 +
 server/routes/model_router.ts                 |  4 +++-
 server/services/model_service.ts              |  1 +
 server/services/utils/model.ts                | 21 +++++++++++++++----
 5 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/public/apis/model.ts b/public/apis/model.ts
index daf0689c..040feb25 100644
--- a/public/apis/model.ts
+++ b/public/apis/model.ts
@@ -30,6 +30,7 @@ export class Model {
     size: number;
     states?: MODEL_STATE[];
     nameOrId?: string;
+    exclude?: 'REMOTE_MODEL';
   }) {
     return InnerHttpProvider.getHttp().get<ModelSearchResponse>(MODEL_API_ENDPOINT, {
       query,
diff --git a/public/components/monitoring/use_monitoring.ts b/public/components/monitoring/use_monitoring.ts
index 35031ad5..34907b5e 100644
--- a/public/components/monitoring/use_monitoring.ts
+++ b/public/components/monitoring/use_monitoring.ts
@@ -47,6 +47,7 @@ const fetchDeployedModels = async (params: Params) => {
         ? [MODEL_STATE.loadFailed, MODEL_STATE.loaded, MODEL_STATE.partiallyLoaded]
         : states,
     sort: [`${params.sort.field}-${params.sort.direction}`],
+    exclude: 'REMOTE_MODEL',
   });
   const totalPages = Math.ceil(result.total_models / params.pageSize);
   return {
diff --git a/server/routes/model_router.ts b/server/routes/model_router.ts
index 822df4c2..94a8512e 100644
--- a/server/routes/model_router.ts
+++ b/server/routes/model_router.ts
@@ -42,11 +42,12 @@ export const modelRouter = (router: IRouter) => {
           ),
           states: schema.maybe(schema.oneOf([schema.arrayOf(modelStateSchema), modelStateSchema])),
           nameOrId: schema.maybe(schema.string()),
+          exclude: schema.maybe(schema.literal('REMOTE_MODEL')),
         }),
       },
     },
     async (context, request) => {
-      const { from, size, sort, states, nameOrId } = request.query;
+      const { from, size, sort, states, nameOrId, exclude } = request.query;
       try {
         const payload = await ModelService.search({
           client: context.core.opensearch.client,
@@ -55,6 +56,7 @@ export const modelRouter = (router: IRouter) => {
           sort: typeof sort === 'string' ? [sort] : sort,
           states: typeof states === 'string' ? [states] : states,
           nameOrId,
+          exclude,
         });
         return opensearchDashboardsResponseFactory.ok({ body: payload });
       } catch (err) {
diff --git a/server/services/model_service.ts b/server/services/model_service.ts
index 2b98dfdf..35480209 100644
--- a/server/services/model_service.ts
+++ b/server/services/model_service.ts
@@ -43,6 +43,7 @@ export class ModelService {
     sort?: ModelSearchSort[];
     states?: MODEL_STATE[];
     nameOrId?: string;
+    exclude?: 'REMOTE_MODEL';
   }) {
     const {
       body: { hits },
diff --git a/server/services/utils/model.ts b/server/services/utils/model.ts
index f7b585c4..e347fd33 100644
--- a/server/services/utils/model.ts
+++ b/server/services/utils/model.ts
@@ -9,9 +9,11 @@ import { generateTermQuery } from './query';
 export const generateModelSearchQuery = ({
   states,
   nameOrId,
+  exclude,
 }: {
   states?: MODEL_STATE[];
   nameOrId?: string;
+  exclude?: 'REMOTE_MODEL';
 }) => ({
   bool: {
     must: [
@@ -33,10 +35,21 @@ export const generateModelSearchQuery = ({
           ]
         : []),
     ],
-    must_not: {
-      exists: {
-        field: 'chunk_number',
+    must_not: [
+      {
+        exists: {
+          field: 'chunk_number',
+        },
       },
-    },
+      ...(exclude === 'REMOTE_MODEL'
+        ? [
+            {
+              term: {
+                algorithm: 'REMOTE',
+              },
+            },
+          ]
+        : []),
+    ],
   },
 });