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
2 type RuntimeAdapter,2 type RuntimeAdapter,
3 type RuntimeExecuteError,3 type RuntimeExecuteError,
4 type RuntimeExecuteInput,4 type RuntimeExecuteInput,
5 type RuntimeRunnableTool,
6} from "@executor-v2/engine";5} from "@executor-v2/engine";
7import { makeLocalInProcessRuntimeAdapter } from "@executor-v2/runtime-local-inproc";6import { makeLocalInProcessRuntimeAdapter } from "@executor-v2/runtime-local-inproc";
8 7
9export type { RuntimeAdapter, RuntimeExecuteError, RuntimeExecuteInput, RuntimeRunnableTool };8export type { RuntimeAdapter, RuntimeExecuteError, RuntimeExecuteInput };
10 9
11export { makeLocalInProcessRuntimeAdapter };10export { makeLocalInProcessRuntimeAdapter };
v2/packages/ai-sdk-adapter/src/index.test.ts+12−6
@@ -1,7 +1,7 @@
1import { describe, expect, it } from "@effect/vitest";1import { describe, expect, it } from "@effect/vitest";
2import {2import {
3 createInMemoryRuntimeRunClient,
4 createRuntimeRunClient,3 createRuntimeRunClient,
4 createStaticToolRegistry,
5} from "@executor-v2/engine";5} from "@executor-v2/engine";
6import { makeDenoSubprocessRuntimeAdapter } from "@executor-v2/runtime-deno-subprocess";6import { makeDenoSubprocessRuntimeAdapter } from "@executor-v2/runtime-deno-subprocess";
7import { makeLocalInProcessRuntimeAdapter } from "@executor-v2/runtime-local-inproc";7import { makeLocalInProcessRuntimeAdapter } from "@executor-v2/runtime-local-inproc";
@@ -92,11 +92,14 @@ describe("toAiSdkTools", () => { · 84 unmodified lines
92 },92 },
93 });93 });
94 94
95 const runClient = createInMemoryRuntimeRunClient({95 const runtimeAdapter = makeLocalInProcessRuntimeAdapter();
96 runtimeAdapter: makeLocalInProcessRuntimeAdapter(),96 const runClient = createRuntimeRunClient({
97 tools: {97 runtimeAdapter,
98 search_docs: searchDocsTool,98 toolRegistry: createStaticToolRegistry({
99 },99 tools: {
100 search_docs: searchDocsTool,
101 },
102 }),
100 defaults: {103 defaults: {
101 timeoutMs: 30_000,104 timeoutMs: 30_000,
102 },105 },
@@ -133,6 +136,9 @@ describe("toAiSdkTools", () => { · 30 unmodified lines
133 136
134 const runClient = createRuntimeRunClient({137 const runClient = createRuntimeRunClient({
135 runtimeAdapter,138 runtimeAdapter,
139 toolRegistry: createStaticToolRegistry({
140 tools: {},
141 }),
136 });142 });
137 143
138 const tools = toAiSdkTools({144 const tools = toAiSdkTools({
v2/packages/engine/src/index.ts+15−4
@@ -9,18 +9,29 @@ export { · 8 unmodified lines
9 type RuntimeAdapterRegistry,9 type RuntimeAdapterRegistry,
10 type RuntimeExecuteError,10 type RuntimeExecuteError,
11 type RuntimeExecuteInput,11 type RuntimeExecuteInput,
12 type RuntimeRunnableTool,
13 type RuntimeToolCallService,12 type RuntimeToolCallService,
14} from "./runtime-adapters";13} from "./runtime-adapters";
15 14
16export {15export {
17 createInMemoryRuntimeRunClient,
18 createRuntimeRunClient,16 createRuntimeRunClient,
19 type CreateInMemoryRuntimeRunClientOptions,
20 type CreateRuntimeRunClientOptions,17 type CreateRuntimeRunClientOptions,
18} from "./run-client";
19
20export {
21 createRuntimeToolCallService,
22 createStaticToolRegistry,
21 type InMemorySandboxTool,23 type InMemorySandboxTool,
22 type InMemorySandboxToolMap,24 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";
24 35
25export {36export {
26 makeOpenApiToolProvider,37 makeOpenApiToolProvider,
v2/packages/engine/src/run-client.ts+6−71
@@ -1,12 +1,8 @@
1import * as Effect from "effect/Effect";1import * as Effect from "effect/Effect";
2import * as Either from "effect/Either";2import * as Either from "effect/Either";
3 3
4import type {4import type { RuntimeAdapter, RuntimeExecuteError } from "./runtime-adapters";
5 RuntimeAdapter,5import { createRuntimeToolCallService, type ToolRegistry } from "./tool-registry";
6 RuntimeExecuteError,
7 RuntimeToolCallService,
8} from "./runtime-adapters";
9import { RuntimeAdapterError } from "./runtime-adapters";
10 6
11export type RuntimeRunClientExecuteInput = {7export type RuntimeRunClientExecuteInput = {
12 code: string;8 code: string;
@@ -28,25 +24,9 @@ export type RuntimeRunClient = { · 15 unmodified lines
28 ) => Promise<RuntimeRunClientExecuteResult>;24 ) => Promise<RuntimeRunClientExecuteResult>;
29};25};
30 26
31export type InMemorySandboxTool = {
32 description?: string | null;
33 execute?: (...args: Array<any>) => Promise<any> | any;
34};
35
36export type InMemorySandboxToolMap = Record<string, InMemorySandboxTool>;
37
38export type CreateRuntimeRunClientOptions = {27export type CreateRuntimeRunClientOptions = {
39 runtimeAdapter: RuntimeAdapter;28 runtimeAdapter: RuntimeAdapter;
40 toolCallService?: RuntimeToolCallService;29 toolRegistry: ToolRegistry;
41 defaults?: {
42 timeoutMs?: number;
43 };
44 makeRunId?: () => string;
45};
46
47export type CreateInMemoryRuntimeRunClientOptions = {
48 runtimeAdapter: RuntimeAdapter;
49 tools: InMemorySandboxToolMap;
50 defaults?: {30 defaults?: {
51 timeoutMs?: number;31 timeoutMs?: number;
52 };32 };
@@ -56,47 +36,11 @@ export type CreateInMemoryRuntimeRunClientOptions = { · 3 unmodified lines
56const formatRuntimeExecuteError = (error: RuntimeExecuteError): string =>36const formatRuntimeExecuteError = (error: RuntimeExecuteError): string =>
57 error.details ? `${error.message}: ${error.details}` : error.message;37 error.details ? `${error.message}: ${error.details}` : error.message;
58 38
59const 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
96export const createRuntimeRunClient = (39export const createRuntimeRunClient = (
97 options: CreateRuntimeRunClientOptions,40 options: CreateRuntimeRunClientOptions,
98): RuntimeRunClient => {41): RuntimeRunClient => {
99 const runIdFactory = options.makeRunId ?? (() => `run_${crypto.randomUUID()}`);42 const runIdFactory = options.makeRunId ?? (() => `run_${crypto.randomUUID()}`);
43 const toolCallService = createRuntimeToolCallService(options.toolRegistry);
100 44
101 return {45 return {
102 execute: async (46 execute: async (
@@ -130,7 +74,7 @@ export const createRuntimeRunClient = ( · 27 unmodified lines
130 runId,74 runId,
131 code: input.code,75 code: input.code,
132 timeoutMs: input.timeoutMs ?? options.defaults?.timeoutMs,76 timeoutMs: input.timeoutMs ?? options.defaults?.timeoutMs,
133 toolCallService: options.toolCallService,77 toolCallService,
134 }),78 }),
135 ),79 ),
136 );80 );
@@ -152,13 +96,4 @@ export const createRuntimeRunClient = ( · 15 unmodified lines
152 };96 };
153};97};
154 98
155export const createInMemoryRuntimeRunClient = (99
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};
v2/packages/engine/src/runtime-adapters.ts+0−9
@@ -3,17 +3,8 @@ import * as Data from "effect/Data"; · 2 unmodified lines
3import * as Effect from "effect/Effect";3import * as Effect from "effect/Effect";
4import * as Layer from "effect/Layer";4import * as Layer from "effect/Layer";
5 5
6import type { Source } from "@executor-v2/schema";
7
8import type { CanonicalToolDescriptor } from "./tool-providers";
9
10export type RuntimeAdapterKind = string;6export type RuntimeAdapterKind = string;
11 7
12export type RuntimeRunnableTool = {
13 descriptor: CanonicalToolDescriptor;
14 source: Source | null;
15};
16
17export type RuntimeToolCallService = {8export type RuntimeToolCallService = {
18 callTool: (input: {9 callTool: (input: {
19 runId: string;10 runId: string;
v2/packages/engine/src/tool-registry.test.ts+100−0
@@ -0,0 +1,100 @@
1import { describe, expect, it } from "@effect/vitest";
2import * as Effect from "effect/Effect";
3
4import {
5 createRuntimeToolCallService,
6 createStaticToolRegistry,
7} from "./tool-registry";
8
9describe("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 @@
1import * as Effect from "effect/Effect";
2
3import {
4 RuntimeAdapterError,
5 type RuntimeToolCallService,
6} from "./runtime-adapters";
7
8export type ToolRegistryCallInput = {
9 runId: string;
10 callId: string;
11 toolPath: string;
12 input?: Record<string, unknown>;
13};
14
15export type ToolRegistryDiscoverInput = {
16 query?: string;
17 limit?: number;
18 compact?: boolean;
19 includeSchemas?: boolean;
20};
21
22export type ToolRegistryCatalogNamespacesInput = {
23 limit?: number;
24};
25
26export type ToolRegistryCatalogToolsInput = {
27 namespace?: string;
28 query?: string;
29 limit?: number;
30 compact?: boolean;
31 includeSchemas?: boolean;
32};
33
34export type ToolRegistryToolSummary = {
35 path: string;
36 source?: string;
37 approval: "auto" | "required";
38 description?: string;
39};
40
41export type ToolRegistryDiscoverOutput = {
42 bestPath: string | null;
43 results: Array<ToolRegistryToolSummary>;
44 total: number;
45};
46
47export type ToolRegistryCatalogNamespacesOutput = {
48 namespaces: Array<{
49 namespace: string;
50 toolCount: number;
51 samplePaths: Array<string>;
52 }>;
53 total: number;
54};
55
56export type ToolRegistryCatalogToolsOutput = {
57 results: Array<ToolRegistryToolSummary>;
58 total: number;
59};
60
61export 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
74export type InMemorySandboxTool = {
75 description?: string | null;
76 execute?: (...args: Array<any>) => Promise<any> | any;
77};
78
79export type InMemorySandboxToolMap = Record<string, InMemorySandboxTool>;
80
81type StaticToolRegistryOptions = {
82 tools: InMemorySandboxToolMap;
83};
84
85const staticToolRegistryRuntimeKind = "static-tool-registry";
86
87const 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
99const 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
107const 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
118const 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
127const 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
139const 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
151const 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
179const 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
203const 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
231const 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
258export 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
297export 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});