pyb-ts
Version:
PYB-CLI - Minimal AI Agent with multi-model support and CLI interface
202 lines (200 loc) • 6.47 kB
JavaScript
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