@effect/ai 0.14.1 → 0.16.0

package/src/AiToolkit.ts

@@ -3,299 +3,259 @@
  */
 import * as Context from "effect/Context"
 import * as Effect from "effect/Effect"
-import * as Effectable from "effect/Effectable"
-import { identity } from "effect/Function"
-import * as HashMap from "effect/HashMap"
-import * as Inspectable from "effect/Inspectable"
+import { CommitPrototype } from "effect/Effectable"
+import type { Inspectable } from "effect/Inspectable"
+import { BaseProto as InspectableProto } from "effect/Inspectable"
 import * as Layer from "effect/Layer"
+import type { ParseError } from "effect/ParseResult"
+import type { Pipeable } from "effect/Pipeable"
 import { pipeArguments } from "effect/Pipeable"
-import type * as Schema from "effect/Schema"
-import type { Scope } from "effect/Scope"
-import type * as Types from "effect/Types"
+import * as Predicate from "effect/Predicate"
+import * as Schema from "effect/Schema"
+import type * as Scope from "effect/Scope"
+import { AiError } from "./AiError.js"
+import * as AiTool from "./AiTool.js"
 
 /**
  * @since 1.0.0
- * @category type ids
+ * @category Type Ids
  */
 export const TypeId: unique symbol = Symbol.for("@effect/ai/AiToolkit")
 
 /**
  * @since 1.0.0
- * @category type ids
+ * @category Type Ids
  */
 export type TypeId = typeof TypeId
 
 /**
+ * An `AiToolkit` represents a set of tools that a large language model can
+ * use to augment its response.
+ *
  * @since 1.0.0
- * @category models
+ * @category Models
  */
-export interface AiToolkit<in out Tools extends Tool.AnySchema>
-  extends Effect.Effect<Handlers<Tools>, never, Tool.Services<Tools> | Registry>, Inspectable.Inspectable
+export interface AiToolkit<in out Tools extends AiTool.Any>
+  extends Effect.Effect<ToHandler<Tools>, never, AiTool.ToHandler<Tools>>, Inspectable, Pipeable
 {
-  readonly [TypeId]: TypeId
-  readonly tools: HashMap.HashMap<string, Tools>
-  readonly add: <S extends Tool.AnySchema>(tool: S) => AiToolkit<Tools | S>
-  readonly addAll: <ToAdd extends ReadonlyArray<Tool.AnySchema>>(
-    ...tools: ToAdd
-  ) => AiToolkit<Tools | ToAdd[number]>
-  readonly concat: <T extends Tool.AnySchema>(that: AiToolkit<T>) => AiToolkit<Tools | T>
-  readonly implement: <R, EX = never, RX = never>(
-    f: (
-      handlers: Handlers<Tools>
-    ) => Handlers<never, R> | Effect.Effect<Handlers<never, R>, EX, RX>
-  ) => Layer.Layer<Tool.ServiceFromTag<Tools["_tag"]> | Registry, EX, Exclude<R | RX, Scope>>
-}
-
-/**
- * @since 1.0.0
- * @category models
- */
-export declare namespace AiToolkit {
-  /**
-   * @since 1.0.0
-   * @category models
-   */
-  export type Tools<A> = A extends AiToolkit<infer Tools> ? Tools : never
-
-  /**
-   * @since 1.0.0
-   * @category models
-   */
-  export type SuccessSchema<A> = A extends AiToolkit<infer Tools> ? Tools["success"] : never
-}
-
-/**
- * @since 1.0.0
- * @category tool
- */
-export declare namespace Tool {
-  /**
-   * @since 1.0.0
-   * @category tool
-   */
-  export interface AnySchema {
-    readonly [Schema.TypeId]: any
-    readonly _tag: string
-    readonly Type: Schema.SerializableWithResult.All
-    readonly success: Schema.Schema.Any
-  }
-
-  /**
-   * @since 1.0.0
-   * @category tool
-   */
-  export type Success<Tool extends AnySchema> = Schema.WithResult.Success<Tool["Type"]>
-
-  /**
-   * @since 1.0.0
-   * @category tool
-   */
-  export type Failure<Tool extends AnySchema> = Schema.WithResult.Failure<Tool["Type"]>
-
-  /**
-   * @since 1.0.0
-   * @category tool
-   */
-  export type Context<Tool extends AnySchema> = Schema.WithResult.Context<Tool["Type"]>
-
-  /**
-   * @since 1.0.0
-   * @category tool
-   */
-  export type Handler<Tool extends AnySchema, R> = (
-    params: Tool["Type"]
-  ) => Effect.Effect<Success<Tool>, Failure<Tool>, R>
+  new(_: never): {}
 
-  /**
-   * @since 1.0.0
-   * @category tool
-   */
-  export type HandlerAny = (params: any) => Effect.Effect<any, any, any>
+  readonly [TypeId]: TypeId
 
   /**
-   * @since 1.0.0
-   * @category tool
+   * A map containing the tools that are part of this toolkit.
    */
-  export interface Service<Tag extends string> {
-    readonly _: unique symbol
-    readonly name: Types.Invariant<Tag>
-  }
+  readonly tools: AiTool.ByName<Tools>
 
   /**
-   * @since 1.0.0
-   * @category tool
+   * Converts this toolkit into a `Context` object containing the handlers for
+   * all tools in the toolkit.
    */
-  export type ServiceFromTag<Tag extends string> = Tag extends infer T ? T extends string ? Service<T> : never : never
+  toContext<Handlers extends HandlersFrom<Tools>, EX = never, RX = never>(
+    build: Handlers | Effect.Effect<Handlers, EX, RX>
+  ): Effect.Effect<Context.Context<AiTool.ToHandler<Tools>>, EX, RX>
 
   /**
-   * @since 1.0.0
-   * @category tool
+   * Converts this toolkit into a `Layer` containing the handlers for all tools
+   * in the toolkit.
    */
-  export type Services<Tools extends AnySchema> = ServiceFromTag<Tools["_tag"]>
+  toLayer<Handlers extends HandlersFrom<Tools>, EX = never, RX = never>(
+    build: Handlers | Effect.Effect<Handlers, EX, RX>
+  ): Layer.Layer<AiTool.ToHandler<Tools>, EX, Exclude<RX, Scope.Scope>>
 }
 
 /**
  * @since 1.0.0
- * @category registry
+ * @category Models
  */
-export class Registry extends Context.Tag("@effect/ai/AiToolkit/Registry")<
-  Registry,
-  Map<Tool.AnySchema, Tool.HandlerAny>
->() {
-  static readonly Live: Layer.Layer<Registry> = Layer.sync(Registry, () => new Map())
-}
-
-class AiToolkitImpl<Tools extends Tool.AnySchema>
-  extends Effectable.Class<Handlers<Tools>, never, Tool.Services<Tools> | Registry>
-  implements AiToolkit<Tools>
-{
+export interface Any {
   readonly [TypeId]: TypeId
-  constructor(readonly tools: HashMap.HashMap<string, Tools>) {
-    super()
-    this[TypeId] = TypeId
-  }
-  toJSON(): unknown {
-    return {
-      _id: "@effect/ai/AiToolkit",
-      tools: [...HashMap.values(this.tools)].map((tool) => tool._tag)
-    }
-  }
-  toString(): string {
-    return Inspectable.format(this)
-  }
-  [Inspectable.NodeInspectSymbol](): string {
-    return Inspectable.format(this)
-  }
-  pipe() {
-    return pipeArguments(this, arguments)
-  }
-  add<S extends Tool.AnySchema>(tool: S): AiToolkit<Tools | S> {
-    return new AiToolkitImpl(HashMap.set(this.tools, tool._tag, tool as any)) as any
-  }
-  addAll<ToAdd extends ReadonlyArray<Tool.AnySchema>>(...tools: ToAdd): AiToolkit<Tools | ToAdd[number]> {
-    let map = this.tools
-    for (const tool of tools) {
-      map = HashMap.set(map, tool._tag, tool as any)
-    }
-    return new AiToolkitImpl(map as any)
-  }
-  concat<T extends Tool.AnySchema>(that: AiToolkit<T>): AiToolkit<Tools | T> {
-    return new AiToolkitImpl(HashMap.union(this.tools, that.tools))
-  }
-  implement<R, EX = never, RX = never>(
-    f: (
-      handlers: Handlers<Tools>
-    ) => Handlers<never, R> | Effect.Effect<Handlers<never, R>, EX, RX>
-  ): Layer.Layer<Tool.ServiceFromTag<Tools["_tag"]> | Registry, EX, Exclude<R | RX, Scope>> {
-    return registerHandlers(this as any, f as any).pipe(Layer.scopedDiscard, Layer.provideMerge(Registry.Live))
-  }
-  commit(): Effect.Effect<Handlers<Tools>, never, Tool.ServiceFromTag<Tools["_tag"]> | Registry> {
-    return Effect.map(Registry, (map) => {
-      let handlers = HashMap.empty<string, Tool.HandlerAny>()
-      for (const [tag, tool] of this.tools) {
-        handlers = HashMap.set(handlers, tag, map.get(tool)!)
-      }
-      return new HandlersImpl(this as any, handlers)
-    }) as any
-  }
+  readonly tools: ReadonlyMap<string, AiTool.Any>
 }
 
-const registerHandlers = (
-  toolkit: AiToolkit<any>,
-  f: (handlers: Handlers<any, any>) => Handlers<any, any> | Effect.Effect<Handlers<any, any>>
-) =>
-  Effect.context<any>().pipe(
-    Effect.bindTo("context"),
-    Effect.bind("handlers", () => {
-      const handlers = f(HandlersImpl.fromToolkit(toolkit))
-      return Effect.isEffect(handlers) ? handlers : Effect.succeed(handlers)
-    }),
-    Effect.tap(({ context, handlers }) => {
-      const registry = Context.unsafeGet(context, Registry)
-      for (const [tag, handler] of handlers.handlers) {
-        const tool = HashMap.unsafeGet(handlers.toolkit.tools, tag)
-        registry.set(tool, function(params: any) {
-          return Effect.withSpan(
-            Effect.mapInputContext(handler(params), (input) => Context.merge(input, context)),
-            "AiToolkit.handler",
-            {
-              captureStackTrace: false,
-              attributes: {
-                tool: tag,
-                parameters: params
-              }
-            }
-          )
-        })
-      }
-    })
-  )
-
 /**
+ * Represents an `AiToolkit` which has been augmented with a handler function
+ * for resolving tool call requests.
+ *
  * @since 1.0.0
- * @category constructors
+ * @category Models
  */
-export const empty: AiToolkit<never> = new AiToolkitImpl(HashMap.empty())
+export interface ToHandler<in out Tool extends AiTool.Any> {
+  readonly tools: ReadonlyArray<Tool>
+  readonly handle: (toolName: AiTool.Name<Tool>, toolParams: AiTool.Parameters<Tool>) => AiTool.HandlerEffect<Tool>
+}
 
 /**
+ * A utility mapped type which associates tool names with their handlers.
+ *
  * @since 1.0.0
- * @category handlers
+ * @category Utility Types
  */
-export const HandlersTypeId: unique symbol = Symbol.for("@effect/ai/AiToolkit/Handlers")
+export type HandlersFrom<Tool extends AiTool.Any> = {
+  [Name in Tool as Tool["name"]]: (params: AiTool.Parameters<Tool>) => AiTool.HandlerEffect<Tool>
+}
 
 /**
+ * A utility type which returns the tools in an `AiToolkit`.
+ *
  * @since 1.0.0
- * @category handlers
+ * @category Utility Types
  */
-export type HandlersTypeId = typeof HandlersTypeId
+export type Tools<Toolkit> = Toolkit extends AiToolkit<infer Tool> ? string extends Tool["name"] ? never : Tool : never
+
+const Proto = {
+  ...CommitPrototype,
+  ...InspectableProto,
+  [TypeId]: TypeId,
+  toContext(this: AiToolkit<any>, build: Effect.Effect<Record<string, (params: any) => any>>) {
+    return Effect.gen(this, function*() {
+      const context = yield* Effect.context<never>()
+      const handlers = Effect.isEffect(build) ? yield* build : build
+      const contextMap = new Map<string, unknown>()
+      for (const [name, handler] of Object.entries(handlers)) {
+        const tool = this.tools[name]!
+        contextMap.set(tool.key, { handler, context })
+      }
+      return Context.unsafeMake(contextMap)
+    })
+  },
+  toLayer(this: AiToolkit<any>, build: Effect.Effect<Record<string, (params: any) => any>>) {
+    return Layer.scopedContext(this.toContext(build))
+  },
+  commit(this: AiToolkit<AiTool.AnyWithProtocol>) {
+    return Effect.gen(this, function*() {
+      const context = yield* Effect.context<never>()
+      const tools = this.tools
+      const schemasCache = new WeakMap<any, {
+        readonly context: Context.Context<never>
+        readonly handler: (params: any) => Effect.Effect<any, any>
+        readonly encodeSuccess: (u: unknown) => Effect.Effect<unknown, ParseError>
+        readonly decodeFailure: (u: unknown) => Effect.Effect<AiTool.Failure<any>, ParseError>
+        readonly decodeParameters: (u: unknown) => Effect.Effect<AiTool.Parameters<any>, ParseError>
+      }>()
+      const getSchemas = (tool: AiTool.AnyWithProtocol) => {
+        let schemas = schemasCache.get(tool)
+        if (Predicate.isUndefined(schemas)) {
+          const handler = context.unsafeMap.get(tool.key)! as AiTool.Handler<any>
+          const encodeSuccess = Schema.encodeUnknown(tool.successSchema) as any
+          const decodeFailure = Schema.decodeUnknown(tool.failureSchema as any) as any
+          const decodeParameters = Schema.decodeUnknown(tool.parametersSchema) as any
+          schemas = {
+            context: handler.context,
+            handler: handler.handler,
+            encodeSuccess,
+            decodeFailure,
+            decodeParameters
+          }
+          schemasCache.set(tool, schemas)
+        }
+        return schemas
+      }
+      const handle = Effect.fn("AiToolkit.handler", { captureStackTrace: false })(
+        function*(toolName: string, toolParams: unknown) {
+          yield* Effect.annotateCurrentSpan({
+            tool: toolName,
+            parameters: toolParams
+          })
+          const tool = tools[toolName]!
+          const schemas = getSchemas(tool)
+          const decodedParams = yield* Effect.mapError(
+            schemas.decodeParameters(toolParams),
+            (cause) =>
+              new AiError({
+                module: "AiToolkit",
+                method: `${toolName}.handle`,
+                description: `Failed to decode tool call parameters for tool '${toolName}' from '${toolParams}'`,
+                cause
+              })
+          )
+          const result = yield* schemas.handler(decodedParams).pipe(
+            Effect.mapInputContext((input) => Context.merge(schemas.context, input)),
+            Effect.catchAll((error) =>
+              schemas.decodeFailure(error).pipe(
+                Effect.mapError((cause) =>
+                  new AiError({
+                    module: "AiToolkit",
+                    method: `${toolName}.handle`,
+                    description: `Failed to decode tool call failure for tool '${toolName}'`,
+                    cause
+                  })
+                ),
+                Effect.flatMap(Effect.fail)
+              )
+            )
+          )
+          const encodedResult = yield* Effect.mapError(
+            schemas.encodeSuccess(result),
+            (cause) =>
+              new AiError({
+                module: "AiToolkit",
+                method: `${toolName}.handle`,
+                description: `Failed to encode tool call result for tool '${toolName}'`,
+                cause
+              })
+          )
+          return {
+            result,
+            encodedResult
+          } satisfies AiTool.HandlerResult<any>
+        }
+      )
+      return {
+        tools: Array.from(Object.values(tools)),
+        handle
+      }
+    })
+  },
+  toJSON(this: AiToolkit<any>): unknown {
+    return {
+      _id: "@effect/ai/AiToolkit",
+      tools: Array.from(Object.values(this.tools)).map((tool) => tool.name)
+    }
+  },
+  pipe() {
+    return pipeArguments(this, arguments)
+  }
+}
+
+const makeProto = <Tools extends AiTool.Any>(tools: Record<string, Tools>): AiToolkit<Tools> =>
+  Object.assign(function() {}, Proto, { tools }) as any
+
+const resolveInput = <Tools extends ReadonlyArray<AiTool.Any>>(
+  ...tools: Tools
+): Record<string, Tools[number]> => {
+  const output = {} as Record<string, Tools[number]>
+  for (const tool of tools) {
+    const value = (Schema.isSchema(tool) ? AiTool.fromTaggedRequest(tool as any) : tool) as any
+    output[tool.name] = value
+  }
+  return output
+}
 
 /**
+ * Constructs a new `AiToolkit` from the specified tools.
+ *
  * @since 1.0.0
- * @category handlers
+ * @category Constructors
  */
-export interface Handlers<in out Tools extends Tool.AnySchema, R = never> {
-  readonly [HandlersTypeId]: Handlers.Variance<Tools>
-  readonly toolkit: AiToolkit<Tools>
-  readonly handlers: HashMap.HashMap<string, Tool.Handler<any, R>>
-  readonly handle: <Tag extends Types.Tags<Tools>, RH>(
-    tag: Tag,
-    f: Tool.Handler<Types.ExtractTag<Tools, Tag>, RH>
-  ) => Handlers<Types.ExcludeTag<Tools, Tag>, R | RH | Tool.Context<Types.ExtractTag<Tools, Tag>>>
-}
+export const make = <const Tools extends ReadonlyArray<AiTool.Any>>(
+  ...tools: Tools
+): AiToolkit<Tools[number]> => makeProto(resolveInput(...tools))
 
 /**
+ * Merges this toolkit with one or more other toolkits.
+ *
  * @since 1.0.0
- * @category handlers
+ * @category Merging
  */
-export declare namespace Handlers {
-  /**
-   * @since 1.0.0
-   * @category handlers
-   */
-  export interface Variance<Tools extends Tool.AnySchema> {
-    readonly _Tools: Types.Invariant<Tools>
-  }
-}
-
-const handlersVariance = {
-  _Tools: identity
-}
-
-class HandlersImpl<Tools extends Tool.AnySchema, R = never> implements Handlers<Tools, R> {
-  readonly [HandlersTypeId]: Handlers.Variance<Tools>
-  constructor(
-    readonly toolkit: AiToolkit<Tools>,
-    readonly handlers: HashMap.HashMap<string, Tool.Handler<any, R>>
-  ) {
-    this[HandlersTypeId] = handlersVariance
-  }
-  static fromToolkit<Tools extends Tool.AnySchema>(toolkit: AiToolkit<Tools>): Handlers<Tools> {
-    return new HandlersImpl(toolkit, HashMap.empty())
-  }
-  handle<Tag extends Types.Tags<Tools>, RH>(
-    tag: Tag,
-    f: Tool.Handler<Types.ExtractTag<Tools, Tag>, RH>
-  ): Handlers<Types.ExcludeTag<Tools, Tag>, R | RH | Tool.Context<Types.ExtractTag<Tools, Tag>>> {
-    return new HandlersImpl(this.toolkit as any, HashMap.set(this.handlers, tag, f as any))
+export const merge = <const Toolkits extends ReadonlyArray<Any>>(
+  ...toolkits: Toolkits
+): AiToolkit<Tools<Toolkits[number]>> => {
+  const tools = {} as Record<string, any>
+  for (const toolkit of toolkits) {
+    for (const [name, tool] of toolkit.tools) {
+      tools[name] = tool
+    }
   }
+  return makeProto(tools) as any
 }

package/src/Tokenizer.ts

@@ -1,17 +1,15 @@
 /**
  * @since 1.0.0
  */
-import * as Chunk from "effect/Chunk"
 import * as Context from "effect/Context"
 import * as Effect from "effect/Effect"
-import * as Option from "effect/Option"
+import * as Predicate from "effect/Predicate"
 import type { AiError } from "./AiError.js"
-import type { Message } from "./AiInput.js"
 import * as AiInput from "./AiInput.js"
 
 /**
  * @since 1.0.0
- * @category tags
+ * @category Tags
  */
 export class Tokenizer extends Context.Tag("@effect/ai/Tokenizer")<
   Tokenizer,
@@ -20,22 +18,21 @@ export class Tokenizer extends Context.Tag("@effect/ai/Tokenizer")<
 
 /**
  * @since 1.0.0
- * @category models
  */
 export declare namespace Tokenizer {
   /**
    * @since 1.0.0
-   * @models
+   * @category Models
    */
   export interface Service {
-    readonly tokenize: (content: AiInput.Input) => Effect.Effect<Array<number>, AiError>
-    readonly truncate: (content: AiInput.Input, tokens: number) => Effect.Effect<AiInput.Input, AiError>
+    readonly tokenize: (input: AiInput.Raw) => Effect.Effect<Array<number>, AiError>
+    readonly truncate: (input: AiInput.Raw, tokens: number) => Effect.Effect<AiInput.AiInput, AiError>
   }
 }
 
 /**
  * @since 1.0.0
- * @category constructors
+ * @category Constructors
  */
 export const make = (options: {
   readonly tokenize: (content: AiInput.AiInput) => Effect.Effect<Array<number>, AiError>
@@ -44,33 +41,32 @@ export const make = (options: {
     tokenize(input) {
       return options.tokenize(AiInput.make(input))
     },
-    truncate(content, tokens) {
-      return truncate(AiInput.make(content), options.tokenize, tokens)
+    truncate(input, tokens) {
+      return truncate(AiInput.make(input), options.tokenize, tokens)
     }
   })
 
 const truncate = (
   self: AiInput.AiInput,
-  tokenize: (content: AiInput.AiInput) => Effect.Effect<Array<number>, AiError>,
+  tokenize: (input: AiInput.AiInput) => Effect.Effect<Array<number>, AiError>,
   maxTokens: number
 ): Effect.Effect<AiInput.AiInput, AiError> =>
   Effect.suspend(() => {
     let count = 0
-    let inParts = self
-    let outParts: Chunk.Chunk<Message> = Chunk.empty()
+    let inputMessages = self.messages
+    let outputMessages: Array<AiInput.Message> = []
     const loop: Effect.Effect<AiInput.AiInput, AiError> = Effect.suspend(() => {
-      const o = Chunk.last(inParts)
-      if (Option.isNone(o)) {
-        return Effect.succeed(AiInput.make(outParts))
+      const message = inputMessages[inputMessages.length - 1]
+      if (Predicate.isUndefined(message)) {
+        return Effect.succeed(AiInput.make(outputMessages))
       }
-      const part = o.value
-      inParts = Chunk.dropRight(inParts, 1)
-      return Effect.flatMap(tokenize(Chunk.of(part)), (tokens) => {
+      inputMessages = inputMessages.slice(0, inputMessages.length - 1)
+      return Effect.flatMap(tokenize(AiInput.make(message)), (tokens) => {
         count += tokens.length
         if (count > maxTokens) {
-          return Effect.succeed(AiInput.make(outParts))
+          return Effect.succeed(AiInput.make(outputMessages))
         }
-        outParts = Chunk.prepend(outParts, part)
+        outputMessages = [message, ...outputMessages]
         return loop
       })
     })

package/src/index.ts

@@ -6,37 +6,37 @@ export * as AiChat from "./AiChat.js"
 /**
  * @since 1.0.0
  */
-export * as AiError from "./AiError.js"
+export * as AiEmbeddingModel from "./AiEmbeddingModel.js"
 
 /**
  * @since 1.0.0
  */
-export * as AiInput from "./AiInput.js"
+export * as AiError from "./AiError.js"
 
 /**
  * @since 1.0.0
  */
-export * as AiModel from "./AiModel.js"
+export * as AiInput from "./AiInput.js"
 
 /**
  * @since 1.0.0
  */
-export * as AiModels from "./AiModels.js"
+export * as AiLanguageModel from "./AiLanguageModel.js"
 
 /**
  * @since 1.0.0
  */
-export * as AiPlan from "./AiPlan.js"
+export * as AiModel from "./AiModel.js"
 
 /**
  * @since 1.0.0
  */
-export * as AiResponse from "./AiResponse.js"
+export * as AiPlan from "./AiPlan.js"
 
 /**
  * @since 1.0.0
  */
-export * as AiRole from "./AiRole.js"
+export * as AiResponse from "./AiResponse.js"
 
 /**
  * @since 1.0.0
@@ -46,17 +46,12 @@ export * as AiTelemetry from "./AiTelemetry.js"
 /**
  * @since 1.0.0
  */
-export * as AiToolkit from "./AiToolkit.js"
-
-/**
- * @since 1.0.0
- */
-export * as Completions from "./Completions.js"
+export * as AiTool from "./AiTool.js"
 
 /**
  * @since 1.0.0
  */
-export * as Embeddings from "./Embeddings.js"
+export * as AiToolkit from "./AiToolkit.js"
 
 /**
  * @since 1.0.0