193 lines · 6.0 KB · TypeScript
RawDownload
import { QuestionTool } from "./question"import { BashTool } from "./bash"import { EditTool } from "./edit"import { GlobTool } from "./glob"import { GrepTool } from "./grep"import { BatchTool } from "./batch"import { ReadTool } from "./read"import { TaskTool } from "./task"import { TodoWriteTool, TodoReadTool } from "./todo"import { PlanWriteTool } from "./planwrite"import { WebFetchTool } from "./webfetch"import { WriteTool } from "./write"import { InvalidTool } from "./invalid"import { SkillTool } from "./skill"import type { Agent } from "../agent/agent"import { Tool } from "./tool"import { Instance } from "../project/instance"import { Config } from "../config/config"import path from "path"import { type ToolContext as PluginToolContext, type ToolDefinition } from "@synsci/plugin"import z from "zod"import { Plugin } from "../plugin"import { WebSearchTool } from "./websearch"import { CodeSearchTool } from "./codesearch"import { Flag } from "@/flag/flag"import { Log } from "@/util/log"import { LspTool } from "./lsp"import { Truncate } from "./truncation"import { PlanExitTool, PlanEnterTool } from "./plan"import { ApplyPatchTool } from "./apply_patch"import { BiologyTools, BIOLOGY_TOOL_IDS } from "./biology"import { ArtifactTool } from "./artifact"import { LearnTool } from "./learn"import { ScienceTools } from "./science"import { ProvenanceTools } from "./provenance"import { NotebookTool } from "./notebook"import { RKernelTool } from "./rkernel" export namespace ToolRegistry {  const log = Log.create({ service: "tool.registry" })   export const state = Instance.state(async () => {    const custom = [] as Tool.Info[]    const glob = new Bun.Glob("{tool,tools}/*.{js,ts}")     for (const dir of await Config.directories()) {      for await (const match of glob.scan({        cwd: dir,        absolute: true,        followSymlinks: true,        dot: true,      })) {        const namespace = path.basename(match, path.extname(match))        const mod = await import(match)        for (const [id, def] of Object.entries<ToolDefinition>(mod)) {          custom.push(fromPlugin(id === "default" ? namespace : `${namespace}_${id}`, def))        }      }    }     const plugins = await Plugin.list()    for (const plugin of plugins) {      for (const [id, def] of Object.entries(plugin.tool ?? {})) {        custom.push(fromPlugin(id, def))      }    }     return { custom }  })   function fromPlugin(id: string, def: ToolDefinition): Tool.Info {    return {      id,      init: async (initCtx) => ({        parameters: z.object(def.args),        description: def.description,        execute: async (args, ctx) => {          const pluginCtx = {            ...ctx,            directory: Instance.directory,            worktree: Instance.worktree,          } as unknown as PluginToolContext          const result = await def.execute(args as any, pluginCtx)          const out = await Truncate.output(result, {}, initCtx?.agent)          return {            title: "",            output: out.truncated ? out.content : result,            metadata: { truncated: out.truncated, outputPath: out.truncated ? out.outputPath : undefined },          }        },      }),    }  }   export async function register(tool: Tool.Info) {    const { custom } = await state()    const idx = custom.findIndex((t) => t.id === tool.id)    if (idx >= 0) {      custom.splice(idx, 1, tool)      return    }    custom.push(tool)  }   async function all(): Promise<Tool.Info[]> {    const custom = await state().then((x) => x.custom)    const config = await Config.get()     return [      InvalidTool,      ...(["app", "cli", "desktop"].includes(Flag.OPENSCIENCE_CLIENT) ? [QuestionTool] : []),      BashTool,      ReadTool,      GlobTool,      GrepTool,      EditTool,      WriteTool,      TaskTool,      WebFetchTool,      ...(config.experimental?.plan_mode === true ? [PlanWriteTool] : [TodoWriteTool]),      TodoReadTool,      WebSearchTool,      CodeSearchTool,      SkillTool,      ApplyPatchTool,      ...(Flag.OPENSCIENCE_EXPERIMENTAL_LSP_TOOL ? [LspTool] : []),      ...(config.experimental?.batch_tool === true ? [BatchTool] : []),      ...(Flag.OPENSCIENCE_EXPERIMENTAL_PLAN_MODE && Flag.OPENSCIENCE_CLIENT === "cli"        ? [PlanExitTool, PlanEnterTool]        : []),      ...BiologyTools,      ...ScienceTools,      ...ProvenanceTools,      NotebookTool,      RKernelTool,      ArtifactTool,      LearnTool,      ...custom,    ]  }   const ARTIFACT_TOOL_ID = "artifact"  const ARTIFACT_AGENTS = ["research", "biology", "ml"]   export async function ids() {    return all().then((x) => x.map((t) => t.id))  }   export async function tools(    model: {      providerID: string      modelID: string    },    agent?: Agent.Info,  ) {    const tools = await all()    const result = await Promise.all(      tools        .filter((t) => {          // Biology-only tools: only available for the biology agent.          if (BIOLOGY_TOOL_IDS.has(t.id)) {            return agent?.name === "biology"          }           // Artifact tool: only for artifact-oriented scientific agents.          if (t.id === ARTIFACT_TOOL_ID) {            return !!agent?.name && ARTIFACT_AGENTS.includes(agent.name)          }           // Enable websearch/codesearch for zen users OR via enable flag          if (t.id === "codesearch" || t.id === "websearch") {            return model.providerID === "synsci" || Flag.OPENSCIENCE_ENABLE_EXA          }           // use apply tool in same format as codex          const usePatch =            model.modelID.includes("gpt-") && !model.modelID.includes("oss") && !model.modelID.includes("gpt-4")          if (t.id === "apply_patch") return usePatch          if (t.id === "edit" || t.id === "write") return !usePatch           return true        })        .map(async (t) => {          using _ = log.time(t.id)          return {            id: t.id,            ...(await t.init({ agent })),          }        }),    )    return result  }}