Feat/admin/auto-add-knowledge-tool (#932)

* feat: add knowledge tool upon knowledge upload in agents Signed-off-by: Ryan Hopper-Lowe <[email protected]> * feat: auto-add knowledge tool upon adding knowledge to a workflow --------- Signed-off-by: Ryan Hopper-Lowe <[email protected]>
obot-platform · Dec 18, 2024 · 932786c · 932786c
1 parent b4c9456
commit 932786c
Show file tree

Hide file tree

Showing 9 changed files with 84 additions and 15 deletions.
diff --git a/ui/admin/app/components/agent/Agent.tsx b/ui/admin/app/components/agent/Agent.tsx
@@ -75,7 +75,8 @@ export function Agent({ className, onRefresh }: AgentProps) {
             updateAgent(updatedAgent);
 
             setAgentUpdates(updatedAgent);
-            setLoadingAgentId(updatedAgent.id);
+
+            if (changes.alias) setLoadingAgentId(changes.alias);
         },
         [agentUpdates, updateAgent, agent]
     );
@@ -138,6 +139,13 @@ export function Agent({ className, onRefresh }: AgentProps) {
                         agentId={agent.id}
                         agent={agent}
                         updateAgent={debouncedSetAgentInfo}
+                        addTool={(tool) => {
+                            if (agent?.tools?.includes(tool)) return;
+
+                            debouncedSetAgentInfo({
+                                tools: [...(agent.tools || []), tool],
+                            });
+                        }}
                     />
                 </div>
             </ScrollArea>

diff --git a/ui/admin/app/components/agent/AgentContext.tsx b/ui/admin/app/components/agent/AgentContext.tsx
@@ -15,7 +15,7 @@ import { useAsync } from "~/hooks/useAsync";
 interface AgentContextType {
     agent: Agent;
     agentId: string;
-    updateAgent: (agent: Agent) => void;
+    updateAgent: (agent: Agent) => Promise<unknown>;
     isUpdating: boolean;
     error?: unknown;
     lastUpdated?: Date;
@@ -59,7 +59,7 @@ export function AgentProvider({
             value={{
                 agentId,
                 agent: getAgent.data ?? agent,
-                updateAgent: updateAgent.execute,
+                updateAgent: updateAgent.executeAsync,
                 isUpdating: updateAgent.isLoading,
                 lastUpdated,
                 error: updateAgent.error,

diff --git a/ui/admin/app/components/agent/ToolForm.tsx b/ui/admin/app/components/agent/ToolForm.tsx
@@ -70,23 +70,33 @@ export function ToolForm({
         resolver: zodResolver(formSchema),
         defaultValues,
     });
+    const { control, handleSubmit, getValues, reset, watch } = form;
+
+    useEffect(() => {
+        const unchanged = compareArrays(
+            defaultValues.tools.map((x) => x.tool),
+            getValues("tools").map((x) => x.tool)
+        );
+
+        if (unchanged) return;
+
+        reset(defaultValues);
+    }, [defaultValues, reset, getValues]);
 
     const toolFields = useFieldArray({
-        control: form.control,
+        control,
         name: "tools",
     });
 
-    const handleSubmit = form.handleSubmit(onSubmit || noop);
-
     useEffect(() => {
-        return form.watch((values) => {
+        return watch((values) => {
             const { data, success } = formSchema.safeParse(values);
 
             if (!success) return;
 
             onChange?.(data);
         }).unsubscribe;
-    }, [form, onChange]);
+    }, [watch, onChange]);
 
     const [allTools, fixedFields, userFields] = useMemo(() => {
         return [
@@ -116,7 +126,10 @@ export function ToolForm({
 
     return (
         <Form {...form}>
-            <form onSubmit={handleSubmit} className="flex flex-col gap-2">
+            <form
+                onSubmit={handleSubmit(onSubmit || noop)}
+                className="flex flex-col gap-2"
+            >
                 <TypographyP className="flex justify-between items-end font-normal">
                     Agent Tools
                 </TypographyP>
@@ -213,3 +226,11 @@ export function ToolForm({
         </Form>
     );
 }
+
+function compareArrays(a: string[], b: string[]) {
+    const aSet = new Set(a);
+
+    if (aSet.size !== b.length) return false;
+
+    return b.every((tool) => aSet.has(tool));
+}
diff --git a/ui/admin/app/components/knowledge/AddSourceModal.tsx b/ui/admin/app/components/knowledge/AddSourceModal.tsx
@@ -1,5 +1,6 @@
 import { FC, useState } from "react";
 
+import { KNOWLEDGE_TOOL } from "~/lib/model/agents";
 import { KnowledgeSourceType } from "~/lib/model/knowledge";
 import { KnowledgeService } from "~/lib/service/api/knowledgeService";
 
@@ -16,6 +17,7 @@ interface AddSourceModalProps {
     isOpen: boolean;
     onOpenChange: (open: boolean) => void;
     onSave: (knowledgeSourceId: string) => void;
+    addTool: (tool: string) => void;
 }
 
 const AddSourceModal: FC<AddSourceModalProps> = ({
@@ -25,6 +27,7 @@ const AddSourceModal: FC<AddSourceModalProps> = ({
     isOpen,
     onOpenChange,
     onSave,
+    addTool,
 }) => {
     const [newWebsite, setNewWebsite] = useState("");
     const [newLink, setNewLink] = useState("");
@@ -68,6 +71,7 @@ const AddSourceModal: FC<AddSourceModalProps> = ({
         } else if (sourceType === KnowledgeSourceType.OneDrive) {
             await handleAddOneDrive();
         }
+        addTool(KNOWLEDGE_TOOL);
         startPolling();
         onOpenChange(false);
     };

diff --git a/ui/admin/app/components/knowledge/AgentKnowledgePanel.tsx b/ui/admin/app/components/knowledge/AgentKnowledgePanel.tsx
@@ -14,7 +14,7 @@ import { useCallback, useEffect, useMemo, useRef, useState } from "react";
 import { $path } from "remix-routes";
 import useSWR, { SWRResponse } from "swr";
 
-import { Agent } from "~/lib/model/agents";
+import { Agent, KNOWLEDGE_TOOL } from "~/lib/model/agents";
 import {
     KnowledgeFile,
     KnowledgeFileState,
@@ -61,12 +61,14 @@ type AgentKnowledgePanelProps = {
     agentId: string;
     agent: Agent;
     updateAgent: (updatedAgent: Agent) => void;
+    addTool: (tool: string) => void;
 };
 
 export default function AgentKnowledgePanel({
     agentId,
     agent,
     updateAgent,
+    addTool,
 }: AgentKnowledgePanelProps) {
     const fileInputRef = useRef<HTMLInputElement>(null);
     const [blockPollingLocalFiles, setBlockPollingLocalFiles] = useState(false);
@@ -210,6 +212,8 @@ export default function AgentKnowledgePanel({
             Array.from(files).map((file) => [file] as const)
         );
 
+        addTool(KNOWLEDGE_TOOL);
+
         if (fileInputRef.current) fileInputRef.current.value = "";
     };
 
@@ -510,10 +514,9 @@ export default function AgentKnowledgePanel({
                                     const res =
                                         await KnowledgeService.createKnowledgeSource(
                                             agentId,
-                                            {
-                                                notionConfig: {},
-                                            }
+                                            { notionConfig: {} }
                                         );
+                                    addTool(KNOWLEDGE_TOOL);
                                     getKnowledgeSources.mutate();
                                     setSelectedKnowledgeSourceId(res.id);
                                     setIsEditKnowledgeSourceModalOpen(true);
@@ -564,6 +567,7 @@ export default function AgentKnowledgePanel({
                     setSelectedKnowledgeSourceId(knowledgeSourceId);
                     setIsEditKnowledgeSourceModalOpen(true);
                 }}
+                addTool={addTool}
             />
             <ErrorDialog
                 error={errorDialogError}

diff --git a/ui/admin/app/components/tools/BasicToolForm.tsx b/ui/admin/app/components/tools/BasicToolForm.tsx
@@ -34,18 +34,30 @@ export function BasicToolForm({
         resolver: zodResolver(formSchema),
         defaultValues: { tools: defaultValues?.tools || [] },
     });
+    const { watch, getValues, reset } = form;
+
+    useEffect(() => {
+        const unchanged = compareArrays(
+            defaultValues?.tools.map(({ value }) => value) || [],
+            getValues().tools.map(({ value }) => value)
+        );
+
+        if (unchanged) return;
+
+        reset({ tools: defaultValues?.tools || [] });
+    }, [defaultValues, getValues, reset]);
 
     const toolArr = useFieldArray({ control: form.control, name: "tools" });
 
     useEffect(() => {
-        return form.watch((values) => {
+        return watch((values) => {
             const { data, success } = formSchema.safeParse(values);
 
             if (!success) return;
 
             onChange?.({ tools: data.tools.map((t) => t.value) });
         }).unsubscribe;
-    }, [form, onChange]);
+    }, [watch, onChange]);
 
     const removeTools = (toolsToRemove: string[]) => {
         const indexes = toolsToRemove
@@ -83,3 +95,11 @@ export function BasicToolForm({
         </Form>
     );
 }
+
+function compareArrays(a: string[], b: string[]) {
+    const aSet = new Set(a);
+
+    if (aSet.size !== b.length) return false;
+
+    return b.every((tool) => aSet.has(tool));
+}
diff --git a/ui/admin/app/components/workflow/Workflow.tsx b/ui/admin/app/components/workflow/Workflow.tsx
@@ -155,6 +155,13 @@ function WorkflowContent({ className, onPersistThreadId }: WorkflowProps) {
                         agent={workflowUpdates}
                         agentId={workflow.id}
                         updateAgent={debouncedSetWorkflowInfo}
+                        addTool={(tool) => {
+                            if (workflow.tools?.includes(tool)) return;
+
+                            debouncedSetWorkflowInfo({
+                                tools: [...(workflow.tools || []), tool],
+                            });
+                        }}
                     />
                 </div>
             </ScrollArea>

diff --git a/ui/admin/app/lib/model/agents.ts b/ui/admin/app/lib/model/agents.ts
@@ -2,6 +2,8 @@ import { EntityMeta } from "~/lib/model/primitives";
 
 // TODO: implement as zod schemas???
 
+export const KNOWLEDGE_TOOL = "knowledge";
+
 export type AgentBase = {
     name: string;
     description: string;

diff --git a/ui/admin/app/lib/utils.ts b/ui/admin/app/lib/utils.ts
@@ -67,3 +67,6 @@ export const getAliasFrom = (text: Nullish<string>) => {
 
     return text.toLowerCase().replace(/[^a-z0-9-]+/g, "-");
 };
+
+export const isNullish = (value: Nullish<unknown>): value is null | undefined =>
+    [null, undefined].includes(value as Todo);