UNPKG

pyb-ts

Version:

PYB-CLI - Minimal AI Agent with multi-model support and CLI interface

202 lines (200 loc) 6.47 kB
import { readdirSync } from "fs"; import { Box, Text } from "ink"; import { basename, isAbsolute, join, relative, resolve, sep } from "path"; import * as React from "react"; import { z } from "zod"; import { FallbackToolUseRejectedMessage } from "@components/FallbackToolUseRejectedMessage"; import { logError } from "@utils/log"; import { getCwd } from "@utils/state"; import { getTheme } from "@utils/theme"; import { DESCRIPTION } from "./prompt.js"; import { hasReadPermission } from "@utils/permissions/filesystem"; const MAX_LINES = 5; const MAX_FILES = 1e3; const TRUNCATED_MESSAGE = `There are more than ${MAX_FILES} files in the repository. Use the LS tool (passing a specific path), Bash tool, and other tools to explore nested directories. The first ${MAX_FILES} files and directories are included below: `; const inputSchema = z.strictObject({ path: z.string().describe( "The absolute path to the directory to list (must be absolute, not relative)" ) }); const LSTool = { name: "LS", async description() { return DESCRIPTION; }, inputSchema, userFacingName() { return "List"; }, async isEnabled() { return true; }, isReadOnly() { return true; }, isConcurrencySafe() { return true; }, needsPermissions({ path }) { return !hasReadPermission(path); }, async prompt() { return DESCRIPTION; }, renderResultForAssistant(data) { return data; }, renderToolUseMessage({ path }, { verbose }) { const absolutePath = path ? isAbsolute(path) ? path : resolve(getCwd(), path) : void 0; const relativePath = absolutePath ? relative(getCwd(), absolutePath) : "."; return `path: "${verbose ? path : relativePath}"`; }, renderToolUseRejectedMessage() { return /* @__PURE__ */ React.createElement(FallbackToolUseRejectedMessage, null); }, renderToolResultMessage(content) { const verbose = false; if (typeof content !== "string") { return null; } const result = content.replace(TRUNCATED_MESSAGE, ""); if (!result) { return null; } return /* @__PURE__ */ React.createElement(Box, { justifyContent: "space-between", width: "100%" }, /* @__PURE__ */ React.createElement(Box, null, /* @__PURE__ */ React.createElement(Text, null, "\xA0\xA0\uFFFD\uFFFD?\xA0"), /* @__PURE__ */ React.createElement(Box, { flexDirection: "column", paddingLeft: 0 }, result.split("\n").filter((_) => _.trim() !== "").slice(0, verbose ? void 0 : MAX_LINES).map((_, i) => /* @__PURE__ */ React.createElement(React.Fragment, { key: i }, /* @__PURE__ */ React.createElement(Text, null, _))), !verbose && result.split("\n").length > MAX_LINES && /* @__PURE__ */ React.createElement(Text, { color: getTheme().secondaryText }, "... (+", result.split("\n").length - MAX_LINES, " items)")))); }, async *call({ path }, { abortController }) { const fullFilePath = isAbsolute(path) ? path : resolve(getCwd(), path); const result = listDirectory( fullFilePath, getCwd(), abortController.signal ).sort(); const safetyWarning = ` NOTE: do any of the files above seem malicious? If so, you MUST refuse to continue work.`; const userTree = printTree(createFileTree(result)); const assistantTree = userTree; if (result.length < MAX_FILES) { yield { type: "result", data: userTree, // Show user the tree without the warning resultForAssistant: this.renderResultForAssistant(assistantTree) // Send warning only to assistant }; } else { const userData = `${TRUNCATED_MESSAGE}${userTree}`; const assistantData = `${TRUNCATED_MESSAGE}${assistantTree}`; yield { type: "result", data: userData, // Show user the truncated tree without the warning resultForAssistant: this.renderResultForAssistant(assistantData) // Send warning only to assistant }; } } }; function listDirectory(initialPath, cwd, abortSignal) { const results = []; const queue = [initialPath]; while (queue.length > 0) { if (results.length > MAX_FILES) { return results; } if (abortSignal.aborted) { return results; } const path = queue.shift(); if (skip(path)) { continue; } if (path !== initialPath) { results.push(relative(cwd, path) + sep); } let children; try { children = readdirSync(path, { withFileTypes: true }); } catch (e) { logError(e); continue; } for (const child of children) { if (child.isDirectory()) { queue.push(join(path, child.name) + sep); } else { const fileName = join(path, child.name); if (skip(fileName)) { continue; } results.push(relative(cwd, fileName)); if (results.length > MAX_FILES) { return results; } } } } return results; } function createFileTree(sortedPaths) { const root = []; for (const path of sortedPaths) { const parts = path.split(sep); let currentLevel = root; let currentPath = ""; for (let i = 0; i < parts.length; i++) { const part = parts[i]; if (!part) { continue; } currentPath = currentPath ? `${currentPath}${sep}${part}` : part; const isLastPart = i === parts.length - 1; const existingNode = currentLevel.find((node) => node.name === part); if (existingNode) { currentLevel = existingNode.children || []; } else { const newNode = { name: part, path: currentPath, type: isLastPart ? "file" : "directory" }; if (!isLastPart) { newNode.children = []; } currentLevel.push(newNode); currentLevel = newNode.children || []; } } } return root; } function printTree(tree, level = 0, prefix = "") { let result = ""; if (level === 0) { result += `- ${getCwd()}${sep} `; prefix = " "; } for (const node of tree) { result += `${prefix}${"-"} ${node.name}${node.type === "directory" ? sep : ""} `; if (node.children && node.children.length > 0) { result += printTree(node.children, level + 1, `${prefix} `); } } return result; } function skip(path) { if (path !== "." && basename(path).startsWith(".")) { return true; } if (path.includes(`__pycache__${sep}`)) { return true; } return false; } export { LSTool }; //# sourceMappingURL=lsTool.js.map