@effect/ai-openai 0.11.0 → 0.11.2

package/src/OpenAiClient.ts

@@ -90,62 +90,86 @@ export const make = (options: {
         Stream.map((event) => JSON.parse(event.data) as A)
       )
     const stream = (request: StreamCompletionRequest) =>
-      streamRequest<RawCompletionChunk>(HttpClientRequest.post("/chat/completions", {
-        body: HttpBody.unsafeJson({
-          ...request,
-          stream: true
-        })
-      })).pipe(
-        Stream.mapAccum(new Map<number, ContentPart | Array<ToolCallPart>>(), (acc, chunk) => {
-          const parts: Array<StreamChunkPart> = []
-          for (let i = 0; i < chunk.choices.length; i++) {
-            const choice = chunk.choices[i]
-            if ("content" in choice.delta && typeof choice.delta.content === "string") {
-              let part = acc.get(choice.index) as ContentPart | undefined
-              part = {
-                _tag: "Content",
-                content: choice.delta.content
-              }
-              acc.set(choice.index, part)
-              parts.push(part)
-            } else if ("tool_calls" in choice.delta && Array.isArray(choice.delta.tool_calls)) {
-              const parts = (acc.get(choice.index) ?? []) as Array<ToolCallPart>
-              for (const toolCall of choice.delta.tool_calls) {
-                const part = parts[toolCall.index]
-                const toolPart = part?._tag === "ToolCall" ?
-                  {
-                    ...part,
-                    arguments: part.arguments + toolCall.function.arguments
-                  } :
-                  {
-                    _tag: "ToolCall",
-                    ...toolCall,
-                    ...toolCall.function,
-                    role: choice.delta.role!
-                  } as any
-                parts[toolCall.index] = toolPart
+      Stream.suspend(() => {
+        const finishReasons: Array<string> = []
+        return streamRequest<RawCompletionChunk>(HttpClientRequest.post("/chat/completions", {
+          body: HttpBody.unsafeJson({
+            ...request,
+            stream: true,
+            stream_options: { include_usage: true }
+          })
+        })).pipe(
+          Stream.mapAccum(new Map<number, ContentPart | Array<ToolCallPart>>(), (acc, chunk) => {
+            const parts: Array<StreamChunkPart> = []
+            if (chunk.usage !== null) {
+              parts.push({
+                _tag: "Usage",
+                id: chunk.id,
+                model: chunk.model,
+                inputTokens: chunk.usage.prompt_tokens,
+                outputTokens: chunk.usage.completion_tokens,
+                finishReasons,
+                systemFingerprint: chunk.system_fingerprint,
+                serviceTier: chunk.service_tier
+              })
+            }
+            for (let i = 0; i < chunk.choices.length; i++) {
+              const choice = chunk.choices[i]
+              if (choice.finish_reason !== null) {
+                finishReasons.push(choice.finish_reason)
               }
-              acc.set(choice.index, parts)
-            } else if (choice.finish_reason === "tool_calls") {
-              const toolParts = acc.get(choice.index) as Array<ToolCallPart>
-              for (const part of toolParts) {
-                try {
-                  const args = JSON.parse(part.arguments as string)
-                  parts.push({
-                    _tag: "ToolCall",
-                    id: part.id,
-                    name: part.name,
-                    arguments: args
-                  })
-                  // eslint-disable-next-line no-empty
-                } catch {}
+              if ("content" in choice.delta && typeof choice.delta.content === "string") {
+                let part = acc.get(choice.index) as ContentPart | undefined
+                part = {
+                  _tag: "Content",
+                  content: choice.delta.content
+                }
+                acc.set(choice.index, part)
+                parts.push(part)
+              } else if ("tool_calls" in choice.delta && Array.isArray(choice.delta.tool_calls)) {
+                const parts = (acc.get(choice.index) ?? []) as Array<ToolCallPart>
+                for (const toolCall of choice.delta.tool_calls) {
+                  const part = parts[toolCall.index]
+                  const toolPart = part?._tag === "ToolCall" ?
+                    {
+                      ...part,
+                      arguments: part.arguments + toolCall.function.arguments
+                    } :
+                    {
+                      _tag: "ToolCall",
+                      ...toolCall,
+                      ...toolCall.function,
+                      role: choice.delta.role!
+                    } as any
+                  parts[toolCall.index] = toolPart
+                }
+                acc.set(choice.index, parts)
+              } else if (choice.finish_reason === "tool_calls") {
+                const toolParts = acc.get(choice.index) as Array<ToolCallPart>
+                for (const part of toolParts) {
+                  try {
+                    const args = JSON.parse(part.arguments as string)
+                    parts.push({
+                      _tag: "ToolCall",
+                      id: part.id,
+                      name: part.name,
+                      arguments: args
+                    })
+                    // eslint-disable-next-line no-empty
+                  } catch {}
+                }
               }
             }
-          }
-          return [acc, parts.length === 0 ? Option.none() : Option.some(new StreamChunk({ parts }))]
-        }),
-        Stream.filterMap(identity)
-      )
+            return [
+              acc,
+              parts.length === 0
+                ? Option.none()
+                : Option.some(new StreamChunk({ parts }))
+            ]
+          }),
+          Stream.filterMap(identity)
+        )
+      })
     return OpenAiClient.of({ client, streamRequest, stream })
   })
 
@@ -202,6 +226,13 @@ interface RawCompletionChunk {
     }
   >
   readonly system_fingerprint: string
+  readonly service_tier: string
+  readonly usage: RawUsage | null
+}
+
+interface RawUsage {
+  readonly prompt_tokens: number
+  readonly completion_tokens: number
 }
 
 type RawDelta = {
@@ -254,19 +285,24 @@ export class StreamChunk extends Data.Class<{
       })
     }
     const part = this.parts[0]
-    return part._tag === "Content" ?
-      AiResponse.AiResponse.fromText({
-        role: AiRole.model,
-        content: part.content
-      }) :
-      new AiResponse.AiResponse({
-        role: AiRole.model,
-        parts: Chunk.of(AiResponse.ToolCallPart.fromUnknown({
-          id: part.id,
-          name: part.name,
-          params: part.arguments
-        }))
-      })
+    switch (part._tag) {
+      case "Content":
+        return AiResponse.AiResponse.fromText({
+          role: AiRole.model,
+          content: part.content
+        })
+      case "ToolCall":
+        return new AiResponse.AiResponse({
+          role: AiRole.model,
+          parts: Chunk.of(AiResponse.ToolCallPart.fromUnknown({
+            id: part.id,
+            name: part.name,
+            params: part.arguments
+          }))
+        })
+      case "Usage":
+        return AiResponse.AiResponse.empty
+    }
   }
 }
 
@@ -274,7 +310,7 @@ export class StreamChunk extends Data.Class<{
  * @since 1.0.0
  * @category models
  */
-export type StreamChunkPart = ContentPart | ToolCallPart
+export type StreamChunkPart = ContentPart | ToolCallPart | UsagePart
 
 /**
  * @since 1.0.0
@@ -296,3 +332,18 @@ export interface ToolCallPart {
   readonly name: string
   readonly arguments: unknown
 }
+
+/**
+ * @since 1.0.0
+ * @category models
+ */
+export interface UsagePart {
+  readonly _tag: "Usage"
+  readonly id: string
+  readonly model: string
+  readonly inputTokens: number
+  readonly outputTokens: number
+  readonly finishReasons: ReadonlyArray<string>
+  readonly systemFingerprint: string
+  readonly serviceTier: string | null
+}

package/src/OpenAiCompletions.ts

@@ -11,14 +11,24 @@ import * as Arr from "effect/Array"
 import * as Effect from "effect/Effect"
 import * as Layer from "effect/Layer"
 import type * as Option from "effect/Option"
+import * as Predicate from "effect/Predicate"
 import * as Stream from "effect/Stream"
+import type { Span } from "effect/Tracer"
 import type * as Generated from "./Generated.js"
+import type { StreamChunk } from "./OpenAiClient.js"
 import { OpenAiClient } from "./OpenAiClient.js"
 import { OpenAiConfig } from "./OpenAiConfig.js"
+import { addGenAIAnnotations } from "./OpenAiTelemetry.js"
 import * as OpenAiTokenizer from "./OpenAiTokenizer.js"
 
+/**
+ * @since 1.0.0
+ * @category models
+ */
+export type Model = typeof Generated.CreateChatCompletionRequestModel.Encoded
+
 const make = (options: {
-  readonly model: string
+  readonly model: (string & {}) | Model
 }) =>
   Effect.gen(function*() {
     const client = yield* OpenAiClient
@@ -66,9 +76,20 @@ const make = (options: {
     }
 
     return yield* Completions.make({
-      create(options) {
+      create({ span, ...options }) {
         return makeRequest(options).pipe(
+          Effect.tap((request) => annotateRequest(span, request)),
           Effect.flatMap(client.client.createChatCompletion),
+          Effect.tap((response) => annotateChatResponse(span, response)),
+          Effect.flatMap((response) =>
+            makeResponse(
+              response,
+              "create",
+              options.tools.length === 1 && options.tools[0].structured
+                ? options.tools[0]
+                : undefined
+            )
+          ),
           Effect.catchAll((cause) =>
             Effect.fail(
               new AiError({
@@ -78,20 +99,19 @@ const make = (options: {
                 cause
               })
             )
-          ),
-          Effect.flatMap((response) =>
-            makeResponse(
-              response,
-              "create",
-              options.tools.length === 1 && options.tools[0].structured ? options.tools[0] : undefined
-            )
           )
         )
       },
-      stream(options) {
+      stream({ span, ...options }) {
         return makeRequest(options).pipe(
+          Effect.tap((request) => annotateRequest(span, request)),
           Effect.map(client.stream),
           Stream.unwrap,
+          Stream.tap((response) => {
+            annotateStreamResponse(span, response)
+            return Effect.void
+          }),
+          Stream.map((response) => response.asAiResponse),
           Stream.catchAll((cause) =>
             Effect.fail(
               new AiError({
@@ -101,8 +121,7 @@ const make = (options: {
                 cause
               })
             )
-          ),
-          Stream.map((response) => response.asAiResponse)
+          )
         )
       }
     })
@@ -113,7 +132,7 @@ const make = (options: {
  * @category layers
  */
 export const layerCompletions = (options: {
-  readonly model: string
+  readonly model: (string & {}) | Model
 }): Layer.Layer<Completions.Completions, never, OpenAiClient> => Layer.effect(Completions.Completions, make(options))
 
 /**
@@ -121,7 +140,7 @@ export const layerCompletions = (options: {
  * @category layers
  */
 export const layer = (options: {
-  readonly model: string
+  readonly model: (string & {}) | Model
 }): Layer.Layer<Completions.Completions | Tokenizer.Tokenizer, never, OpenAiClient> =>
   Layer.merge(layerCompletions(options), OpenAiTokenizer.layer(options))
 
@@ -291,3 +310,78 @@ const makeSystemMessage = (content: string): typeof Generated.ChatCompletionRequ
 }
 
 const safeName = (name: string) => name.replace(/[^a-zA-Z0-9_-]/g, "_").replace(/_+/, "_")
+
+const annotateRequest = (
+  span: Span,
+  request: typeof Generated.CreateChatCompletionRequest.Encoded
+): void => {
+  addGenAIAnnotations(span, {
+    system: "openai",
+    operation: { name: "chat" },
+    request: {
+      model: request.model,
+      temperature: request.temperature,
+      topP: request.top_p,
+      maxTokens: request.max_tokens,
+      stopSequences: Arr.ensure(request.stop).filter(Predicate.isNotNullable),
+      frequencyPenalty: request.frequency_penalty,
+      presencePenalty: request.presence_penalty,
+      seed: request.seed
+    },
+    openai: {
+      request: {
+        responseFormat: request.response_format?.type,
+        serviceTier: request.service_tier
+      }
+    }
+  })
+}
+
+const annotateChatResponse = (
+  span: Span,
+  response: Generated.CreateChatCompletionResponse
+): void => {
+  addGenAIAnnotations(span, {
+    response: {
+      id: response.id,
+      model: response.model,
+      finishReasons: response.choices.map((choice) => choice.finish_reason)
+    },
+    usage: {
+      inputTokens: response.usage?.prompt_tokens,
+      outputTokens: response.usage?.completion_tokens
+    },
+    openai: {
+      response: {
+        systemFingerprint: response.system_fingerprint,
+        serviceTier: response.service_tier
+      }
+    }
+  })
+}
+
+const annotateStreamResponse = (
+  span: Span,
+  response: StreamChunk
+) => {
+  const usage = response.parts.find((part) => part._tag === "Usage")
+  if (Predicate.isNotNullable(usage)) {
+    addGenAIAnnotations(span, {
+      response: {
+        id: usage.id,
+        model: usage.model,
+        finishReasons: usage.finishReasons
+      },
+      usage: {
+        inputTokens: usage.inputTokens,
+        outputTokens: usage.outputTokens
+      },
+      openai: {
+        response: {
+          systemFingerprint: usage.systemFingerprint,
+          serviceTier: usage.serviceTier
+        }
+      }
+    })
+  }
+}

package/src/OpenAiEmbeddings.ts

@@ -0,0 +1,149 @@
+/**
+ * @since 1.0.0
+ */
+import { AiError } from "@effect/ai/AiError"
+import * as Embeddings from "@effect/ai/Embeddings"
+import * as Context from "effect/Context"
+import type * as Duration from "effect/Duration"
+import * as Effect from "effect/Effect"
+import * as Layer from "effect/Layer"
+import type { Simplify } from "effect/Types"
+import type * as Generated from "./Generated.js"
+import { OpenAiClient } from "./OpenAiClient.js"
+
+/**
+ * @since 1.0.0
+ * @category models
+ */
+export type Model = typeof Generated.CreateEmbeddingRequestModel.Encoded
+
+/**
+ * @since 1.0.0
+ * @category tags
+ */
+export class OpenAiEmbeddingsConfig extends Context.Tag("@effect/ai-openai/OpenAiEmbeddings/Config")<
+  OpenAiEmbeddingsConfig,
+  Simplify<
+    Partial<
+      Omit<
+        typeof Generated.CreateEmbeddingRequest.Encoded,
+        "input"
+      >
+    >
+  >
+>() {
+  /**
+   * @since 1.0.0
+   */
+  static readonly getOrUndefined: Effect.Effect<typeof OpenAiEmbeddingsConfig.Service | undefined> = Effect.map(
+    Effect.context<never>(),
+    (context) => context.unsafeMap.get(OpenAiEmbeddingsConfig.key)
+  )
+}
+
+const makeRequest = (
+  client: OpenAiClient.Service,
+  input: ReadonlyArray<string>,
+  parentConfig: typeof OpenAiEmbeddingsConfig.Service | undefined,
+  options: {
+    readonly model: string
+    readonly maxBatchSize?: number
+    readonly cache?: {
+      readonly capacity: number
+      readonly timeToLive: Duration.DurationInput
+    }
+  }
+) =>
+  Effect.context<never>().pipe(
+    Effect.flatMap((context) => {
+      const localConfig = context.unsafeMap.get(OpenAiEmbeddingsConfig.key)
+      return client.client.createEmbedding({
+        input,
+        model: options.model,
+        ...parentConfig,
+        ...localConfig
+      })
+    }),
+    Effect.map((response) =>
+      response.data.map(({ embedding, index }) => ({
+        embeddings: embedding as Array<number>,
+        index
+      }))
+    ),
+    Effect.mapError((cause) => {
+      const common = {
+        module: "OpenAiEmbeddings",
+        method: "embed",
+        cause
+      }
+      if (cause._tag === "ParseError") {
+        return new AiError({
+          description: "Malformed input detected in request",
+          ...common
+        })
+      }
+      return new AiError({
+        description: "An error occurred with the OpenAI API",
+        ...common
+      })
+    })
+  )
+
+const make = Effect.fnUntraced(function*(options: {
+  readonly model: (string & {}) | Model
+  readonly maxBatchSize?: number
+  readonly cache?: {
+    readonly capacity: number
+    readonly timeToLive: Duration.DurationInput
+  }
+}) {
+  const client = yield* OpenAiClient
+  const parentConfig = yield* OpenAiEmbeddingsConfig.getOrUndefined
+  return yield* Embeddings.make({
+    cache: options.cache,
+    maxBatchSize: options.maxBatchSize ?? 2048,
+    embedMany(input) {
+      return makeRequest(client, input, parentConfig, options)
+    }
+  })
+})
+
+const makeDataLoader = Effect.fnUntraced(function*(options: {
+  readonly model: (string & {}) | Model
+  readonly window: Duration.DurationInput
+  readonly maxBatchSize?: number
+}) {
+  const client = yield* OpenAiClient
+  const parentConfig = yield* OpenAiEmbeddingsConfig.getOrUndefined
+  return yield* Embeddings.makeDataLoader({
+    window: options.window,
+    maxBatchSize: options.maxBatchSize ?? 2048,
+    embedMany(input) {
+      return makeRequest(client, input, parentConfig, options)
+    }
+  })
+})
+
+/**
+ * @since 1.0.0
+ * @category layers
+ */
+export const layer = (options: {
+  readonly model: (string & {}) | Model
+  readonly maxBatchSize?: number
+  readonly cache?: {
+    readonly capacity: number
+    readonly timeToLive: Duration.DurationInput
+  }
+}): Layer.Layer<Embeddings.Embeddings, never, OpenAiClient> => Layer.effect(Embeddings.Embeddings, make(options))
+
+/**
+ * @since 1.0.0
+ * @category layers
+ */
+export const layerDataLoader = (options: {
+  readonly model: (string & {}) | Model
+  readonly window: Duration.DurationInput
+  readonly maxBatchSize?: number
+}): Layer.Layer<Embeddings.Embeddings, never, OpenAiClient> =>
+  Layer.scoped(Embeddings.Embeddings, makeDataLoader(options))