Merge pull request #45 from tiwater/revert-44-38-import-custom-agent

Revert 44 38 import custom agent
lyu-ai · Jan 3, 2024 · 14b26fc · 14b26fc
2 parents d90e7f4 + 1aa22ae
commit 14b26fc
Show file tree

Hide file tree

Showing 9 changed files with 62 additions and 29 deletions.
diff --git a/ui/src/app/[locale]/components/Flow.tsx b/ui/src/app/[locale]/components/Flow.tsx
@@ -1,11 +1,17 @@
 import ReactFlow, {
+  Node,
+  Edge,
   Background,
   Controls,
   applyNodeChanges,
   applyEdgeChanges,
+  EdgeChange,
+  NodeChange,
+  SelectionMode,
   useStoreApi,
   useReactFlow,
   addEdge,
+  ConnectionLineType,
 } from 'reactflow';
 import 'reactflow/dist/style.css';
 import { nodeTypes, initialEdges, initialNodes } from '../utils/flow';
@@ -25,8 +31,8 @@ const Autoflow = ({ flowId }: any) => {
   const { flow, updateFlow, isUpdating, isLoading, isError } = useFlow(flowId);
 
   const [mode, setMode] = useState<'main' | 'flow' | 'json' | 'python'>('flow');
-  const [nodes, setNodes] = useState<any[]>(initialNodes);
-  const [edges, setEdges] = useState<any[]>(initialEdges);
+  const [nodes, setNodes] = useState<Node[]>(initialNodes);
+  const [edges, setEdges] = useState<Edge[]>(initialEdges);
   const [isDirty, setIsDirty] = useState<boolean>(false);
   const flowParent = useRef<HTMLDivElement>(null);
   const { fitView, screenToFlowPosition, toObject } = useReactFlow();
@@ -37,7 +43,7 @@ const Autoflow = ({ flowId }: any) => {
   // https://github.com/xyflow/xyflow/issues/3243
   const store = useStoreApi();
   if (process.env.NODE_ENV === 'development') {
-    store.getState().onError = (code: any, message: any) => {
+    store.getState().onError = (code, message) => {
       if (code === '002') {
         return;
       }
@@ -78,7 +84,7 @@ const Autoflow = ({ flowId }: any) => {
   }, [flow, nodes, edges, isDirty, flowId, toObject, debouncedUpdateFlow]);
 
   const onNodesChange = useCallback(
-    (changes: any[]) => {
+    (changes: NodeChange[]) => {
       if (
         !initialLoad.current &&
         changes.some(change => change.type !== 'select')
@@ -90,7 +96,7 @@ const Autoflow = ({ flowId }: any) => {
     [setNodes]
   );
   const onEdgesChange = useCallback(
-    (changes: any[]) => {
+    (changes: EdgeChange[]) => {
       if (
         !initialLoad.current &&
         changes.some(change => change.type !== 'select')
@@ -144,7 +150,7 @@ const Autoflow = ({ flowId }: any) => {
       // Generate a unique node ID
       const newId = genId();
 
-      const newNode = {
+      const newNode: Node = {
         id: newId,
         type: data.type,
         position,
@@ -155,7 +161,7 @@ const Autoflow = ({ flowId }: any) => {
       // Add the new node to the list of nodes in state
       // And clean the previous selections
       setNodes(nds =>
-        nds.map(nd => ({ ...nd, selected: false })).concat(newNode)
+        nds.map(nd => ({ ...nd, selected: false } as Node)).concat(newNode)
       );
     },
     // Specify dependencies for useCallback
@@ -181,7 +187,7 @@ const Autoflow = ({ flowId }: any) => {
     const randInt = (max: number) =>
       Math.floor(Math.random() * Math.floor(max));
 
-    const newNode = {
+    const newNode: Node = {
       id: newId,
       type,
       position: { x: 150 + randInt(100), y: 50 + randInt(50) },
@@ -231,13 +237,13 @@ const Autoflow = ({ flowId }: any) => {
         onEdgesChange={onEdgesChange}
         nodeTypes={nodeTypes}
         onConnect={onConnect}
-        connectionLineType={'smoothstep'}
+        connectionLineType={ConnectionLineType.SmoothStep}
         connectionLineStyle={{ strokeWidth: 2, stroke: 'lightgreen' }}
         onDragOver={onDragOver}
         onDrop={onDrop}
         panOnScroll
         selectionOnDrag
-        selectionMode={'partial'}
+        selectionMode={SelectionMode.Partial}
         fitView
         fitViewOptions={{ maxZoom: 1 }}
         proOptions={{ hideAttribution: true }}

diff --git a/ui/src/app/[locale]/components/node/Assistant.tsx b/ui/src/app/[locale]/components/node/Assistant.tsx
@@ -2,7 +2,7 @@ import clsx from 'clsx';
 import React, { memo, useEffect, useState } from 'react';
 import { FaEye, FaMeta } from 'react-icons/fa6';
 import { GoGear } from 'react-icons/go';
-import { Handle, useReactFlow } from 'reactflow';
+import { Handle, Position, useReactFlow } from 'reactflow';
 import { IconType } from 'react-icons';
 import Toolbar from './Toolbar';
 import { TbBrandOpenai } from 'react-icons/tb';
@@ -175,7 +175,11 @@ function AssistantNode({ id, data, selected }: any) {
         )}
       </div>
 
-      <Handle type="source" position={'right'} className="w-16 !bg-green-600" />
+      <Handle
+        type="source"
+        position={Position.Right}
+        className="w-16 !bg-green-600"
+      />
 
       <AssistantConfig
         show={showOptions}

diff --git a/ui/src/app/[locale]/components/node/CustomConversable.tsx b/ui/src/app/[locale]/components/node/CustomConversable.tsx
@@ -2,7 +2,7 @@ import EditButton from '@/components/EditButton';
 import EditableText from '@/components/EditableText';
 import clsx from 'clsx';
 import { IoExtensionPuzzleSharp } from 'react-icons/io5';
-import { Handle, useReactFlow } from 'reactflow';
+import { Handle, Position, useReactFlow } from 'reactflow';
 import { setNodeData } from '../../utils/flow';
 import Toolbar from './Toolbar';
 import { useState } from 'react';
@@ -75,7 +75,11 @@ const CustomConversable = ({ id, data, selected }: any) => {
         </div>
       </div>
 
-      <Handle type="source" position={'right'} className="w-16 !bg-green-600" />
+      <Handle
+        type="source"
+        position={Position.Right}
+        className="w-16 !bg-green-600"
+      />
     </div>
   );
 };

diff --git a/ui/src/app/[locale]/components/node/GroupChat.tsx b/ui/src/app/[locale]/components/node/GroupChat.tsx
@@ -1,5 +1,5 @@
 import clsx from 'clsx';
-import { Handle, useReactFlow } from 'reactflow';
+import { Handle, Position, useReactFlow } from 'reactflow';
 import { getNodeLabel, setNodeData } from '../../utils/flow';
 import { FaUserGroup } from 'react-icons/fa6';
 import Toolbar from './Toolbar';
@@ -55,8 +55,16 @@ const GroupChatManager = ({ id, selected, data }: any) => {
           />
         </div>
       </div>
-      <Handle type="target" position={'left'} className="w-16 !bg-primary" />
-      <Handle type="source" position={'right'} className="w-16 !bg-primary" />
+      <Handle
+        type="target"
+        position={Position.Left}
+        className="w-16 !bg-primary"
+      />
+      <Handle
+        type="source"
+        position={Position.Right}
+        className="w-16 !bg-primary"
+      />
     </div>
   );
 };

diff --git a/ui/src/app/[locale]/components/node/Toolbar.tsx b/ui/src/app/[locale]/components/node/Toolbar.tsx
@@ -1,7 +1,7 @@
 import clsx from 'clsx';
 import { useTranslations } from 'next-intl';
 import { GoTrash } from 'react-icons/go';
-import { NodeToolbar, useReactFlow } from 'reactflow';
+import { Position, NodeToolbar, useReactFlow } from 'reactflow';
 
 const Toolbar = ({
   nodeId,
@@ -26,7 +26,7 @@ const Toolbar = ({
   return (
     <NodeToolbar
       isVisible={selected}
-      position={'top'}
+      position={Position.Top}
       align="end"
       className={clsx(
         className,

diff --git a/ui/src/app/[locale]/components/node/UserProxy.tsx b/ui/src/app/[locale]/components/node/UserProxy.tsx
@@ -1,6 +1,6 @@
 import clsx from 'clsx';
 import { RiChatSmile2Fill } from 'react-icons/ri';
-import { Handle, useReactFlow } from 'reactflow';
+import { Handle, Position, useReactFlow } from 'reactflow';
 import { getNodeLabel, setNodeData } from '../../utils/flow';
 import Toolbar from './Toolbar';
 import { useState } from 'react';
@@ -126,8 +126,16 @@ const UserProxyAgent = ({ id, selected, data }: any) => {
         </button>
       </div>
 
-      <Handle type="target" position={'left'} className="w-16 !bg-primary" />
-      <Handle type="source" position={'right'} className="w-16 !bg-primary" />
+      <Handle
+        type="target"
+        position={Position.Left}
+        className="w-16 !bg-primary"
+      />
+      <Handle
+        type="source"
+        position={Position.Right}
+        className="w-16 !bg-primary"
+      />
       <UserConfig
         show={showOptions}
         nodeId={id}

diff --git a/ui/src/app/[locale]/gallery/[id]/page.tsx b/ui/src/app/[locale]/gallery/[id]/page.tsx
@@ -15,7 +15,7 @@ const FlowViewer = ({ template, className }: any) => {
   // https://github.com/xyflow/xyflow/issues/3243
   const store = useStoreApi();
   if (process.env.NODE_ENV === 'development') {
-    store.getState().onError = (code: any, message: any) => {
+    store.getState().onError = (code, message) => {
       if (code === '002') {
         return;
       }

diff --git a/ui/src/app/[locale]/utils/flow.ts b/ui/src/app/[locale]/utils/flow.ts
@@ -15,6 +15,7 @@ import {
   RiUser4Line,
   RiUserSearchLine,
 } from 'react-icons/ri';
+import { Edge, Node, ReactFlowInstance } from 'reactflow';
 import { genId } from '@/utils/id';
 import CustomConversable from '../components/node/CustomConversable';
 import { ComponentType } from 'react';
@@ -145,7 +146,7 @@ export const getNodeLabel = (label: string, tNodeMeta: any) => {
 
 // ---------------------
 
-export const initialNodes: any[] = [
+export const initialNodes: Node[] = [
   {
     id: '1',
     type: 'assistant',
@@ -198,7 +199,7 @@ export const initialNodes: any[] = [
   },
 ];
 
-export const initialEdges: any[] = [
+export const initialEdges: Edge[] = [
   {
     id: '1-3',
     source: '1',
@@ -212,7 +213,7 @@ export const initialEdges: any[] = [
 ];
 
 export const setNodeData = (
-  instance: any,
+  instance: ReactFlowInstance,
   id: string,
   dataset: { [key: string]: any },
   scope: string = ''
@@ -233,7 +234,7 @@ export const setNodeData = (
   );
 };
 
-export const getFlowName = (nodes: any[]) => {
+export const getFlowName = (nodes: Node[]) => {
   const configNode = nodes.find((node: any) => node.type === 'config');
   let name = 'flow-unknown';
   if (configNode && configNode?.data.flow_name) {
@@ -244,7 +245,7 @@ export const getFlowName = (nodes: any[]) => {
   return name;
 };
 
-export const getFlowDescription = (nodes: any[]) => {
+export const getFlowDescription = (nodes: Node[]) => {
   const configNode = nodes.find((node: any) => node.type === 'config');
   let description = '';
   if (configNode && configNode?.data.flow_description) {

diff --git a/ui/src/utils/codegen.ts b/ui/src/utils/codegen.ts
@@ -1,3 +1,5 @@
+import { Edge, Node } from 'reactflow';
+
 export enum AgentTypes {
   user_proxy = 'UserProxyAgent',
   assistant = 'AssistantAgent',
@@ -67,7 +69,7 @@ const importDict: Record<string, string> = {
 };
 
 export const genEntry = (
-  data: { nodes: any[]; edges: any[] },
+  data: { nodes: Node[]; edges: Edge[] },
   message: string
 ) => {
   const { nodes, edges } = data;