agent.ts 7.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258
  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 PROMPT_GENERATE from "./generate.txt"
  6. import { SystemPrompt } from "../session/system"
  7. import { Instance } from "../project/instance"
  8. import { mergeDeep } from "remeda"
  9. export namespace Agent {
  10. export const Info = z
  11. .object({
  12. name: z.string(),
  13. description: z.string().optional(),
  14. mode: z.enum(["subagent", "primary", "all"]),
  15. builtIn: z.boolean(),
  16. topP: z.number().optional(),
  17. temperature: z.number().optional(),
  18. color: z.string().optional(),
  19. permission: z.object({
  20. edit: Config.Permission,
  21. bash: z.record(z.string(), Config.Permission),
  22. webfetch: Config.Permission.optional(),
  23. doom_loop: Config.Permission.optional(),
  24. external_directory: Config.Permission.optional(),
  25. }),
  26. model: z
  27. .object({
  28. modelID: z.string(),
  29. providerID: z.string(),
  30. })
  31. .optional(),
  32. prompt: z.string().optional(),
  33. tools: z.record(z.string(), z.boolean()),
  34. options: z.record(z.string(), z.any()),
  35. })
  36. .meta({
  37. ref: "Agent",
  38. })
  39. export type Info = z.infer<typeof Info>
  40. const state = Instance.state(async () => {
  41. const cfg = await Config.get()
  42. const defaultTools = cfg.tools ?? {}
  43. const defaultPermission: Info["permission"] = {
  44. edit: "allow",
  45. bash: {
  46. "*": "allow",
  47. },
  48. webfetch: "allow",
  49. doom_loop: "ask",
  50. external_directory: "ask",
  51. }
  52. const agentPermission = mergeAgentPermissions(defaultPermission, cfg.permission ?? {})
  53. const planPermission = mergeAgentPermissions(
  54. {
  55. edit: "deny",
  56. bash: {
  57. "cut*": "allow",
  58. "diff*": "allow",
  59. "du*": "allow",
  60. "file *": "allow",
  61. "find * -delete*": "ask",
  62. "find * -exec*": "ask",
  63. "find * -fprint*": "ask",
  64. "find * -fls*": "ask",
  65. "find * -fprintf*": "ask",
  66. "find * -ok*": "ask",
  67. "find *": "allow",
  68. "git diff*": "allow",
  69. "git log*": "allow",
  70. "git show*": "allow",
  71. "git status*": "allow",
  72. "git branch": "allow",
  73. "git branch -v": "allow",
  74. "grep*": "allow",
  75. "head*": "allow",
  76. "less*": "allow",
  77. "ls*": "allow",
  78. "more*": "allow",
  79. "pwd*": "allow",
  80. "rg*": "allow",
  81. "sort --output=*": "ask",
  82. "sort -o *": "ask",
  83. "sort*": "allow",
  84. "stat*": "allow",
  85. "tail*": "allow",
  86. "tree -o *": "ask",
  87. "tree*": "allow",
  88. "uniq*": "allow",
  89. "wc*": "allow",
  90. "whereis*": "allow",
  91. "which*": "allow",
  92. "*": "ask",
  93. },
  94. webfetch: "allow",
  95. },
  96. cfg.permission ?? {},
  97. )
  98. const result: Record<string, Info> = {
  99. general: {
  100. name: "general",
  101. description:
  102. "General-purpose agent for researching complex questions, searching for code, and executing multi-step tasks. When you are searching for a keyword or file and are not confident that you will find the right match in the first few tries use this agent to perform the search for you.",
  103. tools: {
  104. todoread: false,
  105. todowrite: false,
  106. ...defaultTools,
  107. },
  108. options: {},
  109. permission: agentPermission,
  110. mode: "subagent",
  111. builtIn: true,
  112. },
  113. build: {
  114. name: "build",
  115. tools: { ...defaultTools },
  116. options: {},
  117. permission: agentPermission,
  118. mode: "primary",
  119. builtIn: true,
  120. },
  121. plan: {
  122. name: "plan",
  123. options: {},
  124. permission: planPermission,
  125. tools: {
  126. ...defaultTools,
  127. },
  128. mode: "primary",
  129. builtIn: true,
  130. },
  131. }
  132. for (const [key, value] of Object.entries(cfg.agent ?? {})) {
  133. if (value.disable) {
  134. delete result[key]
  135. continue
  136. }
  137. let item = result[key]
  138. if (!item)
  139. item = result[key] = {
  140. name: key,
  141. mode: "all",
  142. permission: agentPermission,
  143. options: {},
  144. tools: {},
  145. builtIn: false,
  146. }
  147. const { name, model, prompt, tools, description, temperature, top_p, mode, permission, color, ...extra } = value
  148. item.options = {
  149. ...item.options,
  150. ...extra,
  151. }
  152. if (model) item.model = Provider.parseModel(model)
  153. if (prompt) item.prompt = prompt
  154. if (tools)
  155. item.tools = {
  156. ...item.tools,
  157. ...tools,
  158. }
  159. item.tools = {
  160. ...defaultTools,
  161. ...item.tools,
  162. }
  163. if (description) item.description = description
  164. if (temperature != undefined) item.temperature = temperature
  165. if (top_p != undefined) item.topP = top_p
  166. if (mode) item.mode = mode
  167. if (color) item.color = color
  168. // just here for consistency & to prevent it from being added as an option
  169. if (name) item.name = name
  170. if (permission ?? cfg.permission) {
  171. item.permission = mergeAgentPermissions(cfg.permission ?? {}, permission ?? {})
  172. }
  173. }
  174. return result
  175. })
  176. export async function get(agent: string) {
  177. return state().then((x) => x[agent])
  178. }
  179. export async function list() {
  180. return state().then((x) => Object.values(x))
  181. }
  182. export async function generate(input: { description: string }) {
  183. const defaultModel = await Provider.defaultModel()
  184. const model = await Provider.getModel(defaultModel.providerID, defaultModel.modelID)
  185. const system = SystemPrompt.header(defaultModel.providerID)
  186. system.push(PROMPT_GENERATE)
  187. const existing = await list()
  188. const result = await generateObject({
  189. temperature: 0.3,
  190. prompt: [
  191. ...system.map(
  192. (item): ModelMessage => ({
  193. role: "system",
  194. content: item,
  195. }),
  196. ),
  197. {
  198. role: "user",
  199. 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`,
  200. },
  201. ],
  202. model: model.language,
  203. schema: z.object({
  204. identifier: z.string(),
  205. whenToUse: z.string(),
  206. systemPrompt: z.string(),
  207. }),
  208. })
  209. return result.object
  210. }
  211. }
  212. function mergeAgentPermissions(basePermission: any, overridePermission: any): Agent.Info["permission"] {
  213. if (typeof basePermission.bash === "string") {
  214. basePermission.bash = {
  215. "*": basePermission.bash,
  216. }
  217. }
  218. if (typeof overridePermission.bash === "string") {
  219. overridePermission.bash = {
  220. "*": overridePermission.bash,
  221. }
  222. }
  223. const merged = mergeDeep(basePermission ?? {}, overridePermission ?? {}) as any
  224. let mergedBash
  225. if (merged.bash) {
  226. if (typeof merged.bash === "string") {
  227. mergedBash = {
  228. "*": merged.bash,
  229. }
  230. } else if (typeof merged.bash === "object") {
  231. mergedBash = mergeDeep(
  232. {
  233. "*": "allow",
  234. },
  235. merged.bash,
  236. )
  237. }
  238. }
  239. const result: Agent.Info["permission"] = {
  240. edit: merged.edit ?? "allow",
  241. webfetch: merged.webfetch ?? "allow",
  242. bash: mergedBash ?? { "*": "allow" },
  243. doom_loop: merged.doom_loop,
  244. external_directory: merged.external_directory,
  245. }
  246. return result
  247. }