├── .gitignore ├── .prettierrc ├── LICENSE ├── README.md ├── api ├── execute.py └── requirements.txt ├── docs └── source │ └── _static │ └── logo.png ├── index.html ├── package-lock.json ├── package.json ├── public ├── meta-full.png └── meta.jpg ├── src ├── components │ ├── App.tsx │ ├── NodeInfo.tsx │ ├── Prompt.tsx │ ├── modals │ │ ├── APIKeyModal.tsx │ │ └── SettingsModal.tsx │ ├── nodes │ │ ├── CustomNode.tsx │ │ ├── LabelUpdaterNode.tsx │ │ ├── useAnimatedNodes.tsx │ │ └── useExpandCollapse.tsx │ ├── tree.ts │ └── utils │ │ ├── APIKeyInput.tsx │ │ ├── BigButton.tsx │ │ ├── LabeledInputs.tsx │ │ ├── Markdown.tsx │ │ ├── NavigationBar.tsx │ │ └── Whisper.tsx ├── index.css ├── main.tsx ├── types │ └── highlightjs-solidity.d.ts ├── utils │ ├── apikey.ts │ ├── branchesEdge.ts │ ├── branchesNode.ts │ ├── chakra.tsx │ ├── clipboard.ts │ ├── color.ts │ ├── constants.ts │ ├── debounce.ts │ ├── humanEval.ts │ ├── human_eval_problems.json │ ├── llm.ts │ ├── lstore.ts │ ├── mod.ts │ ├── models.ts │ ├── nodeId.ts │ ├── platform.ts │ ├── prompt.ts │ ├── qparams.ts │ ├── rand.ts │ ├── resize.ts │ ├── tot.ts │ └── types.ts └── vite-env.d.ts ├── tsconfig.json ├── tsconfig.node.json ├── vercel.json └── vite.config.ts /.gitignore: -------------------------------------------------------------------------------- 1 | # Logs 2 | logs 3 | *.log 4 | npm-debug.log* 5 | yarn-debug.log* 6 | yarn-error.log* 7 | pnpm-debug.log* 8 | lerna-debug.log* 9 | 10 | node_modules 11 | dist 12 | dist-ssr 13 | *.local 14 | 15 | # Editor directories and files 16 | .vscode/* 17 | !.vscode/extensions.json 18 | .idea 19 | .DS_Store 20 | *.suo 21 | *.ntvs* 22 | *.njsproj 23 | *.sln 24 | *.sw? 25 | .env 26 | .vercel 27 | -------------------------------------------------------------------------------- /.prettierrc: -------------------------------------------------------------------------------- 1 | { 2 | "printWidth": 90 3 | } 4 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 t11s 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 |
2 | 3 | Branches Logo 4 | 5 | # Branches 6 | 7 | Prototype advanced LLM algorithms for reasoning and planning. 8 | 9 | [Try Online](http://code-gen-tree.vercel.app) • 10 | [Report a Bug](https://github.com/normal-computing/branches/issues) • 11 | [Stay tuned](#stay-tuned-for) 12 | 13 |
14 | 15 | ![Branches in action: Tree-search visualization for code generation with self-correction in the HumanEval benchmark](https://storage.googleapis.com/normal-blog-artifacts/systerm2/tot_demo.gif) 16 | 17 | ***Tree-search visualization during code generation.** We visualize a reasoning algorithm which learns from feedback, automatically correcting itself by analyzing error tracebacks to refine its solutions. In this case, we benchmark Python programming problems from the HumanEval dataset.* 18 | 19 | ## About 20 | 21 | Branches is an AI tool for graph-based prototyping of advanced algorithms for LLM reasoning and planning -- like Tree of Thoughts and Reflexion. Branches is adapted from [Flux](https://github.com/paradigmxyz/flux). 22 | 23 | Designed for researchers and developers, it allows users to directly interact with AI reasoning processes, streamlining the exploration of complex coding challenges and strategic problem-solving. 24 | 25 | ### Code Generation (HumanEval) 26 | 27 | Branches automatically expands decision trees to solve programming problems from the [HumanEval dataset](https://huggingface.co/datasets/openai_humaneval), visualizing reasoning chains and facilitating self-correction through error tracebacks. This is found on the `main` branch and is currently hosted. 28 | 29 | ### Game of 24 30 | Branches includes a specialized evaluation mechanism for the [Game of 24 puzzle](https://en.wikipedia.org/wiki/24_(puzzle)), leveraging a scoring system to enhance breadth-first search (BFS) by prioritizing promising paths. This is found on the `game-of-24` branch. 31 | 32 | ## Features 33 | 34 | - [x] 🌳 **Automated Tree Expansion**: Leveraging Tree of Thoughts for dynamic expansion in problem-solving. 35 | - [x] 🧠 **Pre-loaded Prompts**: Curated for search-based reasoning to solve specific problems. 36 | - [x] 💻 **Code Interpretation**: Instant execution and error analysis for self-correcting AI-generated code. 37 | - [x] 🔍 **Scoring Mechanism**: Advanced BFS for the Game of 24 with node evaluation for search optimization. 38 | - [x] 📊 **Interactive Visualization**: Graphical representation of tree searches for easy analysis and education. Largely adapted from [Flux](https://github.com/paradigmxyz/flux). 39 | 40 | ## Usage 41 | 42 | To get started with Branches, you can either visit [code-gen-tree.vercel.app](https://code-gen-tree.vercel.app) for the hosted version or run it locally by following the instructions below. 43 | 44 | ## Deploy to Vercel 45 | ```sh 46 | npm i -g vercel 47 | vercel 48 | ``` 49 | 50 | ## Stay Tuned For 51 | 52 | Our commitment to enhancing Branches continues, with exciting new developments on the way: 53 | 54 | - More reasoning and planning algorithms beyond the defaults ([#10](https://github.com/normal-computing/branches/issues/10)) 55 | - Node Value Editing and Regenerate Subtree Functionality ([#5](https://github.com/normal-computing/branches/issues/5)) 56 | - UI Color Fixes and Customization Features ([#6](https://github.com/normal-computing/branches/issues/6)) 57 | - Address Model/UI Timeout Issues ([#7](https://github.com/normal-computing/branches/issues/7)) 58 | - Enhance Game of 24 Logic, Model Cost Tracking, and Prompt Engineering ([#8](https://github.com/normal-computing/branches/issues/8)) 59 | 60 | ## Contributing 61 | 62 | Your contributions make Branches better. Whether it’s bug reports, new features, or feedback, we welcome it all! Report bugs or request features by creating an issue [here](https://github.com/normal-computing/Branches/issues). 63 | 64 | ## License 65 | 66 | Branches is open-source and continues to uphold the [MIT license](LICENSE). 67 | -------------------------------------------------------------------------------- /api/execute.py: -------------------------------------------------------------------------------- 1 | from http import HTTPStatus 2 | import json 3 | from concurrent.futures import ThreadPoolExecutor 4 | from human_eval.execution import check_correctness 5 | from flask import Flask, request, jsonify 6 | 7 | 8 | app = Flask(__name__) 9 | 10 | executor = ThreadPoolExecutor(max_workers=5) 11 | 12 | 13 | @app.route("/execute", methods=["POST"]) 14 | def execute(): 15 | data = request.json 16 | 17 | problem = data.get("problem", "") 18 | completion = data.get("completion", "") 19 | timeout = data.get("timeout", 5.0) 20 | args = (problem, completion, timeout) 21 | 22 | if not completion: 23 | response = jsonify({"error": "No completion provided"}) 24 | response.status_code = HTTPStatus.BAD_REQUEST 25 | return response 26 | 27 | try: 28 | future = executor.submit(check_correctness, problem, completion, timeout) 29 | result = future.result() 30 | return jsonify({"result": result}) 31 | except Exception as e: 32 | response = jsonify({"error": str(e)}) 33 | response.status_code = HTTPStatus.INTERNAL_SERVER_ERROR 34 | return response 35 | 36 | 37 | # check if a 500 error code is thrown 38 | @app.errorhandler(500) 39 | def internal_error(error): 40 | return "500 error: {}".format(str(error)), 500 41 | -------------------------------------------------------------------------------- /api/requirements.txt: -------------------------------------------------------------------------------- 1 | openai 2 | git+https://github.com/arunpatro/human-eval.git@pipgit 3 | Flask -------------------------------------------------------------------------------- /docs/source/_static/logo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/normal-computing/branches/eb7111da6edb7a762bfbe4bdad79b6978890657c/docs/source/_static/logo.png -------------------------------------------------------------------------------- /index.html: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | Branches 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 |
31 | 32 | 33 | 34 | -------------------------------------------------------------------------------- /package.json: -------------------------------------------------------------------------------- 1 | { 2 | "name": "branches", 3 | "private": true, 4 | "version": "0.0.0", 5 | "scripts": { 6 | "dev": "vite", 7 | "build": "tsc && vite build", 8 | "preview": "vite preview" 9 | }, 10 | "dependencies": { 11 | "@chakra-ui/icons": "^2.0.17", 12 | "@chakra-ui/react": "^2.5.5", 13 | "@emotion/react": "^11.10.6", 14 | "@emotion/styled": "^11.10.6", 15 | "d3-hierarchy": "^3.1.2", 16 | "d3-timer": "^3.0.1", 17 | "framer-motion": "^9.0.4", 18 | "highlightjs-solidity": "^2.0.6", 19 | "js-tiktoken": "^1.0.7", 20 | "mathjs": "^11.11.0", 21 | "mixpanel-browser": "^2.46.0", 22 | "nunjucks": "^3.2.4", 23 | "openai": "^4.5.0", 24 | "openai-streams": "^4.2.0", 25 | "re-resizable": "^6.9.9", 26 | "react": "^18.2.0", 27 | "react-beforeunload": "^2.5.3", 28 | "react-dom": "^18.2.0", 29 | "react-hotkeys-hook": "^4.3.7", 30 | "react-icons": "^4.11.0", 31 | "react-markdown": "^8.0.6", 32 | "react-textarea-autosize": "^8.4.0", 33 | "reactflow": "^11.9.4", 34 | "rehype-highlight": "^6.0.0", 35 | "yield-stream": "^2.3.0" 36 | }, 37 | "devDependencies": { 38 | "@types/mixpanel-browser": "^2.38.1", 39 | "@types/node": "^18.14.2", 40 | "@types/nunjucks": "^3.2.5", 41 | "@types/react": "^18.0.27", 42 | "@types/react-beforeunload": "^2.1.1", 43 | "@types/react-dom": "^18.0.10", 44 | "@vitejs/plugin-react-swc": "^3.0.0", 45 | "typescript": "^5.1.6", 46 | "vite": "^5.3.3" 47 | } 48 | } 49 | -------------------------------------------------------------------------------- /public/meta-full.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/normal-computing/branches/eb7111da6edb7a762bfbe4bdad79b6978890657c/public/meta-full.png -------------------------------------------------------------------------------- /public/meta.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/normal-computing/branches/eb7111da6edb7a762bfbe4bdad79b6978890657c/public/meta.jpg -------------------------------------------------------------------------------- /src/components/App.tsx: -------------------------------------------------------------------------------- 1 | import { MIXPANEL_TOKEN } from "../main"; 2 | import { isValidAPIKey } from "../utils/apikey"; 3 | import { Column, Row } from "../utils/chakra"; 4 | import { 5 | API_KEY_LOCAL_STORAGE_KEY, 6 | DEFAULT_SETTINGS, 7 | FIT_VIEW_SETTINGS, 8 | MODEL_SETTINGS_LOCAL_STORAGE_KEY, 9 | REACT_FLOW_NODE_TYPES, 10 | REACT_FLOW_LOCAL_STORAGE_KEY, 11 | TOAST_CONFIG, 12 | SAVED_CHAT_SIZE_LOCAL_STORAGE_KEY, 13 | } from "../utils/constants"; 14 | import { useDebouncedEffect } from "../utils/debounce"; 15 | import { newBranchesEdge } from "../utils/branchesEdge"; 16 | import { 17 | getBranchesNode, 18 | newBranchesNode, 19 | appendTextToBranchesNodeAsGPT, 20 | getBranchesNodeLineage, 21 | modifyBranchesNodeText, 22 | markOnlyNodeAsSelected, 23 | getConnectionAllowed, 24 | } from "../utils/branchesNode"; 25 | import { useLocalStorage } from "../utils/lstore"; 26 | import { getAvailableChatModels } from "../utils/models"; 27 | import { generateNodeId, generateStreamId } from "../utils/nodeId"; 28 | import { 29 | explanationMessage, 30 | humanEvalMessageFromNode, 31 | regenMessage, 32 | } from "../utils/prompt"; 33 | import { resetURL } from "../utils/qparams"; 34 | import { useDebouncedWindowResize } from "../utils/resize"; 35 | import { 36 | ToTNodeData, 37 | BranchesNodeType, 38 | Settings, 39 | HumanEvalProblemsType, 40 | } from "../utils/types"; 41 | import { NodeInfo } from "./NodeInfo"; 42 | import { APIKeyModal } from "./modals/APIKeyModal"; 43 | import { SettingsModal } from "./modals/SettingsModal"; 44 | import { NavigationBar } from "./utils/NavigationBar"; 45 | import { CheckCircleIcon } from "@chakra-ui/icons"; 46 | import { Box, useDisclosure, Spinner, useToast } from "@chakra-ui/react"; 47 | import mixpanel from "mixpanel-browser"; 48 | import { OpenAI } from "openai-streams"; 49 | import { Resizable } from "re-resizable"; 50 | import { useEffect, useState, useCallback, useRef } from "react"; 51 | import { useBeforeunload } from "react-beforeunload"; 52 | import rawHumanEvalProblems from "../utils/human_eval_problems.json"; 53 | import useExpandCollapse from "./nodes/useExpandCollapse"; 54 | import useAnimatedNodes from "./nodes/useAnimatedNodes"; 55 | 56 | import ReactFlow, { 57 | addEdge, 58 | Background, 59 | Connection, 60 | Node, 61 | Edge, 62 | NodeMouseHandler, 63 | useEdgesState, 64 | useNodesState, 65 | SelectionMode, 66 | ReactFlowInstance, 67 | ReactFlowJsonObject, 68 | useReactFlow, 69 | updateEdge, 70 | } from "reactflow"; 71 | import "reactflow/dist/style.css"; 72 | import { yieldStream } from "yield-stream"; 73 | import { treeDemo } from "./tree"; 74 | import { getBranchesNodeColor } from "../utils/color"; 75 | import { getEncoding, encodingForModel } from "js-tiktoken"; 76 | 77 | const HUMAN_EVAL_PROBLEMS = rawHumanEvalProblems as HumanEvalProblemsType; 78 | 79 | function App() { 80 | const toast = useToast(); 81 | 82 | /*////////////////////////////////////////////////////////////// 83 | CORE REACT FLOW LOGIC 84 | //////////////////////////////////////////////////////////////*/ 85 | 86 | type NodeWithText = { 87 | node: Node; 88 | text: string; 89 | }; 90 | 91 | const { setViewport, fitView } = useReactFlow(); 92 | 93 | const [reactFlow, setReactFlow] = useState(null); 94 | 95 | const [nodes, setNodes, onNodesChange] = useNodesState([]); 96 | const [edges, setEdges, onEdgesChange] = useEdgesState([]); 97 | 98 | const treeWidth: number = 220; 99 | const treeHeight: number = 150; 100 | const animationDuration: number = 200; 101 | 102 | const { nodes: visibleNodes, edges: visibleEdges } = useExpandCollapse(nodes, edges, { 103 | treeWidth, 104 | treeHeight, 105 | }); 106 | const { nodes: animatedNodes } = useAnimatedNodes(visibleNodes, { animationDuration }); 107 | 108 | const [filteredNodes, setFilteredNodes] = useState([]); 109 | const [showAnswerPathOnly, setShowAnswerPathOnly] = useState(false); 110 | 111 | const [inputTokenCount, setInputTokenCount] = useState(0); 112 | const [outputTokenCount, setOutputTokenCount] = useState(0); 113 | 114 | const edgeUpdateSuccessful = useRef(true); 115 | 116 | const onEdgeUpdateStart = useCallback(() => { 117 | edgeUpdateSuccessful.current = false; 118 | }, []); 119 | 120 | const onEdgeUpdate = (oldEdge: Edge, newConnection: Connection) => { 121 | if ( 122 | !getConnectionAllowed(nodes, edges, { 123 | source: newConnection.source!, 124 | target: newConnection.target!, 125 | }) 126 | ) 127 | return; 128 | 129 | edgeUpdateSuccessful.current = true; 130 | 131 | setEdges((edges) => updateEdge(oldEdge, newConnection, edges)); 132 | }; 133 | 134 | const onEdgeUpdateEnd = (_: unknown, edge: Edge) => { 135 | if (!edgeUpdateSuccessful.current) { 136 | setEdges((edges) => edges.filter((e) => e.id !== edge.id)); 137 | } 138 | 139 | edgeUpdateSuccessful.current = true; 140 | }; 141 | 142 | const onConnect = (connection: Edge | Connection) => { 143 | if ( 144 | !getConnectionAllowed(nodes, edges, { 145 | source: connection.source!, 146 | target: connection.target!, 147 | }) 148 | ) 149 | return; 150 | 151 | setEdges((eds) => addEdge({ ...connection }, eds)); 152 | }; 153 | 154 | const autoZoom = () => setTimeout(() => fitView(FIT_VIEW_SETTINGS), 50); 155 | 156 | const autoZoomIfNecessary = () => { 157 | if (settings.autoZoom) autoZoom(); 158 | }; 159 | 160 | const save = () => { 161 | if (reactFlow) { 162 | localStorage.setItem( 163 | REACT_FLOW_LOCAL_STORAGE_KEY, 164 | JSON.stringify(reactFlow.toObject()) 165 | ); 166 | } 167 | }; 168 | 169 | // Auto save. 170 | const isSavingReactFlow = useDebouncedEffect( 171 | save, 172 | 1000, // 1 second. 173 | [reactFlow, nodes, edges] 174 | ); 175 | 176 | // Auto restore on load. 177 | useEffect(() => { 178 | if (reactFlow) { 179 | // const rawFlow = undefined; 180 | 181 | // const flow: ReactFlowJsonObject = rawFlow ? JSON.parse(rawFlow) : null; 182 | const flow: ReactFlowJsonObject = treeDemo; 183 | 184 | if (flow !== null) { 185 | setEdges(flow.edges || []); 186 | setViewport(flow.viewport); 187 | 188 | const nodes = flow.nodes; // For brevity. 189 | 190 | if (nodes.length > 0) { 191 | // Either the first selected node we find, or the first node in the array. 192 | const toSelect = nodes.find((node) => node.selected)?.id ?? nodes[0].id; 193 | 194 | // Add the nodes to the React Flow array and select the node. 195 | selectNode(toSelect, () => nodes); 196 | 197 | // If there was a newTreeWith query param, create a new tree with that content. 198 | // We pass false for forceAutoZoom because we'll do it 500ms later to avoid lag. 199 | } 200 | } 201 | 202 | setTimeout(() => { 203 | // Do this with a more generous timeout to make sure 204 | // the nodes are rendered and the settings have loaded in. 205 | if (settings.autoZoom) fitView(FIT_VIEW_SETTINGS); 206 | }, 500); 207 | 208 | resetURL(); // Get rid of the query params. 209 | } 210 | }, [reactFlow]); 211 | 212 | /*////////////////////////////////////////////////////////////// 213 | AI PROMPT CALLBACKS 214 | //////////////////////////////////////////////////////////////*/ 215 | 216 | // Takes a prompt, submits it to the GPT API with n responses, 217 | // then creates a child node for each response under the selected node. 218 | const submitPrompt = async () => { 219 | const temp = settings.temp; 220 | const model = settings.model; 221 | const parentNode = selectedNodeLineage[0]; 222 | const submittedNode = getBranchesNode(nodes, parentNode.id)!; 223 | 224 | console.log("current node", submittedNode); 225 | 226 | type SetNodes = React.Dispatch>; 227 | type SetEdges = React.Dispatch>; 228 | type BranchesNodeInput = { 229 | id?: string; 230 | x: number; 231 | y: number; 232 | branchesNodeType: BranchesNodeType; 233 | input: string; 234 | text: string; 235 | streamId?: string; 236 | steps: string[]; 237 | solutions: any[]; 238 | errors: any[]; 239 | style: any; 240 | explanations: any[]; 241 | }; 242 | type BranchesEdgeInput = { 243 | source: string; 244 | target: string; 245 | animated: boolean; 246 | } 247 | 248 | const createNewNodeAndEdge = ( 249 | currentNode: Node, 250 | newBranchesNode: (node: BranchesNodeInput) => Node, 251 | newBranchesEdge: (node: BranchesEdgeInput) => Edge, 252 | setNodes: SetNodes, 253 | setEdges: SetEdges, 254 | streamId: string, 255 | isSolutionNode: boolean, 256 | callback: (newNode: Node) => void 257 | ) => { 258 | const currentChildNodeId = generateNodeId(); 259 | 260 | setNodes((prevNodes: Node[]) => { 261 | const matchingNode = prevNodes.find((n) => n.id === currentNode.id); 262 | if (!matchingNode) { 263 | throw new Error("Node not found"); 264 | } 265 | 266 | // Create a new node using the currentErrors 267 | const newNode = newBranchesNode({ 268 | id: currentChildNodeId, 269 | x: matchingNode.position.x + 10, 270 | y: matchingNode.position.y + 100, 271 | branchesNodeType: BranchesNodeType.GPT, 272 | input: matchingNode.data.input, 273 | text: "", 274 | streamId, 275 | steps: [...matchingNode.data.steps, ""], 276 | solutions: isSolutionNode 277 | ? [...matchingNode.data.solutions, ""] 278 | : [...matchingNode.data.solutions], 279 | style: { background: getBranchesNodeColor(!isSolutionNode, true, false, true, 0) }, 280 | errors: [...matchingNode.data.errors], 281 | explanations: isSolutionNode 282 | ? [...matchingNode.data.explanations] 283 | : [...matchingNode.data.explanations, ""], 284 | }); 285 | 286 | callback(newNode); 287 | 288 | return [...prevNodes, newNode]; 289 | }); 290 | 291 | setEdges((prevEdges) => [ 292 | ...prevEdges, 293 | newBranchesEdge({ 294 | source: currentNode.id, 295 | target: currentChildNodeId, 296 | animated: true, 297 | }), 298 | ]); 299 | 300 | setTimeout(autoZoomIfNecessary, 500); 301 | }; 302 | 303 | const updateNodeColor = ( 304 | nodeId: string, 305 | setNodes: SetNodes, 306 | isExplanation?: boolean 307 | ) => { 308 | setNodes((prevNodes: Node[]) => { 309 | const newNodes = prevNodes.map((node) => { 310 | if (node.id === nodeId) { 311 | console.log(node.data.score) 312 | return { 313 | ...node, 314 | style: { 315 | background: getBranchesNodeColor( 316 | isExplanation || false, 317 | false, 318 | node.data.isTerminal || false, 319 | !node.data.errors || node.data.errors.length == 0, 320 | node.data.score || 0 321 | ), 322 | }, 323 | }; 324 | } 325 | return node; 326 | }); 327 | return newNodes; 328 | }); 329 | }; 330 | 331 | const updatePreviousEdge = (currentChildNodeId: string, setEdges: SetEdges) => { 332 | setEdges((prevEdges: Edge[]) => { 333 | return prevEdges.map((edge) => { 334 | if (edge.target === currentChildNodeId) { 335 | return { ...edge, animated: false }; 336 | } 337 | return edge; 338 | }); 339 | }); 340 | }; 341 | 342 | // need to modify to use model name, currenlty defining enc in function is very slow 343 | const enc = encodingForModel("gpt-3.5-turbo"); 344 | function countTokens(text: string): number { 345 | const tokens = enc.encode(text); 346 | return tokens.length; 347 | } 348 | 349 | const addError = (nodeId: string, error: string, setNodes: SetNodes) => { 350 | setNodes((prevNodes: Node[]) => { 351 | const newNodes = prevNodes.map((node) => { 352 | if (node.id === nodeId) { 353 | const existingErrors = node.data.errors || []; // Initialize to empty array if it doesn't exist 354 | return { 355 | ...node, 356 | data: { 357 | ...node.data, 358 | errors: [...existingErrors, error], // Append the new error to the existing array 359 | }, 360 | }; 361 | } 362 | return node; 363 | }); 364 | return newNodes; 365 | }); 366 | }; 367 | 368 | async function executeInterpreter( 369 | node: Node, 370 | solutionText: string, 371 | finalNode: boolean 372 | ): Promise { 373 | let data = { 374 | problem: HUMAN_EVAL_PROBLEMS[node.data.input], 375 | completion: solutionText, 376 | }; 377 | 378 | console.log("node solution text", solutionText); 379 | console.log("data", JSON.stringify(data)); 380 | 381 | let url = "/execute"; 382 | let response = await fetch(url, { 383 | method: "POST", 384 | headers: { 385 | "Content-Type": "application/json", 386 | }, 387 | body: JSON.stringify(data), 388 | }); 389 | 390 | // Parse JSON response 391 | let jsonResponse = await response.json(); 392 | console.log("json response", jsonResponse); 393 | 394 | const passed = jsonResponse["result"]["passed"]; 395 | console.log("passed", passed); 396 | 397 | if (passed) { 398 | handleFinishedNode(node, true, false); 399 | return null; 400 | } else { 401 | const error = jsonResponse["result"]["result"]; 402 | addError(node.id, error, setNodes); 403 | updateNodeColor(node.id, setNodes); 404 | 405 | if (!finalNode) { 406 | const explanationPromises = Array(settings.N_EXPLANATION_FANOUT) 407 | .fill(null) 408 | .map(async () => { 409 | return await generateChild(node, "explanation", error, false); 410 | }); 411 | 412 | const explanationChildrenWithText: NodeWithText[] = await Promise.all( 413 | explanationPromises 414 | ); 415 | 416 | const regenChildrenPromises: Promise[] = 417 | explanationChildrenWithText.map(async (explanationChildWithText) => { 418 | // Create N_ANSWER_FANOUT number of promises for each explanation child 419 | const regenPromises = Array(settings.N_ANSWER_FANOUT) 420 | .fill(null) 421 | .map(async () => { 422 | // Assuming that `error` is available in the current scope 423 | return await generateChild( 424 | explanationChildWithText.node, 425 | "regen", 426 | error, 427 | true 428 | ); 429 | }); 430 | 431 | // Await all regenPromises for the current explanation child 432 | return await Promise.all(regenPromises); 433 | }); 434 | 435 | // Await all regenChildrenPromises for all explanation children 436 | const regenChildrenArrays: NodeWithText[][] = await Promise.all( 437 | regenChildrenPromises 438 | ); 439 | 440 | // Flatten the array of arrays into a single array 441 | const regenChildrenWithText: NodeWithText[] = ([] as NodeWithText[]).concat(...regenChildrenArrays); 442 | 443 | return regenChildrenWithText; 444 | } 445 | return null; 446 | } 447 | } 448 | 449 | const markAsAnswerPath = ( 450 | targetNodeId: string, 451 | setNodes: SetNodes, 452 | setEdges: SetEdges 453 | ) => { 454 | setEdges((prevEdges) => { 455 | const edges = [...prevEdges]; // Make a shallow copy for reference 456 | console.log("Edges are:", edges); 457 | 458 | setNodes((prevNodes) => { 459 | const markNodeAndAncestors = (nodeId: string, nodes: Node[]) => { 460 | let updatedNodes: Node[] = []; 461 | 462 | const nodeToUpdate = nodes.find((node) => node.id === nodeId); 463 | if (nodeToUpdate) { 464 | const updatedNode = { 465 | ...nodeToUpdate, 466 | data: { ...nodeToUpdate.data, isInAnswerPath: true }, 467 | }; 468 | updatedNodes.push(updatedNode); 469 | } 470 | 471 | edges.forEach((edge) => { 472 | if (edge.target === nodeId) { 473 | updatedNodes = [ 474 | ...updatedNodes, 475 | ...markNodeAndAncestors(edge.source, nodes), 476 | ]; 477 | } 478 | }); 479 | 480 | return updatedNodes; 481 | }; 482 | 483 | const nodesToUpdate = markNodeAndAncestors(targetNodeId, prevNodes); 484 | return prevNodes.map((node) => { 485 | const nodeToUpdate = nodesToUpdate.find((n) => n.id === node.id); 486 | return nodeToUpdate || node; 487 | }); 488 | }); 489 | 490 | return edges; // return the edges as-is since we're not modifying them 491 | }); 492 | }; 493 | 494 | async function handleFinishedNode( 495 | finishedNode: Node, 496 | isTerminal: boolean, 497 | isExplanation: boolean 498 | ): Promise> { 499 | console.log("handling node", finishedNode); 500 | console.log("is explanation?", isExplanation); 501 | let modifiedNode = { ...finishedNode }; 502 | if (isTerminal) { 503 | console.log("found terminal node"); 504 | markAsAnswerPath(finishedNode.id, setNodes, setEdges); 505 | setNodes((prevNodes: Node[]) => { 506 | const newNodes = prevNodes.map((node) => { 507 | if (node.id === finishedNode?.id) { 508 | modifiedNode = { 509 | ...node, 510 | style: { 511 | background: getBranchesNodeColor( 512 | isExplanation, 513 | false, 514 | isTerminal, 515 | true, 516 | finishedNode.data.score || 0 517 | ), 518 | }, 519 | data: { 520 | ...node.data, 521 | isTerminal: true, 522 | }, 523 | }; 524 | return modifiedNode; 525 | } 526 | return node; 527 | }); 528 | return newNodes; 529 | }); 530 | } 531 | 532 | updateNodeColor(finishedNode?.id!, setNodes, isExplanation); 533 | updatePreviousEdge(finishedNode?.id!, setEdges); 534 | 535 | return modifiedNode; 536 | } 537 | 538 | async function generateChild( 539 | node: Node, 540 | nodeType: string, 541 | error: string, 542 | isSolutionNode: boolean 543 | ): Promise { 544 | console.log("generating from this node", node); 545 | const DECODER = new TextDecoder(); 546 | 547 | const abortController = new AbortController(); 548 | 549 | const streamId = generateStreamId(); 550 | console.log("new stream id", streamId); 551 | let isNewNode = true; 552 | 553 | const question = HUMAN_EVAL_PROBLEMS[node.data.input]["prompt"]; 554 | let answer = node.data.steps[0]; 555 | console.log("this is the node we're generating from", node); 556 | let explanation = ""; 557 | if (nodeType == "regen") { 558 | explanation = node.data.steps[1]; 559 | } 560 | 561 | const messages = 562 | nodeType == "explanation" 563 | ? explanationMessage(question, answer, error) 564 | : nodeType == "regen" 565 | ? regenMessage(question, answer, error, explanation) 566 | : humanEvalMessageFromNode(node); 567 | const newInputTokens = countTokens(messages[0]["content"]); 568 | setInputTokenCount((prevCount) => prevCount + newInputTokens); 569 | 570 | const stream = await OpenAI( 571 | "chat", 572 | { 573 | model, 574 | temperature: temp, 575 | messages, 576 | }, 577 | { apiKey: apiKey!, mode: "raw" } 578 | ); 579 | let currentText: string = ""; 580 | let currentChildNode: Node | null = null; 581 | 582 | for await (const chunk of yieldStream(stream, abortController)) { 583 | if (abortController.signal.aborted) break; 584 | 585 | try { 586 | const decoded = JSON.parse(DECODER.decode(chunk)); 587 | const choice = decoded.choices[0]; 588 | 589 | if (choice.delta?.content) { 590 | const chars = choice.delta.content; 591 | const newTokens = countTokens(chars); 592 | setOutputTokenCount((prevCount) => prevCount + newTokens); 593 | 594 | // new node 595 | if (isNewNode) { 596 | createNewNodeAndEdge( 597 | node, 598 | newBranchesNode, 599 | newBranchesEdge, 600 | setNodes, 601 | setEdges, 602 | streamId, 603 | isSolutionNode, 604 | (newNode) => { 605 | currentChildNode = newNode; 606 | } 607 | ); 608 | isNewNode = false; 609 | } 610 | currentText += chars; 611 | 612 | setNodes((prevNodes: Node[]) => { 613 | return appendTextToBranchesNodeAsGPT( 614 | prevNodes, 615 | { 616 | id: currentChildNode?.id!, 617 | text: currentText, 618 | streamId, 619 | }, 620 | isSolutionNode 621 | ); 622 | }); 623 | 624 | // We cannot return within the loop, and we do 625 | // not want to execute the code below, so we break. 626 | if (abortController.signal.aborted) break; 627 | } 628 | } catch (err) { 629 | console.error(err); 630 | } 631 | } 632 | 633 | const finalChild: Node = await handleFinishedNode( 634 | currentChildNode!, 635 | false, 636 | nodeType == "explanation" 637 | ); 638 | return { node: finalChild, text: currentText }; 639 | } 640 | 641 | const promises = Array(settings.N_ANSWER_FANOUT) 642 | .fill(null) 643 | .map(async () => { 644 | return await generateChild(submittedNode, "normal", "", true); 645 | }); 646 | 647 | const childrenWithText = await Promise.all(promises); 648 | 649 | autoZoomIfNecessary(); 650 | 651 | const interpretChildrenPromises = childrenWithText.map(async (childWithText) => { 652 | return await executeInterpreter(childWithText.node, childWithText.text, false); 653 | }); 654 | 655 | const regenChildren = await Promise.all(interpretChildrenPromises); 656 | 657 | autoZoomIfNecessary(); 658 | 659 | if (regenChildren) { 660 | const combinedRegenChildren = regenChildren.flatMap((childArray) => 661 | childArray !== null ? childArray : [] 662 | ); 663 | 664 | combinedRegenChildren.map(async (regenChild: NodeWithText) => { 665 | return await executeInterpreter(regenChild.node, regenChild.text, true); 666 | }); 667 | } 668 | 669 | autoZoomIfNecessary(); 670 | 671 | if (MIXPANEL_TOKEN) mixpanel.track("Submitted Prompt"); // KPI 672 | }; 673 | 674 | /*////////////////////////////////////////////////////////////// 675 | SELECTED NODE LOGIC 676 | //////////////////////////////////////////////////////////////*/ 677 | 678 | const [selectedNodeId, setSelectedNodeId] = useState(null); 679 | 680 | const selectedNodeLineage = 681 | selectedNodeId !== null ? getBranchesNodeLineage(nodes, edges, selectedNodeId) : []; 682 | 683 | /*////////////////////////////////////////////////////////////// 684 | NODE MUTATION CALLBACKS 685 | //////////////////////////////////////////////////////////////*/ 686 | 687 | /*////////////////////////////////////////////////////////////// 688 | NODE SELECTION CALLBACKS 689 | //////////////////////////////////////////////////////////////*/ 690 | 691 | const selectNode = ( 692 | id: string, 693 | computeNewNodes?: (currNodes: Node[]) => Node[] 694 | ) => { 695 | setSelectedNodeId(id); 696 | setNodes((currNodes) => 697 | // If we were passed a computeNewNodes function, use it, otherwise just use the current nodes. 698 | markOnlyNodeAsSelected(computeNewNodes ? computeNewNodes(currNodes) : currNodes, id) 699 | ); 700 | }; 701 | 702 | /*////////////////////////////////////////////////////////////// 703 | SETTINGS MODAL LOGIC 704 | //////////////////////////////////////////////////////////////*/ 705 | 706 | const { 707 | isOpen: isSettingsModalOpen, 708 | onOpen: onOpenSettingsModal, 709 | onClose: onCloseSettingsModal, 710 | } = useDisclosure(); 711 | 712 | const [settings, setSettings] = useState(() => { 713 | const rawSettings = localStorage.getItem(MODEL_SETTINGS_LOCAL_STORAGE_KEY); 714 | 715 | if (rawSettings !== null) { 716 | return JSON.parse(rawSettings) as Settings; 717 | } else { 718 | return DEFAULT_SETTINGS; 719 | } 720 | }); 721 | 722 | const isGPT4 = settings.model.includes("gpt-4"); 723 | 724 | // Auto save. 725 | const isSavingSettings = useDebouncedEffect( 726 | () => { 727 | localStorage.setItem(MODEL_SETTINGS_LOCAL_STORAGE_KEY, JSON.stringify(settings)); 728 | }, 729 | 1000, // 1 second. 730 | [settings] 731 | ); 732 | 733 | /*////////////////////////////////////////////////////////////// 734 | API KEY LOGIC 735 | //////////////////////////////////////////////////////////////*/ 736 | 737 | const [apiKey, setApiKey] = useLocalStorage(API_KEY_LOCAL_STORAGE_KEY); 738 | 739 | const [availableModels, setAvailableModels] = useState(null); 740 | 741 | // modelsLoadCounter lets us discard the results of the requests if a concurrent newer one was made. 742 | const modelsLoadCounter = useRef(0); 743 | useEffect(() => { 744 | if (isValidAPIKey(apiKey)) { 745 | const modelsLoadIndex = modelsLoadCounter.current + 1; 746 | modelsLoadCounter.current = modelsLoadIndex; 747 | 748 | setAvailableModels(null); 749 | 750 | (async () => { 751 | let modelList: string[] = []; 752 | try { 753 | modelList = await getAvailableChatModels(apiKey!); 754 | } catch (e) { 755 | toast({ 756 | title: "Failed to load model list!", 757 | status: "error", 758 | ...TOAST_CONFIG, 759 | }); 760 | } 761 | if (modelsLoadIndex !== modelsLoadCounter.current) return; 762 | 763 | if (modelList.length === 0) modelList.push(settings.model); 764 | 765 | setAvailableModels(modelList); 766 | 767 | if (!modelList.includes(settings.model)) { 768 | const oldModel = settings.model; 769 | const newModel = modelList.includes(DEFAULT_SETTINGS.model) 770 | ? DEFAULT_SETTINGS.model 771 | : modelList[0]; 772 | 773 | setSettings((settings) => ({ ...settings, model: newModel })); 774 | 775 | toast({ 776 | title: `Model "${oldModel}" no longer available!`, 777 | description: `Switched to "${newModel}"`, 778 | status: "warning", 779 | ...TOAST_CONFIG, 780 | }); 781 | } 782 | })(); 783 | } 784 | }, [apiKey]); 785 | 786 | useEffect(() => { 787 | const updatedNodes: Node[] = animatedNodes.filter( 788 | (node) => node.data?.isInAnswerPath 789 | ); 790 | setFilteredNodes(updatedNodes); 791 | }, [nodes]); 792 | 793 | const isAnythingSaving = isSavingReactFlow || isSavingSettings; 794 | const isAnythingLoading = isAnythingSaving || availableModels === null; 795 | 796 | useBeforeunload((event: BeforeUnloadEvent) => { 797 | // Prevent leaving the page before saving. 798 | if (isAnythingSaving) event.preventDefault(); 799 | }); 800 | 801 | /*////////////////////////////////////////////////////////////// 802 | WINDOW RESIZE LOGIC 803 | //////////////////////////////////////////////////////////////*/ 804 | 805 | useDebouncedWindowResize(autoZoomIfNecessary, 100); 806 | 807 | /*////////////////////////////////////////////////////////////// 808 | CHAT RESIZE LOGIC 809 | //////////////////////////////////////////////////////////////*/ 810 | 811 | const [savedChatSize, setSavedChatSize] = useLocalStorage( 812 | SAVED_CHAT_SIZE_LOCAL_STORAGE_KEY 813 | ); 814 | 815 | /*////////////////////////////////////////////////////////////// 816 | APP 817 | //////////////////////////////////////////////////////////////*/ 818 | 819 | const onNodeClick: NodeMouseHandler = useCallback( 820 | (_, node) => { 821 | setSelectedNodeId(node.id); 822 | setNodes((nds) => 823 | nds.map((n) => { 824 | if (n.id === node.id) { 825 | return { 826 | ...n, 827 | data: { ...n.data, expanded: !n.data.expanded }, 828 | }; 829 | } 830 | 831 | return n; 832 | }) 833 | ); 834 | }, 835 | [setNodes] 836 | ); 837 | 838 | return ( 839 | <> 840 | {!isValidAPIKey(apiKey) && } 841 | 842 | 851 | 857 | 858 | { 877 | setSavedChatSize(ref.style.width); 878 | autoZoomIfNecessary(); 879 | 880 | if (MIXPANEL_TOKEN) mixpanel.track("Resized chat window"); 881 | }} 882 | > 883 | 890 | 899 | { 901 | onOpenSettingsModal(); 902 | 903 | if (MIXPANEL_TOKEN) mixpanel.track("Opened Settings Modal"); // KPI 904 | }} 905 | onToggleAnswerFilter={() => { 906 | setShowAnswerPathOnly(!showAnswerPathOnly); 907 | }} 908 | showAnswerPathOnly={showAnswerPathOnly} 909 | /> 910 | 911 | 912 |

Input Token Count: {inputTokenCount}

913 | 914 | 915 | 916 |

Output Token Count: {outputTokenCount}

917 | 918 | 919 | {/* 920 |

Total Cost (GPT-4): ${((inputTokenCount * 0.03 / 1000) + (outputTokenCount * 0.06 / 1000)).toFixed(2)}

921 | */} 922 | 923 | 924 | {isAnythingLoading ? ( 925 | 926 | ) : ( 927 | 928 | )} 929 | 930 | 931 | 932 | 961 | 962 | 963 | 964 | 965 | 966 | 967 | { 972 | setNodes((nodes) => 973 | modifyBranchesNodeText(nodes, { 974 | asHuman: true, 975 | id: selectedNodeId!, 976 | text, 977 | isRunning: false, 978 | }) 979 | ); 980 | }} 981 | apiKey={apiKey} 982 | nodes={nodes} 983 | edges={edges} 984 | /> 985 | 986 | 987 | 988 | 989 | ); 990 | } 991 | 992 | export default App; 993 | -------------------------------------------------------------------------------- /src/components/NodeInfo.tsx: -------------------------------------------------------------------------------- 1 | import { Node, Edge } from "reactflow"; 2 | import rawHumanEvalProblems from "../utils/human_eval_problems.json"; 3 | import { 4 | Box, 5 | Text, 6 | Tag, 7 | TagLeftIcon, 8 | TagLabel, 9 | Textarea, 10 | List, 11 | ListItem, 12 | ListIcon, 13 | Input, 14 | Heading, 15 | Flex, 16 | } from "@chakra-ui/react"; 17 | import { CheckIcon } from "@chakra-ui/icons"; 18 | import { MdCheck, MdQuestionMark, MdClose, MdThumbUpOffAlt } from "react-icons/md"; 19 | import { Settings, ToTNodeData, HumanEvalProblemsType } from "../utils/types"; 20 | import { Prompt } from "./Prompt"; 21 | import { getBranchesNodeParent } from "../utils/branchesNode"; 22 | import { useEffect, useState } from "react"; 23 | import { Markdown } from "./utils/Markdown"; 24 | import { Row } from "../utils/chakra"; 25 | 26 | const HUMAN_EVAL_PROBLEMS = rawHumanEvalProblems as HumanEvalProblemsType; 27 | 28 | function EvalListItem({ item }: { item: string }) { 29 | if (item) { 30 | const lines = item.split("\n"); 31 | const lastLine = lines[lines.length - 1]; 32 | let icon = null; 33 | if (lastLine === "sure") { 34 | icon = ; 35 | } else if (lastLine === "impossible") { 36 | icon = ; 37 | } else if (lastLine === "likely") { 38 | icon = ; 39 | } 40 | 41 | return ( 42 | 43 | {icon} 44 | {lines.map((line, i) => { 45 | return ( 46 | 47 | {line} 48 |
49 | 50 | ); 51 | })} 52 | 53 | ); 54 | } 55 | return ; 56 | } 57 | 58 | export function NodeInfo({ 59 | lineage, 60 | selectNode, 61 | submitPrompt, 62 | apiKey, 63 | onPromptType, 64 | nodes, 65 | edges, 66 | }: { 67 | lineage: Node[] | null; 68 | settings?: Settings; 69 | setSettings?: (settings: Settings) => void; 70 | isGPT4?: boolean; 71 | submitPrompt: () => Promise; 72 | selectNode: (id: string) => void; 73 | apiKey: string | null; 74 | onPromptType: (text: string) => void; 75 | nodes: Node[]; 76 | edges: Edge[]; 77 | }) { 78 | const selectedNode = 79 | lineage && 80 | (lineage.find((n) => n.selected === true) as Node | undefined); 81 | const selectedNodeId = selectedNode?.id ?? null; 82 | const rootNode = lineage ? lineage[lineage.length - 1] : undefined; 83 | 84 | const [selectedNodeParent, setSelectedNodeParent] = useState< 85 | Node | null | undefined 86 | >(null); 87 | 88 | useEffect(() => { 89 | const newSelectedNodeParent = 90 | selectedNodeId !== null ? getBranchesNodeParent(nodes, edges, selectedNodeId) : null; 91 | setSelectedNodeParent(newSelectedNodeParent); 92 | }, [selectedNodeId, nodes, edges]); 93 | 94 | return ( 95 |
96 | {selectedNode?.data.isTerminal ? ( 97 | 103 | 104 | Terminal 105 | 106 | ) : null} 107 | {selectedNode?.data.isValid ? ( 108 | 109 | 110 | Valid 111 | 112 | ) : null} 113 | 114 | {/* 115 | Input 116 | 117 | {selectedNodeParent || selectedNodeId == null ? ( 118 |

{selectedNode?.data.input ?? ""}

119 | ) : ( 120 |