History

refactor: unify v2 runtime around tool registry

Rhys Sullivan committed 8 months ago

Commit 3c7e2c81 parent 68065917 files · +451 −92Browse files

Unified · Split

7 files
  1. v2/apps/convex/convex/runtime_adapter.ts +1 −2
  2. v2/packages/ai-sdk-adapter/src/index.test.ts +12 −6
  3. v2/packages/engine/src/index.ts +15 −4
  4. v2/packages/engine/src/run-client.ts +6 −71
  5. v2/packages/engine/src/runtime-adapters.ts +0 −9
  6. v2/packages/engine/src/tool-registry.test.ts +100 −0
  7. v2/packages/engine/src/tool-registry.ts +317 −0
v2/apps/convex/convex/runtime_adapter.ts+1−2
@@ -2,10 +2,9 @@ import { · 1 unmodified line
22 type RuntimeAdapter,
33 type RuntimeExecuteError,
44 type RuntimeExecuteInput,
5− type RuntimeRunnableTool,
65 } from "@executor-v2/engine";
76 import { makeLocalInProcessRuntimeAdapter } from "@executor-v2/runtime-local-inproc";
87
9−export type { RuntimeAdapter, RuntimeExecuteError, RuntimeExecuteInput, RuntimeRunnableTool };
8+export type { RuntimeAdapter, RuntimeExecuteError, RuntimeExecuteInput };
109
1110 export { makeLocalInProcessRuntimeAdapter };
v2/packages/ai-sdk-adapter/src/index.test.ts+12−6
@@ -1,7 +1,7 @@
11 import { describe, expect, it } from "@effect/vitest";
22 import {
3− createInMemoryRuntimeRunClient,
43 createRuntimeRunClient,
4+ createStaticToolRegistry,
55 } from "@executor-v2/engine";
66 import { makeDenoSubprocessRuntimeAdapter } from "@executor-v2/runtime-deno-subprocess";
77 import { makeLocalInProcessRuntimeAdapter } from "@executor-v2/runtime-local-inproc";
@@ -92,11 +92,14 @@ describe("toAiSdkTools", () => { · 84 unmodified lines
9292 },
9393 });
9494
95− const runClient = createInMemoryRuntimeRunClient({
96− runtimeAdapter: makeLocalInProcessRuntimeAdapter(),
97− tools: {
98− search_docs: searchDocsTool,
99− },
95+ const runtimeAdapter = makeLocalInProcessRuntimeAdapter();
96+ const runClient = createRuntimeRunClient({
97+ runtimeAdapter,
98+ toolRegistry: createStaticToolRegistry({
99+ tools: {
100+ search_docs: searchDocsTool,
101+ },
102+ }),
100103 defaults: {
101104 timeoutMs: 30_000,
102105 },
@@ -133,6 +136,9 @@ describe("toAiSdkTools", () => { · 30 unmodified lines
133136
134137 const runClient = createRuntimeRunClient({
135138 runtimeAdapter,
139+ toolRegistry: createStaticToolRegistry({
140+ tools: {},
141+ }),
136142 });
137143
138144 const tools = toAiSdkTools({
v2/packages/engine/src/index.ts+15−4
@@ -9,18 +9,29 @@ export { · 8 unmodified lines
99 type RuntimeAdapterRegistry,
1010 type RuntimeExecuteError,
1111 type RuntimeExecuteInput,
12− type RuntimeRunnableTool,
1312 type RuntimeToolCallService,
1413 } from "./runtime-adapters";
1514
1615 export {
17− createInMemoryRuntimeRunClient,
1816 createRuntimeRunClient,
19− type CreateInMemoryRuntimeRunClientOptions,
2017 type CreateRuntimeRunClientOptions,
18+} from "./run-client";
19+
20+export {
21+ createRuntimeToolCallService,
22+ createStaticToolRegistry,
2123 type InMemorySandboxTool,
2224 type InMemorySandboxToolMap,
23−} from "./run-client";
25+ type ToolRegistry,
26+ type ToolRegistryCallInput,
27+ type ToolRegistryCatalogNamespacesInput,
28+ type ToolRegistryCatalogNamespacesOutput,
29+ type ToolRegistryCatalogToolsInput,
30+ type ToolRegistryCatalogToolsOutput,
31+ type ToolRegistryDiscoverInput,
32+ type ToolRegistryDiscoverOutput,
33+ type ToolRegistryToolSummary,
34+} from "./tool-registry";
2435
2536 export {
2637 makeOpenApiToolProvider,
v2/packages/engine/src/run-client.ts+6−71
@@ -1,12 +1,8 @@
11 import * as Effect from "effect/Effect";
22 import * as Either from "effect/Either";
33
4−import type {
5− RuntimeAdapter,
6− RuntimeExecuteError,
7− RuntimeToolCallService,
8−} from "./runtime-adapters";
9−import { RuntimeAdapterError } from "./runtime-adapters";
4+import type { RuntimeAdapter, RuntimeExecuteError } from "./runtime-adapters";
5+import { createRuntimeToolCallService, type ToolRegistry } from "./tool-registry";
106
117 export type RuntimeRunClientExecuteInput = {
128 code: string;
@@ -28,25 +24,9 @@ export type RuntimeRunClient = { · 15 unmodified lines
2824 ) => Promise<RuntimeRunClientExecuteResult>;
2925 };
3026
31−export type InMemorySandboxTool = {
32− description?: string | null;
33− execute?: (...args: Array<any>) => Promise<any> | any;
34−};
35−
36−export type InMemorySandboxToolMap = Record<string, InMemorySandboxTool>;
37−
3827 export type CreateRuntimeRunClientOptions = {
3928 runtimeAdapter: RuntimeAdapter;
40− toolCallService?: RuntimeToolCallService;
41− defaults?: {
42− timeoutMs?: number;
43− };
44− makeRunId?: () => string;
45−};
46−
47−export type CreateInMemoryRuntimeRunClientOptions = {
48− runtimeAdapter: RuntimeAdapter;
49− tools: InMemorySandboxToolMap;
29+ toolRegistry: ToolRegistry;
5030 defaults?: {
5131 timeoutMs?: number;
5232 };
@@ -56,47 +36,11 @@ export type CreateInMemoryRuntimeRunClientOptions = { · 3 unmodified lines
5636 const formatRuntimeExecuteError = (error: RuntimeExecuteError): string =>
5737 error.details ? `${error.message}: ${error.details}` : error.message;
5838
59−const inMemoryToolCallService = (
60− runtimeKind: string,
61− tools: InMemorySandboxToolMap,
62−): RuntimeToolCallService => ({
63− callTool: (input) => {
64− const implementation = tools[input.toolPath];
65− if (!implementation) {
66− return new RuntimeAdapterError({
67− operation: "call_tool",
68− runtimeKind,
69− message: `Unknown in-memory tool: ${input.toolPath}`,
70− details: null,
71− });
72− }
73−
74− if (!implementation.execute) {
75− return new RuntimeAdapterError({
76− operation: "call_tool",
77− runtimeKind,
78− message: `In-memory tool '${input.toolPath}' has no execute function`,
79− details: null,
80− });
81− }
82−
83− return Effect.tryPromise({
84− try: () => implementation.execute!(input.input ?? {}, undefined),
85− catch: (cause) =>
86− new RuntimeAdapterError({
87− operation: "call_tool",
88− runtimeKind,
89− message: `In-memory tool invocation failed: ${input.toolPath}`,
90− details: String(cause),
91− }),
92− });
93− },
94−});
95−
9639 export const createRuntimeRunClient = (
9740 options: CreateRuntimeRunClientOptions,
9841 ): RuntimeRunClient => {
9942 const runIdFactory = options.makeRunId ?? (() => `run_${crypto.randomUUID()}`);
43+ const toolCallService = createRuntimeToolCallService(options.toolRegistry);
10044
10145 return {
10246 execute: async (
@@ -130,7 +74,7 @@ export const createRuntimeRunClient = ( · 27 unmodified lines
13074 runId,
13175 code: input.code,
13276 timeoutMs: input.timeoutMs ?? options.defaults?.timeoutMs,
133− toolCallService: options.toolCallService,
77+ toolCallService,
13478 }),
13579 ),
13680 );
@@ -152,13 +96,4 @@ export const createRuntimeRunClient = ( · 15 unmodified lines
15296 };
15397 };
15498
155−export const createInMemoryRuntimeRunClient = (
156− options: CreateInMemoryRuntimeRunClientOptions,
157−): RuntimeRunClient => {
158− return createRuntimeRunClient({
159− runtimeAdapter: options.runtimeAdapter,
160− toolCallService: inMemoryToolCallService(options.runtimeAdapter.kind, options.tools),
161− defaults: options.defaults,
162− makeRunId: options.makeRunId,
163− });
164−};
99+
v2/packages/engine/src/runtime-adapters.ts+0−9
@@ -3,17 +3,8 @@ import * as Data from "effect/Data"; · 2 unmodified lines
33 import * as Effect from "effect/Effect";
44 import * as Layer from "effect/Layer";
55
6−import type { Source } from "@executor-v2/schema";
7−
8−import type { CanonicalToolDescriptor } from "./tool-providers";
9−
106 export type RuntimeAdapterKind = string;
117
12−export type RuntimeRunnableTool = {
13− descriptor: CanonicalToolDescriptor;
14− source: Source | null;
15−};
16−
178 export type RuntimeToolCallService = {
189 callTool: (input: {
1910 runId: string;
v2/packages/engine/src/tool-registry.test.ts+100−0
@@ -0,0 +1,100 @@
1+import { describe, expect, it } from "@effect/vitest";
2+import * as Effect from "effect/Effect";
3+
4+import {
5+ createRuntimeToolCallService,
6+ createStaticToolRegistry,
7+} from "./tool-registry";
8+
9+describe("tool registry", () => {
10+ it.effect("supports callTool plus discover and catalog", () =>
11+ Effect.gen(function* () {
12+ const registry = createStaticToolRegistry({
13+ tools: {
14+ search_docs: {
15+ description: "Search docs",
16+ execute: (input: { query: string }) => ({ hits: [input.query] }),
17+ },
18+ "github.repos.get": {
19+ description: "Get repository",
20+ execute: (input: { owner: string; repo: string }) => ({
21+ full_name: `${input.owner}/${input.repo}`,
22+ }),
23+ },
24+ },
25+ });
26+
27+ const discovered = yield* registry.discover({ query: "github", limit: 5 });
28+ expect(discovered.bestPath).toBe("github.repos.get");
29+ expect(discovered.results.some((entry) => entry.path === "github.repos.get")).toBe(
30+ true,
31+ );
32+
33+ const namespaces = yield* registry.catalogNamespaces({});
34+ expect(namespaces.namespaces.map((namespace) => namespace.namespace)).toEqual([
35+ "github",
36+ "search_docs",
37+ ]);
38+
39+ const catalogTools = yield* registry.catalogTools({ namespace: "github" });
40+ expect(catalogTools.results).toHaveLength(1);
41+ expect(catalogTools.results[0]?.path).toBe("github.repos.get");
42+
43+ const callResult = yield* registry.callTool({
44+ runId: "run_1",
45+ callId: "call_1",
46+ toolPath: "github.repos.get",
47+ input: {
48+ owner: "octocat",
49+ repo: "hello-world",
50+ },
51+ });
52+ expect(callResult).toEqual({
53+ full_name: "octocat/hello-world",
54+ });
55+ }),
56+ );
57+
58+ it.effect("maps runtime tool paths to registry methods", () =>
59+ Effect.gen(function* () {
60+ const calls: Array<string> = [];
61+
62+ const registry = createStaticToolRegistry({
63+ tools: {
64+ ping: {
65+ execute: () => "pong",
66+ },
67+ },
68+ });
69+
70+ const runtimeToolCallService = createRuntimeToolCallService({
71+ ...registry,
72+ callTool: (input) => {
73+ calls.push(input.toolPath);
74+ return registry.callTool(input);
75+ },
76+ });
77+
78+ const discoverResult = yield* runtimeToolCallService.callTool({
79+ runId: "run_2",
80+ callId: "call_2",
81+ toolPath: "discover",
82+ input: {
83+ query: "ping",
84+ },
85+ });
86+
87+ expect((discoverResult as { bestPath: string | null }).bestPath).toBe("ping");
88+ expect(calls).toEqual([]);
89+
90+ const toolResult = yield* runtimeToolCallService.callTool({
91+ runId: "run_2",
92+ callId: "call_3",
93+ toolPath: "ping",
94+ });
95+
96+ expect(toolResult).toBe("pong");
97+ expect(calls).toEqual(["ping"]);
98+ }),
99+ );
100+});
v2/packages/engine/src/tool-registry.ts+317−0
@@ -0,0 +1,317 @@
1+import * as Effect from "effect/Effect";
2+
3+import {
4+ RuntimeAdapterError,
5+ type RuntimeToolCallService,
6+} from "./runtime-adapters";
7+
8+export type ToolRegistryCallInput = {
9+ runId: string;
10+ callId: string;
11+ toolPath: string;
12+ input?: Record<string, unknown>;
13+};
14+
15+export type ToolRegistryDiscoverInput = {
16+ query?: string;
17+ limit?: number;
18+ compact?: boolean;
19+ includeSchemas?: boolean;
20+};
21+
22+export type ToolRegistryCatalogNamespacesInput = {
23+ limit?: number;
24+};
25+
26+export type ToolRegistryCatalogToolsInput = {
27+ namespace?: string;
28+ query?: string;
29+ limit?: number;
30+ compact?: boolean;
31+ includeSchemas?: boolean;
32+};
33+
34+export type ToolRegistryToolSummary = {
35+ path: string;
36+ source?: string;
37+ approval: "auto" | "required";
38+ description?: string;
39+};
40+
41+export type ToolRegistryDiscoverOutput = {
42+ bestPath: string | null;
43+ results: Array<ToolRegistryToolSummary>;
44+ total: number;
45+};
46+
47+export type ToolRegistryCatalogNamespacesOutput = {
48+ namespaces: Array<{
49+ namespace: string;
50+ toolCount: number;
51+ samplePaths: Array<string>;
52+ }>;
53+ total: number;
54+};
55+
56+export type ToolRegistryCatalogToolsOutput = {
57+ results: Array<ToolRegistryToolSummary>;
58+ total: number;
59+};
60+
61+export type ToolRegistry = {
62+ callTool: (input: ToolRegistryCallInput) => Effect.Effect<unknown, RuntimeAdapterError>;
63+ discover: (
64+ input: ToolRegistryDiscoverInput,
65+ ) => Effect.Effect<ToolRegistryDiscoverOutput, RuntimeAdapterError>;
66+ catalogNamespaces: (
67+ input: ToolRegistryCatalogNamespacesInput,
68+ ) => Effect.Effect<ToolRegistryCatalogNamespacesOutput, RuntimeAdapterError>;
69+ catalogTools: (
70+ input: ToolRegistryCatalogToolsInput,
71+ ) => Effect.Effect<ToolRegistryCatalogToolsOutput, RuntimeAdapterError>;
72+};
73+
74+export type InMemorySandboxTool = {
75+ description?: string | null;
76+ execute?: (...args: Array<any>) => Promise<any> | any;
77+};
78+
79+export type InMemorySandboxToolMap = Record<string, InMemorySandboxTool>;
80+
81+type StaticToolRegistryOptions = {
82+ tools: InMemorySandboxToolMap;
83+};
84+
85+const staticToolRegistryRuntimeKind = "static-tool-registry";
86+
87+const toRuntimeAdapterError = (
88+ operation: string,
89+ message: string,
90+ details: string | null,
91+): RuntimeAdapterError =>
92+ new RuntimeAdapterError({
93+ runtimeKind: staticToolRegistryRuntimeKind,
94+ operation,
95+ message,
96+ details,
97+ });
98+
99+const normalizeObjectInput = (input: unknown): Record<string, unknown> => {
100+ if (!input || typeof input !== "object" || Array.isArray(input)) {
101+ return {};
102+ }
103+
104+ return input as Record<string, unknown>;
105+};
106+
107+const normalizeDiscoverInput = (input: unknown): ToolRegistryDiscoverInput => {
108+ const value = normalizeObjectInput(input);
109+ return {
110+ query: typeof value.query === "string" ? value.query : undefined,
111+ limit: typeof value.limit === "number" ? value.limit : undefined,
112+ compact: typeof value.compact === "boolean" ? value.compact : undefined,
113+ includeSchemas:
114+ typeof value.includeSchemas === "boolean" ? value.includeSchemas : undefined,
115+ };
116+};
117+
118+const normalizeCatalogNamespacesInput = (
119+ input: unknown,
120+): ToolRegistryCatalogNamespacesInput => {
121+ const value = normalizeObjectInput(input);
122+ return {
123+ limit: typeof value.limit === "number" ? value.limit : undefined,
124+ };
125+};
126+
127+const normalizeCatalogToolsInput = (input: unknown): ToolRegistryCatalogToolsInput => {
128+ const value = normalizeObjectInput(input);
129+ return {
130+ namespace: typeof value.namespace === "string" ? value.namespace : undefined,
131+ query: typeof value.query === "string" ? value.query : undefined,
132+ limit: typeof value.limit === "number" ? value.limit : undefined,
133+ compact: typeof value.compact === "boolean" ? value.compact : undefined,
134+ includeSchemas:
135+ typeof value.includeSchemas === "boolean" ? value.includeSchemas : undefined,
136+ };
137+};
138+
139+const asToolSummaries = (
140+ tools: InMemorySandboxToolMap,
141+): Array<ToolRegistryToolSummary> =>
142+ Object.entries(tools)
143+ .sort(([leftPath], [rightPath]) => leftPath.localeCompare(rightPath))
144+ .map(([path, tool]) => ({
145+ path,
146+ source: "in-memory",
147+ approval: "auto" as const,
148+ description: tool.description ?? undefined,
149+ }));
150+
151+const scorePath = (entry: ToolRegistryToolSummary, query: string): number => {
152+ if (query.length === 0) {
153+ return 1;
154+ }
155+
156+ const lowerPath = entry.path.toLowerCase();
157+ const lowerDescription = (entry.description ?? "").toLowerCase();
158+ const lowerQuery = query.toLowerCase();
159+
160+ if (lowerPath === lowerQuery) {
161+ return 100;
162+ }
163+
164+ if (lowerPath.startsWith(lowerQuery)) {
165+ return 80;
166+ }
167+
168+ if (lowerPath.includes(lowerQuery)) {
169+ return 60;
170+ }
171+
172+ if (lowerDescription.includes(lowerQuery)) {
173+ return 40;
174+ }
175+
176+ return 0;
177+};
178+
179+const inMemoryDiscover = (
180+ tools: InMemorySandboxToolMap,
181+ input: ToolRegistryDiscoverInput,
182+): ToolRegistryDiscoverOutput => {
183+ const limit = Math.max(1, Math.min(50, input.limit ?? 8));
184+ const query = (input.query ?? "").trim().toLowerCase();
185+
186+ const ranked = asToolSummaries(tools)
187+ .map((entry) => ({
188+ entry,
189+ score: scorePath(entry, query),
190+ }))
191+ .filter((item) => item.score > 0)
192+ .sort((left, right) => right.score - left.score)
193+ .slice(0, limit)
194+ .map((item) => item.entry);
195+
196+ return {
197+ bestPath: ranked[0]?.path ?? null,
198+ results: ranked,
199+ total: ranked.length,
200+ };
201+};
202+
203+const inMemoryCatalogNamespaces = (
204+ tools: InMemorySandboxToolMap,
205+ input: ToolRegistryCatalogNamespacesInput,
206+): ToolRegistryCatalogNamespacesOutput => {
207+ const limit = Math.max(1, Math.min(200, input.limit ?? 50));
208+ const grouped = new Map<string, Array<string>>();
209+
210+ for (const path of Object.keys(tools)) {
211+ const namespace = path.split(".")[0] ?? "default";
212+ const list = grouped.get(namespace) ?? [];
213+ list.push(path);
214+ grouped.set(namespace, list);
215+ }
216+
217+ const namespaces = [...grouped.entries()]
218+ .sort((left, right) => left[0].localeCompare(right[0]))
219+ .map(([namespace, paths]) => ({
220+ namespace,
221+ toolCount: paths.length,
222+ samplePaths: [...paths].sort((left, right) => left.localeCompare(right)).slice(0, 3),
223+ }));
224+
225+ return {
226+ namespaces: namespaces.slice(0, limit),
227+ total: namespaces.length,
228+ };
229+};
230+
231+const inMemoryCatalogTools = (
232+ tools: InMemorySandboxToolMap,
233+ input: ToolRegistryCatalogToolsInput,
234+): ToolRegistryCatalogToolsOutput => {
235+ const limit = Math.max(1, Math.min(200, input.limit ?? 50));
236+ const query = (input.query ?? "").trim().toLowerCase();
237+ const namespace = (input.namespace ?? "").trim().toLowerCase();
238+
239+ const filtered = asToolSummaries(tools)
240+ .filter((entry) => {
241+ if (namespace.length > 0) {
242+ const entryNamespace = entry.path.split(".")[0]?.toLowerCase() ?? "";
243+ if (entryNamespace !== namespace) {
244+ return false;
245+ }
246+ }
247+
248+ return scorePath(entry, query) > 0;
249+ })
250+ .slice(0, limit);
251+
252+ return {
253+ results: filtered,
254+ total: filtered.length,
255+ };
256+};
257+
258+export const createStaticToolRegistry = (
259+ options: StaticToolRegistryOptions,
260+): ToolRegistry => ({
261+ callTool: (input) => {
262+ const implementation = options.tools[input.toolPath];
263+ if (!implementation) {
264+ return new RuntimeAdapterError({
265+ operation: "call_tool",
266+ runtimeKind: staticToolRegistryRuntimeKind,
267+ message: `Unknown in-memory tool: ${input.toolPath}`,
268+ details: null,
269+ });
270+ }
271+
272+ if (!implementation.execute) {
273+ return new RuntimeAdapterError({
274+ operation: "call_tool",
275+ runtimeKind: staticToolRegistryRuntimeKind,
276+ message: `In-memory tool '${input.toolPath}' has no execute function`,
277+ details: null,
278+ });
279+ }
280+
281+ return Effect.tryPromise({
282+ try: () => Promise.resolve(implementation.execute!(input.input ?? {}, undefined)),
283+ catch: (cause) =>
284+ toRuntimeAdapterError(
285+ "call_tool",
286+ `In-memory tool invocation failed: ${input.toolPath}`,
287+ String(cause),
288+ ),
289+ });
290+ },
291+ discover: (input) => Effect.succeed(inMemoryDiscover(options.tools, input)),
292+ catalogNamespaces: (input) =>
293+ Effect.succeed(inMemoryCatalogNamespaces(options.tools, input)),
294+ catalogTools: (input) => Effect.succeed(inMemoryCatalogTools(options.tools, input)),
295+});
296+
297+export const createRuntimeToolCallService = (
298+ toolRegistry: ToolRegistry,
299+): RuntimeToolCallService => ({
300+ callTool: (input) => {
301+ if (input.toolPath === "discover") {
302+ return toolRegistry.discover(normalizeDiscoverInput(input.input));
303+ }
304+
305+ if (input.toolPath === "catalog.namespaces") {
306+ return toolRegistry.catalogNamespaces(
307+ normalizeCatalogNamespacesInput(input.input),
308+ );
309+ }
310+
311+ if (input.toolPath === "catalog.tools") {
312+ return toolRegistry.catalogTools(normalizeCatalogToolsInput(input.input));
313+ }
314+
315+ return toolRegistry.callTool(input);
316+ },
317+});