From 9d324c5ed6d7b09706663038b9154394628fb929 Mon Sep 17 00:00:00 2001
From: Ramiro Medina <64783088+ramedina86@users.noreply.github.com>
Date: Tue, 29 Oct 2024 09:57:29 +0000
Subject: [PATCH] feat: Support tools field

---
 src/ui/src/builder/BuilderApp.vue             |  4 +-
 src/ui/src/builder/BuilderFieldsTools.vue     | 43 ++++++++++--------
 .../workflows/WorkflowsWorkflow.vue           | 18 ++++----
 src/ui/src/renderer/useEvaluator.ts           |  2 +-
 src/writer/workflows_blocks/writerchat.py     | 45 +++++++++----------
 5 files changed, 56 insertions(+), 56 deletions(-)

diff --git a/src/ui/src/builder/BuilderApp.vue b/src/ui/src/builder/BuilderApp.vue
index fee3e0d26..eab0913c7 100644
--- a/src/ui/src/builder/BuilderApp.vue
+++ b/src/ui/src/builder/BuilderApp.vue
@@ -319,8 +319,8 @@ onMounted(() => {
 	--builderSubtleHighlightColor: rgba(0, 0, 0, 0.05);
 	--builderSubtleHighlightColorSolid: #f2f2f2;
 	--builderDisabledColor: rgb(180, 180, 180);
-	--builderSidebarWidth: max(265px, 27vh);
-	--builderSettingsWidth: max(265px, 27vh);
+	--builderSidebarWidth: 265px;
+	--builderSettingsWidth: 265px;
 	--builderActionOngoingColor: #333333;
 	--builderTopBarHeight: 48px;
 	--builderWarningTextColor: white;
diff --git a/src/ui/src/builder/BuilderFieldsTools.vue b/src/ui/src/builder/BuilderFieldsTools.vue
index 67ab803e7..326e23ece 100644
--- a/src/ui/src/builder/BuilderFieldsTools.vue
+++ b/src/ui/src/builder/BuilderFieldsTools.vue
@@ -17,11 +17,7 @@
 			</div>
 		</div>
 
-		<WdsButton
-			variant="builder"
-			size="small"
-			@click="resetAndShowToolFormModal"
-		>
+		<WdsButton size="small" @click="resetAndShowToolFormModal">
 			<i class="material-symbols-outlined">add</i>
 			Add tool</WdsButton
 		>
@@ -45,16 +41,14 @@
 				<WdsTextInput v-model="toolForm.graphId"></WdsTextInput>
 			</template>
 			<div>
-				<WdsButton variant="builder" @click="saveToolForm"
-					>Save</WdsButton
-				>
+				<WdsButton @click="saveToolForm">Save</WdsButton>
 			</div>
 		</BuilderModal>
 	</div>
 </template>
 
 <script setup lang="ts">
-import { toRefs, inject, computed, ref, Ref, ComputedRef } from "vue";
+import { toRefs, inject, computed, ref, Ref } from "vue";
 import { Component, FieldControl } from "@/writerTypes";
 import { useComponentActions } from "./useComponentActions";
 import injectionKeys from "../injectionKeys";
@@ -87,6 +81,7 @@ type Tool = FunctionTool | GraphTool;
 type ToolForm = {
 	isShown: boolean;
 	type: "function" | "graph";
+	originalName?: string;
 	name: string;
 	code: string;
 	graphId: string;
@@ -98,10 +93,10 @@ const { setContentValue } = useComponentActions(wf, ssbm);
 
 const initFunctionToolCode = `
 {
-	"description": "An example tool",
-	"parameters": {
-		"city": string
-	}
+	"description": "Gets info for an employee, given an employee id",
+    "parameters": {
+      "id": {"type": "string", "description": "Id of the employee"}
+    }
 }
 `.trim();
 
@@ -122,7 +117,7 @@ const props = defineProps<{
 const { componentId, fieldKey } = toRefs(props);
 const component = computed(() => wf.getComponentById(componentId.value));
 
-const tools: ComputedRef<Record<string, Tool>> = computed(() => {
+const tools = computed<Record<string, Tool>>(() => {
 	let value = {};
 	try {
 		value = JSON.parse(component.value.content[fieldKey.value]);
@@ -158,11 +153,14 @@ function getToolFromForm(): Tool {
 
 function validateToolForm(form: ToolForm): string[] {
 	let errors = [];
-	const { name } = form;
+	const { originalName, name } = form;
 	if (!name) {
 		errors.push("The name cannot be empty.");
 	}
-	if (Object.keys(tools.value).includes(name)) {
+	if (
+		Object.keys(tools.value).includes(name) &&
+		!(originalName || originalName == name)
+	) {
 		errors.push("An existing tool with the specified name already exists.");
 	}
 	return errors;
@@ -183,6 +181,9 @@ function saveToolForm() {
 	}
 	const newFieldValue = JSON.stringify({
 		...tools.value,
+		...(toolForm.value.originalName
+			? { [toolForm.value.originalName]: undefined }
+			: {}),
 		[toolForm.value.name]: toolFromForm,
 	});
 	setContentValue(component.value.id, fieldKey.value, newFieldValue);
@@ -195,6 +196,7 @@ function getFormFromToolEntry(toolName: string, tool: Tool): ToolForm {
 		return {
 			isShown: true,
 			type: "function",
+			originalName: toolName,
 			name: toolName,
 			graphId: "",
 			code: JSON.stringify(tool, undefined, 2),
@@ -204,6 +206,7 @@ function getFormFromToolEntry(toolName: string, tool: Tool): ToolForm {
 		return {
 			isShown: true,
 			type: "graph",
+			originalName: toolName,
 			name: toolName,
 			graphId: tool.graph_ids?.[0],
 			code: "",
@@ -234,10 +237,12 @@ const customHandlerModalCloseAction: ModalAction = {
 };
 </script>
 
-<style>
-@import "@/renderer/colorTransformations.css";
+<style scoped>
+@import "./sharedStyles.css";
 
-.BuilderFieldsTool {
+.BuilderFieldsTools {
+	--separatorColor: var(--builderSeparatorColor);
+	--intensifiedButtonColor: red;
 }
 
 .tools {
diff --git a/src/ui/src/components/workflows/WorkflowsWorkflow.vue b/src/ui/src/components/workflows/WorkflowsWorkflow.vue
index f3cea9160..edcf51884 100644
--- a/src/ui/src/components/workflows/WorkflowsWorkflow.vue
+++ b/src/ui/src/components/workflows/WorkflowsWorkflow.vue
@@ -154,23 +154,23 @@ const activeCanvasMove: Ref<{
 } | null> = ref(null);
 
 function refreshArrows() {
-	const a = [];
+	const newArrows = [];
 	nodes.value
 		.filter((node) => node.outs?.length > 0)
 		.forEach((node) => {
 			const fromNodeId = node.id;
 			node.outs.forEach((out) => {
-				a.push(
-					calculateArrow(
-						fromNodeId,
-						out.outId,
-						undefined,
-						out.toNodeId,
-					),
+				const arrow = calculateArrow(
+					fromNodeId,
+					out.outId,
+					undefined,
+					out.toNodeId,
 				);
+				if (!arrow) return;
+				newArrows.push(arrow);
 			});
 		});
-	arrows.value = a;
+	arrows.value = newArrows;
 }
 
 const isUnselectable = computed(() => {
diff --git a/src/ui/src/renderer/useEvaluator.ts b/src/ui/src/renderer/useEvaluator.ts
index ca2a1a1a1..5c28f8e8f 100644
--- a/src/ui/src/renderer/useEvaluator.ts
+++ b/src/ui/src/renderer/useEvaluator.ts
@@ -173,7 +173,7 @@ export function useEvaluator(wf: Core) {
 			typeof evaluated == "undefined" ||
 			evaluated === null ||
 			evaluated === "";
-		if (fieldType == FieldType.Object || fieldType == FieldType.KeyValue) {
+		if (fieldType == FieldType.Object || fieldType == FieldType.KeyValue || fieldType == FieldType.Tools) {
 			if (!evaluated) {
 				return JSON.parse(defaultValue ?? null);
 			}
diff --git a/src/writer/workflows_blocks/writerchat.py b/src/writer/workflows_blocks/writerchat.py
index eed2a931b..ac30703d7 100644
--- a/src/writer/workflows_blocks/writerchat.py
+++ b/src/writer/workflows_blocks/writerchat.py
@@ -5,15 +5,6 @@
 
 DEFAULT_MODEL = "palmyra-x-004"
 
-function_tools_init = """{
-  "get_employee_info": {
-    "parameters": {
-      "id": {"type": "float", "description": "Id of the employee"},
-    }
-  }
-
-}"""
-
 class WriterChat(WorkflowBlock):
 
     @classmethod
@@ -41,12 +32,6 @@ def register(cls, type: str):
                         "type": "Number",
                         "default": "0.7"
                     },
-                    "functionTools": {
-                        "name": "Function tools",
-                        "type": "Object",
-                        "default": "{}",
-                        "init": function_tools_init
-                    },
                     "tools": {
                         "name": "Tools",
                         "type": "Tools",
@@ -56,7 +41,7 @@ def register(cls, type: str):
                 },
                 "outs": {
                     "$dynamic": {
-                        "field": "functionTools"
+                        "field": "tools"
                     },
                     "success": {
                         "name": "Success",
@@ -100,17 +85,27 @@ def run(self):
             temperature = float(self._get_field("temperature", False, "0.7"))
             model_id = self._get_field("modelId", False, default_field_value=DEFAULT_MODEL)
             config = { "temperature": temperature, "model": model_id}
-            function_tools_raw = self._get_field("functionTools", True)
+            tools_raw = self._get_field("tools", True)
             tools = []
 
-            for tool_name, tool_raw in function_tools_raw.items():
-                tool = writer.ai.FunctionTool(
-                    type="function",
-                    name=tool_name,
-                    description=tool_raw.get("description"),
-                    callable=self._make_callable(tool_name),
-                    parameters=tool_raw.get("parameters")
-                )
+            for tool_name, tool_raw in tools_raw.items():
+                tool_type = tool_raw.get("type")
+                tool = None
+                if tool_type == "function":
+                    tool = writer.ai.FunctionTool(
+                        type="function",
+                        name=tool_name,
+                        description=tool_raw.get("description"),
+                        callable=self._make_callable(tool_name),
+                        parameters=tool_raw.get("parameters")
+                    )
+                elif tool_type == "graph":
+                    tool = {
+                        "type": "graph",
+                        "graph_ids": tool_raw.get("graph_ids")
+                    }
+                else:
+                    continue
                 tools.append(tool)
 
             conversation = self.evaluator.evaluate_expression(conversation_state_element, self.instance_path, self.execution_env)