@effect/ai 0.8.1 → 0.8.2

package/src/AiTelemetry.ts

@@ -0,0 +1,297 @@
+/**
+ * @since 1.0.0
+ */
+
+import { dual } from "effect/Function"
+import * as Predicate from "effect/Predicate"
+import * as String from "effect/String"
+import type { Span } from "effect/Tracer"
+import type { Simplify } from "effect/Types"
+
+/**
+ * The attributes used to describe telemetry in the context of Generative
+ * Artificial Intelligence (GenAI) Models requests and responses.
+ *
+ * {@see https://opentelemetry.io/docs/specs/semconv/attributes-registry/gen-ai/}
+ *
+ * @since 1.0.0
+ * @category models
+ */
+export type GenAITelemetryAttributes = Simplify<
+  & GenAI.AttributesWithPrefix<GenAI.BaseAttributes, "gen_ai">
+  & GenAI.AttributesWithPrefix<GenAI.OperationAttributes, "gen_ai.operation">
+  & GenAI.AttributesWithPrefix<GenAI.TokenAttributes, "gen_ai.token">
+  & GenAI.AttributesWithPrefix<GenAI.UsageAttributes, "gen_ai.usage">
+  & GenAI.AttributesWithPrefix<GenAI.RequestAttributes, "gen_ai.request">
+  & GenAI.AttributesWithPrefix<GenAI.ResponseAttributes, "gen_ai.response">
+>
+
+/**
+ * @since 1.0.0
+ * @category models
+ */
+export declare namespace GenAI {
+  /**
+   * All telemetry attributes which are part of the GenAI specification.
+   *
+   * @since 1.0.0
+   * @category models
+   */
+  export type AllAttributes =
+    & BaseAttributes
+    & OperationAttributes
+    & TokenAttributes
+    & UsageAttributes
+    & RequestAttributes
+    & ResponseAttributes
+
+  /**
+   * Telemetry attributes which are part of the GenAI specification and are
+   * namespaced by `gen_ai`.
+   *
+   * @since 1.0.0
+   * @category models
+   */
+  export interface BaseAttributes {
+    /**
+     * The Generative AI product as identified by the client or server
+     * instrumentation.
+     */
+    readonly system?: (string & {}) | WellKnownSystem | null | undefined
+  }
+
+  /**
+   * Telemetry attributes which are part of the GenAI specification and are
+   * namespaced by `gen_ai.operation`.
+   *
+   * @since 1.0.0
+   * @category models
+   */
+  export interface OperationAttributes {
+    readonly name?: (string & {}) | WellKnownOperationName | null | undefined
+  }
+
+  /**
+   * Telemetry attributes which are part of the GenAI specification and are
+   * namespaced by `gen_ai.token`.
+   *
+   * @since 1.0.0
+   * @category models
+   */
+  export interface TokenAttributes {
+    readonly type?: string | null | undefined
+  }
+
+  /**
+   * Telemetry attributes which are part of the GenAI specification and are
+   * namespaced by `gen_ai.usage`.
+   *
+   * @since 1.0.0
+   * @category models
+   */
+  export interface UsageAttributes {
+    readonly inputTokens?: number | null | undefined
+    readonly outputTokens?: number | null | undefined
+  }
+
+  /**
+   * Telemetry attributes which are part of the GenAI specification and are
+   * namespaced by `gen_ai.request`.
+   *
+   * @since 1.0.0
+   * @category models
+   */
+  export interface RequestAttributes {
+    /**
+     * The name of the GenAI model a request is being made to.
+     */
+    readonly model?: string | null | undefined
+    /**
+     * The temperature setting for the GenAI request.
+     */
+    readonly temperature?: number | null | undefined
+    /**
+     * The temperature setting for the GenAI request.
+     */
+    readonly topK?: number | null | undefined
+    /**
+     * The top_k sampling setting for the GenAI request.
+     */
+    readonly topP?: number | null | undefined
+    /**
+     * The top_p sampling setting for the GenAI request.
+     */
+    readonly maxTokens?: number | null | undefined
+    /**
+     * The encoding formats requested in an embeddings operation, if specified.
+     */
+    readonly encodingFormats?: ReadonlyArray<string> | null | undefined
+    /**
+     * List of sequences that the model will use to stop generating further
+     * tokens.
+     */
+    readonly stopSequences?: ReadonlyArray<string> | null | undefined
+    /**
+     * The frequency penalty setting for the GenAI request.
+     */
+    readonly frequencyPenalty?: number | null | undefined
+    /**
+     * The presence penalty setting for the GenAI request.
+     */
+    readonly presencePenalty?: number | null | undefined
+    /**
+     * The seed setting for the GenAI request. Requests with same seed value
+     * are more likely to return same result.
+     */
+    readonly seed?: number | null | undefined
+  }
+
+  /**
+   * Telemetry attributes which are part of the GenAI specification and are
+   * namespaced by `gen_ai.response`.
+   *
+   * @since 1.0.0
+   * @category models
+   */
+  export interface ResponseAttributes {
+    /**
+     * The unique identifier for the completion.
+     */
+    readonly id?: string | null | undefined
+    /**
+     * The name of the model that generated the response.
+     */
+    readonly model?: string | null | undefined
+    /**
+     * Array of reasons the model stopped generating tokens, corresponding to
+     * each generation received.
+     */
+    readonly finishReasons?: ReadonlyArray<string> | null | undefined
+  }
+
+  /**
+   * The `gen_ai.operation.name` attribute has the following list of well-known
+   * values.
+   *
+   * If one of them applies, then the respective value **MUST** be used;
+   * otherwise, a custom value **MAY** be used.
+   *
+   * @since 1.0.0
+   * @category models
+   */
+  export type WellKnownOperationName = "chat" | "embeddings" | "text_completion"
+
+  /**
+   * The `gen_ai.system` attribute has the following list of well-known values.
+   *
+   * If one of them applies, then the respective value **MUST** be used;
+   * otherwise, a custom value **MAY** be used.
+   *
+   * @since 1.0.0
+   * @category models
+   */
+  export type WellKnownSystem =
+    | "anthropic"
+    | "aws.bedrock"
+    | "az.ai.inference"
+    | "az.ai.openai"
+    | "cohere"
+    | "deepseek"
+    | "gemini"
+    | "groq"
+    | "ibm.watsonx.ai"
+    | "mistral_ai"
+    | "openai"
+    | "perplexity"
+    | "vertex_ai"
+    | "xai"
+
+  /**
+   * @since 1.0.0
+   * @category models
+   */
+  export type AttributesWithPrefix<Attributes extends Record<string, any>, Prefix extends string> = {
+    [Name in keyof Attributes as `${Prefix}.${FormatAttributeName<Name>}`]: Attributes[Name]
+  }
+
+  /**
+   * @since 1.0.0
+   * @category models
+   */
+  export type FormatAttributeName<T extends string | number | symbol> = T extends string ?
+    T extends `${infer First}${infer Rest}`
+      ? `${First extends Uppercase<First> ? "_" : ""}${Lowercase<First>}${FormatAttributeName<Rest>}`
+    : T :
+    never
+}
+
+/**
+ * @since 1.0.0
+ * @category utilities
+ */
+export const addSpanAttributes = (
+  keyPrefix: string,
+  transformKey: (key: string) => string
+) =>
+<Attributes extends Record<string, any>>(span: Span, attributes: Attributes): void => {
+  for (const [key, value] of Object.entries(attributes)) {
+    if (Predicate.isNotNullable(value)) {
+      span.attribute(`${keyPrefix}.${transformKey(key)}`, value)
+    }
+  }
+}
+
+const addSpanBaseAttributes = addSpanAttributes("gen_ai", String.camelToSnake)<GenAI.BaseAttributes>
+const addSpanOperationAttributes = addSpanAttributes("gen_ai.operation", String.camelToSnake)<GenAI.OperationAttributes>
+const addSpanRequestAttributes = addSpanAttributes("gen_ai.request", String.camelToSnake)<GenAI.RequestAttributes>
+const addSpanResponseAttributes = addSpanAttributes("gen_ai.response", String.camelToSnake)<GenAI.ResponseAttributes>
+const addSpanTokenAttributes = addSpanAttributes("gen_ai.token", String.camelToSnake)<GenAI.TokenAttributes>
+const addSpanUsageAttributes = addSpanAttributes("gen_ai.usage", String.camelToSnake)<GenAI.UsageAttributes>
+
+/**
+ * @since 1.0.0
+ * @since models
+ */
+export type GenAITelemetryAttributeOptions = GenAI.BaseAttributes & {
+  readonly operation?: GenAI.OperationAttributes | undefined
+  readonly request?: GenAI.RequestAttributes | undefined
+  readonly response?: GenAI.ResponseAttributes | undefined
+  readonly token?: GenAI.TokenAttributes | undefined
+  readonly usage?: GenAI.UsageAttributes | undefined
+}
+
+/**
+ * Applies the specified GenAI telemetry attributes to the provided `Span`.
+ *
+ * **NOTE**: This method will mutate the `Span` **in-place**.
+ *
+ * @since 1.0.0
+ * @since utilities
+ */
+export const addGenAIAnnotations = dual<
+  /**
+   * Applies the specified GenAI telemetry attributes to the provided `Span`.
+   *
+   * **NOTE**: This method will mutate the `Span` **in-place**.
+   *
+   * @since 1.0.0
+   * @since utilities
+   */
+  (options: GenAITelemetryAttributeOptions) => (span: Span) => void,
+  /**
+   * Applies the specified GenAI telemetry attributes to the provided `Span`.
+   *
+   * **NOTE**: This method will mutate the `Span` **in-place**.
+   *
+   * @since 1.0.0
+   * @since utilities
+   */
+  (span: Span, options: GenAITelemetryAttributeOptions) => void
+>(2, (span, options) => {
+  addSpanBaseAttributes(span, { system: options.system })
+  if (Predicate.isNotNullable(options.operation)) addSpanOperationAttributes(span, options.operation)
+  if (Predicate.isNotNullable(options.request)) addSpanRequestAttributes(span, options.request)
+  if (Predicate.isNotNullable(options.response)) addSpanResponseAttributes(span, options.response)
+  if (Predicate.isNotNullable(options.token)) addSpanTokenAttributes(span, options.token)
+  if (Predicate.isNotNullable(options.usage)) addSpanUsageAttributes(span, options.usage)
+})

package/src/Completions.ts

@@ -10,6 +10,7 @@ import * as Option from "effect/Option"
 import * as Schema from "effect/Schema"
 import * as AST from "effect/SchemaAST"
 import * as Stream from "effect/Stream"
+import type { Span } from "effect/Tracer"
 import type { Concurrency } from "effect/Types"
 import { AiError } from "./AiError.js"
 import type { Message } from "./AiInput.js"
@@ -34,26 +35,44 @@ export class Completions extends Context.Tag("@effect/ai/Completions")<
 export declare namespace Completions {
   /**
    * @since 1.0.0
-   * @models
+   * @category models
    */
-  export interface StructuredSchema<A, I, R> extends Schema.Schema<A, I, R> {
-    readonly _tag?: string
+  export type StructuredSchema<A, I, R> = TaggedSchema<A, I, R> | IdentifiedSchema<A, I, R>
+
+  /**
+   * @since 1.0.0
+   * @category models
+   */
+  export interface TaggedSchema<A, I, R> extends Schema.Schema<A, I, R> {
+    readonly _tag: string
+  }
+
+  /**
+   * @since 1.0.0
+   * @category models
+   */
+  export interface IdentifiedSchema<A, I, R> extends Schema.Schema<A, I, R> {
     readonly identifier: string
   }
 
   /**
    * @since 1.0.0
-   * @models
+   * @category 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 structured: {
+      <A, I, R>(options: {
         readonly input: AiInput.Input
         readonly schema: StructuredSchema<A, I, R>
-      }
-    ) => Effect.Effect<WithResolved<A>, AiError, R>
+      }): Effect.Effect<WithResolved<A>, AiError, R>
+      <A, I, R>(options: {
+        readonly input: AiInput.Input
+        readonly schema: Schema.Schema<A, I, R>
+        readonly toolCallId: string
+      }): Effect.Effect<WithResolved<A>, AiError, R>
+    }
     readonly toolkit: <Tools extends AiToolkit.Tool.AnySchema>(
       options: {
         readonly input: AiInput.Input
@@ -114,6 +133,7 @@ export const make = (options: {
       readonly structured: boolean
     }>
     readonly required: boolean | string
+    readonly span: Span
   }) => Effect.Effect<AiResponse, AiError>
   readonly stream: (options: {
     readonly system: Option.Option<string>
@@ -125,91 +145,103 @@ export const make = (options: {
       readonly structured: boolean
     }>
     readonly required: boolean | string
+    readonly span: Span
   }) => 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 })
+        return Effect.useSpan(
+          "Completions.create",
+          { captureStackTrace: false },
+          (span) =>
+            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,
+                  span
+                })
+              )
+            )
         )
       },
       stream(input_) {
         const input = AiInput.make(input_)
-        return Effect.serviceOption(AiInput.SystemInstruction).pipe(
-          Effect.map((system) =>
+        return Effect.makeSpanScoped("Completions.stream", { captureStackTrace: false }).pipe(
+          Effect.zip(Effect.serviceOption(AiInput.SystemInstruction)),
+          Effect.map(([span, system]) =>
             options.stream({
               input: input as Chunk.NonEmptyChunk<Message>,
               system: Option.orElse(system, () => parentSystem),
               tools: [],
-              required: false
+              required: false,
+              span
             })
           ),
-          Stream.unwrap,
-          Stream.withSpan("Completions.stream", { captureStackTrace: false })
+          Stream.unwrapScoped
         )
       },
       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, true)],
-              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]])
+        const decode = Schema.decodeUnknown(opts.schema)
+        const toolId = "toolCallId" in opts
+          ? opts.toolCallId
+          : "_tag" in opts.schema
+          ? opts.schema._tag
+          : opts.schema.identifier
+        return Effect.useSpan(
+          "Completions.structured",
+          { attributes: { toolId }, captureStackTrace: false },
+          (span) =>
+            Effect.serviceOption(AiInput.SystemInstruction).pipe(
+              Effect.flatMap((system) =>
+                options.create({
+                  input: input as Chunk.NonEmptyChunk<Message>,
+                  system: Option.orElse(system, () => parentSystem),
+                  tools: [convertTool(toolId, opts.schema, true)],
+                  required: true,
+                  span
+                })
+              ),
+              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 }) {
@@ -221,26 +253,25 @@ export const make = (options: {
           structured: boolean
         }> = []
         for (const [, tool] of tools.toolkit.tools) {
-          toolArr.push(convertTool(tool as any))
+          toolArr.push(convertTool(tool._tag, 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
+        return Effect.useSpan(
+          "Completions.toolkit",
+          { attributes: { concurrency, required }, captureStackTrace: false },
+          (span) =>
+            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,
+                  span
+                })
+              ),
+              Effect.flatMap((response) => resolveParts({ response, tools, concurrency, method: "toolkit" }))
+            ) as any
+        )
       },
       toolkitStream({ concurrency, input, required = false, tools }) {
         const toolArr: Array<{
@@ -250,38 +281,40 @@ export const make = (options: {
           structured: boolean
         }> = []
         for (const [, tool] of tools.toolkit.tools) {
-          toolArr.push(convertTool(tool as any))
+          toolArr.push(convertTool(tool._tag, tool as any))
         }
-        return Effect.serviceOption(AiInput.SystemInstruction).pipe(
-          Effect.map((system) =>
+        return Effect.makeSpanScoped("Completions.stream", {
+          captureStackTrace: false,
+          attributes: { required, concurrency }
+        }).pipe(
+          Effect.zip(Effect.serviceOption(AiInput.SystemInstruction)),
+          Effect.map(([span, system]) =>
             options.stream({
               input: AiInput.make(input) as Chunk.NonEmptyChunk<Message>,
               system: Option.orElse(system, () => parentSystem),
               tools: toolArr,
-              required: required as any
+              required: required as any,
+              span
             })
           ),
-          Stream.unwrap,
+          Stream.unwrapScoped,
           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>, structured = false) => ({
-  name: tool._tag ?? tool.identifier,
-  description: getDescription(tool.ast),
-  parameters: makeJsonSchema(tool.ast),
+const convertTool = <A, I, R>(
+  name: string,
+  schema: Schema.Schema<A, I, R>,
+  structured = false
+) => ({
+  name,
+  description: getDescription(schema.ast),
+  parameters: makeJsonSchema(schema.ast),
   structured
 })