1/**
2 * Groups AI tools together with their handlers.
3 *
4 * A toolkit connects `Tool` schemas to the handler functions an application
5 * provides for a language model workflow. It can build a handler context or
6 * layer and execute tool calls by name. Execution validates parameters, runs the
7 * handler, encodes the result, supports preliminary streamed results, and
8 * applies the tool's failure mode.
9 *
10 * @stability unstable
11 * @since 4.0.0
12 */
13import * as Cause from "../Cause.ts"
14import * as Context from "../Context.ts"
15import * as Effect from "../Effect.ts"
16import * as Effectable from "../Effectable.ts"
17import * as Fiber from "../Fiber.ts"
18import { identity } from "../Function.ts"
19import * as InternalRecord from "../internal/record.ts"
20import * as Layer from "../Layer.ts"
21import * as Predicate from "../Predicate.ts"
22import * as Queue from "../Queue.ts"
23import * as Result from "../Result.ts"
24import * as Schema from "../Schema.ts"
25import type * as SchemaAST from "../SchemaAST.ts"
26import type * as Scope from "../Scope.ts"
27import * as Stream from "../Stream.ts"
28import * as AiError from "./AiError.ts"
29import * as Tool from "./Tool.ts"
30
31const TypeId = "~effect/ai/Toolkit" as const
32
33/**
34 * Represents a collection of tools which can be used to enhance the
35 * capabilities of a large language model.
36 *
37 * **Example** (Defining AI toolkits)
38 *
39 * ```ts import.meta.vitest
40 * import { Effect, Schema } from "effect"
41 * import { Tool, Toolkit } from "effect/ai"
42 *
43 * const SearchDocs = Tool.make("SearchDocs", {
44 * description: "Search project documentation",
45 * parameters: Schema.Struct({ query: Schema.String }),
46 * success: Schema.Array(Schema.String)
47 * })
48 *
49 * const SummarizeText = Tool.make("SummarizeText", {
50 * description: "Summarize text",
51 * parameters: Schema.Struct({ text: Schema.String }),
52 * success: Schema.String
53 * })
54 *
55 * const AiToolkit = Toolkit.make(SearchDocs, SummarizeText)
56 *
57 * const ready = AiToolkit.pipe(Effect.provide(AiToolkit.toLayer({
58 * SearchDocs: ({ query }) => Effect.succeed([query]),
59 * SummarizeText: ({ text }) => Effect.succeed(text)
60 * })))
61 *
62 * Object.keys((await Effect.runPromise(ready)).tools) // => ["SearchDocs", "SummarizeText"]
63 * ```
64 *
65 * @stability unstable
66 * @category models
67 * @since 4.0.0
68 */
69export interface Toolkit<in out Tools extends Record<string, Tool.Any>> extends
70 Effect.Effect<
71 WithHandler<Tools>,
72 never,
73 Tool.HandlersFor<Tools>
74 >
75{
76 new(_: never): {}
77
78 readonly [TypeId]: typeof TypeId
79
80 /**
81 * A record containing all tools in this toolkit.
82 */
83 readonly tools: Tools
84
85 /**
86 * A helper method which can be used for type-safe handler declarations.
87 */
88 of<Handlers extends HandlersFrom<Tools>>(handlers: Handlers): Handlers
89
90 /**
91 * Converts a toolkit into a `Context` containing handlers for each tool
92 * in the toolkit.
93 */
94 toHandlers<Handlers extends HandlersFrom<Tools>, EX = never, RX = never>(
95 build: Handlers | Effect.Effect<Handlers, EX, RX>
96 ): Effect.Effect<Context.Context<Tool.HandlersFor<Tools>>, EX, RX>
97
98 /**
99 * Converts a toolkit into a `Layer` containing handlers for each tool in the
100 * toolkit.
101 */
102 toLayer<Handlers extends HandlersFrom<Tools>, EX = never, RX = never>(
103 /**
104 * Handler functions or Effect that produces handlers.
105 */
106 build: Handlers | Effect.Effect<Handlers, EX, RX>
107 ): Layer.Layer<Tool.HandlersFor<Tools>, EX, Exclude<RX, Scope.Scope>>
108}
109
110/**
111 * Context provided to tool handlers during execution.
112 *
113 * @stability unstable
114 * @category models
115 * @since 4.0.0
116 */
117export interface HandlerContext<Tool extends Tool.Any> {
118 /**
119 * The unique identifier of the tool call, when available.
120 */
121 readonly toolCallId?: string | undefined
122 /**
123 * Emit a preliminary result during long-running tool calls.
124 *
125 * **Details**
126 *
127 * Preliminary results are streamed to the caller before the handler completes,
128 * enabling real-time progress updates for lengthy operations.
129 */
130 readonly preliminary: (result: Tool.Success<Tool>) => Effect.Effect<void>
131}
132
133/**
134 * Represents any `Toolkit` instance, used for generic constraints.
135 *
136 * @stability unstable
137 * @category utility types
138 * @since 4.0.0
139 */
140export interface Any {
141 readonly [TypeId]: typeof TypeId
142 readonly tools: Record<string, Tool.Any>
143}
144
145/**
146 * A utility type which can be used to extract the tool definitions from a
147 * toolkit.
148 *
149 * @stability unstable
150 * @category utility types
151 * @since 4.0.0
152 */
153export type Tools<T> = T extends Toolkit<infer Tools> ? Tools : never
154
155/**
156 * A utility type which transforms either a record or an array of tools into
157 * a record where keys are tool names and values are the tool instances.
158 *
159 * @stability unstable
160 * @category utility types
161 * @since 4.0.0
162 */
163export type ToolsByName<Tools> = Tools extends Record<string, Tool.Any> ?
164 { readonly [Name in keyof Tools]: Tools[Name] }
165 : Tools extends ReadonlyArray<Tool.Any> ? { readonly [Tool in Tools[number] as Tool["name"]]: Tool }
166 : never
167
168/**
169 * A utility type that maps tool names to their required handler functions.
170 *
171 * **Details**
172 *
173 * Handlers can return either the tool's custom failure type, an `AiErrorReason`
174 * (which will be wrapped in `AiError`), or a full `AiError`.
175 *
176 * @stability unstable
177 * @category utility types
178 * @since 4.0.0
179 */
180export type HandlersFrom<Tools extends Record<string, Tool.Any>> = {
181 readonly [Name in keyof Tools as Tool.RequiresHandler<Tools[Name]> extends true ? Name : never]: (
182 params: Tool.Parameters<Tools[Name]>,
183 context: HandlerContext<Tools[Name]>
184 ) => Effect.Effect<
185 Tool.Success<Tools[Name]>,
186 Tool.Failure<Tools[Name]> | AiError.AiError | AiError.AiErrorReason,
187 Tool.HandlerServices<Tools[Name]>
188 >
189}
190
191/**
192 * A toolkit instance with registered handlers ready for tool execution.
193 *
194 * @stability unstable
195 * @category models
196 * @since 4.0.0
197 */
198export interface WithHandler<in out Tools extends Record<string, Tool.Any>> {
199 /**
200 * The tools available in this toolkit instance.
201 */
202 readonly tools: Tools
203
204 /**
205 * Executes a tool call by name.
206 *
207 * **Details**
208 *
209 * Validates the input parameters, executes the corresponding handler, and
210 * streams back both the typed result and encoded result. Streaming allows
211 * handlers to emit preliminary results before completion.
212 */
213 readonly handle: <Name extends keyof Tools>(
214 /**
215 * The name of the tool to execute.
216 */
217 name: Name,
218 /**
219 * Encoded parameters to decode and pass to the tool handler.
220 */
221 params: Tool.ParametersEncoded<Tools[Name]>,
222 /**
223 * The unique identifier of the tool call.
224 */
225 toolCallId?: string,
226 /**
227 * Schema parse options for tool parameters.
228 */
229 options?: SchemaAST.ParseOptions
230 ) => Effect.Effect<
231 Stream.Stream<
232 Tool.HandlerResult<Tools[Name]>,
233 Tool.HandlerError<Tools[Name]> | AiError.AiError,
234 Tool.HandlerServices<Tools[Name]>
235 >,
236 AiError.AiError,
237 Tool.HandlerServices<Tools[Name]>
238 >
239}
240
241/**
242 * A utility type which can be used to extract the tools from a toolkit with
243 * handlers.
244 *
245 * @stability unstable
246 * @category utility types
247 * @since 4.0.0
248 */
249export type WithHandlerTools<T> = T extends WithHandler<infer Tools> ? Tools : never
250
251/**
252 * Cause annotation identifying the phase of a typed `Toolkit.handle` failure.
253 *
254 * **Details**
255 *
256 * Typed failures raised by `Toolkit.handle` carry one of these origins:
257 *
258 * - `"parameters"`: the tool call arguments failed to decode
259 * - `"handler"`: the tool handler itself failed
260 * - `"result"`: the handler's output failed to validate or encode
261 *
262 * Unannotated causes default to `"result"`, allowing consumers to treat them
263 * as internal failures.
264 *
265 * Returned failures expose the same origin in `Tool.HandlerResult.failureOrigin`.
266 *
267 * **Example** (Reading a handler failure's origin)
268 *
269 * ```ts import.meta.vitest
270 * import { Cause, Context, Effect, Schema, Stream } from "effect"
271 * import { Tool, Toolkit } from "effect/ai"
272 *
273 * const toolkit = Toolkit.make(Tool.make("Lookup", {
274 * failure: Schema.String,
275 * failureMode: "error"
276 * }))
277 *
278 * const program = Effect.gen(function*() {
279 * const handlers = yield* toolkit
280 * yield* handlers.handle("Lookup", {}).pipe(Effect.flatMap(Stream.runDrain))
281 * }).pipe(
282 * Effect.catchCause((cause) =>
283 * Effect.succeed(Context.get(Cause.annotations(cause), Toolkit.FailureOrigin))
284 * ),
285 * Effect.provide(toolkit.toLayer({ Lookup: () => Effect.fail("Not found") }))
286 * )
287 *
288 * await Effect.runPromise(program) // => "handler"
289 * ```
290 *
291 * @stability unstable
292 * @category services
293 * @since 4.0.0
294 */
295export const FailureOrigin = Context.Reference<Tool.FailureOrigin>("effect/ai/Toolkit/FailureOrigin", {
296 defaultValue: () => "result"
297})
298
299const failureCause = <E>(error: E, origin: Tool.FailureOrigin): Cause.Cause<E> =>
300 Cause.annotate(Cause.fail(error), Context.make(FailureOrigin, origin))
301
302const Proto = {
303 ...Effectable.Prototype({
304 label: "Toolkit",
305 evaluate: Effect.fnUntraced(function*(this: Toolkit<Record<string, Tool.Any>>, parent) {
306 const tools = this.tools
307 const services = parent.context
308 const schemasCache = new WeakMap<any, {
309 readonly context: Context.Context<never>
310 readonly handler: Tool.Handler<any>["handler"]
311 readonly decodeParameters: (
312 u: unknown,
313 options?: SchemaAST.ParseOptions
314 ) => Effect.Effect<unknown, Schema.SchemaError>
315 readonly encodeResult: (u: unknown, isFailure: boolean) => Effect.Effect<unknown, Schema.SchemaError>
316 }>()
317
318 const getSchemas = (tool: Tool.Any) => {
319 let schemas = schemasCache.get(tool)
320 if (Predicate.isUndefined(schemas)) {
321 const handler = services.mapUnsafe.get(tool.id)! as Tool.Handler<any>
322 const decodeParameters = Schema.isSchema(tool.parametersSchema)
323 ? Schema.decodeUnknownEffect(tool.parametersSchema) as any
324 : (u: unknown) => Effect.succeed(u)
325 const encodeSuccess = Schema.encodeUnknownEffect(tool.successSchema) as any
326 const encodeFailure = Schema.encodeUnknownEffect(Tool.failureResultSchema(tool)) as any
327 const encodeResult = (u: unknown, isFailure: boolean) => isFailure ? encodeFailure(u) : encodeSuccess(u)
328 schemas = {
329 context: handler.context,
330 handler: handler.handler,
331 decodeParameters,
332 encodeResult
333 }
334 schemasCache.set(tool, schemas)
335 }
336 return schemas
337 }
338
339 const handle = Effect.fnUntraced(function*(
340 name: string,
341 params: unknown,
342 toolCallId?: string,
343 options?: SchemaAST.ParseOptions
344 ) {
345 const tool = Object.hasOwn(tools, name) ? tools[name] : undefined
346
347 yield* Effect.annotateCurrentSpan({
348 tool: name,
349 parameters: params
350 })
351
352 if (Predicate.isUndefined(tool)) {
353 return yield* AiError.make({
354 module: "Toolkit",
355 method: `${name}.handle`,
356 reason: new AiError.ToolNotFoundError({
357 toolName: name,
358 availableTools: Object.keys(tools)
359 })
360 })
361 }
362
363 const schemas = getSchemas(tool)
364
365 const encodeResult = (result: any, isFailure: boolean) =>
366 schemas.encodeResult(result, isFailure).pipe(
367 Effect.mapError((cause) =>
368 AiError.make({
369 module: "Toolkit",
370 method: `${name}.handle`,
371 reason: new AiError.ToolResultEncodingError({
372 toolName: name,
373 toolResult: result,
374 description: cause.message
375 })
376 })
377 )
378 )
379
380 const decodedParamsResult = yield* Effect.result(schemas.decodeParameters(params, options))
381 if (Result.isFailure(decodedParamsResult)) {
382 const error = AiError.make({
383 module: "Toolkit",
384 method: `${name}.handle`,
385 reason: new AiError.ToolParameterValidationError({
386 toolName: name,
387 description: decodedParamsResult.failure.message
388 })
389 })
390 if (tool.failureMode === "error") {
391 return yield* Effect.failCause(failureCause(error, "parameters"))
392 }
393 return Stream.fromEffect(
394 Effect.map(encodeResult(error, true), (encodedResult) => ({
395 result: error,
396 isFailure: true,
397 failureOrigin: "parameters" as const,
398 preliminary: false,
399 encodedResult
400 }))
401 ) satisfies Stream.Stream<Tool.HandlerResult<any>, any>
402 }
403 const decodedParams = decodedParamsResult.success
404
405 const queue = yield* Queue.make<{
406 readonly result: any
407 readonly isFailure: boolean
408 readonly preliminary: boolean
409 }, Cause.Done>()
410 const context: HandlerContext<any> = {
411 toolCallId,
412 preliminary: (result) =>
413 Effect.asVoid(Queue.offer(queue, {
414 result,
415 isFailure: false,
416 preliminary: true
417 }))
418 }
419
420 const fiber = yield* schemas.handler(decodedParams, context).pipe(
421 Effect.flatMap((result) => Queue.offer(queue, { result, isFailure: false, preliminary: false })),
422 Effect.updateContext((input) => Context.merge(schemas.context, input)),
423 Effect.matchCauseEffect({
424 onFailure: (cause) => Queue.failCause(queue, cause),
425 onSuccess: () => Queue.end(queue)
426 }),
427 Effect.forkChild
428 )
429
430 const normalizeError = (error: unknown) => {
431 const normalizedError = Schema.isSchemaError(error)
432 ? AiError.make({
433 module: "Toolkit",
434 method: `${name}.handle`,
435 reason: new AiError.InvalidToolResultError({
436 toolName: name,
437 description: `Tool handler returned invalid result: ${error.message}`
438 })
439 })
440 : AiError.isAiErrorReason(error)
441 ? AiError.make({
442 module: "Toolkit",
443 method: `${name}.handle`,
444 reason: error
445 })
446 : error
447 return normalizedError
448 }
449
450 return Stream.fromQueue(queue).pipe(
451 Stream.catch((error) => {
452 const normalizedError = normalizeError(error)
453 const failureOrigin: Tool.FailureOrigin = Schema.isSchemaError(error) ? "result" : "handler"
454 return tool.failureMode === "error"
455 ? Stream.failCause(failureCause(normalizedError, failureOrigin))
456 : Stream.succeed({ result: normalizedError, isFailure: true, failureOrigin, preliminary: false })
457 }),
458 Stream.mapEffect(Effect.fnUntraced(function*(output) {
459 const encodedResult = yield* encodeResult(output.result, output.isFailure)
460 return { ...output, encodedResult }
461 })),
462 Stream.onEnd(Fiber.interrupt(fiber))
463 ) satisfies Stream.Stream<Tool.HandlerResult<any>, any>
464 })
465
466 return {
467 tools,
468 handle: handle as any
469 } satisfies WithHandler<Record<string, any>>
470 })
471 }),
472 [TypeId]: TypeId,
473 of: identity,
474 toHandlers(
475 this: Toolkit<Record<string, Tool.Any>>,
476 build: Record<string, (params: any) => any> | Effect.Effect<Record<string, (params: any) => any>>
477 ) {
478 return Effect.gen({ self: this }, function*() {
479 const services = yield* Effect.context<never>()
480 const handlers = Effect.isEffect(build) ? yield* build : build
481 const context = new Map<string, unknown>()
482 for (const [name, handler] of Object.entries(handlers)) {
483 const tool = Object.hasOwn(this.tools, name) ? this.tools[name] : undefined
484 if (tool !== undefined) {
485 context.set(tool.id, { name, handler, context: services })
486 }
487 }
488 return Context.makeUnsafe(context)
489 })
490 },
491 toLayer(
492 this: Toolkit<Record<string, Tool.Any>>,
493 build: Record<string, (params: any) => any> | Effect.Effect<Record<string, (params: any) => any>>
494 ) {
495 return Layer.effectContext(this.toHandlers(build))
496 },
497 toJSON(this: Toolkit<any>): unknown {
498 return {
499 _id: "effect/ai/Toolkit",
500 tools: Array.from(Object.values(this.tools)).map((tool) => (tool as Tool.Any).name)
501 }
502 }
503}
504
505const makeProto = <Tools extends Record<string, Tool.Any>>(tools: Tools): Toolkit<Tools> =>
506 Object.assign(function() {}, Proto, { tools }) as any
507
508const resolveInput = <Tools extends ReadonlyArray<Tool.Any>>(
509 ...tools: Tools
510): Record<string, Tools[number]> => {
511 const output = {} as Record<string, Tools[number]>
512 for (const tool of tools) {
513 InternalRecord.assignProperty(output, tool.name, tool)
514 }
515 return output
516}
517
518/**
519 * An empty toolkit with no tools.
520 *
521 * **When to use**
522 *
523 * Use when you need an empty starting point for building toolkits or a default
524 * toolkit value that can be extended with `merge`.
525 *
526 * @stability unstable
527 * @category constructors
528 * @since 4.0.0
529 */
530export const empty: Toolkit<{}> = makeProto({})
531
532/**
533 * Creates a new toolkit from the specified tools.
534 *
535 * **Details**
536 *
537 * This is the primary constructor for creating toolkits. It accepts multiple
538 * tools and organizes them into a toolkit that can be provided to AI language
539 * models.
540 *
541 * **Example** (Creating a toolkit)
542 *
543 * ```ts import.meta.vitest
544 * import { Effect, Schema } from "effect"
545 * import { Tool, Toolkit } from "effect/ai"
546 *
547 * const GetCurrentTime = Tool.make("GetCurrentTime", {
548 * description: "Get the current timestamp",
549 * success: Schema.Number
550 * })
551 *
552 * const GetWeather = Tool.make("get_weather", {
553 * description: "Get weather information",
554 * parameters: Schema.Struct({ location: Schema.String }),
555 * success: Schema.Struct({
556 * temperature: Schema.Number,
557 * condition: Schema.String
558 * })
559 * })
560 *
561 * const toolkit = Toolkit.make(GetCurrentTime, GetWeather)
562 * const ready = toolkit.pipe(Effect.provide(toolkit.toLayer({
563 * GetCurrentTime: () => Effect.succeed(0),
564 * get_weather: () => Effect.succeed({ temperature: 20, condition: "clear" })
565 * })))
566 *
567 * Object.keys((await Effect.runPromise(ready)).tools) // => ["GetCurrentTime", "get_weather"]
568 * ```
569 *
570 * @stability unstable
571 * @category constructors
572 * @since 4.0.0
573 */
574export const make = <Tools extends ReadonlyArray<Tool.Any>>(
575 ...tools: Tools
576): Toolkit<ToolsByName<Tools>> => makeProto(resolveInput(...tools)) as any
577
578/**
579 * A utility type which flattens a record type for improved IDE display.
580 *
581 * @stability unstable
582 * @category utility types
583 * @since 4.0.0
584 */
585export type SimplifyRecord<T> = { [K in keyof T]: T[K] } & {}
586
587/**
588 * A utility type which merges a union of tool records into a single record.
589 *
590 * @stability unstable
591 * @category utility types
592 * @since 4.0.0
593 */
594export type MergeRecords<U> = {
595 readonly [K in Extract<U extends unknown ? keyof U : never, string>]: Extract<
596 U extends Record<K, infer V> ? V : never,
597 Tool.Any
598 >
599}
600
601/**
602 * A utility type which merges the tools from multiple toolkits into a single
603 * record.
604 *
605 * @stability unstable
606 * @category utility types
607 * @since 4.0.0
608 */
609export type MergedTools<Toolkits extends ReadonlyArray<Any>> = SimplifyRecord<
610 MergeRecords<Tools<Toolkits[number]>>
611>
612
613/**
614 * Merges multiple toolkits into a single toolkit.
615 *
616 * **Details**
617 *
618 * Combines all tools from the provided toolkits into one unified toolkit.
619 * If there are naming conflicts, tools from later toolkits will override
620 * tools from earlier ones.
621 *
622 * **Example** (Merging toolkits)
623 *
624 * ```ts import.meta.vitest
625 * import { Effect, Schema } from "effect"
626 * import { Tool, Toolkit } from "effect/ai"
627 *
628 * const mathToolkit = Toolkit.make(
629 * Tool.make("add", { success: Schema.Number }),
630 * Tool.make("subtract", { success: Schema.Number })
631 * )
632 *
633 * const utilityToolkit = Toolkit.make(
634 * Tool.make("get_time", { success: Schema.Number }),
635 * Tool.make("get_weather", { success: Schema.String })
636 * )
637 *
638 * const combined = Toolkit.merge(mathToolkit, utilityToolkit)
639 * const ready = combined.pipe(Effect.provide(combined.toLayer({
640 * add: () => Effect.succeed(1),
641 * subtract: () => Effect.succeed(0),
642 * get_time: () => Effect.succeed(0),
643 * get_weather: () => Effect.succeed("clear")
644 * })))
645 *
646 * Object.keys((await Effect.runPromise(ready)).tools) // => ["add", "subtract", "get_time", "get_weather"]
647 * ```
648 *
649 * @stability unstable
650 * @category constructors
651 * @since 4.0.0
652 */
653export const merge = <const Toolkits extends ReadonlyArray<Any>>(
654 /**
655 * The toolkits to merge together.
656 */
657 ...toolkits: Toolkits
658): Toolkit<MergedTools<Toolkits>> => {
659 const tools = {} as Record<string, any>
660 for (const toolkit of toolkits) {
661 for (const [name, tool] of Object.entries(toolkit.tools)) {
662 InternalRecord.assignProperty(tools, name, tool)
663 }
664 }
665 return makeProto(tools) as any
666}