@effect/ai 0.1.0

package/src/AiToolkit.ts

@@ -0,0 +1,314 @@
+/**
+ * @since 1.0.0
+ */
+import type * as Schema from "@effect/schema/Schema"
+import type * as Serializable from "@effect/schema/Serializable"
+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 * as Layer from "effect/Layer"
+import { pipeArguments } from "effect/Pipeable"
+import type { Scope } from "effect/Scope"
+import type * as Types from "effect/Types"
+
+/**
+ * @since 1.0.0
+ * @category type ids
+ */
+export const TypeId: unique symbol = Symbol.for("@effect/ai/AiToolkit")
+
+/**
+ * @since 1.0.0
+ * @category type ids
+ */
+export type TypeId = typeof TypeId
+
+/**
+ * @since 1.0.0
+ * @category models
+ */
+export interface AiToolkit<in out Tools extends Tool.AnySchema>
+  extends Effect.Effect<Handlers<Tools>, never, Tool.Services<Tools> | Registry>, Inspectable.Inspectable
+{
+  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, R | RX>
+  readonly implementScoped: <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: Serializable.SerializableWithResult.All
+    readonly success: Schema.Schema.Any
+  }
+
+  /**
+   * @since 1.0.0
+   * @category tool
+   */
+  export type Success<Tool extends AnySchema> = Serializable.WithResult.Success<Tool["Type"]>
+
+  /**
+   * @since 1.0.0
+   * @category tool
+   */
+  export type Failure<Tool extends AnySchema> = Serializable.WithResult.Failure<Tool["Type"]>
+
+  /**
+   * @since 1.0.0
+   * @category tool
+   */
+  export type Context<Tool extends AnySchema> = Serializable.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>
+
+  /**
+   * @since 1.0.0
+   * @category tool
+   */
+  export type HandlerAny = (params: any) => Effect.Effect<any, any, any>
+
+  /**
+   * @since 1.0.0
+   * @category tool
+   */
+  export interface Service<Tag extends string> {
+    readonly _: unique symbol
+    readonly name: Types.Invariant<Tag>
+  }
+
+  /**
+   * @since 1.0.0
+   * @category tool
+   */
+  export type ServiceFromTag<Tag extends string> = Tag extends infer T ? T extends string ? Service<T> : never : never
+
+  /**
+   * @since 1.0.0
+   * @category tool
+   */
+  export type Services<Tools extends AnySchema> = ServiceFromTag<Tools["_tag"]>
+}
+
+/**
+ * @since 1.0.0
+ * @category registry
+ */
+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>
+{
+  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, R | RX> {
+    return registerHandlers(this as any, f as any).pipe(Layer.effectDiscard, Layer.provideMerge(Registry.Live))
+  }
+  implementScoped<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
+  }
+}
+
+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
+              }
+            }
+          )
+        })
+      }
+    })
+  )
+
+/**
+ * @since 1.0.0
+ * @category constructors
+ */
+export const empty: AiToolkit<never> = new AiToolkitImpl(HashMap.empty())
+
+/**
+ * @since 1.0.0
+ * @category handlers
+ */
+export const HandlersTypeId: unique symbol = Symbol.for("@effect/ai/AiToolkit/Handlers")
+
+/**
+ * @since 1.0.0
+ * @category handlers
+ */
+export type HandlersTypeId = typeof HandlersTypeId
+
+/**
+ * @since 1.0.0
+ * @category handlers
+ */
+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>>>
+}
+
+/**
+ * @since 1.0.0
+ * @category handlers
+ */
+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))
+  }
+}

package/src/Completions.ts

@@ -0,0 +1,354 @@
+/**
+ * @since 1.0.0
+ */
+import * as JsonSchema from "@effect/platform/OpenApiJsonSchema"
+import * as AST from "@effect/schema/AST"
+import * as Schema from "@effect/schema/Schema"
+import * as Chunk from "effect/Chunk"
+import * as Context from "effect/Context"
+import * as Effect from "effect/Effect"
+import * as HashMap from "effect/HashMap"
+import * as Option from "effect/Option"
+import * as Stream from "effect/Stream"
+import type { Concurrency } from "effect/Types"
+import { AiError } from "./AiError.js"
+import type { Message } from "./AiInput.js"
+import * as AiInput from "./AiInput.js"
+import type { AiResponse, ToolCallId, ToolCallPart } from "./AiResponse.js"
+import { WithResolved } from "./AiResponse.js"
+import type * as AiToolkit from "./AiToolkit.js"
+
+/**
+ * @since 1.0.0
+ * @category tags
+ */
+export class Completions extends Context.Tag("@effect/ai/Completions")<
+  Completions,
+  Completions.Service
+>() {}
+
+/**
+ * @since 1.0.0
+ * @category models
+ */
+export declare namespace Completions {
+  /**
+   * @since 1.0.0
+   * @models
+   */
+  export interface StructuredSchema<A, I, R> extends Schema.Schema<A, I, R> {
+    readonly _tag?: string
+    readonly identifier: string
+  }
+
+  /**
+   * @since 1.0.0
+   * @models
+   */
+  export interface Service {
+    readonly create: (input: AiInput.Input) => Effect.Effect<AiResponse, AiError>
+    readonly stream: (input: AiInput.Input) => Stream.Stream<AiResponse, AiError>
+    readonly structured: <A, I, R>(
+      options: {
+        readonly input: AiInput.Input
+        readonly schema: StructuredSchema<A, I, R>
+      }
+    ) => Effect.Effect<WithResolved<A>, AiError, R>
+    readonly toolkit: <Tools extends AiToolkit.Tool.AnySchema>(
+      options: {
+        readonly input: AiInput.Input
+        readonly tools: AiToolkit.Handlers<Tools>
+        readonly required?: Tools["_tag"] | boolean | undefined
+        readonly concurrency?: Concurrency | undefined
+      }
+    ) => Effect.Effect<
+      WithResolved<AiToolkit.Tool.Success<Tools>>,
+      AiError | AiToolkit.Tool.Failure<Tools>,
+      AiToolkit.Tool.Context<Tools>
+    >
+    readonly toolkitStream: <Tools extends AiToolkit.Tool.AnySchema>(
+      options: {
+        readonly input: AiInput.Input
+        readonly tools: AiToolkit.Handlers<Tools>
+        readonly required?: Tools["_tag"] | boolean | undefined
+        readonly concurrency?: Concurrency | undefined
+      }
+    ) => Stream.Stream<
+      WithResolved<AiToolkit.Tool.Success<Tools>>,
+      AiError | AiToolkit.Tool.Failure<Tools>,
+      AiToolkit.Tool.Context<Tools>
+    >
+  }
+}
+
+const constEmptyMap = new Map<never, never>()
+
+/**
+ * @since 1.0.0
+ * @category models
+ */
+export interface CompletionOptions {
+  readonly system: Option.Option<string>
+  readonly input: Chunk.NonEmptyChunk<Message>
+  readonly tools: Array<{
+    readonly name: string
+    readonly description: string
+    readonly parameters: JsonSchema.JsonSchema
+  }>
+  readonly required: boolean | string
+}
+
+/**
+ * @since 1.0.0
+ * @category constructors
+ */
+export const make = (options: {
+  readonly create: (options: {
+    readonly system: Option.Option<string>
+    readonly input: Chunk.NonEmptyChunk<Message>
+    readonly tools: Array<{
+      readonly name: string
+      readonly description: string
+      readonly parameters: JsonSchema.JsonSchema
+    }>
+    readonly required: boolean | string
+  }) => Effect.Effect<AiResponse, AiError>
+  readonly stream: (options: {
+    readonly system: Option.Option<string>
+    readonly input: Chunk.NonEmptyChunk<Message>
+    readonly tools: Array<{
+      readonly name: string
+      readonly description: string
+      readonly parameters: JsonSchema.JsonSchema
+    }>
+    readonly required: boolean | string
+  }) => Stream.Stream<AiResponse, AiError>
+}): Effect.Effect<Completions.Service> =>
+  Effect.map(Effect.serviceOption(AiInput.SystemInstruction), (parentSystem) => {
+    return Completions.of({
+      create(input) {
+        return Effect.serviceOption(AiInput.SystemInstruction).pipe(
+          Effect.flatMap((system) =>
+            options.create({
+              input: AiInput.make(input) as Chunk.NonEmptyChunk<Message>,
+              system: Option.orElse(system, () => parentSystem),
+              tools: [],
+              required: false
+            })
+          ),
+          Effect.withSpan("Completions.create", { captureStackTrace: false })
+        )
+      },
+      stream(input_) {
+        const input = AiInput.make(input_)
+        return Effect.serviceOption(AiInput.SystemInstruction).pipe(
+          Effect.map((system) =>
+            options.stream({
+              input: input as Chunk.NonEmptyChunk<Message>,
+              system: Option.orElse(system, () => parentSystem),
+              tools: [],
+              required: false
+            })
+          ),
+          Stream.unwrap,
+          Stream.withSpan("Completions.stream", { captureStackTrace: false })
+        )
+      },
+      structured(opts) {
+        const input = AiInput.make(opts.input)
+        const schema = opts.schema
+        const decode = Schema.decodeUnknown(schema)
+        const toolId = schema._tag ?? schema.identifier
+        return Effect.serviceOption(AiInput.SystemInstruction).pipe(
+          Effect.flatMap((system) =>
+            options.create({
+              input: input as Chunk.NonEmptyChunk<Message>,
+              system: Option.orElse(system, () => parentSystem),
+              tools: [convertTool(schema)],
+              required: true
+            })
+          ),
+          Effect.flatMap((response) =>
+            Chunk.findFirst(
+              response.parts,
+              (part): part is ToolCallPart => part._tag === "ToolCall" && part.name === toolId
+            ).pipe(
+              Option.match({
+                onNone: () =>
+                  Effect.fail(
+                    new AiError({
+                      module: "Completions",
+                      method: "structured",
+                      description: `Tool call '${toolId}' not found in response`
+                    })
+                  ),
+                onSome: (toolCall) =>
+                  Effect.matchEffect(decode(toolCall.params), {
+                    onFailure: (cause) =>
+                      new AiError({
+                        module: "Completions",
+                        method: "structured",
+                        description: `Failed to decode tool call '${toolId}' parameters`,
+                        cause
+                      }),
+                    onSuccess: (resolved) =>
+                      Effect.succeed(
+                        new WithResolved({
+                          response,
+                          resolved: new Map([[toolCall.id, resolved]]),
+                          encoded: new Map([[toolCall.id, toolCall.params]])
+                        })
+                      )
+                  })
+              }),
+              Effect.withSpan("Completions.structured", {
+                attributes: { tool: toolId },
+                captureStackTrace: false
+              })
+            )
+          )
+        )
+      },
+      toolkit({ concurrency, input: inputInput, required = false, tools }) {
+        const input = AiInput.make(inputInput)
+        const toolArr: Array<{ name: string; description: string; parameters: JsonSchema.JsonSchema }> = []
+        for (const [, tool] of tools.toolkit.tools) {
+          toolArr.push(convertTool(tool as any))
+        }
+        return Effect.serviceOption(AiInput.SystemInstruction).pipe(
+          Effect.flatMap((system) =>
+            options.create({
+              input: input as Chunk.NonEmptyChunk<Message>,
+              system: Option.orElse(system, () => parentSystem),
+              tools: toolArr,
+              required: required as any
+            })
+          ),
+          Effect.flatMap((response) => resolveParts({ response, tools, concurrency, method: "toolkit" })),
+          Effect.withSpan("Completions.toolkit", {
+            captureStackTrace: false,
+            attributes: {
+              concurrency,
+              required
+            }
+          })
+        ) as any
+      },
+      toolkitStream({ concurrency, input, required = false, tools }) {
+        const toolArr: Array<{ name: string; description: string; parameters: JsonSchema.JsonSchema }> = []
+        for (const [, tool] of tools.toolkit.tools) {
+          toolArr.push(convertTool(tool as any))
+        }
+        return Effect.serviceOption(AiInput.SystemInstruction).pipe(
+          Effect.map((system) =>
+            options.stream({
+              input: AiInput.make(input) as Chunk.NonEmptyChunk<Message>,
+              system: Option.orElse(system, () => parentSystem),
+              tools: toolArr,
+              required: required as any
+            })
+          ),
+          Stream.unwrap,
+          Stream.mapEffect(
+            (chunk) => resolveParts({ response: chunk, tools, concurrency, method: "toolkitStream" }),
+            { concurrency: "unbounded" }
+          ),
+          Stream.withSpan("Completions.toolkitStream", {
+            captureStackTrace: false,
+            attributes: {
+              concurrency,
+              required
+            }
+          })
+        ) as any
+      }
+    })
+  })
+
+const convertTool = <A, I, R>(tool: Completions.StructuredSchema<A, I, R>) => ({
+  name: tool._tag ?? tool.identifier,
+  description: getDescription(tool.ast),
+  parameters: JsonSchema.make(tool)
+})
+
+const getDescription = (ast: AST.AST): string => {
+  const annotations = ast._tag === "Transformation" ?
+    {
+      ...ast.to.annotations,
+      ...ast.annotations
+    } :
+    ast.annotations
+  return AST.DescriptionAnnotationId in annotations ? annotations[AST.DescriptionAnnotationId] as string : ""
+}
+
+const resolveParts = (
+  options: {
+    readonly response: AiResponse
+    readonly tools: AiToolkit.Handlers<any>
+    readonly concurrency: Concurrency | undefined
+    readonly method: string
+  }
+) => {
+  const toolNames: Array<string> = []
+  const toolParts = Chunk.filter(
+    options.response.parts,
+    (part): part is ToolCallPart => {
+      if (part._tag === "ToolCall") {
+        toolNames.push(part.name)
+        return true
+      }
+      return false
+    }
+  )
+  if (Chunk.isEmpty(toolParts)) {
+    return Effect.succeed(
+      new WithResolved({
+        response: options.response,
+        resolved: constEmptyMap,
+        encoded: constEmptyMap
+      })
+    )
+  }
+  const resolved = new Map<ToolCallId, AiToolkit.Tool.Success<any>>()
+  const encoded = new Map<ToolCallId, unknown>()
+  return Effect.annotateCurrentSpan("toolCalls", toolNames).pipe(
+    Effect.zipRight(Effect.forEach(
+      toolParts,
+      (part) => {
+        const tool = HashMap.unsafeGet(options.tools.toolkit.tools, part.name)
+        const handler = HashMap.unsafeGet(options.tools.handlers, part.name)
+        const decodeParams = Schema.decodeUnknown(tool as any)
+        const encodeSuccess = Schema.encode(tool.success)
+        return decodeParams(part.params).pipe(
+          Effect.mapError((cause) =>
+            new AiError({
+              module: "Completions",
+              method: options.method,
+              description: `Failed to decode tool call '${part.name}' parameters`,
+              cause
+            })
+          ),
+          Effect.flatMap(handler),
+          Effect.tap((value) => {
+            return encodeSuccess(value).pipe(
+              Effect.mapError((cause) =>
+                new AiError({
+                  module: "Completions",
+                  method: options.method,
+                  description: `Failed to encode tool call '${part.name}' result`,
+                  cause
+                })
+              ),
+              Effect.map((encodedValue) => {
+                resolved.set(part.id, value)
+                encoded.set(part.id, encodedValue)
+              })
+            )
+          })
+        )
+      },
+      { concurrency: options.concurrency, discard: true }
+    )),
+    Effect.as(new WithResolved({ response: options.response, resolved, encoded }))
+  )
+}