Preserve reported token usage presence

This commit is contained in:
2026-09-23 16:39:43 +00:00
parent ea94d66fc2
commit d31a563ca1
8 changed files with 166 additions and 16 deletions

View File

@@ -125,8 +125,22 @@ func toDomainGenerateResponse(resp *GenerateResponse) *domain.GenerateResponse {
return nil return nil
} }
return &domain.GenerateResponse{ return &domain.GenerateResponse{
Content: resp.Content, Content: resp.Content,
Usage: toDomainTokenUsage(resp.Usage), Usage: toDomainTokenUsage(resp.Usage),
UsagePresence: toDomainTokenUsagePresence(resp.UsagePresence),
}
}
func toDomainTokenUsagePresence(presence *TokenUsagePresence) *domain.TokenUsagePresence {
if presence == nil {
return nil
}
return &domain.TokenUsagePresence{
PromptTokens: presence.PromptTokens,
CompletionTokens: presence.CompletionTokens,
TotalTokens: presence.TotalTokens,
CachedTokens: presence.CachedTokens,
CacheWriteTokens: presence.CacheWriteTokens,
} }
} }

View File

@@ -234,6 +234,8 @@ Run `go test ./...`. Commit when mappings and existing tests pass.
## Stage 2 — Preserve Usage Presence Through The Adapter ## Stage 2 — Preserve Usage Presence Through The Adapter
**Completed:** 2026-09-23.
Read wire response decoding, domain/public response types, facade conversion, Read wire response decoding, domain/public response types, facade conversion,
and JSON contract tests. Add the presence types and optional public response and JSON contract tests. Add the presence types and optional public response
field above. Convert wire usage counts to pointer-aware decoding, retaining field above. Convert wire usage counts to pointer-aware decoding, retaining

View File

@@ -306,8 +306,18 @@ type StructuredOutputJSONSpec struct {
// GenerateResponse is the response received from the LLM client. // GenerateResponse is the response received from the LLM client.
type GenerateResponse struct { type GenerateResponse struct {
Content string Content string
Usage TokenUsage Usage TokenUsage
UsagePresence *TokenUsagePresence
}
// TokenUsagePresence records which token fields a provider reported.
type TokenUsagePresence struct {
PromptTokens bool
CompletionTokens bool
TotalTokens bool
CachedTokens bool
CacheWriteTokens bool
} }
// TokenUsage tracks token consumption. // TokenUsage tracks token consumption.

View File

@@ -187,15 +187,29 @@ func (c *OpenAICompatibleClient) Generate(ctx context.Context, req domain.Genera
return &domain.GenerateResponse{ return &domain.GenerateResponse{
Content: *content, Content: *content,
Usage: domain.TokenUsage{ Usage: domain.TokenUsage{
PromptTokens: wireResp.Usage.PromptTokens, PromptTokens: tokenCount(wireResp.Usage.PromptTokens),
CompletionTokens: wireResp.Usage.CompletionTokens, CompletionTokens: tokenCount(wireResp.Usage.CompletionTokens),
TotalTokens: wireResp.Usage.TotalTokens, TotalTokens: tokenCount(wireResp.Usage.TotalTokens),
CachedTokens: wireResp.Usage.PromptTokensDetails.CachedTokens, CachedTokens: tokenCount(wireResp.Usage.PromptTokensDetails.CachedTokens),
CacheWriteTokens: wireResp.Usage.CacheWriteTokens, CacheWriteTokens: tokenCount(wireResp.Usage.CacheWriteTokens),
},
UsagePresence: &domain.TokenUsagePresence{
PromptTokens: wireResp.Usage.PromptTokens != nil,
CompletionTokens: wireResp.Usage.CompletionTokens != nil,
TotalTokens: wireResp.Usage.TotalTokens != nil,
CachedTokens: wireResp.Usage.PromptTokensDetails.CachedTokens != nil,
CacheWriteTokens: wireResp.Usage.CacheWriteTokens != nil,
}, },
}, nil }, nil
} }
func tokenCount(value *int) int {
if value == nil {
return 0
}
return *value
}
func decodeOpenAIChatResponse(body io.Reader) (openAIChatResponse, error) { func decodeOpenAIChatResponse(body io.Reader) (openAIChatResponse, error) {
limited := &io.LimitedReader{ limited := &io.LimitedReader{
R: body, R: body,
@@ -389,13 +403,13 @@ type openAIChatResponse struct {
Message openAIChatResponseMessage `json:"message"` Message openAIChatResponseMessage `json:"message"`
} `json:"choices"` } `json:"choices"`
Usage struct { Usage struct {
PromptTokens int `json:"prompt_tokens"` PromptTokens *int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"` CompletionTokens *int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"` TotalTokens *int `json:"total_tokens"`
PromptTokensDetails struct { PromptTokensDetails struct {
CachedTokens int `json:"cached_tokens"` CachedTokens *int `json:"cached_tokens"`
} `json:"prompt_tokens_details"` } `json:"prompt_tokens_details"`
CacheWriteTokens int `json:"cache_write_tokens"` CacheWriteTokens *int `json:"cache_write_tokens"`
} `json:"usage"` } `json:"usage"`
} }

View File

@@ -881,6 +881,44 @@ func checkCacheUsageMapping(t *testing.T) {
} }
} }
func TestOpenAICompatibleUsagePresence(t *testing.T) {
for _, tc := range []struct {
name string
usage string
want domain.TokenUsage
presence domain.TokenUsagePresence
}{
{name: "absent"},
{name: "null", usage: `"usage":null,`},
{name: "partial nested", usage: `"usage":{"prompt_tokens":7,"prompt_tokens_details":{"cached_tokens":0}},`,
want: domain.TokenUsage{PromptTokens: 7}, presence: domain.TokenUsagePresence{PromptTokens: true, CachedTokens: true}},
{name: "explicit null fields", usage: `"usage":{"prompt_tokens":null,"completion_tokens":0,"prompt_tokens_details":{"cached_tokens":null}},`,
presence: domain.TokenUsagePresence{CompletionTokens: true}},
{name: "all reported", usage: `"usage":{"prompt_tokens":1,"completion_tokens":2,"total_tokens":3,"prompt_tokens_details":{"cached_tokens":4},"cache_write_tokens":5},`,
want: domain.TokenUsage{PromptTokens: 1, CompletionTokens: 2, TotalTokens: 3, CachedTokens: 4, CacheWriteTokens: 5},
presence: domain.TokenUsagePresence{PromptTokens: true, CompletionTokens: true, TotalTokens: true, CachedTokens: true, CacheWriteTokens: true}},
} {
t.Run(tc.name, func(t *testing.T) {
provider := newRecordingProvider(t)
provider.respond(http.StatusOK, `{`+tc.usage+`"choices":[{"message":{"content":"ok"}}]}`)
client := newProviderClient(t, provider, OpenAICompatibleConfig{Model: "model"})
resp, err := client.Generate(context.Background(), ordinaryGenerateRequest())
if err != nil {
t.Fatal(err)
}
if resp.Usage != tc.want || resp.UsagePresence == nil || *resp.UsagePresence != tc.presence {
t.Fatalf("usage=%+v, presence=%+v", resp.Usage, resp.UsagePresence)
}
})
}
provider := newRecordingProvider(t)
provider.respond(http.StatusOK, `{"usage":{"prompt_tokens":"seven"},"choices":[{"message":{"content":"ok"}}]}`)
client := newProviderClient(t, provider, OpenAICompatibleConfig{Model: "model"})
if resp, err := client.Generate(context.Background(), ordinaryGenerateRequest()); resp != nil || !errors.Is(err, ErrMalformedResponse) {
t.Fatalf("invalid numeric usage response=%+v, error=%v", resp, err)
}
}
func checkOptionalResponseFieldOmission(t *testing.T) { func checkOptionalResponseFieldOmission(t *testing.T) {
provider := newRecordingProvider(t) provider := newRecordingProvider(t)
client := newProviderClient(t, provider, OpenAICompatibleConfig{}) client := newProviderClient(t, provider, OpenAICompatibleConfig{})

View File

@@ -12,6 +12,36 @@ import (
"gitea.maximumdirect.net/eric/promptkit" "gitea.maximumdirect.net/eric/promptkit"
) )
func TestGenerateResponseUsagePresenceJSON(t *testing.T) {
value := promptkit.GenerateResponse{
Content: "ok",
Usage: promptkit.TokenUsage{TotalTokens: 5},
UsagePresence: &promptkit.TokenUsagePresence{
PromptTokens: true, CompletionTokens: true, TotalTokens: true,
CachedTokens: true, CacheWriteTokens: true,
},
}
payload, err := json.Marshal(value)
if err != nil {
t.Fatal(err)
}
object := decodeJSONObject(t, payload)
if _, ok := object["usage_presence"]; !ok {
t.Fatalf("missing usage_presence: %s", payload)
}
var decoded promptkit.GenerateResponse
if err := json.Unmarshal(payload, &decoded); err != nil || !reflect.DeepEqual(decoded, value) {
t.Fatalf("round trip = %+v, %v", decoded, err)
}
payload, err = json.Marshal(promptkit.GenerateResponse{})
if err != nil {
t.Fatal(err)
}
if _, ok := decodeJSONObject(t, payload)["usage_presence"]; ok {
t.Fatalf("nil usage_presence was serialized: %s", payload)
}
}
func TestPreparedRunJSONContractRoundTripsAllFields(t *testing.T) { func TestPreparedRunJSONContractRoundTripsAllFields(t *testing.T) {
value := fullyPopulatedPreparedRun() value := fullyPopulatedPreparedRun()
payload, err := json.Marshal(value) payload, err := json.Marshal(value)

View File

@@ -47,6 +47,30 @@ func TestPublicLLMClientAdapterGivesClientOwnedNestedValues(t *testing.T) {
} }
} }
func TestPublicLLMClientAdapterCopiesUsagePresence(t *testing.T) {
presence := &TokenUsagePresence{PromptTokens: true}
client := &usagePresenceClient{response: &GenerateResponse{
Content: "ok",
Usage: TokenUsage{PromptTokens: 0, TotalTokens: 5},
UsagePresence: presence,
}}
adapter := publicLLMClientAdapter{client: client}
got, err := adapter.Generate(context.Background(), domain.GenerateRequest{})
if err != nil {
t.Fatal(err)
}
presence.PromptTokens = false
if got.UsagePresence == nil || !got.UsagePresence.PromptTokens || got.UsagePresence.TotalTokens || got.Usage.TotalTokens != 5 {
t.Fatalf("copied response = %+v", got)
}
}
type usagePresenceClient struct{ response *GenerateResponse }
func (c *usagePresenceClient) Generate(context.Context, GenerateRequest) (*GenerateResponse, error) {
return c.response, nil
}
type retainingMutatingLLMClient struct { type retainingMutatingLLMClient struct {
retained GenerateRequest retained GenerateRequest
} }

View File

@@ -709,8 +709,13 @@ type StructuredOutputJSONSpec struct {
// //
// An arbitrary returned error makes Run or RunPrepared return ErrLLMGenerate // An arbitrary returned error makes Run or RunPrepared return ErrLLMGenerate
// while preserving the client error through errors.Is rather than translating // while preserving the client error through errors.Is rather than translating
// it. A nil response with a nil error also produces ErrLLMGenerate. Promptkit // it. A nil response with a nil error also produces ErrLLMGenerate and
// copies the non-nil response before returning from either method. // ErrMalformedResponse. Promptkit copies the non-nil response before returning
// from either method. When UsagePresence is nil, nonzero usage counts are
// treated as known and zero counts as unknown for observation only; the legacy
// interface cannot distinguish omitted zero from measured zero. An explicit
// UsagePresence can report known zero counts. Result usage sums the supplied
// Usage values independently of presence.
type LLMClient interface { type LLMClient interface {
Generate(context.Context, GenerateRequest) (*GenerateResponse, error) Generate(context.Context, GenerateRequest) (*GenerateResponse, error)
} }
@@ -741,6 +746,19 @@ type GenerateResponse struct {
Content string `json:"content"` Content string `json:"content"`
// Usage is the client's token accounting. // Usage is the client's token accounting.
Usage TokenUsage `json:"usage"` Usage TokenUsage `json:"usage"`
// UsagePresence states which usage fields were reported, including zeros.
// Nil uses the legacy observation fallback described by LLMClient.
UsagePresence *TokenUsagePresence `json:"usage_presence,omitempty"`
}
// TokenUsagePresence identifies provider-reported token counts for observation.
// It does not change TokenUsage values or successful-result accounting.
type TokenUsagePresence struct {
PromptTokens bool `json:"prompt_tokens"`
CompletionTokens bool `json:"completion_tokens"`
TotalTokens bool `json:"total_tokens"`
CachedTokens bool `json:"cached_tokens"`
CacheWriteTokens bool `json:"cache_write_tokens"`
} }
// File returns a file-backed artifact reference whose URI is path. // File returns a file-backed artifact reference whose URI is path.