diff --git a/packages/fold-core/src/EventLog/Usage.ts b/packages/fold-core/src/EventLog/Usage.ts index 630af3b..db5af2e 100644 --- a/packages/fold-core/src/EventLog/Usage.ts +++ b/packages/fold-core/src/EventLog/Usage.ts @@ -1,5 +1,5 @@ import { Schema } from 'effect' -import { Response } from 'effect/unstable/ai' +import type { Response } from 'effect/unstable/ai' /** Best-effort token count reported by a model provider. Providers may omit any usage field. */ export const UsageTokenCount = Schema.Finite.check(Schema.isGreaterThanOrEqualTo(0)).annotate({ @@ -41,12 +41,26 @@ export const UsageEncoded = Schema.Struct({ }).annotate({ identifier: 'UsageEncoded' }) export type UsageEncoded = typeof UsageEncoded.Type -const encodeResponseUsage = Schema.encodeUnknownSync(Response.Usage) const decodeUsageEncoded = Schema.decodeUnknownSync(UsageEncoded) +const nonNegativeTokenCount = (value: number | undefined): number | undefined => + value !== undefined && Number.isFinite(value) && value >= 0 ? value : undefined + /** Convert Effect AI usage into fold's tolerant durable usage shape. */ export const usageFromResponseUsage = (usage: Response.Usage): UsageEncoded => - decodeUsageEncoded(encodeResponseUsage(usage)) + decodeUsageEncoded({ + inputTokens: { + uncached: nonNegativeTokenCount(usage.inputTokens.uncached), + total: nonNegativeTokenCount(usage.inputTokens.total), + cacheRead: nonNegativeTokenCount(usage.inputTokens.cacheRead), + cacheWrite: nonNegativeTokenCount(usage.inputTokens.cacheWrite), + }, + outputTokens: { + total: nonNegativeTokenCount(usage.outputTokens.total), + text: nonNegativeTokenCount(usage.outputTokens.text), + reasoning: nonNegativeTokenCount(usage.outputTokens.reasoning), + }, + }) /** Best estimate of total input tokens from whatever fields the provider reported. */ export const usageInputTotal = (usage: UsageEncoded): number => { diff --git a/packages/fold-core/test/EventLog/Usage.vi.test.ts b/packages/fold-core/test/EventLog/Usage.vi.test.ts new file mode 100644 index 0000000..25a1b54 --- /dev/null +++ b/packages/fold-core/test/EventLog/Usage.vi.test.ts @@ -0,0 +1,36 @@ +import { describe, expect, it } from '@effect/vitest' +import type { Response } from 'effect/unstable/ai' + +import { usageFromResponseUsage } from '../../src/EventLog/Usage' + +describe('usageFromResponseUsage', () => { + it('omits invalid derived token details while preserving provider totals', () => { + const usage: Response.Usage = { + inputTokens: { + uncached: 10, + total: 10, + cacheRead: 0, + cacheWrite: 0, + }, + outputTokens: { + total: 8, + text: -4, + reasoning: 12, + }, + } + + expect(usageFromResponseUsage(usage)).toEqual({ + inputTokens: { + uncached: 10, + total: 10, + cacheRead: 0, + cacheWrite: 0, + }, + outputTokens: { + total: 8, + text: undefined, + reasoning: 12, + }, + }) + }) +}) diff --git a/packages/fold-xai/src/XaiModel.ts b/packages/fold-xai/src/XaiModel.ts index 931a223..0170485 100644 --- a/packages/fold-xai/src/XaiModel.ts +++ b/packages/fold-xai/src/XaiModel.ts @@ -1,9 +1,14 @@ /** FoldModel factory for xAI's OpenAI-compatible inference API authenticated with OAuth. */ import { OpenAiClient, OpenAiLanguageModel } from '@effect/ai-openai-compat' +import type { + ChatCompletionChunk, + CreateResponse200, + CreateResponse200Sse, +} from '@effect/ai-openai-compat/OpenAiClient' import * as NodeFileSystem from '@effect/platform-node/NodeFileSystem' import { customModel, resolveOpenAiReasoning } from '@humanlayer/fold-core' import type { FoldModel, ReasoningLevel } from '@humanlayer/fold-core' -import { Context, Effect, Layer } from 'effect' +import { Context, Effect, Layer, Option, Predicate, Schema, Stream } from 'effect' import type { Scope } from 'effect' import type { LanguageModel } from 'effect/unstable/ai' import { FetchHttpClient, HttpClient } from 'effect/unstable/http' @@ -14,6 +19,48 @@ import { DEFAULT_XAI_MODEL_ID } from './XaiModelCatalog' export const XAI_API_URL = 'https://api.x.ai/v1' +const TokenCount = Schema.Finite.check(Schema.isInt(), Schema.isGreaterThanOrEqualTo(0)) +const XaiCompletionTokenDetails = Schema.Struct({ reasoning_tokens: TokenCount }) +const decodeXaiCompletionTokenDetails = Schema.decodeUnknownOption(XaiCompletionTokenDetails) + +type XaiUsage = NonNullable + +/** + * xAI reports `completion_tokens` as text-only while putting reasoning tokens in + * `completion_tokens_details`. OpenAI-compatible clients expect `completion_tokens` to include both. + */ +export const normalizeXaiChatCompletionUsage = (usage: XaiUsage): XaiUsage => { + const details = decodeXaiCompletionTokenDetails(usage.completion_tokens_details) + return Option.match(details, { + onNone: () => usage, + onSome: ({ reasoning_tokens: reasoningTokens }) => ({ + ...usage, + completion_tokens: usage.completion_tokens + reasoningTokens, + }), + }) +} + +const normalizeXaiResponse = ( + response: Response, +): Response => { + if (Predicate.isNullish(response.usage)) return response + return { ...response, usage: normalizeXaiChatCompletionUsage(response.usage) } +} + +const normalizeXaiStreamResponse = (response: CreateResponse200Sse): CreateResponse200Sse => + typeof response === 'string' || '_tag' in response ? response : normalizeXaiResponse(response) + +/** Normalize xAI's token semantics before the stock OpenAI-compatible model derives usage details. */ +export const decorateXaiClient = (inner: OpenAiClient.Service): OpenAiClient.Service => ({ + ...inner, + createResponse: (options) => + inner.createResponse(options).pipe(Effect.map(([body, response]) => [normalizeXaiResponse(body), response])), + createResponseStream: (options) => + inner + .createResponseStream(options) + .pipe(Effect.map(([response, stream]) => [response, stream.pipe(Stream.map(normalizeXaiStreamResponse))])), +}) + export type XaiModelOptions = { readonly model?: string readonly reasoning?: ReasoningLevel @@ -35,8 +82,9 @@ export const makeXaiLanguageModel = ( const clientContext = yield* Layer.build(OpenAiClient.layer({ apiUrl: options.apiUrl ?? XAI_API_URL })).pipe( Effect.provideService(HttpClient.HttpClient, withXaiAuth(base, auth)), ) + const client = decorateXaiClient(Context.get(clientContext, OpenAiClient.OpenAiClient)) return yield* OpenAiLanguageModel.make({ model: options.model ?? DEFAULT_XAI_MODEL_ID }).pipe( - Effect.provideService(OpenAiClient.OpenAiClient, Context.get(clientContext, OpenAiClient.OpenAiClient)), + Effect.provideService(OpenAiClient.OpenAiClient, client), ) }).pipe(Effect.provide(NodeFileSystem.layer)) diff --git a/packages/fold-xai/test/Xai.vi.test.ts b/packages/fold-xai/test/Xai.vi.test.ts index f86359b..43874bc 100644 --- a/packages/fold-xai/test/Xai.vi.test.ts +++ b/packages/fold-xai/test/Xai.vi.test.ts @@ -8,6 +8,7 @@ import { buildXaiAuthorizeUrl, DEFAULT_XAI_MODEL_ID, makeXaiAuthStore, + normalizeXaiChatCompletionUsage, XAI_FRONTIER_MODELS, XAI_BROWSER_REDIRECT_URI, XAI_CLIENT_ID, @@ -50,6 +51,37 @@ describe('xAI OAuth', () => { }) describe('xaiModel', () => { + it('normalizes xAI text-only completion tokens for the OpenAI-compatible adapter', () => { + expect( + normalizeXaiChatCompletionUsage({ + prompt_tokens: 641, + completion_tokens: 1, + total_tokens: 889, + prompt_tokens_details: { cached_tokens: 512 }, + completion_tokens_details: { reasoning_tokens: 247 }, + }), + ).toMatchObject({ + prompt_tokens: 641, + completion_tokens: 248, + total_tokens: 889, + completion_tokens_details: { reasoning_tokens: 247 }, + }) + }) + + it('normalizes xAI usage independently of the aggregate total', () => { + const usage = { + prompt_tokens: 10, + completion_tokens: 8, + total_tokens: 18, + completion_tokens_details: { reasoning_tokens: 3 }, + } + + expect(normalizeXaiChatCompletionUsage(usage)).toEqual({ + ...usage, + completion_tokens: 11, + }) + }) + it('exports the supported frontier catalog and defaults to its newest model', () => { expect(XAI_FRONTIER_MODELS).toEqual([ { modelId: 'grok-4.5', label: 'Grok 4.5' },