pyb-ts
Version:
PYB-CLI - Minimal AI Agent with multi-model support and CLI interface
355 lines (351 loc) • 13.2 kB
JavaScript
import { last, memoize } from "lodash-es";
import React from "react";
import { Box, Text } from "ink";
import { z } from "zod";
import { FallbackToolUseRejectedMessage } from "@components/FallbackToolUseRejectedMessage";
import { getAgentPrompt } from "@constants/prompts";
import { getContext } from "@context";
import { hasPermissionsToUseTool } from "@permissions";
import { query } from "@query";
import { formatDuration, formatNumber } from "@utils/format";
import {
getMessagesPath,
getNextAvailableLogSidechainNumber,
overwriteLog
} from "@utils/log";
import {
createAssistantMessage,
createUserMessage,
getLastAssistantMessageId,
INTERRUPT_MESSAGE,
normalizeMessages
} from "@utils/messages";
import { getModelManager } from "@utils/model";
import { getMaxThinkingTokens } from "@utils/thinking";
import { getTheme } from "@utils/theme";
import { generateAgentId } from "@utils/agentStorage";
import { globalMemoryHook } from "@utils/memoryRecorder";
import { getTaskTools, getPrompt } from "./prompt.js";
import { TOOL_NAME } from "./constants.js";
import { getAgentByType, getAvailableAgentTypes } from "@utils/agentLoader";
const inputSchema = z.object({
description: z.string().describe("A short (3-5 word) description of the task"),
prompt: z.string().describe("The task for the agent to perform"),
model_name: z.string().optional().describe(
"Optional: Specific model name to use for this task. If not provided, uses the default task model pointer."
),
subagent_type: z.string().optional().describe(
"The type of specialized agent to use for this task"
)
});
const TaskTool = {
async prompt({ safeMode }) {
return await getPrompt(safeMode);
},
name: TOOL_NAME,
async description() {
return "Launch a new task";
},
inputSchema,
async *call({ description, prompt, model_name, subagent_type }, {
abortController,
options: { safeMode = false, forkNumber, messageLogName, verbose },
readFileTimestamps
}) {
const startTime = Date.now();
const agentType = subagent_type || "general-purpose";
let effectivePrompt = prompt;
let effectiveModel = model_name || "task";
let toolFilter = null;
let temperature = void 0;
if (agentType) {
const agentConfig = await getAgentByType(agentType);
if (!agentConfig) {
const availableTypes = await getAvailableAgentTypes();
const helpMessage = `Agent type '${agentType}' not found.
Available agents:
${availableTypes.map((t) => ` \uFFFD?${t}`).join("\n")}
Use /agents command to manage agent configurations.`;
yield {
type: "result",
data: [{ type: "text", text: helpMessage }],
resultForAssistant: helpMessage
};
return;
}
if (agentConfig.systemPrompt) {
effectivePrompt = `${agentConfig.systemPrompt}
${prompt}`;
}
if (!model_name && agentConfig.model_name) {
if (agentConfig.model_name !== "inherit") {
effectiveModel = agentConfig.model_name;
}
}
toolFilter = agentConfig.tools;
}
const messages = [createUserMessage(effectivePrompt)];
let tools = await getTaskTools(safeMode);
if (toolFilter) {
const isAllArray = Array.isArray(toolFilter) && toolFilter.length === 1 && toolFilter[0] === "*";
if (toolFilter === "*" || isAllArray) {
} else if (Array.isArray(toolFilter)) {
tools = tools.filter((tool) => toolFilter.includes(tool.name));
}
}
const modelToUse = effectiveModel;
yield {
type: "progress",
content: createAssistantMessage(`Starting agent: ${agentType}`),
normalizedMessages: normalizeMessages(messages),
tools
};
yield {
type: "progress",
content: createAssistantMessage(`Using model: ${modelToUse}`),
normalizedMessages: normalizeMessages(messages),
tools
};
yield {
type: "progress",
content: createAssistantMessage(`Task: ${description}`),
normalizedMessages: normalizeMessages(messages),
tools
};
yield {
type: "progress",
content: createAssistantMessage(`Prompt: ${prompt.length > 150 ? prompt.substring(0, 150) + "..." : prompt}`),
normalizedMessages: normalizeMessages(messages),
tools
};
const [taskPrompt, context, maxThinkingTokens] = await Promise.all([
getAgentPrompt(),
getContext(),
getMaxThinkingTokens(messages)
]);
taskPrompt.push(`
IMPORTANT: You are currently running as ${modelToUse}. You do not need to consult ${modelToUse} via AskExpertModel since you ARE ${modelToUse}. Complete tasks directly using your capabilities.`);
let toolUseCount = 0;
const getSidechainNumber = memoize(
() => getNextAvailableLogSidechainNumber(messageLogName, forkNumber)
);
const taskId = generateAgentId();
await globalMemoryHook.recordTaskStart(taskId, agentType, prompt, agentType);
const queryOptions = {
safeMode,
forkNumber,
messageLogName,
tools,
commands: [],
verbose,
maxThinkingTokens,
model: modelToUse
};
if (temperature !== void 0) {
queryOptions["temperature"] = temperature;
}
for await (const message of query(
messages,
taskPrompt,
context,
hasPermissionsToUseTool,
{
abortController,
options: queryOptions,
messageId: getLastAssistantMessageId(messages),
agentId: taskId,
readFileTimestamps,
setToolJSX: () => {
}
// No-op implementation for TaskTool
}
)) {
messages.push(message);
overwriteLog(
getMessagesPath(messageLogName, forkNumber, getSidechainNumber()),
messages.filter((_) => _.type !== "progress")
);
if (message.type !== "assistant") {
continue;
}
const normalizedMessages2 = normalizeMessages(messages);
for (const content of message.message.content) {
if (content.type === "text" && content.text && content.text !== INTERRUPT_MESSAGE) {
const preview = content.text.length > 200 ? content.text.substring(0, 200) + "..." : content.text;
yield {
type: "progress",
content: createAssistantMessage(`${preview}`),
normalizedMessages: normalizedMessages2,
tools
};
} else if (content.type === "tool_use") {
toolUseCount++;
const toolMessage = normalizedMessages2.find(
(_) => _.type === "assistant" && _.message.content[0]?.type === "tool_use" && _.message.content[0].id === content.id
);
if (toolMessage) {
const modifiedMessage = {
...toolMessage,
message: {
...toolMessage.message,
content: toolMessage.message.content.map((c) => {
if (c.type === "tool_use" && c.id === content.id) {
return {
...c,
name: c.name
// Keep original name, UI will handle display
};
}
return c;
})
}
};
yield {
type: "progress",
content: modifiedMessage,
normalizedMessages: normalizedMessages2,
tools
};
}
}
}
}
const normalizedMessages = normalizeMessages(messages);
const lastMessage = last(messages);
if (lastMessage?.type !== "assistant") {
throw new Error("Last message was not an assistant message");
}
if (lastMessage.message.content.some(
(_) => _.type === "text" && _.text === INTERRUPT_MESSAGE
)) {
} else {
const result = [
toolUseCount === 1 ? "1 tool use" : `${toolUseCount} tool uses`,
formatNumber(
(lastMessage.message.usage.cache_creation_input_tokens ?? 0) + (lastMessage.message.usage.cache_read_input_tokens ?? 0) + lastMessage.message.usage.input_tokens + lastMessage.message.usage.output_tokens
) + " tokens",
formatDuration(Date.now() - startTime)
];
yield {
type: "progress",
content: createAssistantMessage(`Task completed (${result.join(" \xB7 ")})`),
normalizedMessages,
tools
};
}
const data = lastMessage.message.content.filter((_) => _.type === "text");
await globalMemoryHook.recordTaskCompletion(taskId, agentType, "completed");
yield {
type: "result",
data,
resultForAssistant: this.renderResultForAssistant(data)
};
},
isReadOnly() {
return true;
},
isConcurrencySafe() {
return true;
},
async validateInput(input, context) {
if (!input.description || typeof input.description !== "string") {
return {
result: false,
message: "Description is required and must be a string"
};
}
if (!input.prompt || typeof input.prompt !== "string") {
return {
result: false,
message: "Prompt is required and must be a string"
};
}
if (input.model_name) {
const modelManager = getModelManager();
const availableModels = modelManager.getAllAvailableModelNames();
if (!availableModels.includes(input.model_name)) {
return {
result: false,
message: `Model '${input.model_name}' does not exist. Available models: ${availableModels.join(", ")}`,
meta: {
model_name: input.model_name,
availableModels
}
};
}
}
if (input.subagent_type) {
const availableTypes = await getAvailableAgentTypes();
if (!availableTypes.includes(input.subagent_type)) {
return {
result: false,
message: `Agent type '${input.subagent_type}' does not exist. Available types: ${availableTypes.join(", ")}`,
meta: {
subagent_type: input.subagent_type,
availableTypes
}
};
}
}
return { result: true };
},
async isEnabled() {
return true;
},
userFacingName(input) {
const agentType = input?.subagent_type || "general-purpose";
return `agent-${agentType}`;
},
needsPermissions() {
return false;
},
renderResultForAssistant(data) {
return data.map((block) => block.type === "text" ? block.text : "").join("\n");
},
renderToolUseMessage({ description, prompt, model_name, subagent_type }, { verbose }) {
if (!description || !prompt) return null;
const modelManager = getModelManager();
const defaultTaskModel = modelManager.getModelName("task");
const actualModel = model_name || defaultTaskModel;
const agentType = subagent_type || "general-purpose";
const promptPreview = prompt.length > 80 ? prompt.substring(0, 80) + "..." : prompt;
const theme = getTheme();
if (verbose) {
return /* @__PURE__ */ React.createElement(Box, { flexDirection: "column" }, /* @__PURE__ */ React.createElement(Text, null, "[", agentType, "] ", actualModel, ": ", description), /* @__PURE__ */ React.createElement(
Box,
{
paddingLeft: 2,
borderLeftStyle: "single",
borderLeftColor: theme.secondaryBorder
},
/* @__PURE__ */ React.createElement(Text, { color: theme.secondaryText }, promptPreview)
));
}
return `[${agentType}] ${actualModel}: ${description}`;
},
renderToolUseRejectedMessage() {
return /* @__PURE__ */ React.createElement(FallbackToolUseRejectedMessage, null);
},
renderToolResultMessage(content) {
const theme = getTheme();
if (Array.isArray(content)) {
const textBlocks = content.filter((block) => block.type === "text");
const totalLength = textBlocks.reduce(
(sum, block) => sum + block.text.length,
0
);
const isInterrupted = content.some(
(block) => block.type === "text" && block.text === INTERRUPT_MESSAGE
);
if (isInterrupted) {
return /* @__PURE__ */ React.createElement(Box, { flexDirection: "row" }, /* @__PURE__ */ React.createElement(Text, null, "\xA0\xA0\uFFFD?\xA0"), /* @__PURE__ */ React.createElement(Text, { color: theme.error }, "Interrupted by user"));
}
return /* @__PURE__ */ React.createElement(Box, { flexDirection: "column" }, /* @__PURE__ */ React.createElement(Box, { justifyContent: "space-between", width: "100%" }, /* @__PURE__ */ React.createElement(Box, { flexDirection: "row" }, /* @__PURE__ */ React.createElement(Text, null, "\xA0\xA0\uFFFD?\xA0"), /* @__PURE__ */ React.createElement(Text, null, "Task completed"), textBlocks.length > 0 && /* @__PURE__ */ React.createElement(Text, { color: theme.secondaryText }, " ", "(", totalLength, " characters)"))));
}
return /* @__PURE__ */ React.createElement(Box, { flexDirection: "row" }, /* @__PURE__ */ React.createElement(Text, null, "\xA0\xA0\uFFFD?\xA0"), /* @__PURE__ */ React.createElement(Text, { color: theme.secondaryText }, "Task completed"));
}
};
export {
TaskTool
};
//# sourceMappingURL=TaskTool.js.map