From 85520247c70ec85bc19cdbbff4fb5bbf62508029 Mon Sep 17 00:00:00 2001
From: David Gauldie <dgauldie@uncharted.software>
Date: Mon, 18 Nov 2024 14:26:06 -0500
Subject: [PATCH 1/6] add autocomplete to dataset transformer

---
 .../dataset/tera-dataset-jupyter-panel.vue    |  6 +-
 .../src/components/llm/tera-beaker-input.vue  | 68 ++++++++++++++++++-
 .../tera-dataset-transformer.vue              |  8 +++
 3 files changed, 78 insertions(+), 4 deletions(-)

diff --git a/packages/client/hmi-client/src/components/dataset/tera-dataset-jupyter-panel.vue b/packages/client/hmi-client/src/components/dataset/tera-dataset-jupyter-panel.vue
index ac7c32a645..65067bad16 100644
--- a/packages/client/hmi-client/src/components/dataset/tera-dataset-jupyter-panel.vue
+++ b/packages/client/hmi-client/src/components/dataset/tera-dataset-jupyter-panel.vue
@@ -14,6 +14,7 @@
 			<tera-beaker-input
 				class="tera-beaker-input"
 				:kernel-is-busy="kernelStatus !== KernelState.idle"
+				:default-options="sampleAgentQuestions"
 				context="dataset"
 				@submitQuery="onSubmitQuery"
 				@add-code-cell="onAddCodeCell"
@@ -135,8 +136,8 @@ import {
 	getServerSettings,
 	getSessionManager,
 	JupyterMessage,
-	newSession,
-	KernelState
+	KernelState,
+	newSession
 } from '@/services/jupyter';
 import { SessionContext } from '@jupyterlab/apputils/lib/sessioncontext';
 import { createMessage } from '@jupyterlab/services/lib/kernel/messages';
@@ -179,6 +180,7 @@ const props = defineProps<{
 	programmingLanguage?: string;
 	kernelState: any;
 	selectedDataset: string | null;
+	sampleAgentQuestions: string[];
 }>();
 
 const languages = programmingLanguageOptions();
diff --git a/packages/client/hmi-client/src/components/llm/tera-beaker-input.vue b/packages/client/hmi-client/src/components/llm/tera-beaker-input.vue
index c34a49ae85..5f47cb7f9d 100644
--- a/packages/client/hmi-client/src/components/llm/tera-beaker-input.vue
+++ b/packages/client/hmi-client/src/components/llm/tera-beaker-input.vue
@@ -3,6 +3,7 @@
 		<div class="chat-input-container">
 			<div class="chat-input-elements">
 				<textarea
+					v-if="hideAutoComplete || props.defaultOptions.length === 0"
 					class="input"
 					ref="inputElement"
 					v-model="queryString"
@@ -13,6 +14,19 @@
 					@input="autoGrow"
 					rows="1"
 				/>
+				<AutoComplete
+					v-else
+					ref="autoComplete"
+					v-model="queryString"
+					class="auto-complete"
+					:suggestions="filteredOptions"
+					:completeOnFocus="true"
+					@complete="searchOptions"
+					@change="onInputChange"
+					placeholder="What do you want to do?"
+					emptySearchMessage="No suggestions"
+					:disabled="props.kernelIsBusy"
+				/>
 				<Button
 					v-if="queryString.length !== 0"
 					text
@@ -57,12 +71,13 @@
 </template>
 
 <script setup lang="ts">
-import { onBeforeUnmount, onMounted, ref } from 'vue';
+import { nextTick, onBeforeUnmount, onMounted, ref } from 'vue';
 import ProgressBar from 'primevue/progressbar';
 import Button from 'primevue/button';
 import * as EventService from '@/services/event';
 import { EventType } from '@/types/Types';
 import { useProjects } from '@/composables/project';
+import AutoComplete from 'primevue/autocomplete';
 
 const emit = defineEmits(['submit-query', 'add-code-cell', 'run-all-cells']);
 
@@ -84,7 +99,15 @@ const props = defineProps({
 		type: String,
 		default: 'black'
 	},
-	kernelIsBusy: Boolean
+	kernelIsBusy: Boolean,
+	defaultOptions: {
+		type: Array,
+		default: () => []
+	},
+	maxChars: {
+		type: Number,
+		default: 75
+	}
 });
 
 const queryString = ref('');
@@ -93,6 +116,9 @@ const containerElement = ref<HTMLElement | null>(null);
 const inputElement = ref<HTMLInputElement | null>(null);
 const addCodeCellButton = ref<HTMLElement | null>(null);
 const fixedDivWidth = ref(0);
+const hideAutoComplete = ref(false);
+const filteredOptions = ref<any[]>([]);
+const autoComplete = ref();
 
 const submitQuery = () => {
 	EventService.create(EventType.TransformPrompt, useProjects().activeProject.value?.id, queryString.value);
@@ -112,6 +138,30 @@ const runAllCells = () => {
 	emit('run-all-cells');
 };
 
+const searchOptions = () => {
+	const query = queryString.value.toLowerCase();
+	filteredOptions.value = props.defaultOptions.filter((option) => option.toLowerCase().includes(query));
+};
+
+const onInputChange = async () => {
+	// If there are no search suggestions to show, the auto complete element is never shown
+	if (props.defaultOptions.length === 0) return;
+
+	// If the number of characters entered exceeds the maximum, switch from auto complete to text area
+	// Force focus for a smoother transition
+	const numChars = queryString.value.length;
+	if (numChars <= props.maxChars && numChars >= 0) {
+		hideAutoComplete.value = false;
+		await nextTick();
+		const inputEl = autoComplete.value?.$el.querySelector('input');
+		inputEl.focus();
+	} else {
+		hideAutoComplete.value = true;
+		await nextTick();
+		textArea.value?.$el.focus();
+	}
+};
+
 onMounted(() => {
 	if (props.focusInput) {
 		inputElement.value?.focus();
@@ -204,4 +254,18 @@ const autoGrow = () => {
 	height: 31px;
 	width: 10rem;
 }
+.auto-complete {
+	width: 100%;
+}
+
+.auto-complete:deep(input),
+.text-area {
+	width: 100%;
+	background-image: url('@assets/svg/icons/message.svg');
+	background-size: 1rem;
+	background-position: var(--gap-2) 9px;
+	background-repeat: no-repeat;
+	padding-right: 2rem;
+	padding-left: 2rem;
+}
 </style>
diff --git a/packages/client/hmi-client/src/components/workflow/ops/dataset-transformer/tera-dataset-transformer.vue b/packages/client/hmi-client/src/components/workflow/ops/dataset-transformer/tera-dataset-transformer.vue
index 202e411ac2..6edf51c89d 100644
--- a/packages/client/hmi-client/src/components/workflow/ops/dataset-transformer/tera-dataset-transformer.vue
+++ b/packages/client/hmi-client/src/components/workflow/ops/dataset-transformer/tera-dataset-transformer.vue
@@ -31,6 +31,7 @@
 					@update-kernel-state="updateKernelState"
 					:kernelState="kernelState"
 					:selected-dataset="selectedDataset"
+					:sample-agent-questions="sampleAgentQuestions"
 				/>
 			</Suspense>
 		</div>
@@ -85,6 +86,13 @@ const updateKernelState = (newKernelState: any) => {
 	}
 };
 
+const sampleAgentQuestions = [
+	'View the dataset to see the columns and data types.',
+	'Join the two datasets together on date',
+	'Add a new column named timestamp',
+	"Let's assume that average time to recover is 14 days and average time to exit the hospital is 10 days. Can you convert this data into prevalence data? Ideally please map it to SIRHD. Assume a population of 150 million."
+];
+
 const notebookSession = ref(<NotebookSession | undefined>undefined);
 
 onMounted(async () => {

From cd4381bb88d477316882631bfb2b679774669bcb Mon Sep 17 00:00:00 2001
From: David Gauldie <dgauldie@uncharted.software>
Date: Mon, 18 Nov 2024 15:31:36 -0500
Subject: [PATCH 2/6] update prompts

---
 .../dataset-transformer/tera-dataset-transformer.vue   | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/packages/client/hmi-client/src/components/workflow/ops/dataset-transformer/tera-dataset-transformer.vue b/packages/client/hmi-client/src/components/workflow/ops/dataset-transformer/tera-dataset-transformer.vue
index 6edf51c89d..047b8c9566 100644
--- a/packages/client/hmi-client/src/components/workflow/ops/dataset-transformer/tera-dataset-transformer.vue
+++ b/packages/client/hmi-client/src/components/workflow/ops/dataset-transformer/tera-dataset-transformer.vue
@@ -87,10 +87,12 @@ const updateKernelState = (newKernelState: any) => {
 };
 
 const sampleAgentQuestions = [
-	'View the dataset to see the columns and data types.',
-	'Join the two datasets together on date',
-	'Add a new column named timestamp',
-	"Let's assume that average time to recover is 14 days and average time to exit the hospital is 10 days. Can you convert this data into prevalence data? Ideally please map it to SIRHD. Assume a population of 150 million."
+	'I have two dataframes d1, d2. Join them on the column named "date". Name the joined dataframe d3.',
+	'I have a dataframe d1. Show me the data types by column.',
+	"I have three dataframes d1, d2, d3. d1 is incident case counts. d2 is incident hospitalization counts. d3 is cumulative death counts. Let's assume that average time to recover is 14 days and average time to exit the hospital is 10 days. Can you convert this data into prevalence data? Ideally please map it to SIRHD. Assume a population of 150 million.",
+	'Add a new column to the dataframe d1 that indexes the rows from 0, 1, 2 to N.',
+	'Download geojson of US counties from the Plotly GitHub repo using urlopen.',
+	'I have a geopandas dataframe d1. Use Matplotlib to create a chloropleth map using the column R0. Use the cividis colormap. Add a colorbar.'
 ];
 
 const notebookSession = ref(<NotebookSession | undefined>undefined);

From 36392a5264680cea2712a232b548c21bd1d97804 Mon Sep 17 00:00:00 2001
From: David Gauldie <dgauldie@uncharted.software>
Date: Tue, 19 Nov 2024 09:55:52 -0500
Subject: [PATCH 3/6] fix typechecking

---
 .../hmi-client/src/components/llm/tera-beaker-input.vue       | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/packages/client/hmi-client/src/components/llm/tera-beaker-input.vue b/packages/client/hmi-client/src/components/llm/tera-beaker-input.vue
index 5f47cb7f9d..f3e950cc51 100644
--- a/packages/client/hmi-client/src/components/llm/tera-beaker-input.vue
+++ b/packages/client/hmi-client/src/components/llm/tera-beaker-input.vue
@@ -101,7 +101,7 @@ const props = defineProps({
 	},
 	kernelIsBusy: Boolean,
 	defaultOptions: {
-		type: Array,
+		type: Array<String>,
 		default: () => []
 	},
 	maxChars: {
@@ -158,7 +158,7 @@ const onInputChange = async () => {
 	} else {
 		hideAutoComplete.value = true;
 		await nextTick();
-		textArea.value?.$el.focus();
+		inputElement.value?.focus();
 	}
 };
 

From a202e97197039f93ef82fccb987a99a441c153e1 Mon Sep 17 00:00:00 2001
From: David Gauldie <dgauldie@uncharted.software>
Date: Tue, 19 Nov 2024 10:54:52 -0500
Subject: [PATCH 4/6] minor fix for mit-tr

---
 .../hmiserver/controller/knowledge/KnowledgeController.java     | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/packages/server/src/main/java/software/uncharted/terarium/hmiserver/controller/knowledge/KnowledgeController.java b/packages/server/src/main/java/software/uncharted/terarium/hmiserver/controller/knowledge/KnowledgeController.java
index fcd5a96e92..7a8d927ef6 100644
--- a/packages/server/src/main/java/software/uncharted/terarium/hmiserver/controller/knowledge/KnowledgeController.java
+++ b/packages/server/src/main/java/software/uncharted/terarium/hmiserver/controller/knowledge/KnowledgeController.java
@@ -679,7 +679,7 @@ public ResponseEntity<Dataset> postProfileDataset(
 				if (groundings.getIdentifiers() == null) {
 					groundings.setIdentifiers(new ArrayList<>());
 				}
-				groundings.getIdentifiers().add(new DKG(g.get(0).asText(), g.get(1).asText(), null));
+				groundings.getIdentifiers().add(new DKG(g.get(0).asText(), g.get(1).asText(), ""));
 			}
 
 			// remove groundings from annotation object

From 8c14321ab8cf244a1b001f0fbbda3102f862ebfe Mon Sep 17 00:00:00 2001
From: David Gauldie <dgauldie@uncharted.software>
Date: Tue, 19 Nov 2024 11:07:30 -0500
Subject: [PATCH 5/6] fix broken tests

---
 .../terarium/hmiserver/configuration/TestReBACService.java      | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/packages/server/src/test/java/software/uncharted/terarium/hmiserver/configuration/TestReBACService.java b/packages/server/src/test/java/software/uncharted/terarium/hmiserver/configuration/TestReBACService.java
index 10a73a0278..fdce52e24a 100644
--- a/packages/server/src/test/java/software/uncharted/terarium/hmiserver/configuration/TestReBACService.java
+++ b/packages/server/src/test/java/software/uncharted/terarium/hmiserver/configuration/TestReBACService.java
@@ -160,7 +160,7 @@ public void removeRelationship(
 		final SchemaObject who,
 		final SchemaObject what,
 		final Schema.Relationship relationship
-	) throws Exception, RelationshipAlreadyExistsException {}
+	) throws Exception {}
 
 	@Override
 	public List<RebacPermissionRelationship> getRelationships(final SchemaObject what) throws Exception {

From 506dcad927b23166516ddb60f2882620b1fafc16 Mon Sep 17 00:00:00 2001
From: David Gauldie <dgauldie@uncharted.software>
Date: Tue, 19 Nov 2024 11:13:26 -0500
Subject: [PATCH 6/6] update tests

---
 .../hmiserver/service/data/DatasetServiceTests.java         | 2 +-
 .../hmiserver/service/data/DocumentServiceTests.java        | 2 +-
 .../service/data/TerariumAssetCloneServiceTests.java        | 6 +++---
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/packages/server/src/test/java/software/uncharted/terarium/hmiserver/service/data/DatasetServiceTests.java b/packages/server/src/test/java/software/uncharted/terarium/hmiserver/service/data/DatasetServiceTests.java
index 7b8c399503..82dc7ce6e2 100644
--- a/packages/server/src/test/java/software/uncharted/terarium/hmiserver/service/data/DatasetServiceTests.java
+++ b/packages/server/src/test/java/software/uncharted/terarium/hmiserver/service/data/DatasetServiceTests.java
@@ -60,7 +60,7 @@ static Grounding createGrounding(final String key) {
 		final Grounding grounding = new Grounding();
 		grounding.setContext(mapper.createObjectNode().put("hello", "world-" + key).put("foo", "bar-" + key));
 		grounding.setIdentifiers(new ArrayList<>());
-		grounding.getIdentifiers().add(new DKG("curie", "maria", null));
+		grounding.getIdentifiers().add(new DKG("curie", "maria", ""));
 		return grounding;
 	}
 
diff --git a/packages/server/src/test/java/software/uncharted/terarium/hmiserver/service/data/DocumentServiceTests.java b/packages/server/src/test/java/software/uncharted/terarium/hmiserver/service/data/DocumentServiceTests.java
index 49400fbe7b..400065cdfb 100644
--- a/packages/server/src/test/java/software/uncharted/terarium/hmiserver/service/data/DocumentServiceTests.java
+++ b/packages/server/src/test/java/software/uncharted/terarium/hmiserver/service/data/DocumentServiceTests.java
@@ -63,7 +63,7 @@ static Grounding createGrounding(final String key) {
 		final Grounding grounding = new Grounding();
 		grounding.setContext(mapper.createObjectNode().put("hello", "world-" + key).put("foo", "bar-" + key));
 		grounding.setIdentifiers(new ArrayList<>());
-		grounding.getIdentifiers().add(new DKG("curie", "maria", null));
+		grounding.getIdentifiers().add(new DKG("curie", "maria", ""));
 		return grounding;
 	}
 
diff --git a/packages/server/src/test/java/software/uncharted/terarium/hmiserver/service/data/TerariumAssetCloneServiceTests.java b/packages/server/src/test/java/software/uncharted/terarium/hmiserver/service/data/TerariumAssetCloneServiceTests.java
index f8426a4bba..7580cee1e6 100644
--- a/packages/server/src/test/java/software/uncharted/terarium/hmiserver/service/data/TerariumAssetCloneServiceTests.java
+++ b/packages/server/src/test/java/software/uncharted/terarium/hmiserver/service/data/TerariumAssetCloneServiceTests.java
@@ -74,7 +74,7 @@ static Grounding createGrounding(final String key) {
 		final Grounding grounding = new Grounding();
 		grounding.setContext(objectMapper.createObjectNode().put("hello", "world-" + key).put("foo", "bar-" + key));
 		grounding.setIdentifiers(new ArrayList<>());
-		grounding.getIdentifiers().add(new DKG("curie", "maria", null));
+		grounding.getIdentifiers().add(new DKG("curie", "maria", ""));
 		return grounding;
 	}
 
@@ -133,7 +133,7 @@ public void testItCanCloneAndPersistAssets() throws Exception {
 					after.getId(),
 					filename,
 					ContentType.TEXT_PLAIN,
-					new String("This is my sample file containing" + filename).getBytes()
+					("This is my sample file containing" + filename).getBytes()
 				);
 			}
 		}
@@ -182,7 +182,7 @@ public void testItCanExportImportProject() throws Exception {
 					after.getId(),
 					filename,
 					ContentType.TEXT_PLAIN,
-					new String("This is my sample file containing" + filename).getBytes()
+					("This is my sample file containing" + filename).getBytes()
 				);
 			}
 		}