From 4916a2a67251fc4ebd3e451f08c5512b4ef14b2b Mon Sep 17 00:00:00 2001
From: Jaehwan Ryu <jryu@uncharted.software>
Date: Mon, 3 Jun 2024 12:18:29 -0400
Subject: [PATCH] 3595 notebook summary - model stratify operator Part 1
 (#3732)

---
 .../components/drilldown/tera-drilldown.vue   | 12 +++
 .../llm/tera-notebook-jupyter-input.vue       |  3 +-
 .../ops/model-edit/tera-model-edit.vue        |  8 +-
 .../stratify-mira/stratify-mira-operation.ts  | 12 +++
 .../ops/stratify-mira/tera-stratify-mira.vue  | 18 ++--
 .../src/components/workflow/tera-workflow.vue | 26 +++++-
 .../client/hmi-client/src/services/beaker.ts  | 11 +++
 .../hmi-client/src/services/notebook.ts       | 83 ++++++++++++++++++-
 .../hmi-client/src/services/workflow.ts       | 27 ++++++
 .../client/hmi-client/src/types/workflow.ts   |  4 +
 10 files changed, 191 insertions(+), 13 deletions(-)
 create mode 100644 packages/client/hmi-client/src/services/beaker.ts

diff --git a/packages/client/hmi-client/src/components/drilldown/tera-drilldown.vue b/packages/client/hmi-client/src/components/drilldown/tera-drilldown.vue
index d350bd399f..bce9893bcf 100644
--- a/packages/client/hmi-client/src/components/drilldown/tera-drilldown.vue
+++ b/packages/client/hmi-client/src/components/drilldown/tera-drilldown.vue
@@ -39,6 +39,16 @@
 						:state="node.state"
 						@update-state="(state: any) => emit('update-state', state)"
 					/>
+					<!-- TODO: Uncomment and enable this when summary UI is ready -->
+					<!-- <tera-operator-annotation
+						v-if="activeOutputSummary === undefined"
+						:state="node.state"
+						@update-state="(state: any) => emit('update-state', state)"
+					/> -->
+					<!-- TODO: Make this as a summary component and render for the operator supporting AI generated summary-->
+					<!-- <div v-else>
+						{{ activeOutputSummary === '' ? 'Generating AI summary...' : activeOutputSummary }}
+					</div> -->
 				</template>
 			</tera-drilldown-header>
 			<tera-columnar-panel>
@@ -75,6 +85,7 @@ import { isEmpty } from 'lodash';
 import Menu from 'primevue/menu';
 import Button from 'primevue/button';
 import TeraOutputDropdown from '@/components/drilldown/tera-output-dropdown.vue';
+// import { getActiveOutputSummary } from '@/services/workflow';
 
 const props = defineProps<{
 	node: WorkflowNode<any>;
@@ -126,6 +137,7 @@ const outputOptions = computed(() => {
 	}
 	return [];
 });
+// const activeOutputSummary = computed(() => getActiveOutputSummary(props.node));
 
 const toggle = (event) => {
 	menu.value.toggle(event);
diff --git a/packages/client/hmi-client/src/components/llm/tera-notebook-jupyter-input.vue b/packages/client/hmi-client/src/components/llm/tera-notebook-jupyter-input.vue
index fbfa6e6c02..04e026fe98 100644
--- a/packages/client/hmi-client/src/components/llm/tera-notebook-jupyter-input.vue
+++ b/packages/client/hmi-client/src/components/llm/tera-notebook-jupyter-input.vue
@@ -73,7 +73,8 @@ const submitQuestion = () => {
 	const message = props.kernelManager.sendMessage('llm_request', {
 		request: questionString.value
 	});
-	emit('question-asked');
+	emit('question-asked', questionString.value);
+
 	// May prefer to use a manual status rather than following this. TBD. Both options work for now
 	message.register('status', (data) => {
 		kernelStatus.value = data.content.execution_state;
diff --git a/packages/client/hmi-client/src/components/workflow/ops/model-edit/tera-model-edit.vue b/packages/client/hmi-client/src/components/workflow/ops/model-edit/tera-model-edit.vue
index 88ee1c04db..4d6e066d75 100644
--- a/packages/client/hmi-client/src/components/workflow/ops/model-edit/tera-model-edit.vue
+++ b/packages/client/hmi-client/src/components/workflow/ops/model-edit/tera-model-edit.vue
@@ -25,7 +25,7 @@
 							:context-language="contextLanguage"
 							@llm-output="(data: any) => appendCode(data, 'code')"
 							@llm-thought-output="(data: any) => llmThoughts.push(data)"
-							@question-asked="llmThoughts = []"
+							@question-asked="updateLlmQuery"
 						>
 							<template #toolbar-right-side>
 								<Button
@@ -166,6 +166,7 @@ const contextLanguage = ref<string>('python3');
 const defaultCodeText =
 	'# This environment contains the variable "model" \n# which is displayed on the right';
 const codeText = ref(defaultCodeText);
+const llmQuery = ref('');
 const llmThoughts = ref<any[]>([]);
 
 const executeResponse = ref({
@@ -185,6 +186,11 @@ const menuItems = computed(() => [
 	}
 ]);
 
+const updateLlmQuery = (query: string) => {
+	llmThoughts.value = [];
+	llmQuery.value = query;
+};
+
 const appendCode = (data: any, property: string) => {
 	const newCode = data.content[property] as string;
 	if (newCode) {
diff --git a/packages/client/hmi-client/src/components/workflow/ops/stratify-mira/stratify-mira-operation.ts b/packages/client/hmi-client/src/components/workflow/ops/stratify-mira/stratify-mira-operation.ts
index ab22aef416..9683bc2fda 100644
--- a/packages/client/hmi-client/src/components/workflow/ops/stratify-mira/stratify-mira-operation.ts
+++ b/packages/client/hmi-client/src/components/workflow/ops/stratify-mira/stratify-mira-operation.ts
@@ -1,3 +1,5 @@
+import { getModel } from '@/services/model';
+import { createNotebookFromCode } from '@/services/notebook';
 import type { Operation, BaseState } from '@/types/workflow';
 import { WorkflowOperationTypes } from '@/types/workflow';
 
@@ -18,6 +20,8 @@ export interface StratifyGroup {
 }
 
 export interface StratifyCode {
+	llmQuery: string;
+	llmThoughts: any[];
 	code: string;
 	timestamp: number;
 }
@@ -70,5 +74,13 @@ export const StratifyMiraOperation: Operation = {
 			hasCodeBeenRun: false
 		};
 		return init;
+	},
+	createNotebook: async (state: StratifyOperationStateMira, value?: any[] | null) => {
+		const modelIdToLoad = value?.[0];
+		const outputModel = await getModel(modelIdToLoad);
+		const code = state.strataCodeHistory?.[0].code ?? '';
+		// TODO: Add llm query and thought to the notebook
+		const notebook = createNotebookFromCode(code, 'python3', { 'application/json': outputModel });
+		return notebook;
 	}
 };
diff --git a/packages/client/hmi-client/src/components/workflow/ops/stratify-mira/tera-stratify-mira.vue b/packages/client/hmi-client/src/components/workflow/ops/stratify-mira/tera-stratify-mira.vue
index 11fe2d3e79..445dd00324 100644
--- a/packages/client/hmi-client/src/components/workflow/ops/stratify-mira/tera-stratify-mira.vue
+++ b/packages/client/hmi-client/src/components/workflow/ops/stratify-mira/tera-stratify-mira.vue
@@ -48,7 +48,7 @@
 						:context-language="'python3'"
 						@llm-output="(data: any) => processLLMOutput(data)"
 						@llm-thought-output="(data: any) => llmThoughts.push(data)"
-						@question-asked="llmThoughts = []"
+						@question-asked="updateLlmQuery"
 					>
 						<template #toolbar-right-side>
 							<Button label="Run" size="small" icon="pi pi-play" @click="runCodeStratify" />
@@ -116,6 +116,7 @@ import TeraNotebookJupyterInput from '@/components/llm/tera-notebook-jupyter-inp
 import teraNotebookJupyterThoughtOutput from '@/components/llm/tera-notebook-jupyter-thought-output.vue';
 
 import { createModel, getModel } from '@/services/model';
+
 import { WorkflowNode, OperatorStatus } from '@/types/workflow';
 import { logger } from '@/utils/logger';
 import Button from 'primevue/button';
@@ -180,6 +181,7 @@ const kernelManager = new KernelSessionManager();
 
 let editor: VAceEditorInstance['_editor'] | null;
 const codeText = ref('');
+const llmQuery = ref('');
 const llmThoughts = ref<any[]>([]);
 
 const sampleAgentQuestions = [
@@ -188,6 +190,11 @@ const sampleAgentQuestions = [
 	'What is cartesian_control in stratify?'
 ];
 
+const updateLlmQuery = (query: string) => {
+	llmThoughts.value = [];
+	llmQuery.value = query;
+};
+
 const updateStratifyGroupForm = (config: StratifyGroup) => {
 	const state = _.cloneDeep(props.node.state);
 	state.strataGroup = config;
@@ -421,20 +428,19 @@ const runCodeStratify = () => {
 	});
 };
 
-// FIXME: Copy pasted in 3 locations, could be written cleaner and in a service
+// FIXME: Copy pasted in 3 locations, could be written cleaner and in a service. Migrate it to use saveCodeToState from @/services/notebook
 const saveCodeToState = (code: string, hasCodeBeenRun: boolean) => {
 	const state = _.cloneDeep(props.node.state);
 	state.hasCodeBeenRun = hasCodeBeenRun;
-
 	// for now only save the last code executed, may want to save all code executed in the future
 	const codeHistoryLength = props.node.state.strataCodeHistory.length;
 	const timestamp = Date.now();
+	const llm = { llmQuery: llmQuery.value, llmThoughts: llmThoughts.value };
 	if (codeHistoryLength > 0) {
-		state.strataCodeHistory[0] = { code, timestamp };
+		state.strataCodeHistory[0] = { code, timestamp, ...llm };
 	} else {
-		state.strataCodeHistory.push({ code, timestamp });
+		state.strataCodeHistory.push({ code, timestamp, ...llm });
 	}
-
 	emit('update-state', state);
 };
 
diff --git a/packages/client/hmi-client/src/components/workflow/tera-workflow.vue b/packages/client/hmi-client/src/components/workflow/tera-workflow.vue
index 2bf631c25c..fdde24f0b7 100644
--- a/packages/client/hmi-client/src/components/workflow/tera-workflow.vue
+++ b/packages/client/hmi-client/src/components/workflow/tera-workflow.vue
@@ -360,7 +360,13 @@ function appendInputPort(
  * */
 function appendOutput(
 	node: WorkflowNode<any> | null,
-	port: { type: string; label?: string; value: any; state?: any; isSelected?: boolean }
+	port: {
+		type: string;
+		label?: string;
+		value: any;
+		state?: any;
+		isSelected?: boolean;
+	}
 ) {
 	if (!node) return;
 
@@ -389,10 +395,26 @@ function appendOutput(
 	node.outputs = node.outputs.filter((d) => d.value);
 
 	selectOutput(node, uuid);
-
+	generateSummary(node, outputPort);
 	workflowDirty = true;
 }
 
+async function generateSummary(node: WorkflowNode<any>, outputPort: WorkflowOutput<any>) {
+	outputPort.summary = ''; // Indicating that the summary generation is initiated
+	const result = await workflowService.generateSummary(
+		node,
+		outputPort,
+		registry.getOperation(node.operationType)?.createNotebook ?? null
+	);
+	if (!result) return;
+	const updateNode = wf.value.nodes.find((n) => n.id === node.id);
+	const updateOutput = (updateNode?.outputs ?? []).find((o) => o.id === outputPort.id);
+	if (!updateNode || !updateOutput) return;
+	updateOutput.summary = result.summary;
+	updateOutput.label = result.title;
+	updateOutputPort(updateNode, updateOutput);
+}
+
 function updateWorkflowNodeState(node: WorkflowNode<any> | null, state: any) {
 	if (!node) return;
 	workflowService.updateNodeState(wf.value, node.id, state);
diff --git a/packages/client/hmi-client/src/services/beaker.ts b/packages/client/hmi-client/src/services/beaker.ts
new file mode 100644
index 0000000000..4c4b22a0d8
--- /dev/null
+++ b/packages/client/hmi-client/src/services/beaker.ts
@@ -0,0 +1,11 @@
+import axios, { AxiosHeaders } from 'axios';
+
+const BeakerAPI = axios.create({
+	baseURL: '/beaker',
+	headers: new AxiosHeaders()
+});
+
+export const summarizeNotebook = async (notebook: any) => {
+	const { data } = await BeakerAPI.post('/summary', { notebook });
+	return data;
+};
diff --git a/packages/client/hmi-client/src/services/notebook.ts b/packages/client/hmi-client/src/services/notebook.ts
index d54e3c6e21..5f8294d7fd 100644
--- a/packages/client/hmi-client/src/services/notebook.ts
+++ b/packages/client/hmi-client/src/services/notebook.ts
@@ -3,13 +3,21 @@ import { WorkflowNode } from '@/types/workflow';
 
 export interface NotebookHistory {
 	code: string;
+	llmQuery: string;
+	llmThoughts: any[];
 	timestamp: number;
 }
 
 // A common pattern used to save code from a notebook within an operator
 // This is ready to be ported to nodes such as tera-model-config and tera-stratify-mira
 // Not ported yet since this will ruin the states of the nodes that already exist due to their differently named properties
-export const saveCodeToState = (node: WorkflowNode<any>, code: string, hasCodeRun: boolean) => {
+export const saveCodeToState = (
+	node: WorkflowNode<any>,
+	code: string,
+	hasCodeRun: boolean,
+	llmQuery: string = '',
+	llmThoughts: any[] = []
+) => {
 	const state = cloneDeep(node.state);
 
 	if (!('notebookHistory' in state) || !('hasCodeRun' in state)) return state;
@@ -19,9 +27,78 @@ export const saveCodeToState = (node: WorkflowNode<any>, code: string, hasCodeRu
 	const notebookHistoryLength = state.notebookHistory.length;
 	const timestamp = Date.now();
 	if (notebookHistoryLength > 0) {
-		state.notebookHistory[0] = { code, timestamp };
+		state.notebookHistory[0] = { code, timestamp, llmQuery, llmThoughts };
 	} else {
-		state.notebookHistory.push({ code, timestamp });
+		state.notebookHistory.push({ code, timestamp, llmQuery, llmThoughts });
 	}
 	return state;
 };
+
+/**
+ * Create a notebook from a code
+ * @param code code to be added to the notebook
+ * @param executeResult output of the code execution if any
+ * @param language language of the code
+ * @param llmQuery llm query used to generate the code if any
+ * @param thought llm thought generated from the query if any
+ * @returns
+ */
+export const createNotebookFromCode = (
+	code: string,
+	language: string,
+	executionResult?: any,
+	llmQuery?: string,
+	llmThoughts: any[] = []
+) => {
+	// TODO: Consider using jataware/beaker-kernel library to generate notebook (https://github.com/jataware/beaker-kernel/blob/886b2b3913ca1460f0301a5cd97cbcf15de609bc/beaker-ts/src/notebook.ts#L414)
+	const notebook = {
+		nbformat: 4,
+		nbformat_minor: 5,
+		cells: [] as any[],
+		metadata: {
+			kernelspec: {
+				display_name: 'Beaker Kernel',
+				name: 'beaker',
+				language: 'beaker'
+			},
+			language_info: {
+				name: language,
+				display_name: language
+			}
+		}
+	};
+	if (llmQuery) {
+		const beakerQueryCell = {
+			cell_type: 'query',
+			events: llmThoughts.map((thought) => ({ type: 'thought', content: thought.content })),
+			metadata: {},
+			source: llmQuery,
+			status: 'idle'
+		};
+		notebook.cells.push(beakerQueryCell);
+	}
+	const beakerCodeCell = {
+		cell_type: 'code',
+		execution_count: 1,
+		metadata: {},
+		outputs: [] as any[],
+		source: code,
+		status: 'idle'
+	};
+	if (executionResult) {
+		// Make a shallow copy of the execution result
+		const data = { ...executionResult };
+		// Make sure the values of the data is stringified as the beaker summary endpoint seem to have issue with object json value
+		Object.keys(data).forEach((type) => {
+			if (typeof data[type] !== 'string') {
+				data[type] = JSON.stringify(data[type]);
+			}
+		});
+		beakerCodeCell.outputs.push({
+			output_type: 'execute_result',
+			data
+		});
+	}
+	notebook.cells.push(beakerCodeCell);
+	return notebook;
+};
diff --git a/packages/client/hmi-client/src/services/workflow.ts b/packages/client/hmi-client/src/services/workflow.ts
index f9b7f3aaa5..78d4e90505 100644
--- a/packages/client/hmi-client/src/services/workflow.ts
+++ b/packages/client/hmi-client/src/services/workflow.ts
@@ -15,6 +15,7 @@ import type {
 	WorkflowOutput
 } from '@/types/workflow';
 import { WorkflowPortStatus, OperatorStatus } from '@/types/workflow';
+import { summarizeNotebook } from './beaker';
 
 /**
  * Captures common actions performed on workflow nodes/edges. The functions here are
@@ -432,12 +433,38 @@ export function selectOutput(
 	cascadeInvalidateDownstream(operator, nodeCache);
 }
 
+export function getActiveOutputSummary(node: WorkflowNode<any>) {
+	const output = node.outputs.find((o) => o.id === node.active);
+	return output?.summary;
+}
+
 export function updateOutputPort(node: WorkflowNode<any>, updatedOutputPort: WorkflowOutput<any>) {
 	let outputPort = node.outputs.find((port) => port.id === updatedOutputPort.id);
 	if (!outputPort) return;
 	outputPort = Object.assign(outputPort, updatedOutputPort);
 }
 
+// Keep track of the summary generation requests to prevent multiple requests for the same workflow output
+const summaryGenerationRequestIds = new Set<string>();
+
+export async function generateSummary(
+	node: WorkflowNode<any>,
+	outputPort: WorkflowOutput<any>,
+	createNotebookFn: ((state: any, value: WorkflowPort['value']) => Promise<any>) | null
+) {
+	if (!node || !createNotebookFn || summaryGenerationRequestIds.has(outputPort.id)) return null;
+	try {
+		summaryGenerationRequestIds.add(outputPort.id);
+		const notebook = await createNotebookFn(outputPort.state, outputPort.value);
+		const result = await summarizeNotebook(notebook);
+		return result;
+	} catch {
+		return { title: outputPort.label, summary: 'Generating AI summary has failed.' };
+	} finally {
+		summaryGenerationRequestIds.delete(outputPort.id);
+	}
+}
+
 // Check if the current-state matches that of the output-state.
 // Note operatorState subsumes the keys of the outputState
 export const isOperatorStateInSync = (
diff --git a/packages/client/hmi-client/src/types/workflow.ts b/packages/client/hmi-client/src/types/workflow.ts
index e6dea5730e..1fbdf0bb4b 100644
--- a/packages/client/hmi-client/src/types/workflow.ts
+++ b/packages/client/hmi-client/src/types/workflow.ts
@@ -68,6 +68,8 @@ export interface Operation {
 
 	inputs: OperationData[];
 	outputs: OperationData[];
+
+	createNotebook?: (state: any, value: WorkflowPort['value']) => Promise<any>;
 }
 
 // Defines the data-exchange between WorkflowNode
@@ -87,6 +89,8 @@ export interface WorkflowPort {
 export interface WorkflowOutput<S> extends WorkflowPort {
 	isSelected?: boolean;
 	operatorStatus?: OperatorStatus;
+	summary?: string;
+	summaryHasBeenEdited?: boolean;
 	state?: Partial<S>;
 	timestamp?: Date;
 }