From c0836cc7f47932173d27c468b79498597e154b94 Mon Sep 17 00:00:00 2001
From: Google AI Edge <noreply@google.com>
Date: Wed, 25 Sep 2024 11:41:23 -0700
Subject: [PATCH] Internal cleanup

PiperOrigin-RevId: 678786345
---
 src/ui/src/common/consts.ts                   |  3 +++
 src/ui/src/common/types.ts                    |  4 ++++
 .../model_source_input/model_source_input.ts  | 20 ++++++++++++++++++-
 .../services/internal_adapter_extensions.ts   |  8 ++++++++
 4 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/src/ui/src/common/consts.ts b/src/ui/src/common/consts.ts
index 97bb289c..43f5b59c 100644
--- a/src/ui/src/common/consts.ts
+++ b/src/ui/src/common/consts.ts
@@ -18,3 +18,6 @@
 
 /** Prefix for graphs json model sources. */
 export const GRAPHS_MODEL_SOURCE_PREFIX = 'graphs://';
+
+/** Prefix for data nexus model sources. */
+export const DATA_NEXUS_MODEL_SOURCE_PREFIX = 'data-nexus://';
diff --git a/src/ui/src/common/types.ts b/src/ui/src/common/types.ts
index 7187f09a..31859d1f 100644
--- a/src/ui/src/common/types.ts
+++ b/src/ui/src/common/types.ts
@@ -35,6 +35,7 @@ export enum InternalAdapterExtId {
   GRAPHDEF = 'builtin_graphdef',
   MLIR = 'builtin_mlir',
   JSON_LOADER = 'builtin_json',
+  DATA_NEXUS = 'builtin_data_nexus',
 }
 
 /** Extension types. */
@@ -87,6 +88,9 @@ export enum ModelItemType {
   // Remote path/url.
   REMOTE = 'remote',
 
+  // Data nexus models.
+  DATA_NEXUS = 'data_nexus',
+
   // External only
   //
   // Graphs json specified when starting server.
diff --git a/src/ui/src/components/model_source_input/model_source_input.ts b/src/ui/src/components/model_source_input/model_source_input.ts
index b744253d..04beedb3 100644
--- a/src/ui/src/components/model_source_input/model_source_input.ts
+++ b/src/ui/src/components/model_source_input/model_source_input.ts
@@ -51,7 +51,10 @@ import {MatProgressSpinnerModule} from '@angular/material/progress-spinner';
 import {MatSelectModule} from '@angular/material/select';
 import {MatTooltipModule} from '@angular/material/tooltip';
 
-import {GRAPHS_MODEL_SOURCE_PREFIX} from '../../common/consts';
+import {
+  DATA_NEXUS_MODEL_SOURCE_PREFIX,
+  GRAPHS_MODEL_SOURCE_PREFIX,
+} from '../../common/consts';
 import {IS_EXTERNAL} from '../../common/flags';
 import {type ModelLoaderServiceInterface} from '../../common/model_loader_service_interface';
 import {
@@ -184,6 +187,21 @@ export class ModelSourceInput {
           selectedAdapter: ext,
         };
       }
+      // Data nexus model source.
+      else if (modelSource.url.startsWith(DATA_NEXUS_MODEL_SOURCE_PREFIX)) {
+        const ext = this.adapterExtensionService.getExtensionById(
+          InternalAdapterExtId.DATA_NEXUS,
+        );
+        const adapterCandidates = ext == null ? [] : [ext];
+        return {
+          path: modelSource.url,
+          type: ModelItemType.DATA_NEXUS,
+          status: signal<ModelItemStatus>(ModelItemStatus.NOT_STARTED),
+          selected: true,
+          adapterCandidates,
+          selectedAdapter: ext,
+        };
+      }
       // Other typical use cases.
       else {
         const adapterCandidates = getAdapterCandidates(
diff --git a/src/ui/src/services/internal_adapter_extensions.ts b/src/ui/src/services/internal_adapter_extensions.ts
index 8a472534..e87ba521 100644
--- a/src/ui/src/services/internal_adapter_extensions.ts
+++ b/src/ui/src/services/internal_adapter_extensions.ts
@@ -81,6 +81,13 @@ const jsonAdapterExtension: AdapterExtension = {
   description:
     'Loads JSON graphs data file or tfjs model from the given url or uploaded file and convert them to Model Explorer format.',
 };
+const dataNexusAdapterExtension: AdapterExtension = {
+  type: ExtensionType.ADAPTER,
+  fileExts: ['data_nexus'],
+  id: InternalAdapterExtId.DATA_NEXUS,
+  name: 'Data Nexus adapter',
+  description: 'Loads data from Data Nexus.',
+};
 
 /** All internal extensions. */
 export const INTERNAL_ADAPTER_EXTENSIONS: AdapterExtension[] = [
@@ -91,4 +98,5 @@ export const INTERNAL_ADAPTER_EXTENSIONS: AdapterExtension[] = [
   graphdefAdapterExtension,
   mlirAdapterExtension,
   jsonAdapterExtension,
+  dataNexusAdapterExtension,
 ];