diff --git a/AGENTS.md b/AGENTS.md index 3110195..3503a64 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -317,7 +317,8 @@ The active checklist is in `docs/30_DAY_MVP_REPORT.md`; detailed long-term works ## Known Risks - `core.internal` implementation classes are package-private by design. Cross-module construction should continue through public factory/configuration APIs. -- Current usage totals and cached/reasoning breakdown semantics can double-count; fix the domain invariant before expanding pricing. +- `TokenUsage` now enforces normalized inclusive totals, optional cache-read/cache-creation/reasoning details, and explicit usage provenance. `DefaultCostCalculator` partitions overlapping totals into disjoint billable amounts before applying rates. +- Spring AI usage extraction converts map/JSON-compatible native usage objects into the normalized core model. Real-provider compatibility fixtures remain required because provider and Spring AI usage shapes can change independently. - Current budget flow is check-then-add, is not an atomic reservation, and may not enforce `BLOCK` before provider invocation. - Current Micrometer `ai.token.*` metrics may duplicate Spring AI Observability; preserve compatibility while deciding default suppression or replacement. - Spring AI 2.0.0 documentation cannot be assumed to describe the current 1.1.4 runtime exactly; capability detection and supported-version tests are release requirements. @@ -382,6 +383,8 @@ Stage and deploy a Central release: ### 2026-07-20 +- Replaced generic cached input/cached output details with optional cache-read input, cache-creation input, and reasoning-output breakdowns; `null` now means unreported and `0` means reported zero. +- Added `UsageSource`, provider-specific total normalization in the Spring AI adapter, and disjoint cost partitioning to prevent details from being charged twice. - Positioned Token Pilot as a framework-independent Java LLM control and accounting core with Spring AI as an optional adapter. - Chose to reuse Spring AI Observability for standard latency, trace, and token telemetry while keeping Token Pilot metrics focused on cost, policy, budget, and reconciliation. - Kept the existing starter as a Spring AI convenience distribution rather than the sole product identity. diff --git a/README.md b/README.md index daaba51..42460d0 100644 --- a/README.md +++ b/README.md @@ -89,8 +89,8 @@ This configuration describes the current starter path. The 30-day MVP will exten | Capability | Status | MVP decision | | --- | --- | --- | -| Provider usage -> `TokenUsage` normalization | Basic implementation | Fix total/breakdown invariants first | -| `BigDecimal` cost calculation and pricing registry | Basic implementation | Define precision, rounding, and missing-price policy | +| Provider usage -> `TokenUsage` normalization | Inclusive totals, optional cache/reasoning breakdown, and source implemented | Add supported-provider compatibility fixtures | +| `BigDecimal` cost calculation and pricing registry | Disjoint billable partition implemented | Define precision, rounding, and missing-price policy | | `LedgerManager` event publication | Basic implementation | Add idempotent estimate/actual reconciliation contract | | Spring AI `LedgerAdvisor` | Basic implementation | Keep as an adapter; enforce preflight decisions before the call | | Micrometer metrics | Basic implementation | Avoid duplicating Spring AI token/latency telemetry by default | @@ -99,6 +99,10 @@ This configuration describes the current starter path. The 30-day MVP will exten | Exact byte-level BPE and JMH optimization | Post-MVP | Keep the estimator SPI so it can be added without API churn | | Provider routing, retry, fallback, gateway runtime | Future | Build only after accounting correctness is demonstrated | +`TokenUsage` represents normalized input/output as inclusive totals. `TokenUsageDetails` stores optional cache-read input, cache-creation input, and reasoning-output breakdowns; `null` means the provider did not report the value, while `0` means it reported zero usage. `UsageSource` distinguishes reported, provider-derived, locally calculated, estimated, and unavailable values. Cost calculation partitions these overlapping totals into disjoint billable amounts before applying rates, so details are never charged twice. + +Before the first `0.1.0` release, the former `Map` record component was replaced by `inputTokens`, `outputTokens`, `TokenUsageDetails`, `UsageSource`, and `metadata`. `CACHED_PROMPT`/`CACHED_COMPLETION` were replaced by `CACHE_READ_PROMPT`/`CACHE_CREATION_PROMPT`; consumers of the pre-release source API must migrate constructor calls and pricing keys accordingly. + ## Modules | Module | Purpose | Status | diff --git a/token-pilot-core/src/main/java/io/tokenpilot/core/domain/PricingPlan.java b/token-pilot-core/src/main/java/io/tokenpilot/core/domain/PricingPlan.java index 7da178d..fb962ca 100644 --- a/token-pilot-core/src/main/java/io/tokenpilot/core/domain/PricingPlan.java +++ b/token-pilot-core/src/main/java/io/tokenpilot/core/domain/PricingPlan.java @@ -70,8 +70,7 @@ public BigDecimal completionPricePerK() { /** * 특정 토큰 타입의 단가를 가져옵니다. 없을 시 계층 구조에 따라 대체값을 반환합니다. * REASONING -> COMPLETION - * CACHED_PROMPT -> PROMPT - * CACHED_COMPLETION -> COMPLETION + * CACHE_READ_PROMPT, CACHE_CREATION_PROMPT -> PROMPT */ public BigDecimal getRate(TokenType type) { if (rates.containsKey(type)) { @@ -80,8 +79,9 @@ public BigDecimal getRate(TokenType type) { // Fallback Logic return switch (type) { - case REASONING, CACHED_COMPLETION -> rates.getOrDefault(TokenType.COMPLETION, BigDecimal.ZERO); - case CACHED_PROMPT -> rates.getOrDefault(TokenType.PROMPT, BigDecimal.ZERO); + case REASONING -> rates.getOrDefault(TokenType.COMPLETION, BigDecimal.ZERO); + case CACHE_READ_PROMPT, CACHE_CREATION_PROMPT -> + rates.getOrDefault(TokenType.PROMPT, BigDecimal.ZERO); default -> rates.getOrDefault(type, BigDecimal.ZERO); }; } diff --git a/token-pilot-core/src/main/java/io/tokenpilot/core/domain/TokenType.java b/token-pilot-core/src/main/java/io/tokenpilot/core/domain/TokenType.java index 24abfe8..94ddd9b 100644 --- a/token-pilot-core/src/main/java/io/tokenpilot/core/domain/TokenType.java +++ b/token-pilot-core/src/main/java/io/tokenpilot/core/domain/TokenType.java @@ -10,22 +10,26 @@ public enum TokenType { COMPLETION, /** 추론 (Reasoning) - 주로 출력 계열 */ REASONING, - /** 캐시된 입력 (Cached Prompt) */ - CACHED_PROMPT, - /** 캐시된 출력 (Cached Completion) */ - CACHED_COMPLETION; + /** 캐시에서 읽은 입력 (Cache Read Prompt) */ + CACHE_READ_PROMPT, + /** 캐시에 새로 저장한 입력 (Cache Creation Prompt) */ + CACHE_CREATION_PROMPT; /** * 해당 토큰 타입이 입력(Prompt) 계열인지 확인합니다. + * + * @return 입력 계열이면 {@code true} */ public boolean isPrompt() { - return this == PROMPT || this == CACHED_PROMPT; + return this == PROMPT || this == CACHE_READ_PROMPT || this == CACHE_CREATION_PROMPT; } /** * 해당 토큰 타입이 출력(Completion) 계열인지 확인합니다. + * + * @return 출력 계열이면 {@code true} */ public boolean isCompletion() { - return this == COMPLETION || this == REASONING || this == CACHED_COMPLETION; + return this == COMPLETION || this == REASONING; } } diff --git a/token-pilot-core/src/main/java/io/tokenpilot/core/domain/TokenUsage.java b/token-pilot-core/src/main/java/io/tokenpilot/core/domain/TokenUsage.java index 05861aa..6c82eaf 100644 --- a/token-pilot-core/src/main/java/io/tokenpilot/core/domain/TokenUsage.java +++ b/token-pilot-core/src/main/java/io/tokenpilot/core/domain/TokenUsage.java @@ -1,77 +1,170 @@ package io.tokenpilot.core.domain; -import java.util.Collections; -import java.util.EnumMap; import java.util.Map; +import java.util.Objects; /** * AI 모델 호출 시 발생하는 토큰 사용량 정보. - * 세부적인 {@link TokenType} 별로 사용량을 관리합니다. + * 전체 입력/출력 토큰과 세부 사용량을 관리합니다. * - * @param tokenCounts 토큰 타입별 사용량 (Map) - * @param metadata 추가 메타데이터 (예: 모델 정보 등) + * @param inputTokens 전체 입력 토큰 + * @param outputTokens 전체 출력 토큰 + * @param details 세부 토큰 사용량 + * @param source 사용량 값의 출처 + * @param metadata 추가 메타데이터 (예: 모델 정보 등) */ public record TokenUsage( - Map tokenCounts, + long inputTokens, + long outputTokens, + TokenUsageDetails details, + UsageSource source, Map metadata ) { + /** + * 토큰 총량과 세부량의 포함 관계를 검증하고 metadata를 불변 복사합니다. + * + * @throws IllegalArgumentException 토큰 수가 음수이거나 세부량이 총량을 초과한 경우 + * @throws NullPointerException details/source가 null이거나 metadata에 null key/value가 있는 경우 + */ public TokenUsage { - tokenCounts = Collections.unmodifiableMap(new EnumMap<>(tokenCounts)); - metadata = (metadata != null) ? Collections.unmodifiableMap(metadata) : Map.of(); + if (inputTokens < 0) { + throw new IllegalArgumentException( + "inputTokens must be non-negative" + ); + } + + if (outputTokens < 0) { + throw new IllegalArgumentException( + "outputTokens must be non-negative" + ); + } + + details = Objects.requireNonNull( + details, + "details must not be null" + ); + source = Objects.requireNonNull( + source, + "source must not be null" + ); + + long cacheRead = countOrZero(details.cacheReadInputTokens()); + long cacheCreation = countOrZero(details.cacheCreationInputTokens()); + if (cacheRead > inputTokens + || cacheCreation > inputTokens - cacheRead) { + throw new IllegalArgumentException( + "Input details must not exceed inputTokens" + ); + } + + long reasoning = countOrZero(details.reasoningOutputTokens()); + if (reasoning > outputTokens) { + throw new IllegalArgumentException( + "reasoningOutputTokens must not exceed outputTokens" + ); + } + + metadata = metadata == null ? Map.of() : Map.copyOf(metadata); } /** * 기본 입력/출력 토큰을 사용하는 {@link TokenUsage}를 생성합니다. + * + * @param prompt 전체 입력 토큰 + * @param completion 전체 출력 토큰 + * @return 세부량과 metadata가 비어 있는 사용량 */ public static TokenUsage from(long prompt, long completion) { - Map counts = new EnumMap<>(TokenType.class); - counts.put(TokenType.PROMPT, prompt); - counts.put(TokenType.COMPLETION, completion); - return new TokenUsage(counts, Map.of()); + return new TokenUsage( + prompt, + completion, + TokenUsageDetails.unreported(), + UsageSource.PROVIDER_REPORTED, + Map.of() + ); } /** * 입력/출력/추론 토큰을 포함하는 {@link TokenUsage}를 생성합니다. + * + * @param prompt 전체 입력 토큰 + * @param completion 전체 출력 토큰 + * @param reasoning 전체 출력에 포함된 reasoning 토큰 + * @return reasoning 세부량을 포함하는 사용량 */ public static TokenUsage from(long prompt, long completion, long reasoning) { - Map counts = new EnumMap<>(TokenType.class); - counts.put(TokenType.PROMPT, prompt); - counts.put(TokenType.COMPLETION, completion); - counts.put(TokenType.REASONING, reasoning); - return new TokenUsage(counts, Map.of()); + return new TokenUsage( + prompt, + completion, + new TokenUsageDetails(null, null, reasoning), + UsageSource.PROVIDER_REPORTED, + Map.of() + ); + } + + /** + * provider 응답에서 사용량 정보를 얻지 못한 상태를 생성합니다. + * + * @param metadata 보존할 응답 메타데이터 + * @return 출처가 {@link UsageSource#UNAVAILABLE}인 0 토큰 사용량 + */ + public static TokenUsage unavailable(Map metadata) { + return new TokenUsage( + 0, + 0, + TokenUsageDetails.unreported(), + UsageSource.UNAVAILABLE, + metadata + ); } /** * 모든 종류의 입력/출력 토큰 수의 합계를 반환합니다. + * + * @return 전체 입력 토큰 */ public long promptTokens() { - return tokenCounts.entrySet().stream() - .filter(e -> e.getKey().isPrompt()) - .mapToLong(Map.Entry::getValue) - .sum(); + return inputTokens; } /** * 모든 종류의 출력(추론 포함) 토큰 수의 합계를 반환합니다. + * + * @return 전체 출력 토큰 */ public long completionTokens() { - return tokenCounts.entrySet().stream() - .filter(e -> e.getKey().isCompletion()) - .mapToLong(Map.Entry::getValue) - .sum(); + return outputTokens; } /** * 전체 사용 토큰 수의 합계를 반환합니다. + * + * @return 전체 입력과 출력 토큰의 합 + * @throws ArithmeticException 합계가 {@code long} 범위를 초과한 경우 */ public long totalTokens() { - return tokenCounts.values().stream().mapToLong(Long::longValue).sum(); + return Math.addExact(inputTokens, outputTokens); } /** - * 특정 토큰 타입의 사용량을 가져옵니다. 없을 시 0을 반환합니다. + * 특정 토큰 타입의 전체 또는 세부 사용량을 반환합니다. + * 보고되지 않은 세부량은 이 호환 projection에서 {@code 0}으로 반환되며, + * 미보고 여부는 {@link #details()}의 nullable 필드로 확인해야 합니다. + * + * @param type 조회할 토큰 타입 + * @return 해당 토큰 타입의 사용량 */ public long getCount(TokenType type) { - return tokenCounts.getOrDefault(type, 0L); + return switch (type) { + case PROMPT -> inputTokens; + case COMPLETION -> outputTokens; + case REASONING -> countOrZero(details.reasoningOutputTokens()); + case CACHE_READ_PROMPT -> countOrZero(details.cacheReadInputTokens()); + case CACHE_CREATION_PROMPT -> countOrZero(details.cacheCreationInputTokens()); + }; + } + + private static long countOrZero(Long count) { + return count == null ? 0L : count; } } diff --git a/token-pilot-core/src/main/java/io/tokenpilot/core/domain/TokenUsageDetails.java b/token-pilot-core/src/main/java/io/tokenpilot/core/domain/TokenUsageDetails.java new file mode 100644 index 0000000..58bac48 --- /dev/null +++ b/token-pilot-core/src/main/java/io/tokenpilot/core/domain/TokenUsageDetails.java @@ -0,0 +1,42 @@ +package io.tokenpilot.core.domain; + +/** + * 정규화된 전체 입력/출력 토큰에 포함되는 선택적 세부 사용량. + * {@code null}은 provider가 값을 보고하지 않았음을, {@code 0}은 값을 + * 보고했으나 실제 사용량이 없었음을 뜻합니다. + * + * @param cacheReadInputTokens 전체 입력에 포함된 cache read 토큰 또는 미보고 시 {@code null} + * @param cacheCreationInputTokens 전체 입력에 포함된 cache creation 토큰 또는 미보고 시 {@code null} + * @param reasoningOutputTokens 전체 출력에 포함된 reasoning 토큰 또는 미보고 시 {@code null} + */ +public record TokenUsageDetails( + Long cacheReadInputTokens, + Long cacheCreationInputTokens, + Long reasoningOutputTokens +) { + /** + * 모든 세부 토큰 수가 0 이상인지 검증합니다. + * + * @throws IllegalArgumentException 세부 토큰 수가 음수인 경우 + */ + public TokenUsageDetails { + if (cacheReadInputTokens != null && cacheReadInputTokens < 0) { + throw new IllegalArgumentException("cacheReadInputTokens must be non-negative"); + } + if (cacheCreationInputTokens != null && cacheCreationInputTokens < 0) { + throw new IllegalArgumentException("cacheCreationInputTokens must be non-negative"); + } + if (reasoningOutputTokens != null && reasoningOutputTokens < 0) { + throw new IllegalArgumentException("reasoningOutputTokens must be non-negative"); + } + } + + /** + * provider가 세부 사용량을 보고하지 않은 상태를 생성합니다. + * + * @return 모든 세부량이 미보고 상태인 객체 + */ + public static TokenUsageDetails unreported() { + return new TokenUsageDetails(null, null, null); + } +} diff --git a/token-pilot-core/src/main/java/io/tokenpilot/core/domain/UsageSource.java b/token-pilot-core/src/main/java/io/tokenpilot/core/domain/UsageSource.java new file mode 100644 index 0000000..b958e7b --- /dev/null +++ b/token-pilot-core/src/main/java/io/tokenpilot/core/domain/UsageSource.java @@ -0,0 +1,17 @@ +package io.tokenpilot.core.domain; + +/** + * 토큰 사용량 값의 출처와 생성 방식을 나타냅니다. + */ +public enum UsageSource { + /** provider가 포괄 총량을 직접 보고한 사용량 */ + PROVIDER_REPORTED, + /** provider가 보고한 여러 필드를 어댑터가 정규화해 만든 사용량 */ + PROVIDER_DERIVED, + /** 로컬 tokenizer가 계산한 사용량 */ + LOCAL_TOKENIZER, + /** 휴리스틱으로 근사 추정한 사용량 */ + HEURISTIC_ESTIMATE, + /** 사용량 정보를 얻을 수 없음 */ + UNAVAILABLE +} diff --git a/token-pilot-core/src/main/java/io/tokenpilot/core/internal/DefaultCostCalculator.java b/token-pilot-core/src/main/java/io/tokenpilot/core/internal/DefaultCostCalculator.java index a2d8c79..079daf8 100644 --- a/token-pilot-core/src/main/java/io/tokenpilot/core/internal/DefaultCostCalculator.java +++ b/token-pilot-core/src/main/java/io/tokenpilot/core/internal/DefaultCostCalculator.java @@ -9,11 +9,10 @@ import java.math.BigDecimal; import java.math.RoundingMode; -import java.util.Map; - /** * 기본 비용 계산기 구현체. - * 각 {@link TokenType} 별 단가를 적용하여 정밀하게 계산합니다. + * 포괄 총량에서 cache/reasoning 세부량을 분리한 배타적 구간별로 + * {@link TokenType} 단가를 적용하여 중복 없이 계산합니다. * 1K 토큰당 가격 정보를 사용하여 소수점 10자리까지 중간 계산 후 6자리로 최종 반올림합니다. */ class DefaultCostCalculator implements CostCalculator { @@ -21,21 +20,32 @@ class DefaultCostCalculator implements CostCalculator { @Override public Cost calculate(TokenUsage usage, PricingPlan plan) { - BigDecimal totalCostValue = BigDecimal.ZERO; - - // 사용된 모든 토큰 타입에 대해 각각의 단가를 적용하여 합산 - for (Map.Entry entry : usage.tokenCounts().entrySet()) { - TokenType type = entry.getKey(); - Long count = entry.getValue(); - - if (count > 0) { - BigDecimal rate = plan.getRate(type); - BigDecimal typeCost = rate.multiply(BigDecimal.valueOf(count)) - .divide(THOUSAND, 10, RoundingMode.HALF_UP); - totalCostValue = totalCostValue.add(typeCost); - } - } + long cacheReadInput = countOrZero(usage.details().cacheReadInputTokens()); + long cacheCreationInput = countOrZero(usage.details().cacheCreationInputTokens()); + long reasoningOutput = countOrZero(usage.details().reasoningOutputTokens()); + + long regularInput = usage.inputTokens() - cacheReadInput - cacheCreationInput; + long regularOutput = usage.outputTokens() - reasoningOutput; + + BigDecimal totalCostValue = costFor(regularInput, plan.getRate(TokenType.PROMPT)) + .add(costFor(cacheReadInput, plan.getRate(TokenType.CACHE_READ_PROMPT))) + .add(costFor(cacheCreationInput, plan.getRate(TokenType.CACHE_CREATION_PROMPT))) + .add(costFor(regularOutput, plan.getRate(TokenType.COMPLETION))) + .add(costFor(reasoningOutput, plan.getRate(TokenType.REASONING))); return new Cost(totalCostValue, plan.currency()); } + + private BigDecimal costFor(long count, BigDecimal rate) { + if (count == 0) { + return BigDecimal.ZERO; + } + + return rate.multiply(BigDecimal.valueOf(count)) + .divide(THOUSAND, 10, RoundingMode.HALF_UP); + } + + private long countOrZero(Long count) { + return count == null ? 0L : count; + } } diff --git a/token-pilot-core/src/test/java/io/tokenpilot/core/domain/TokenTypeTest.java b/token-pilot-core/src/test/java/io/tokenpilot/core/domain/TokenTypeTest.java new file mode 100644 index 0000000..dc13e22 --- /dev/null +++ b/token-pilot-core/src/test/java/io/tokenpilot/core/domain/TokenTypeTest.java @@ -0,0 +1,26 @@ +package io.tokenpilot.core.domain; + +import static org.assertj.core.api.Assertions.assertThat; + +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +class TokenTypeTest { + + @Test + @DisplayName("completion 계열에는 전체 출력과 reasoning breakdown이 포함된다") + void shouldClassifyCompletionTokenTypes() { + assertThat(TokenType.COMPLETION.isCompletion()).isTrue(); + assertThat(TokenType.REASONING.isCompletion()).isTrue(); + assertThat(TokenType.PROMPT.isCompletion()).isFalse(); + } + + @Test + @DisplayName("prompt 계열에는 전체 입력과 cache read/create breakdown이 포함된다") + void shouldClassifyPromptTokenTypes() { + assertThat(TokenType.PROMPT.isPrompt()).isTrue(); + assertThat(TokenType.CACHE_READ_PROMPT.isPrompt()).isTrue(); + assertThat(TokenType.CACHE_CREATION_PROMPT.isPrompt()).isTrue(); + assertThat(TokenType.COMPLETION.isPrompt()).isFalse(); + } +} diff --git a/token-pilot-core/src/test/java/io/tokenpilot/core/domain/TokenUsageDetailsTest.java b/token-pilot-core/src/test/java/io/tokenpilot/core/domain/TokenUsageDetailsTest.java new file mode 100644 index 0000000..aa5a401 --- /dev/null +++ b/token-pilot-core/src/test/java/io/tokenpilot/core/domain/TokenUsageDetailsTest.java @@ -0,0 +1,48 @@ +package io.tokenpilot.core.domain; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +class TokenUsageDetailsTest { + + @Test + @DisplayName("cache read input 토큰은 음수일 수 없다") + void shouldRejectNegativeCacheReadInputTokens() { + assertThatThrownBy(() -> + new TokenUsageDetails(-1L, 0L, 0L) + ).isInstanceOf(IllegalArgumentException.class); + } + + @Test + @DisplayName("cache creation input 토큰은 음수일 수 없다") + void shouldRejectNegativeCacheCreationInputTokens() { + assertThatThrownBy(() -> + new TokenUsageDetails(0L, -1L, 0L) + ).isInstanceOf(IllegalArgumentException.class); + } + + @Test + @DisplayName("reasoning output 토큰은 음수일 수 없다") + void shouldRejectNegativeReasoningOutputTokens() { + assertThatThrownBy(() -> + new TokenUsageDetails(0L, 0L, -1L) + ).isInstanceOf(IllegalArgumentException.class); + } + + @Test + @DisplayName("미보고 세부량과 보고된 0을 구분한다") + void shouldDistinguishUnreportedDetailsFromReportedZero() { + TokenUsageDetails unreported = TokenUsageDetails.unreported(); + TokenUsageDetails reportedZero = new TokenUsageDetails(0L, 0L, 0L); + + assertThat(unreported.cacheReadInputTokens()).isNull(); + assertThat(unreported.cacheCreationInputTokens()).isNull(); + assertThat(unreported.reasoningOutputTokens()).isNull(); + assertThat(reportedZero.cacheReadInputTokens()).isZero(); + assertThat(reportedZero.cacheCreationInputTokens()).isZero(); + assertThat(reportedZero.reasoningOutputTokens()).isZero(); + } +} diff --git a/token-pilot-core/src/test/java/io/tokenpilot/core/domain/TokenUsageTest.java b/token-pilot-core/src/test/java/io/tokenpilot/core/domain/TokenUsageTest.java new file mode 100644 index 0000000..7eab1cf --- /dev/null +++ b/token-pilot-core/src/test/java/io/tokenpilot/core/domain/TokenUsageTest.java @@ -0,0 +1,243 @@ +package io.tokenpilot.core.domain; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +import java.util.HashMap; +import java.util.Map; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +class TokenUsageTest { + + @Test + @DisplayName("세부량이 보고되지 않으면 입력과 출력 총량만 사용한다") + void shouldUseInputAndOutputTotalsWithoutReportedDetails() { + TokenUsage usage = TokenUsage.from(100, 200); + + assertThat(usage.inputTokens()).isEqualTo(100); + assertThat(usage.outputTokens()).isEqualTo(200); + assertThat(usage.totalTokens()).isEqualTo(300); + assertThat(usage.details()).isEqualTo(TokenUsageDetails.unreported()); + assertThat(usage.source()).isEqualTo(UsageSource.PROVIDER_REPORTED); + } + + @Test + @DisplayName("입력과 출력 토큰은 모두 0일 수 있다") + void shouldAllowZeroInputAndOutputTokens() { + assertThat(TokenUsage.from(0, 0).totalTokens()).isZero(); + } + + @Test + @DisplayName("reasoning 토큰은 output total에 중복 합산되지 않는다") + void shouldNotAddReasoningTokensToOutputTotalAgain() { + TokenUsage usage = TokenUsage.from(100, 200, 150); + + assertThat(usage.inputTokens()).isEqualTo(100); + assertThat(usage.outputTokens()).isEqualTo(200); + assertThat(usage.totalTokens()).isEqualTo(300); + assertThat(usage.details().reasoningOutputTokens()).isEqualTo(150); + assertThat(usage.getCount(TokenType.REASONING)).isEqualTo(150); + } + + @Test + @DisplayName("cache read와 creation 토큰은 input total에 중복 합산되지 않는다") + void shouldNotAddCacheBreakdownToInputTotalAgain() { + TokenUsage usage = reportedUsage( + 100, + 200, + new TokenUsageDetails(40L, 20L, null), + Map.of() + ); + + assertThat(usage.totalTokens()).isEqualTo(300); + assertThat(usage.getCount(TokenType.CACHE_READ_PROMPT)).isEqualTo(40); + assertThat(usage.getCount(TokenType.CACHE_CREATION_PROMPT)).isEqualTo(20); + } + + @Test + @DisplayName("input 토큰은 음수일 수 없다") + void shouldRejectNegativeInputTokens() { + assertThatThrownBy(() -> reportedUsage(-1, 0, TokenUsageDetails.unreported(), Map.of())) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("inputTokens must be non-negative"); + } + + @Test + @DisplayName("output 토큰은 음수일 수 없다") + void shouldRejectNegativeOutputTokens() { + assertThatThrownBy(() -> reportedUsage(0, -1, TokenUsageDetails.unreported(), Map.of())) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("outputTokens must be non-negative"); + } + + @Test + @DisplayName("토큰 세부 정보는 null일 수 없다") + void shouldRejectNullTokenUsageDetails() { + assertThatThrownBy(() -> reportedUsage(100, 200, null, Map.of())) + .isInstanceOf(NullPointerException.class) + .hasMessage("details must not be null"); + } + + @Test + @DisplayName("사용량 출처는 null일 수 없다") + void shouldRejectNullUsageSource() { + assertThatThrownBy(() -> new TokenUsage( + 100, + 200, + TokenUsageDetails.unreported(), + null, + Map.of() + )) + .isInstanceOf(NullPointerException.class) + .hasMessage("source must not be null"); + } + + @Test + @DisplayName("cache read 토큰은 전체 input보다 클 수 없다") + void shouldRejectCacheReadInputTokensGreaterThanInputTokens() { + assertThatThrownBy(() -> reportedUsage( + 100, + 200, + new TokenUsageDetails(101L, null, null), + Map.of() + )) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("Input details must not exceed inputTokens"); + } + + @Test + @DisplayName("cache creation 토큰은 전체 input보다 클 수 없다") + void shouldRejectCacheCreationInputTokensGreaterThanInputTokens() { + assertThatThrownBy(() -> reportedUsage( + 100, + 200, + new TokenUsageDetails(null, 101L, null), + Map.of() + )) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("Input details must not exceed inputTokens"); + } + + @Test + @DisplayName("cache read와 creation의 합은 전체 input보다 클 수 없다") + void shouldRejectCacheBreakdownGreaterThanInputTokens() { + assertThatThrownBy(() -> reportedUsage( + 100, + 200, + new TokenUsageDetails(60L, 41L, null), + Map.of() + )) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("Input details must not exceed inputTokens"); + } + + @Test + @DisplayName("cache breakdown 검증은 long overflow 없이 수행한다") + void shouldRejectCacheBreakdownWithoutOverflow() { + assertThatThrownBy(() -> reportedUsage( + Long.MAX_VALUE, + 0, + new TokenUsageDetails(Long.MAX_VALUE, 1L, null), + Map.of() + )) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("Input details must not exceed inputTokens"); + } + + @Test + @DisplayName("reasoning output은 전체 output보다 클 수 없다") + void shouldRejectReasoningOutputTokensGreaterThanOutputTokens() { + assertThatThrownBy(() -> reportedUsage( + 100, + 200, + new TokenUsageDetails(null, null, 201L), + Map.of() + )) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("reasoningOutputTokens must not exceed outputTokens"); + } + + @Test + @DisplayName("metadata는 원본 Map의 변경에 영향을 받지 않는다") + void shouldDefensivelyCopyMetadata() { + Map metadata = new HashMap<>(); + metadata.put("model", "gpt-4o"); + TokenUsage usage = reportedUsage(100, 200, TokenUsageDetails.unreported(), metadata); + + metadata.put("model", "changed"); + + assertThat(usage.metadata()).containsEntry("model", "gpt-4o"); + } + + @Test + @DisplayName("metadata가 null이면 빈 Map을 사용한다") + void shouldUseEmptyMetadataWhenMetadataIsNull() { + TokenUsage usage = reportedUsage(100, 200, TokenUsageDetails.unreported(), null); + + assertThat(usage.metadata()).isEmpty(); + } + + @Test + @DisplayName("metadata의 key는 null일 수 없다") + void shouldRejectMetadataWithNullKey() { + Map metadata = new HashMap<>(); + metadata.put(null, "value"); + + assertThatThrownBy(() -> reportedUsage(100, 200, TokenUsageDetails.unreported(), metadata)) + .isInstanceOf(NullPointerException.class); + } + + @Test + @DisplayName("metadata의 value는 null일 수 없다") + void shouldRejectMetadataWithNullValue() { + Map metadata = new HashMap<>(); + metadata.put("key", null); + + assertThatThrownBy(() -> reportedUsage(100, 200, TokenUsageDetails.unreported(), metadata)) + .isInstanceOf(NullPointerException.class); + } + + @Test + @DisplayName("breakdown 합계는 전체 토큰과 같을 수 있다") + void shouldAllowDetailsEqualToTotalTokens() { + TokenUsage usage = reportedUsage( + 100, + 200, + new TokenUsageDetails(70L, 30L, 200L), + Map.of() + ); + + assertThat(usage.inputTokens()).isEqualTo(100); + assertThat(usage.outputTokens()).isEqualTo(200); + } + + @Test + @DisplayName("전체 토큰 합계가 long 범위를 넘으면 예외가 발생한다") + void shouldRejectTotalTokenOverflow() { + TokenUsage usage = reportedUsage( + Long.MAX_VALUE, + 1, + TokenUsageDetails.unreported(), + Map.of() + ); + + assertThatThrownBy(usage::totalTokens) + .isInstanceOf(ArithmeticException.class); + } + + private TokenUsage reportedUsage( + long inputTokens, + long outputTokens, + TokenUsageDetails details, + Map metadata + ) { + return new TokenUsage( + inputTokens, + outputTokens, + details, + UsageSource.PROVIDER_REPORTED, + metadata + ); + } +} diff --git a/token-pilot-core/src/test/java/io/tokenpilot/core/internal/DefaultCostCalculatorTest.java b/token-pilot-core/src/test/java/io/tokenpilot/core/internal/DefaultCostCalculatorTest.java index ac07762..aff726f 100644 --- a/token-pilot-core/src/test/java/io/tokenpilot/core/internal/DefaultCostCalculatorTest.java +++ b/token-pilot-core/src/test/java/io/tokenpilot/core/internal/DefaultCostCalculatorTest.java @@ -4,12 +4,13 @@ import io.tokenpilot.core.domain.PricingPlan; import io.tokenpilot.core.domain.TokenType; import io.tokenpilot.core.domain.TokenUsage; -import org.junit.jupiter.api.BeforeEach; +import io.tokenpilot.core.domain.TokenUsageDetails; +import io.tokenpilot.core.domain.UsageSource; import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; import java.math.BigDecimal; -import java.util.EnumMap; +import java.util.Currency; import java.util.Map; import static org.assertj.core.api.Assertions.assertThat; @@ -34,45 +35,31 @@ void calculateStandardTokens() { } @Test - @DisplayName("추론 토큰 요율이 다를 때 각각 정밀하게 계산되어야 한다") - void calculateReasoningTokensWithDifferentRate() { - // Given: Prompt $0.015, Completion $0.060, Reasoning $0.045 - Map rates = new EnumMap<>(TokenType.class); - rates.put(TokenType.PROMPT, new BigDecimal("0.015")); - rates.put(TokenType.COMPLETION, new BigDecimal("0.060")); - rates.put(TokenType.REASONING, new BigDecimal("0.045")); - PricingPlan plan = new PricingPlan("o1-preview", rates, Cost.DEFAULT_CURRENCY); + @DisplayName("포괄 총량을 배타적인 과금 구간으로 나누어 중복 없이 계산한다") + void calculateInclusiveTotalsWithoutDoubleChargingDetails() { + PricingPlan plan = new PricingPlan( + "provider-model", + Map.of( + TokenType.PROMPT, new BigDecimal("0.01"), + TokenType.CACHE_READ_PROMPT, new BigDecimal("0.002"), + TokenType.CACHE_CREATION_PROMPT, new BigDecimal("0.0125"), + TokenType.COMPLETION, new BigDecimal("0.03"), + TokenType.REASONING, new BigDecimal("0.05") + ), + Currency.getInstance("USD") + ); + TokenUsage usage = new TokenUsage( + 1_000, + 1_000, + new TokenUsageDetails(400L, 100L, 200L), + UsageSource.PROVIDER_REPORTED, + Map.of() + ); - // Usage: Prompt 1000, Completion 500, Reasoning 1500 (Total 3000) - Map counts = new EnumMap<>(TokenType.class); - counts.put(TokenType.PROMPT, 1000L); - counts.put(TokenType.COMPLETION, 500L); - counts.put(TokenType.REASONING, 1500L); - TokenUsage usage = new TokenUsage(counts, Map.of()); - - // When - Cost cost = calculator.calculate(usage, plan); - - // Then - // 1000 * 0.015 / 1000 = 0.015 - // 500 * 0.060 / 1000 = 0.030 - // 1500 * 0.045 / 1000 = 0.0675 - // Total = 0.015 + 0.030 + 0.0675 = 0.1125 - assertThat(cost.value()).isEqualByComparingTo("0.112500"); - } - - @Test - @DisplayName("추론 토큰 요율이 명시되지 않으면 일반 출력 요율을 따라야 한다") - void calculateReasoningTokensWithFallback() { - // Given: Prompt $0.01, Completion $0.03 (Reasoning 없음) - PricingPlan plan = new PricingPlan("gpt-4o", new BigDecimal("0.01"), new BigDecimal("0.03")); - // Usage: Prompt 1000, Completion 0, Reasoning 1000 - TokenUsage usage = TokenUsage.from(1000, 0, 1000); - - // When Cost cost = calculator.calculate(usage, plan); - // Then: 1000 * 0.01 / 1000 + 1000 * 0.03 / 1000 = 0.04 - assertThat(cost.value()).isEqualByComparingTo("0.040000"); + // 500 normal input + 400 cache read + 100 cache creation + // 800 normal output + 200 reasoning output + assertThat(cost.value()).isEqualByComparingTo("0.041050"); } } diff --git a/token-pilot-core/src/test/java/io/tokenpilot/core/internal/InMemoryPricingRegistryTest.java b/token-pilot-core/src/test/java/io/tokenpilot/core/internal/InMemoryPricingRegistryTest.java index 01d3033..8c27125 100644 --- a/token-pilot-core/src/test/java/io/tokenpilot/core/internal/InMemoryPricingRegistryTest.java +++ b/token-pilot-core/src/test/java/io/tokenpilot/core/internal/InMemoryPricingRegistryTest.java @@ -1,7 +1,6 @@ package io.tokenpilot.core.internal; import io.tokenpilot.core.domain.PricingPlan; -import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; diff --git a/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/internal/MicroCostMetricsPublisher.java b/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/internal/MicroCostMetricsPublisher.java index a2fb9dc..1744581 100644 --- a/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/internal/MicroCostMetricsPublisher.java +++ b/token-pilot-micrometer/src/main/java/io/tokenpilot/micrometer/internal/MicroCostMetricsPublisher.java @@ -5,8 +5,9 @@ import io.micrometer.core.instrument.MeterRegistry; import io.micrometer.core.instrument.Tag; import io.micrometer.core.instrument.Tags; -import io.tokenpilot.core.domain.CostRecordedEvent; import io.tokenpilot.core.LedgerListener; +import io.tokenpilot.core.domain.CostRecordedEvent; +import io.tokenpilot.core.domain.TokenType; import io.tokenpilot.micrometer.MetricsOptions; import java.util.Map; @@ -43,7 +44,8 @@ public void onRecord(CostRecordedEvent event) { final Tags finalTags = commonTags; - event.usage().tokenCounts().forEach(((tokenType, count) -> { + for (TokenType tokenType : TokenType.values()) { + long count = event.usage().getCount(tokenType); if (count > 0) { String typeName = tokenType.name().toLowerCase(); DistributionSummary.builder("ai.token.usage.distribution") @@ -51,7 +53,7 @@ public void onRecord(CostRecordedEvent event) { .baseUnit("tokens") .tags(finalTags.and("token_type", typeName)) .register(meterRegistry) - .record(count.doubleValue()); + .record(count); Counter.builder("ai.token.usage.total") .description("Total number of AI tokens recorded") .baseUnit("tokens") @@ -59,7 +61,7 @@ public void onRecord(CostRecordedEvent event) { .register(meterRegistry) .increment(count); } - })); + } Counter.builder("ai.token.cost.total") .description("Total estimated AI token cost") .baseUnit("currency") diff --git a/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/DefaultUsageExtractor.java b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/DefaultUsageExtractor.java index bd0c682..4bf9681 100644 --- a/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/DefaultUsageExtractor.java +++ b/token-pilot-spring-ai/src/main/java/io/tokenpilot/springai/internal/DefaultUsageExtractor.java @@ -1,80 +1,160 @@ package io.tokenpilot.springai.internal; -import io.tokenpilot.core.domain.TokenType; import io.tokenpilot.core.domain.TokenUsage; +import io.tokenpilot.core.domain.TokenUsageDetails; +import io.tokenpilot.core.domain.UsageSource; import io.tokenpilot.springai.UsageExtractor; import org.springframework.ai.chat.client.ChatClientResponse; import org.springframework.ai.chat.metadata.ChatResponseMetadata; +import org.springframework.ai.chat.metadata.EmptyUsage; import org.springframework.ai.chat.metadata.Usage; import org.springframework.ai.chat.model.ChatResponse; +import org.springframework.ai.model.ModelOptionsUtils; -import java.util.EnumMap; import java.util.HashMap; import java.util.Map; /** * 기본 {@link UsageExtractor} 구현체. * Spring AI의 {@link Usage} 정보를 {@link TokenUsage}로 변환하며, - * 메타데이터에 포함된 추론(Reasoning) 토큰 등을 식별하여 세분화된 사용량을 추출합니다. + * provider native usage와 메타데이터에서 cache/reasoning breakdown을 식별합니다. + * Provider가 비포괄 총량을 반환하면 Token Pilot의 포괄 총량 계약에 맞게 정규화합니다. */ public class DefaultUsageExtractor implements UsageExtractor { @Override public TokenUsage extract(ChatClientResponse response) { if (response == null) { - return TokenUsage.from(0, 0); + return TokenUsage.unavailable(Map.of()); } ChatResponse chatResponse = response.chatResponse(); if (chatResponse == null || chatResponse.getMetadata() == null) { - return TokenUsage.from(0, 0); + return TokenUsage.unavailable(Map.of()); } ChatResponseMetadata metadata = chatResponse.getMetadata(); Usage usage = metadata.getUsage(); - if (usage == null) { - return new TokenUsage(defaultCounts(), copyMetadata(metadata, null)); + if (usage == null || usage instanceof EmptyUsage) { + return TokenUsage.unavailable(copyMetadata(metadata, null)); } - Map counts = new EnumMap<>(TokenType.class); - long prompt = (usage.getPromptTokens() != null) ? usage.getPromptTokens() : 0L; long completion = (usage.getCompletionTokens() != null) ? usage.getCompletionTokens() : 0L; - - counts.put(TokenType.PROMPT, prompt); - counts.put(TokenType.COMPLETION, completion); + Object nativeUsage = usage.getNativeUsage(); + Map nativeUsageFields = usageFields(nativeUsage); + + Long cacheRead = extractCacheReadTokens(metadata, nativeUsageFields); + Long cacheCreation = extractCacheCreationTokens(metadata, nativeUsageFields); + Long reasoning = extractReasoningTokens(metadata, nativeUsageFields); + UsageSource source = UsageSource.PROVIDER_REPORTED; - Long reasoning = extractReasoningTokens(metadata); - if (reasoning > 0) { - counts.put(TokenType.REASONING, reasoning); + Long nativeInput = firstNonNull( + numberAt(nativeUsageFields, "input_tokens"), + numberAt(nativeUsageFields, "inputTokens") + ); + if (cacheCreation != null && nativeInput != null && prompt == nativeInput) { + prompt = Math.addExact(prompt, countOrZero(cacheRead)); + prompt = Math.addExact(prompt, cacheCreation); + source = UsageSource.PROVIDER_DERIVED; } - return new TokenUsage(counts, copyMetadata(metadata, usage.getNativeUsage())); + Long nativeCandidates = firstNonNull( + numberAt(nativeUsageFields, "candidatesTokenCount"), + numberAt(nativeUsageFields, "candidates_token_count") + ); + if (reasoning != null && nativeCandidates != null && completion == nativeCandidates) { + completion = Math.addExact(completion, reasoning); + source = UsageSource.PROVIDER_DERIVED; + } + + return new TokenUsage( + prompt, + completion, + new TokenUsageDetails(cacheRead, cacheCreation, reasoning), + source, + copyMetadata(metadata, nativeUsage) + ); + } + + private Long extractCacheReadTokens(ChatResponseMetadata metadata, Object nativeUsage) { + return firstNonNull( + numberAt(metadata.get("prompt_tokens_details"), "cached_tokens"), + numberAt(metadata.get("input_tokens_details"), "cached_tokens"), + numberAt(metadata, "cache_read_input_tokens"), + numberAt(metadata, "cachedContentTokenCount"), + numberAt(nativeUsage, "prompt_tokens_details", "cached_tokens"), + numberAt(nativeUsage, "input_tokens_details", "cached_tokens"), + numberAt(nativeUsage, "cache_read_input_tokens"), + numberAt(nativeUsage, "cacheReadInputTokens"), + numberAt(nativeUsage, "cachedContentTokenCount"), + numberAt(nativeUsage, "cached_content_token_count") + ); + } + + private Long extractCacheCreationTokens(ChatResponseMetadata metadata, Object nativeUsage) { + return firstNonNull( + numberAt(metadata, "cache_creation_input_tokens"), + numberAt(metadata, "cacheCreationInputTokens"), + numberAt(nativeUsage, "cache_creation_input_tokens"), + numberAt(nativeUsage, "cacheCreationInputTokens") + ); + } + + private Long extractReasoningTokens(ChatResponseMetadata metadata, Object nativeUsage) { + return firstNonNull( + numberAt(metadata.get("completion_tokens_details"), "reasoning_tokens"), + numberAt(metadata.get("output_tokens_details"), "reasoning_tokens"), + numberAt(metadata, "reasoning_tokens"), + numberAt(metadata, "thoughtsTokenCount"), + numberAt(nativeUsage, "completion_tokens_details", "reasoning_tokens"), + numberAt(nativeUsage, "output_tokens_details", "reasoning_tokens"), + numberAt(nativeUsage, "reasoning_tokens"), + numberAt(nativeUsage, "reasoningTokens"), + numberAt(nativeUsage, "thoughtsTokenCount"), + numberAt(nativeUsage, "thoughts_token_count") + ); } - private Long extractReasoningTokens(ChatResponseMetadata metadata) { - // metadata는 직접 Map이 아닐 수 있으므로 get 메서드로 개별 접근 - Object details = metadata.get("completion_tokens_details"); - if (details instanceof Map detailsMap) { - Object reasoning = detailsMap.get("reasoning_tokens"); - if (reasoning instanceof Number num) { - return num.longValue(); + private Long numberAt(Object source, String... path) { + Object current = source; + for (String key : path) { + if (current instanceof ChatResponseMetadata responseMetadata) { + current = responseMetadata.get(key); + } + else if (current instanceof Map map) { + current = map.get(key); + } + else { + return null; } } - - Object directReasoning = metadata.get("reasoning_tokens"); - if (directReasoning instanceof Number num) { - return num.longValue(); + return current instanceof Number number ? number.longValue() : null; + } + + private Long firstNonNull(Long... values) { + for (Long value : values) { + if (value != null) { + return value; + } } + return null; + } - return 0L; + private long countOrZero(Long count) { + return count == null ? 0L : count; } - private Map defaultCounts() { - Map counts = new EnumMap<>(TokenType.class); - counts.put(TokenType.PROMPT, 0L); - counts.put(TokenType.COMPLETION, 0L); - return counts; + private Map usageFields(Object nativeUsage) { + if (nativeUsage == null) { + return Map.of(); + } + try { + return ModelOptionsUtils.objectToMap(nativeUsage); + } + catch (RuntimeException ignored) { + return Map.of(); + } } private Map copyMetadata(ChatResponseMetadata metadata, Object nativeUsage) { diff --git a/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/DefaultUsageExtractorTest.java b/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/DefaultUsageExtractorTest.java index 5b7679d..0bcb005 100644 --- a/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/DefaultUsageExtractorTest.java +++ b/token-pilot-spring-ai/src/test/java/io/tokenpilot/springai/internal/DefaultUsageExtractorTest.java @@ -2,6 +2,8 @@ import io.tokenpilot.core.domain.TokenType; import io.tokenpilot.core.domain.TokenUsage; +import io.tokenpilot.core.domain.TokenUsageDetails; +import io.tokenpilot.core.domain.UsageSource; import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; import org.springframework.ai.chat.client.ChatClientResponse; @@ -40,6 +42,8 @@ void extractStandardUsage() { assertThat(result.getCount(TokenType.PROMPT)).isEqualTo(100L); assertThat(result.getCount(TokenType.COMPLETION)).isEqualTo(200L); + assertThat(result.details()).isEqualTo(TokenUsageDetails.unreported()); + assertThat(result.source()).isEqualTo(UsageSource.PROVIDER_REPORTED); assertThat(result.metadata()).containsEntry("nativeUsage", Map.of("provider", "openai")); } @@ -63,6 +67,74 @@ void extractReasoningUsage() { assertThat(result.getCount(TokenType.PROMPT)).isEqualTo(100L); assertThat(result.getCount(TokenType.COMPLETION)).isEqualTo(200L); assertThat(result.getCount(TokenType.REASONING)).isEqualTo(150L); + assertThat(result.details().cacheReadInputTokens()).isNull(); + assertThat(result.details().cacheCreationInputTokens()).isNull(); + } + + @Test + @DisplayName("OpenAI의 cached input과 reasoning breakdown을 포괄 총량과 분리한다") + void extractOpenAiUsageBreakdown() { + Usage usage = mock(Usage.class); + when(usage.getPromptTokens()).thenReturn(100); + when(usage.getCompletionTokens()).thenReturn(200); + + ChatResponseMetadata metadata = ChatResponseMetadata.builder() + .usage(usage) + .keyValue("prompt_tokens_details", Map.of("cached_tokens", 40)) + .keyValue("completion_tokens_details", Map.of("reasoning_tokens", 150)) + .build(); + + TokenUsage result = extractor.extract(responseWith(metadata)); + + assertThat(result.inputTokens()).isEqualTo(100); + assertThat(result.outputTokens()).isEqualTo(200); + assertThat(result.details()).isEqualTo(new TokenUsageDetails(40L, null, 150L)); + assertThat(result.source()).isEqualTo(UsageSource.PROVIDER_REPORTED); + } + + @Test + @DisplayName("Anthropic의 일반 입력과 cache read/create를 포괄 입력 총량으로 정규화한다") + void normalizeAnthropicInputUsage() { + AnthropicNativeUsage nativeUsage = new AnthropicNativeUsage(50, 100, 25, 60); + Usage usage = mock(Usage.class); + when(usage.getPromptTokens()).thenReturn(50); + when(usage.getCompletionTokens()).thenReturn(60); + when(usage.getNativeUsage()).thenReturn(nativeUsage); + ChatResponseMetadata metadata = ChatResponseMetadata.builder() + .usage(usage) + .build(); + + TokenUsage result = extractor.extract(responseWith(metadata)); + + assertThat(result.inputTokens()).isEqualTo(175); + assertThat(result.outputTokens()).isEqualTo(60); + assertThat(result.details()).isEqualTo(new TokenUsageDetails(100L, 25L, null)); + assertThat(result.source()).isEqualTo(UsageSource.PROVIDER_DERIVED); + } + + @Test + @DisplayName("Gemini의 candidates와 thoughts를 포괄 출력 총량으로 정규화한다") + void normalizeGeminiOutputUsage() { + Map nativeUsage = Map.of( + "promptTokenCount", 100, + "cachedContentTokenCount", 40, + "candidatesTokenCount", 80, + "thoughtsTokenCount", 20 + ); + Usage usage = mock(Usage.class); + when(usage.getPromptTokens()).thenReturn(100); + when(usage.getCompletionTokens()).thenReturn(80); + when(usage.getNativeUsage()).thenReturn(nativeUsage); + ChatResponseMetadata metadata = ChatResponseMetadata.builder() + .usage(usage) + .build(); + + TokenUsage result = extractor.extract(responseWith(metadata)); + + assertThat(result.inputTokens()).isEqualTo(100); + assertThat(result.outputTokens()).isEqualTo(100); + assertThat(result.details()).isEqualTo(new TokenUsageDetails(40L, null, 20L)); + assertThat(result.source()).isEqualTo(UsageSource.PROVIDER_DERIVED); } @Test @@ -72,6 +144,7 @@ void extractNullResponse() { assertThat(result.promptTokens()).isZero(); assertThat(result.completionTokens()).isZero(); + assertThat(result.source()).isEqualTo(UsageSource.UNAVAILABLE); assertThat(result.metadata()).isEmpty(); } @@ -86,6 +159,7 @@ void extractNullMetadata() { assertThat(result.promptTokens()).isZero(); assertThat(result.completionTokens()).isZero(); + assertThat(result.source()).isEqualTo(UsageSource.UNAVAILABLE); assertThat(result.metadata()).isEmpty(); } @@ -102,6 +176,23 @@ void preserveMetadataWithoutUsage() { assertThat(result.promptTokens()).isZero(); assertThat(result.completionTokens()).isZero(); + assertThat(result.source()).isEqualTo(UsageSource.UNAVAILABLE); assertThat(result.metadata()).containsEntry("model_region", "ap-northeast-2"); } + + private ChatClientResponse responseWith(ChatResponseMetadata metadata) { + ChatResponse chatResponse = new ChatResponse( + List.of(new Generation(new org.springframework.ai.chat.messages.AssistantMessage("test"))), + metadata + ); + return new ChatClientResponse(chatResponse, Map.of()); + } + + private record AnthropicNativeUsage( + Integer inputTokens, + Integer cacheReadInputTokens, + Integer cacheCreationInputTokens, + Integer outputTokens + ) { + } }