agent.ts 9.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335
  1. import { Config } from "../config/config"
  2. import z from "zod"
  3. import { Provider } from "../provider/provider"
  4. import { generateObject, type ModelMessage } from "ai"
  5. import { SystemPrompt } from "../session/system"
  6. import { Instance } from "../project/instance"
  7. import { mergeDeep } from "remeda"
  8. import PROMPT_GENERATE from "./generate.txt"
  9. import PROMPT_COMPACTION from "./prompt/compaction.txt"
  10. import PROMPT_EXPLORE from "./prompt/explore.txt"
  11. import PROMPT_SUMMARY from "./prompt/summary.txt"
  12. import PROMPT_TITLE from "./prompt/title.txt"
  13. export namespace Agent {
  14. export const Info = z
  15. .object({
  16. name: z.string(),
  17. description: z.string().optional(),
  18. mode: z.enum(["subagent", "primary", "all"]),
  19. native: z.boolean().optional(),
  20. hidden: z.boolean().optional(),
  21. topP: z.number().optional(),
  22. temperature: z.number().optional(),
  23. color: z.string().optional(),
  24. permission: z.object({
  25. edit: Config.Permission,
  26. bash: z.record(z.string(), Config.Permission),
  27. webfetch: Config.Permission.optional(),
  28. doom_loop: Config.Permission.optional(),
  29. external_directory: Config.Permission.optional(),
  30. }),
  31. model: z
  32. .object({
  33. modelID: z.string(),
  34. providerID: z.string(),
  35. })
  36. .optional(),
  37. prompt: z.string().optional(),
  38. tools: z.record(z.string(), z.boolean()),
  39. options: z.record(z.string(), z.any()),
  40. maxSteps: z.number().int().positive().optional(),
  41. })
  42. .meta({
  43. ref: "Agent",
  44. })
  45. export type Info = z.infer<typeof Info>
  46. const state = Instance.state(async () => {
  47. const cfg = await Config.get()
  48. const defaultTools = cfg.tools ?? {}
  49. const defaultPermission: Info["permission"] = {
  50. edit: "allow",
  51. bash: {
  52. "*": "allow",
  53. },
  54. webfetch: "allow",
  55. doom_loop: "ask",
  56. external_directory: "ask",
  57. }
  58. const agentPermission = mergeAgentPermissions(defaultPermission, cfg.permission ?? {})
  59. const planPermission = mergeAgentPermissions(
  60. {
  61. edit: "deny",
  62. bash: {
  63. "cut*": "allow",
  64. "diff*": "allow",
  65. "du*": "allow",
  66. "file *": "allow",
  67. "find * -delete*": "ask",
  68. "find * -exec*": "ask",
  69. "find * -fprint*": "ask",
  70. "find * -fls*": "ask",
  71. "find * -fprintf*": "ask",
  72. "find * -ok*": "ask",
  73. "find *": "allow",
  74. "git diff*": "allow",
  75. "git log*": "allow",
  76. "git show*": "allow",
  77. "git status*": "allow",
  78. "git branch": "allow",
  79. "git branch -v": "allow",
  80. "grep*": "allow",
  81. "head*": "allow",
  82. "less*": "allow",
  83. "ls*": "allow",
  84. "more*": "allow",
  85. "pwd*": "allow",
  86. "rg*": "allow",
  87. "sort --output=*": "ask",
  88. "sort -o *": "ask",
  89. "sort*": "allow",
  90. "stat*": "allow",
  91. "tail*": "allow",
  92. "tree -o *": "ask",
  93. "tree*": "allow",
  94. "uniq*": "allow",
  95. "wc*": "allow",
  96. "whereis*": "allow",
  97. "which*": "allow",
  98. "*": "ask",
  99. },
  100. webfetch: "allow",
  101. },
  102. cfg.permission ?? {},
  103. )
  104. const result: Record<string, Info> = {
  105. build: {
  106. name: "build",
  107. tools: { ...defaultTools },
  108. options: {},
  109. permission: agentPermission,
  110. mode: "primary",
  111. native: true,
  112. },
  113. plan: {
  114. name: "plan",
  115. options: {},
  116. permission: planPermission,
  117. tools: {
  118. ...defaultTools,
  119. },
  120. mode: "primary",
  121. native: true,
  122. },
  123. general: {
  124. name: "general",
  125. description: `General-purpose agent for researching complex questions and executing multi-step tasks. Use this agent to execute multiple units of work in parallel.`,
  126. tools: {
  127. todoread: false,
  128. todowrite: false,
  129. ...defaultTools,
  130. },
  131. options: {},
  132. permission: agentPermission,
  133. mode: "subagent",
  134. native: true,
  135. hidden: true,
  136. },
  137. explore: {
  138. name: "explore",
  139. tools: {
  140. todoread: false,
  141. todowrite: false,
  142. edit: false,
  143. write: false,
  144. ...defaultTools,
  145. },
  146. description: `Fast agent specialized for exploring codebases. Use this when you need to quickly find files by patterns (eg. "src/components/**/*.tsx"), search code for keywords (eg. "API endpoints"), or answer questions about the codebase (eg. "how do API endpoints work?"). When calling this agent, specify the desired thoroughness level: "quick" for basic searches, "medium" for moderate exploration, or "very thorough" for comprehensive analysis across multiple locations and naming conventions.`,
  147. prompt: PROMPT_EXPLORE,
  148. options: {},
  149. permission: agentPermission,
  150. mode: "subagent",
  151. native: true,
  152. },
  153. compaction: {
  154. name: "compaction",
  155. mode: "primary",
  156. native: true,
  157. hidden: true,
  158. prompt: PROMPT_COMPACTION,
  159. tools: {
  160. "*": false,
  161. },
  162. options: {},
  163. permission: agentPermission,
  164. },
  165. title: {
  166. name: "title",
  167. mode: "primary",
  168. options: {},
  169. native: true,
  170. hidden: true,
  171. permission: agentPermission,
  172. prompt: PROMPT_TITLE,
  173. tools: {},
  174. },
  175. summary: {
  176. name: "summary",
  177. mode: "primary",
  178. options: {},
  179. native: true,
  180. hidden: true,
  181. permission: agentPermission,
  182. prompt: PROMPT_SUMMARY,
  183. tools: {},
  184. },
  185. }
  186. for (const [key, value] of Object.entries(cfg.agent ?? {})) {
  187. if (value.disable) {
  188. delete result[key]
  189. continue
  190. }
  191. let item = result[key]
  192. if (!item)
  193. item = result[key] = {
  194. name: key,
  195. mode: "all",
  196. permission: agentPermission,
  197. options: {},
  198. tools: {},
  199. native: false,
  200. }
  201. const {
  202. name,
  203. model,
  204. prompt,
  205. tools,
  206. description,
  207. temperature,
  208. top_p,
  209. mode,
  210. permission,
  211. color,
  212. maxSteps,
  213. ...extra
  214. } = value
  215. item.options = {
  216. ...item.options,
  217. ...extra,
  218. }
  219. if (model) item.model = Provider.parseModel(model)
  220. if (prompt) item.prompt = prompt
  221. if (tools)
  222. item.tools = {
  223. ...item.tools,
  224. ...tools,
  225. }
  226. item.tools = {
  227. ...defaultTools,
  228. ...item.tools,
  229. }
  230. if (description) item.description = description
  231. if (temperature != undefined) item.temperature = temperature
  232. if (top_p != undefined) item.topP = top_p
  233. if (mode) item.mode = mode
  234. if (color) item.color = color
  235. // just here for consistency & to prevent it from being added as an option
  236. if (name) item.name = name
  237. if (maxSteps != undefined) item.maxSteps = maxSteps
  238. if (permission ?? cfg.permission) {
  239. item.permission = mergeAgentPermissions(cfg.permission ?? {}, permission ?? {})
  240. }
  241. }
  242. return result
  243. })
  244. export async function get(agent: string) {
  245. return state().then((x) => x[agent])
  246. }
  247. export async function list() {
  248. return state().then((x) => Object.values(x))
  249. }
  250. export async function generate(input: { description: string; model?: { providerID: string; modelID: string } }) {
  251. const cfg = await Config.get()
  252. const defaultModel = input.model ?? (await Provider.defaultModel())
  253. const model = await Provider.getModel(defaultModel.providerID, defaultModel.modelID)
  254. const language = await Provider.getLanguage(model)
  255. const system = SystemPrompt.header(defaultModel.providerID)
  256. system.push(PROMPT_GENERATE)
  257. const existing = await list()
  258. const result = await generateObject({
  259. experimental_telemetry: {
  260. isEnabled: cfg.experimental?.openTelemetry,
  261. metadata: {
  262. userId: cfg.username ?? "unknown",
  263. },
  264. },
  265. temperature: 0.3,
  266. messages: [
  267. ...system.map(
  268. (item): ModelMessage => ({
  269. role: "system",
  270. content: item,
  271. }),
  272. ),
  273. {
  274. role: "user",
  275. content: `Create an agent configuration based on this request: \"${input.description}\".\n\nIMPORTANT: The following identifiers already exist and must NOT be used: ${existing.map((i) => i.name).join(", ")}\n Return ONLY the JSON object, no other text, do not wrap in backticks`,
  276. },
  277. ],
  278. model: language,
  279. schema: z.object({
  280. identifier: z.string(),
  281. whenToUse: z.string(),
  282. systemPrompt: z.string(),
  283. }),
  284. })
  285. return result.object
  286. }
  287. }
  288. function mergeAgentPermissions(basePermission: any, overridePermission: any): Agent.Info["permission"] {
  289. if (typeof basePermission.bash === "string") {
  290. basePermission.bash = {
  291. "*": basePermission.bash,
  292. }
  293. }
  294. if (typeof overridePermission.bash === "string") {
  295. overridePermission.bash = {
  296. "*": overridePermission.bash,
  297. }
  298. }
  299. const merged = mergeDeep(basePermission ?? {}, overridePermission ?? {}) as any
  300. let mergedBash
  301. if (merged.bash) {
  302. if (typeof merged.bash === "string") {
  303. mergedBash = {
  304. "*": merged.bash,
  305. }
  306. } else if (typeof merged.bash === "object") {
  307. mergedBash = mergeDeep(
  308. {
  309. "*": "allow",
  310. },
  311. merged.bash,
  312. )
  313. }
  314. }
  315. const result: Agent.Info["permission"] = {
  316. edit: merged.edit ?? "allow",
  317. webfetch: merged.webfetch ?? "allow",
  318. bash: mergedBash ?? { "*": "allow" },
  319. doom_loop: merged.doom_loop,
  320. external_directory: merged.external_directory,
  321. }
  322. return result
  323. }