agent.ts 7.3 KB

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