@llm4ts/flow 0.1.0

package/src/CostLedger.ts

@@ -0,0 +1,171 @@
+import * as Effect from "effect/Effect"
+import * as Schema from "effect/Schema"
+import { BudgetExceeded, PlanParseError, type FlowError } from "./FlowError.ts"
+import type { PlainFileStoreShape } from "./Persistence.ts"
+import { stableHash } from "./Plan.ts"
+import { PricesAsOf } from "./PriceList.ts"
+
+export class CostCell extends Schema.Class<CostCell>("CostCell")({
+  stage: Schema.String,
+  agent: Schema.String,
+  model: Schema.String,
+  prompt: Schema.Int,
+  completion: Schema.Int,
+  cached: Schema.optionalKey(Schema.Int),
+  total: Schema.Int,
+  costUsd: Schema.optionalKey(Schema.Number)
+}) {}
+
+export class CostRecord extends Schema.Class<CostRecord>("CostRecord")({
+  schemaVersion: Schema.Int,
+  runId: Schema.String,
+  at: Schema.String,
+  repo: Schema.String,
+  label: Schema.optionalKey(Schema.String),
+  promptHead: Schema.String,
+  promptHash: Schema.String,
+  pricesAsOf: Schema.String,
+  cells: Schema.Array(CostCell),
+  totalPrompt: Schema.Int,
+  totalCompletion: Schema.Int,
+  totalCostUsd: Schema.optionalKey(Schema.Number)
+}) {}
+
+export const CurrentCostRecordSchema = 1
+export const CostLedgerFileName = "costs.jsonl"
+
+export interface CostRecordInput {
+  readonly runId: string
+  readonly at: string
+  readonly repo: string
+  readonly label?: string
+  readonly prompt: string
+  readonly cells: ReadonlyArray<CostCell>
+}
+
+export const makeCostRecord = (input: CostRecordInput): CostRecord => {
+  const costs = input.cells.flatMap((cell) => (cell.costUsd === undefined ? [] : [cell.costUsd]))
+  return CostRecord.make({
+    schemaVersion: CurrentCostRecordSchema,
+    runId: input.runId,
+    at: input.at,
+    repo: input.repo,
+    ...(input.label === undefined ? {} : { label: input.label }),
+    promptHead: input.prompt.trim().replaceAll(/\s+/gu, " ").slice(0, 120),
+    promptHash: stableHash(input.prompt),
+    pricesAsOf: PricesAsOf,
+    cells: input.cells,
+    totalPrompt: input.cells.reduce((sum, cell) => sum + cell.prompt, 0),
+    totalCompletion: input.cells.reduce((sum, cell) => sum + cell.completion, 0),
+    ...(costs.length === 0 ? {} : { totalCostUsd: costs.reduce((sum, cost) => sum + cost, 0) })
+  })
+}
+
+export class CostBudget extends Schema.Class<CostBudget>("CostBudget")({
+  maximumTokens: Schema.optionalKey(Schema.Int),
+  maximumCostUsd: Schema.optionalKey(Schema.Number)
+}) {}
+
+export const checkCostBudget = (
+  record: CostRecord,
+  budget: CostBudget
+): Effect.Effect<void, BudgetExceeded> => {
+  const totalTokens = record.totalPrompt + record.totalCompletion
+  if (budget.maximumTokens !== undefined && totalTokens > budget.maximumTokens) {
+    return Effect.fail(
+      BudgetExceeded.make({
+        metric: "tokens",
+        limit: budget.maximumTokens,
+        actual: totalTokens
+      })
+    )
+  }
+  if (budget.maximumCostUsd !== undefined && (record.totalCostUsd ?? 0) > budget.maximumCostUsd) {
+    return Effect.fail(
+      BudgetExceeded.make({
+        metric: "costUsd",
+        limit: budget.maximumCostUsd,
+        actual: record.totalCostUsd ?? 0
+      })
+    )
+  }
+  return Effect.void
+}
+
+const costRecordCodec = Schema.fromJsonString(CostRecord)
+
+export const appendCostRecord = Effect.fn("@llm4ts/flow/CostLedger.append")(function* (
+  files: PlainFileStoreShape,
+  path: string,
+  record: CostRecord
+): Effect.fn.Return<void, FlowError> {
+  const encoded = yield* Schema.encodeEffect(costRecordCodec)(record).pipe(
+    Effect.mapError((error) =>
+      PlanParseError.make({
+        message: `failed to encode cost record: ${String(error)}`
+      })
+    )
+  )
+  yield* files.append(path, `${encoded}\n`)
+})
+
+export const loadCostLedger = Effect.fn("@llm4ts/flow/CostLedger.load")(function* (
+  files: PlainFileStoreShape,
+  path: string
+): Effect.fn.Return<ReadonlyArray<CostRecord>, FlowError> {
+  const contents = yield* files.read(path)
+  if (contents === undefined) {
+    return []
+  }
+  const parsed = yield* Effect.forEach(
+    contents
+      .split(/\r?\n/u)
+      .map((line) => line.trim())
+      .filter((line) => line.length > 0),
+    (line) =>
+      Schema.decodeUnknownEffect(costRecordCodec)(line).pipe(
+        Effect.option,
+        Effect.map((option) => (option._tag === "Some" ? [option.value] : []))
+      ),
+    { concurrency: 1 }
+  )
+  return parsed.flat()
+})
+
+const sumOptional = (values: ReadonlyArray<number | undefined>): number | undefined => {
+  const present = values.flatMap((value) => (value === undefined ? [] : [value]))
+  return present.length === 0 ? undefined : present.reduce((sum, value) => sum + value, 0)
+}
+
+export const mergeCostCells = (records: ReadonlyArray<CostRecord>): ReadonlyArray<CostCell> => {
+  const groups = new Map<string, Array<CostCell>>()
+  for (const cell of records.flatMap((record) => record.cells)) {
+    const key = JSON.stringify([cell.stage, cell.agent, cell.model])
+    groups.set(key, [...(groups.get(key) ?? []), cell])
+  }
+  return [...groups.values()]
+    .map((cells) => {
+      const first = cells[0]
+      if (first === undefined) {
+        return undefined
+      }
+      const cached = sumOptional(cells.map((cell) => cell.cached))
+      const costUsd = sumOptional(cells.map((cell) => cell.costUsd))
+      return CostCell.make({
+        stage: first.stage,
+        agent: first.agent,
+        model: first.model,
+        prompt: cells.reduce((sum, cell) => sum + cell.prompt, 0),
+        completion: cells.reduce((sum, cell) => sum + cell.completion, 0),
+        total: cells.reduce((sum, cell) => sum + cell.total, 0),
+        ...(cached === undefined ? {} : { cached }),
+        ...(costUsd === undefined ? {} : { costUsd })
+      })
+    })
+    .filter((cell): cell is CostCell => cell !== undefined)
+    .sort((left, right) =>
+      [left.stage, left.agent, left.model]
+        .join("\0")
+        .localeCompare([right.stage, right.agent, right.model].join("\0"))
+    )
+}

package/src/CostTracker.ts

@@ -0,0 +1,170 @@
+import * as Effect from "effect/Effect"
+import * as Ref from "effect/Ref"
+import type * as Scope from "effect/Scope"
+import * as Stream from "effect/Stream"
+import { TokenUsage } from "@llm4ts/core/Models"
+import { CostCell } from "./CostLedger.ts"
+import type { FlowEvent, FlowEventHub } from "./FlowEvents.ts"
+import { estimateCostUsd, PricesAsOf } from "./PriceList.ts"
+
+interface CostState {
+  readonly byAgent: ReadonlyMap<string, TokenUsage>
+  readonly byModel: ReadonlyMap<string, TokenUsage>
+  readonly byCell: ReadonlyMap<string, TokenUsage>
+  readonly cellKeys: ReadonlyMap<string, readonly [stage: string, agent: string, model: string]>
+  readonly openStages: ReadonlyArray<string>
+}
+
+export interface CostTracker {
+  readonly record: (event: FlowEvent) => Effect.Effect<void>
+  readonly cells: Effect.Effect<ReadonlyArray<CostCell>>
+  readonly summary: Effect.Effect<string>
+  readonly consume: (hub: FlowEventHub) => Effect.Effect<void, never, Scope.Scope>
+  readonly awaitDrained: (hub: FlowEventHub) => Effect.Effect<void>
+}
+
+const emptyState: CostState = {
+  byAgent: new Map(),
+  byModel: new Map(),
+  byCell: new Map(),
+  cellKeys: new Map(),
+  openStages: []
+}
+
+const mergeUsage = (previous: TokenUsage | undefined, usage: TokenUsage): TokenUsage => {
+  const cachedValues = [previous?.cached, usage.cached].flatMap((value) =>
+    value === undefined ? [] : [value]
+  )
+  return TokenUsage.make({
+    prompt: (previous?.prompt ?? 0) + usage.prompt,
+    completion: (previous?.completion ?? 0) + usage.completion,
+    total: (previous?.total ?? 0) + usage.total,
+    ...(cachedValues.length === 0
+      ? {}
+      : {
+          cached: cachedValues.reduce((sum, value) => sum + value, 0)
+        })
+  })
+}
+
+const addUsage = (
+  values: ReadonlyMap<string, TokenUsage>,
+  key: string,
+  usage: TokenUsage
+): ReadonlyMap<string, TokenUsage> => new Map(values).set(key, mergeUsage(values.get(key), usage))
+
+const updateCostState = (state: CostState, event: FlowEvent): CostState => {
+  switch (event._tag) {
+    case "TokensUsed": {
+      const model = event.model ?? "(unknown)"
+      const stage = state.openStages[0] ?? "(no stage)"
+      const cellKey = JSON.stringify([stage, event.agent, model])
+      return {
+        ...state,
+        byAgent: addUsage(state.byAgent, event.agent, event.usage),
+        byModel: addUsage(state.byModel, model, event.usage),
+        byCell: addUsage(state.byCell, cellKey, event.usage),
+        cellKeys: new Map(state.cellKeys).set(cellKey, [stage, event.agent, model])
+      }
+    }
+    case "StageStarted":
+      return { ...state, openStages: [event.stage, ...state.openStages] }
+    case "StageCompleted":
+    case "StageFailed":
+      return { ...state, openStages: state.openStages.slice(1) }
+    default:
+      return state
+  }
+}
+
+const tokenLine = (label: string, usage: TokenUsage, cost: number | undefined): string => {
+  const cached = usage.cached === undefined || usage.cached <= 0 ? "" : ` (${usage.cached} cached)`
+  const money = cost === undefined ? "" : ` ($${cost.toFixed(4)}*)`
+  return `  ${label}: ${usage.prompt} in${cached}, ${usage.completion} out${money}`
+}
+
+const renderSummary = (state: CostState): string => {
+  const agents = [...state.byAgent.entries()]
+    .sort(([left], [right]) => left.localeCompare(right))
+    .map(([agent, usage]) => tokenLine(agent, usage, undefined))
+  const models = [...state.byModel.entries()]
+    .sort(([left], [right]) => left.localeCompare(right))
+    .map(([model, usage]) => {
+      const cost = model === "(unknown)" ? undefined : estimateCostUsd(model, usage)
+      return { line: tokenLine(model, usage, cost), cost }
+    })
+  const total = models.reduce((sum, entry) => sum + (entry.cost ?? 0), 0)
+  const footnote =
+    total <= 0
+      ? ""
+      : `\n\n* estimated from the pricing table (rates as of ${PricesAsOf} — may be stale)`
+  return `By agent:\n${agents.join("\n")}\n\nBy model:\n${models
+    .map(({ line }) => line)
+    .join("\n")}\n\nTotal: ${total <= 0 ? "$0.00" : `$${total.toFixed(4)}`}${footnote}`
+}
+
+export const makeCostTracker = Effect.fn("@llm4ts/flow/CostTracker.make")(
+  function* (): Effect.fn.Return<CostTracker> {
+    const state = yield* Ref.make<CostState>(emptyState)
+    const consumed = yield* Ref.make(0)
+    const record = (event: FlowEvent): Effect.Effect<void> =>
+      Ref.update(state, (current) => updateCostState(current, event))
+    return {
+      record,
+      cells: Ref.get(state).pipe(
+        Effect.map((current) =>
+          [...current.byCell.entries()]
+            .flatMap(([key, usage]) => {
+              const parts = current.cellKeys.get(key)
+              if (parts === undefined) {
+                return []
+              }
+              const [stage, agent, model] = parts
+              const costUsd = model === "(unknown)" ? undefined : estimateCostUsd(model, usage)
+              return [
+                CostCell.make({
+                  stage,
+                  agent,
+                  model,
+                  prompt: usage.prompt,
+                  completion: usage.completion,
+                  total: usage.total,
+                  ...(usage.cached === undefined ? {} : { cached: usage.cached }),
+                  ...(costUsd === undefined ? {} : { costUsd })
+                })
+              ]
+            })
+            .sort((left, right) =>
+              [left.stage, left.agent, left.model]
+                .join("\0")
+                .localeCompare([right.stage, right.agent, right.model].join("\0"))
+            )
+        )
+      ),
+      summary: Ref.get(state).pipe(Effect.map(renderSummary)),
+      consume: (hub) =>
+        Effect.gen(function* () {
+          const subscription = yield* hub.subscribe
+          yield* Stream.fromSubscription(subscription).pipe(
+            Stream.runForEach((event) =>
+              record(event).pipe(Effect.andThen(Ref.update(consumed, (count) => count + 1)))
+            ),
+            Effect.forkScoped,
+            Effect.asVoid
+          )
+        }),
+      awaitDrained: (hub) =>
+        Effect.gen(function* () {
+          const target = yield* hub.publishedCount
+          const drain: Effect.Effect<void> = Effect.suspend(() =>
+            Ref.get(consumed).pipe(
+              Effect.flatMap((count) =>
+                count >= target ? Effect.void : Effect.yieldNow.pipe(Effect.andThen(drain))
+              )
+            )
+          )
+          yield* drain
+        })
+    }
+  }
+)