mirror of
https://github.com/BaaRai/ARTEX.git
synced 2026-11-16 05:00:53 +08:00
Compare commits
7
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
f0be3ae51f | ||
|
|
3b755228d6 | ||
|
|
d147dce344 | ||
|
|
60728092e1 | ||
|
|
ff599c65d3 | ||
|
|
ce096cdba1 | ||
|
|
50f6d07640 |
+17
-4
@@ -39,10 +39,17 @@ func enableNoa(opts *agentcore.Options, enabled func() bool, archiveRoot, sessio
|
||||
if opts.OnWarn == nil {
|
||||
opts.OnWarn = onWarn
|
||||
}
|
||||
if err := noaadapter.Enable(opts, noaadapter.Options{
|
||||
ArchiveBaseDir: filepath.Join(archiveRoot, "noa"),
|
||||
SessionID: sessionID,
|
||||
OnWarn: onWarn,
|
||||
// noa sizes its thresholds to the window this agent is configured for, read
|
||||
// before Compaction is cleared below; without it noa assumes its own default.
|
||||
window := defaultWindowK * 1000
|
||||
if opts.Compaction != nil && opts.Compaction.ContextWindow > 0 {
|
||||
window = opts.Compaction.ContextWindow
|
||||
}
|
||||
if err := noaEnable(opts, noaadapter.Options{
|
||||
ArchiveBaseDir: filepath.Join(archiveRoot, "noa"),
|
||||
SessionID: sessionID,
|
||||
OnWarn: onWarn,
|
||||
ModelContextLimit: window,
|
||||
}); err != nil {
|
||||
if onWarn != nil {
|
||||
onWarn("failed to enable noa compression, falling back to built-in compaction: " + err.Error())
|
||||
@@ -52,4 +59,10 @@ func enableNoa(opts *agentcore.Options, enabled func() bool, archiveRoot, sessio
|
||||
// Compactor overrides Compaction, but agentcore warns on every run when both are
|
||||
// set; clear it explicitly.
|
||||
opts.Compaction = nil
|
||||
// noa can only trim what it may drop; a request whose protected part alone
|
||||
// exceeds the window is refused here instead of by the provider.
|
||||
opts.Provider = &noaBudgetProvider{Provider: opts.Provider, window: window}
|
||||
}
|
||||
|
||||
// noaEnable is noaadapter.Enable; tests replace it to observe the options.
|
||||
var noaEnable = noaadapter.Enable
|
||||
|
||||
@@ -0,0 +1,44 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"iter"
|
||||
|
||||
"github.com/Autumn-27/artex/sidequestion"
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
)
|
||||
|
||||
// noaRequestOverhead covers request framing the input estimate does not count.
|
||||
const noaRequestOverhead = 512
|
||||
|
||||
// noaBudgetProvider refuses, before any transport, a request that cannot fit the
|
||||
// window noa was configured with. noa trims what it is allowed to drop; when
|
||||
// the protected part (system prompt, recent turns, the output reservation)
|
||||
// alone is over the window, sending it only buys a provider rejection.
|
||||
type noaBudgetProvider struct {
|
||||
llm.Provider
|
||||
window int
|
||||
}
|
||||
|
||||
func (p *noaBudgetProvider) check(req llm.CompletionRequest) error {
|
||||
estimate := sidequestion.EstimateInputTokens(req) + max(req.MaxTokens, 0) + noaRequestOverhead
|
||||
if estimate > p.window {
|
||||
return fmt.Errorf("Norma NOA request estimate %d exceeds configured budget %d; protected context cannot fit, increase the configured window or start a new session", estimate, p.window)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *noaBudgetProvider) Complete(ctx context.Context, req llm.CompletionRequest) (llm.Message, string, llm.Usage, error) {
|
||||
if err := p.check(req); err != nil {
|
||||
return llm.Message{}, "", llm.Usage{}, err
|
||||
}
|
||||
return p.Provider.Complete(ctx, req)
|
||||
}
|
||||
|
||||
func (p *noaBudgetProvider) Stream(ctx context.Context, req llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
|
||||
if err := p.check(req); err != nil {
|
||||
return func(yield func(llm.StreamEvent, error) bool) { yield(llm.StreamEvent{}, err) }
|
||||
}
|
||||
return p.Provider.Stream(ctx, req)
|
||||
}
|
||||
@@ -0,0 +1,134 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"iter"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Autumn-27/norma/agentcore"
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
"github.com/Autumn-27/norma/noaadapter"
|
||||
)
|
||||
|
||||
// countingProvider answers with one text turn and counts the calls it receives.
|
||||
type countingProvider struct{ calls int }
|
||||
|
||||
func (p *countingProvider) Stream(context.Context, llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
|
||||
p.calls++
|
||||
return func(yield func(llm.StreamEvent, error) bool) {
|
||||
if yield(llm.StreamEvent{Type: llm.SETextDelta, Text: "ok"}, nil) {
|
||||
yield(llm.StreamEvent{Type: llm.SEMessageStop}, nil)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (p *countingProvider) Complete(context.Context, llm.CompletionRequest) (llm.Message, string, llm.Usage, error) {
|
||||
p.calls++
|
||||
return llm.Message{Role: llm.RoleAssistant, Content: []llm.ContentBlock{llm.TextBlock("ok")}}, "end_turn", llm.Usage{}, nil
|
||||
}
|
||||
|
||||
// withNoaEnable swaps the noaadapter.Enable seam for the duration of a test.
|
||||
func withNoaEnable(t *testing.T, fn func(*agentcore.Options, noaadapter.Options) error) {
|
||||
t.Helper()
|
||||
prev := noaEnable
|
||||
noaEnable = fn
|
||||
t.Cleanup(func() { noaEnable = prev })
|
||||
}
|
||||
|
||||
func TestEnableNoaUsesConfiguredWindow(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
compaction func() *agentcore.Options
|
||||
want int
|
||||
}{
|
||||
{"configured", func() *agentcore.Options { return &agentcore.Options{Compaction: compactionConfig(32000)} }, 32000},
|
||||
{"default", func() *agentcore.Options { return &agentcore.Options{} }, defaultWindowK * 1000},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
var got noaadapter.Options
|
||||
withNoaEnable(t, func(opts *agentcore.Options, o noaadapter.Options) error {
|
||||
got = o
|
||||
return noaadapter.Enable(opts, o)
|
||||
})
|
||||
inner := &countingProvider{}
|
||||
opts := tc.compaction()
|
||||
opts.Provider = inner
|
||||
enableNoa(opts, func() bool { return true }, t.TempDir(), "noa-window", nil)
|
||||
|
||||
if got.ModelContextLimit != tc.want {
|
||||
t.Fatalf("noaadapter.Options.ModelContextLimit=%d, want %d", got.ModelContextLimit, tc.want)
|
||||
}
|
||||
bp, ok := opts.Provider.(*noaBudgetProvider)
|
||||
if !ok || bp.window != tc.want || bp.Provider != inner {
|
||||
t.Fatalf("provider %T %+v, want *noaBudgetProvider{window %d} around the original", opts.Provider, opts.Provider, tc.want)
|
||||
}
|
||||
if opts.Compaction != nil || opts.Compactor == nil {
|
||||
t.Fatalf("Compaction=%v Compactor=%v; noa must replace the built-in compaction", opts.Compaction, opts.Compactor)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func irreducibleRequest() llm.CompletionRequest {
|
||||
return llm.CompletionRequest{System: []string{strings.Repeat("protected system prompt ", 8000)}, MaxTokens: 8000}
|
||||
}
|
||||
|
||||
func TestNoaBudgetRefusesIrreducibleRequestBeforeTransport(t *testing.T) {
|
||||
t.Run("atomic", func(t *testing.T) {
|
||||
inner := &countingProvider{}
|
||||
p := &noaBudgetProvider{Provider: inner, window: 32000}
|
||||
_, _, _, err := p.Complete(t.Context(), irreducibleRequest())
|
||||
if err == nil || !strings.Contains(err.Error(), "exceeds configured budget 32000") || inner.calls != 0 {
|
||||
t.Fatalf("err=%v provider calls=%d", err, inner.calls)
|
||||
}
|
||||
if _, _, _, err := p.Complete(t.Context(), llm.CompletionRequest{System: []string{"short"}, MaxTokens: 1000}); err != nil || inner.calls != 1 {
|
||||
t.Fatalf("a request that fits was refused or not sent: err=%v calls=%d", err, inner.calls)
|
||||
}
|
||||
})
|
||||
t.Run("stream", func(t *testing.T) {
|
||||
inner := &countingProvider{}
|
||||
p := &noaBudgetProvider{Provider: inner, window: 32000}
|
||||
var events int
|
||||
var err error
|
||||
for ev, e := range p.Stream(t.Context(), irreducibleRequest()) {
|
||||
if e != nil {
|
||||
err = e
|
||||
break
|
||||
}
|
||||
_ = ev
|
||||
events++
|
||||
}
|
||||
if err == nil || !strings.Contains(err.Error(), "exceeds configured budget 32000") || events != 0 || inner.calls != 0 {
|
||||
t.Fatalf("err=%v events=%d provider calls=%d", err, events, inner.calls)
|
||||
}
|
||||
for _, e := range p.Stream(t.Context(), llm.CompletionRequest{System: []string{"short"}, MaxTokens: 1000}) {
|
||||
if e != nil {
|
||||
t.Fatal(e)
|
||||
}
|
||||
}
|
||||
if inner.calls != 1 {
|
||||
t.Fatalf("a request that fits was not streamed: calls=%d", inner.calls)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestEnableNoaFailureKeepsBuiltInCompaction(t *testing.T) {
|
||||
withNoaEnable(t, func(*agentcore.Options, noaadapter.Options) error { return errors.New("archive dir unusable") })
|
||||
inner := &countingProvider{}
|
||||
compaction := compactionConfig(32000)
|
||||
opts := &agentcore.Options{Compaction: compaction, Provider: inner}
|
||||
var warns []string
|
||||
enableNoa(opts, func() bool { return true }, t.TempDir(), "noa-fail", func(m string) { warns = append(warns, m) })
|
||||
|
||||
if len(warns) != 1 || !strings.Contains(warns[0], "falling back to built-in compaction") || !strings.Contains(warns[0], "archive dir unusable") {
|
||||
t.Fatalf("warnings %q", warns)
|
||||
}
|
||||
if opts.Compaction != compaction || opts.Compactor != nil {
|
||||
t.Fatalf("Compaction=%v Compactor=%v; the built-in compaction must stay", opts.Compaction, opts.Compactor)
|
||||
}
|
||||
if opts.Provider != llm.Provider(inner) {
|
||||
t.Fatalf("provider %T was wrapped although noa is off", opts.Provider)
|
||||
}
|
||||
}
|
||||
+6
-5
@@ -341,11 +341,12 @@ func applyCostEur(d *DB, out []ProfileUsage) error {
|
||||
if !ok || (pr.in == 0 && pr.outp == 0 && pr.cr == 0 && pr.cw == 0) {
|
||||
continue
|
||||
}
|
||||
// InputTokens includes the cached subset (Anthropic and OpenAI usage
|
||||
// semantics), so the fresh share is input minus cache reads; charging
|
||||
// both buckets at full input price double-counts the cache. Clamped at
|
||||
// zero for gateways that report the buckets disjointly.
|
||||
fresh := out[i].InputTokens - out[i].CacheReadTokens
|
||||
// InputTokens already includes both cached subsets: Norma normalizes it to
|
||||
// input + cache read + cache write (OpenAI prompt_tokens semantics). The
|
||||
// fresh share is therefore input minus cache reads minus cache writes;
|
||||
// charging either subset at full input price as well bills it twice.
|
||||
// Clamped at zero for gateways that report the buckets disjointly.
|
||||
fresh := out[i].InputTokens - out[i].CacheReadTokens - out[i].CacheWriteTokens
|
||||
if fresh < 0 {
|
||||
fresh = 0
|
||||
}
|
||||
|
||||
@@ -0,0 +1,125 @@
|
||||
package db
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// costProfiles saves a priced profile (EUR per million tokens: input 2.00,
|
||||
// output 4.00, cache read 0.20, cache write 2.50) and an unpriced one.
|
||||
func costProfiles(t *testing.T, d *DB) (priced, unpriced string) {
|
||||
t.Helper()
|
||||
suffix := time.Now().UnixNano()
|
||||
priced, unpriced = fmt.Sprintf("t-cost-priced-%d", suffix), fmt.Sprintf("t-cost-unpriced-%d", suffix)
|
||||
for _, p := range []*LLMProfile{
|
||||
{Name: priced, Format: "anthropic", Model: "m", APIKey: "k",
|
||||
PriceInputPerM: 2, PriceOutputPerM: 4, PriceCacheReadPerM: 0.2, PriceCacheWritePerM: 2.5},
|
||||
{Name: unpriced, Format: "anthropic", Model: "m", APIKey: "k"},
|
||||
} {
|
||||
id, err := d.SaveProfile(p)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { d.Exec(`DELETE FROM llm_profiles WHERE id=$1`, id) })
|
||||
}
|
||||
return priced, unpriced
|
||||
}
|
||||
|
||||
func assertCost(t *testing.T, got *float64, want float64) {
|
||||
t.Helper()
|
||||
if got == nil {
|
||||
t.Fatalf("cost=nil, want %.2f", want)
|
||||
}
|
||||
if math.Abs(*got-want) > 1e-9 {
|
||||
t.Fatalf("cost=%.4f, want %.2f", *got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// InputTokens already includes the cache read and cache write subsets, so only
|
||||
// the remainder is billed at the input price.
|
||||
func TestCostSeparatesFreshReadAndWriteBuckets(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer d.Close()
|
||||
priced, unpriced := costProfiles(t, d)
|
||||
// Every case adds 100k output tokens: 0.1 × 4.00 = 0.40.
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
in, read, write int
|
||||
want float64
|
||||
}{
|
||||
{"fresh", 1_000_000, 0, 0, 2.40}, // 1M×2.00 + 0.40
|
||||
{"read_only", 1_000_000, 1_000_000, 0, 0.60}, // 0 fresh + 1M×0.20 + 0.40
|
||||
{"write_only", 1_000_000, 0, 1_000_000, 2.90}, // 0 fresh + 1M×2.50 + 0.40
|
||||
{"mixed", 1_000_000, 200_000, 300_000, 2.19}, // 500k×2.00 + 200k×0.20 + 300k×2.50 + 0.40
|
||||
{"inconsistent_gateway", 100_000, 200_000, 300_000, 1.19}, // fresh clamped to 0 + 0.04 + 0.75 + 0.40
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
out := []ProfileUsage{
|
||||
{ProfileName: priced, InputTokens: tc.in, CacheReadTokens: tc.read, CacheWriteTokens: tc.write, OutputTokens: 100_000},
|
||||
{ProfileName: unpriced, InputTokens: tc.in, CacheReadTokens: tc.read, CacheWriteTokens: tc.write, OutputTokens: 100_000},
|
||||
{ProfileName: "t-cost-unknown-profile", InputTokens: tc.in, OutputTokens: 100_000},
|
||||
}
|
||||
if err := applyCostEur(d, out); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
assertCost(t, out[0].CostEur, tc.want)
|
||||
if out[1].CostEur != nil || out[2].CostEur != nil {
|
||||
t.Fatalf("unpriced usage got a cost: %v %v", out[1].CostEur, out[2].CostEur)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Several ledger rows of one profile are summed before pricing; with no clamp
|
||||
// involved the profile cost equals the sum of the per-row costs.
|
||||
func TestUsageByProfileCostAggregatesRows(t *testing.T) {
|
||||
d, err := Open(testDSN(t))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer d.Close()
|
||||
priced, unpriced := costProfiles(t, d)
|
||||
t.Cleanup(func() { d.Exec(`DELETE FROM llm_usage WHERE profile_name IN ($1,$2)`, priced, unpriced) })
|
||||
rows := []LLMUsage{
|
||||
// fresh 400k×2.00 + 200k×0.20 + 40k×4.00 = 0.80 + 0.04 + 0.16 = 1.00
|
||||
{InputTokens: 600_000, CacheRead: 200_000, OutputTokens: 40_000},
|
||||
// fresh 100k×2.00 + 300k×2.50 + 60k×4.00 = 0.20 + 0.75 + 0.24 = 1.19
|
||||
{InputTokens: 400_000, CacheWrite: 300_000, OutputTokens: 60_000},
|
||||
}
|
||||
for _, profile := range []string{priced, unpriced} {
|
||||
for _, r := range rows {
|
||||
r.ProfileName, r.Model, r.Status = profile, "m", "ok"
|
||||
if err := d.InsertLLMUsage(&r); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
usage, err := d.UsageByProfile()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
found := 0
|
||||
for _, p := range usage {
|
||||
switch p.ProfileName {
|
||||
case priced:
|
||||
found++
|
||||
if p.Calls != 2 || p.InputTokens != 1_000_000 || p.CacheReadTokens != 200_000 || p.CacheWriteTokens != 300_000 || p.OutputTokens != 100_000 {
|
||||
t.Fatalf("aggregate %+v", p)
|
||||
}
|
||||
assertCost(t, p.CostEur, 2.19) // = 1.00 + 1.19, the mixed case
|
||||
case unpriced:
|
||||
found++
|
||||
if p.Calls != 2 || p.CostEur != nil {
|
||||
t.Fatalf("unpriced aggregate %+v", p)
|
||||
}
|
||||
}
|
||||
}
|
||||
if found != 2 {
|
||||
t.Fatalf("profiles missing from UsageByProfile: found %d of 2", found)
|
||||
}
|
||||
}
|
||||
@@ -3,7 +3,7 @@ module github.com/Autumn-27/artex
|
||||
go 1.26.10
|
||||
|
||||
require (
|
||||
github.com/Autumn-27/norma v0.4.3
|
||||
github.com/Autumn-27/norma v0.4.9
|
||||
github.com/golang-jwt/jwt/v5 v5.3.1
|
||||
github.com/google/uuid v1.6.0
|
||||
github.com/jackc/pgx/v5 v5.10.0
|
||||
@@ -61,4 +61,4 @@ require (
|
||||
modernc.org/memory v1.11.0 // indirect
|
||||
)
|
||||
|
||||
replace github.com/Autumn-27/norma => github.com/BaaRai/Norma v0.4.8
|
||||
replace github.com/Autumn-27/norma => github.com/BaaRai/Norma v0.4.9
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
github.com/BaaRai/Norma v0.4.8 h1:k3at2aYSJ6aIDjEo7tC4VDSLBHXw2HobYjpKChVe6hI=
|
||||
github.com/BaaRai/Norma v0.4.8/go.mod h1:BudWrf7iHl96FgHqFr/HwN3mKYmGjJDt+xL2ZHXfFZw=
|
||||
github.com/BaaRai/Norma v0.4.9 h1:veWwo9e+OgmszutYEuglRfnvMDKThX9ilaFaXmOEGxY=
|
||||
github.com/BaaRai/Norma v0.4.9/go.mod h1:BudWrf7iHl96FgHqFr/HwN3mKYmGjJDt+xL2ZHXfFZw=
|
||||
github.com/andybalholm/brotli v1.2.1 h1:R+f5xP285VArJDRgowrfb9DqL18yVK0gKAW/F+eTWro=
|
||||
github.com/andybalholm/brotli v1.2.1/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY=
|
||||
github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2 h1:DklsrG3dyBCFEj5IhUbnKptjxatkF07cF2ak3yi77so=
|
||||
|
||||
+6
-6
@@ -195,11 +195,11 @@ func (r *Recorder) Stream(ctx context.Context, req llm.CompletionRequest) iter.S
|
||||
return
|
||||
}
|
||||
// Always track usage (cheap); accumulate text/thinking only for bodies.
|
||||
// Anthropic (and the other providers) split token usage across events:
|
||||
// message_start carries ONLY the input side (input + cache), message_delta
|
||||
// ONLY the output. They must be FOLDED with Add — overwriting on delta
|
||||
// would zero out the input already counted at start (mirrors the SDK's own
|
||||
// llm.Accumulator; see norma/llm/accumulate.go).
|
||||
// message_start carries the input side (input + cache) and is added.
|
||||
// Since Norma v0.4.7 message_delta also echoes that input side next to the
|
||||
// output, so it is folded with llm.Usage.MergeEcho (max on echoed fields,
|
||||
// sum on output), as llm.Accumulator does; Add would count the input and
|
||||
// cache twice. An output-only delta folds the same either way.
|
||||
switch ev.Type {
|
||||
case llm.SETextDelta:
|
||||
if recordBodies {
|
||||
@@ -215,7 +215,7 @@ func (r *Recorder) Stream(ctx context.Context, req llm.CompletionRequest) iter.S
|
||||
if ev.StopReason != "" {
|
||||
stopReason = ev.StopReason
|
||||
}
|
||||
usage.Add(ev.Usage)
|
||||
usage.MergeEcho(ev.Usage)
|
||||
}
|
||||
if !yield(ev, err) {
|
||||
return
|
||||
|
||||
@@ -0,0 +1,118 @@
|
||||
package llmrec
|
||||
|
||||
import (
|
||||
"context"
|
||||
"iter"
|
||||
"testing"
|
||||
|
||||
"github.com/Autumn-27/artex/db"
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// echoStreamProvider streams one message_start and one message_delta with the
|
||||
// given (already normalized) usage, as Norma's providers emit them.
|
||||
type echoStreamProvider struct {
|
||||
completeProvider
|
||||
start, delta llm.Usage
|
||||
}
|
||||
|
||||
func (p echoStreamProvider) Stream(context.Context, llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
|
||||
return func(yield func(llm.StreamEvent, error) bool) {
|
||||
for _, ev := range []llm.StreamEvent{
|
||||
{Type: llm.SEMessageStart, Usage: p.start},
|
||||
{Type: llm.SETextDelta, Text: "answer"},
|
||||
{Type: llm.SEMessageDelta, StopReason: "end_turn", Usage: p.delta},
|
||||
{Type: llm.SEMessageStop},
|
||||
} {
|
||||
if !yield(ev, nil) {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Norma normalizes Anthropic usage so InputTokens includes the cache read and
|
||||
// write subsets, and since v0.4.7 message_delta echoes that input side.
|
||||
var usageEchoCases = []struct {
|
||||
name string
|
||||
start, delta llm.Usage
|
||||
want llm.Usage
|
||||
}{
|
||||
{
|
||||
name: "full_echo",
|
||||
start: llm.Usage{InputTokens: 23, CacheReadTokens: 5, CacheWriteTokens: 3},
|
||||
delta: llm.Usage{InputTokens: 23, CacheReadTokens: 5, CacheWriteTokens: 3, OutputTokens: 5},
|
||||
want: llm.Usage{InputTokens: 23, OutputTokens: 5, CacheReadTokens: 5, CacheWriteTokens: 3},
|
||||
},
|
||||
{
|
||||
// Raw start: input 500, cache read 200, cache creation 40 → normalized 740.
|
||||
// Raw delta: input 500, output 77, no cache fields → normalized 500.
|
||||
name: "partial_echo",
|
||||
start: llm.Usage{InputTokens: 740, CacheReadTokens: 200, CacheWriteTokens: 40},
|
||||
delta: llm.Usage{InputTokens: 500, OutputTokens: 77},
|
||||
want: llm.Usage{InputTokens: 740, OutputTokens: 77, CacheReadTokens: 200, CacheWriteTokens: 40},
|
||||
},
|
||||
{
|
||||
name: "no_echo",
|
||||
start: llm.Usage{InputTokens: 23, CacheReadTokens: 5, CacheWriteTokens: 3},
|
||||
delta: llm.Usage{OutputTokens: 5},
|
||||
want: llm.Usage{InputTokens: 23, OutputTokens: 5, CacheReadTokens: 5, CacheWriteTokens: 3},
|
||||
},
|
||||
}
|
||||
|
||||
func drainRecorderStream(t *testing.T, ctx context.Context, r *Recorder) {
|
||||
t.Helper()
|
||||
for _, err := range r.Stream(ctx, llm.CompletionRequest{}) {
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The metered usage of a streamed call counts an echoed input side once.
|
||||
func TestStreamUsageEchoCountedOnce(t *testing.T) {
|
||||
for _, tc := range usageEchoCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
var got []llm.Usage
|
||||
ctx := WithUsageSink(t.Context(), func(u llm.Usage) { got = append(got, u) })
|
||||
drainRecorderStream(t, ctx, Wrap(echoStreamProvider{start: tc.start, delta: tc.delta}, nil, "fixture", "profile", "", "", func() bool { return false }))
|
||||
if len(got) != 1 || got[0] != tc.want {
|
||||
t.Fatalf("metered %+v, want one call with %+v", got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// The llm_usage ledger gets one row per streamed call, with the echo counted once.
|
||||
func TestStreamUsageEchoLedgerRow(t *testing.T) {
|
||||
dsn, _, err := db.DSN()
|
||||
if err != nil {
|
||||
t.Skipf("no database config (%v)", err)
|
||||
}
|
||||
pg, err := db.Open(dsn)
|
||||
if err != nil {
|
||||
t.Skipf("postgres unavailable (%v)", err)
|
||||
}
|
||||
defer pg.Close()
|
||||
if err := pg.EnsureLLMUsageTable(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, tc := range usageEchoCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
profile := "usage-echo-" + uuid.NewString()
|
||||
t.Cleanup(func() { _, _ = pg.Exec(`DELETE FROM llm_usage WHERE profile_name=$1`, profile) })
|
||||
drainRecorderStream(t, t.Context(), Wrap(echoStreamProvider{start: tc.start, delta: tc.delta}, pg, "fixture", profile, "", "", func() bool { return false }))
|
||||
var rows int
|
||||
var got llm.Usage
|
||||
err := pg.QueryRow(`SELECT count(*), COALESCE(sum(input_tokens),0), COALESCE(sum(output_tokens),0), COALESCE(sum(cache_read),0), COALESCE(sum(cache_write),0)
|
||||
FROM llm_usage WHERE profile_name=$1`, profile).Scan(&rows, &got.InputTokens, &got.OutputTokens, &got.CacheReadTokens, &got.CacheWriteTokens)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if rows != 1 || got != tc.want {
|
||||
t.Fatalf("ledger rows=%d usage=%+v, want 1 row with %+v", rows, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -32,8 +32,11 @@ func (s SideQuestionService) Answer(ctx context.Context, req llm.CompletionReque
|
||||
out.Text += ev.Text
|
||||
case llm.SEToolUseStart:
|
||||
out.ToolUse = true
|
||||
case llm.SEMessageStart, llm.SEMessageDelta:
|
||||
case llm.SEMessageStart:
|
||||
out.Usage.Add(ev.Usage)
|
||||
case llm.SEMessageDelta:
|
||||
// message_delta may echo message_start's input side: MergeEcho, not Add.
|
||||
out.Usage.MergeEcho(ev.Usage)
|
||||
case llm.SEMessageStop:
|
||||
complete = true
|
||||
}
|
||||
|
||||
@@ -0,0 +1,52 @@
|
||||
package sidequestion
|
||||
|
||||
import (
|
||||
"context"
|
||||
"iter"
|
||||
"testing"
|
||||
|
||||
"github.com/Autumn-27/norma/llm"
|
||||
)
|
||||
|
||||
// A streamed side answer counts message_delta's echo of the input side once.
|
||||
func TestAnswerStreamUsageEchoCountedOnce(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
start, delta llm.Usage
|
||||
want llm.Usage
|
||||
}{
|
||||
{"full_echo",
|
||||
llm.Usage{InputTokens: 23, CacheReadTokens: 5, CacheWriteTokens: 3},
|
||||
llm.Usage{InputTokens: 23, CacheReadTokens: 5, CacheWriteTokens: 3, OutputTokens: 5},
|
||||
llm.Usage{InputTokens: 23, OutputTokens: 5, CacheReadTokens: 5, CacheWriteTokens: 3}},
|
||||
{"partial_echo",
|
||||
llm.Usage{InputTokens: 740, CacheReadTokens: 200, CacheWriteTokens: 40},
|
||||
llm.Usage{InputTokens: 500, OutputTokens: 77},
|
||||
llm.Usage{InputTokens: 740, OutputTokens: 77, CacheReadTokens: 200, CacheWriteTokens: 40}},
|
||||
{"no_echo",
|
||||
llm.Usage{InputTokens: 23, CacheReadTokens: 5, CacheWriteTokens: 3},
|
||||
llm.Usage{OutputTokens: 5},
|
||||
llm.Usage{InputTokens: 23, OutputTokens: 5, CacheReadTokens: 5, CacheWriteTokens: 3}},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
p := fakeProvider{stream: func(context.Context, llm.CompletionRequest) iter.Seq2[llm.StreamEvent, error] {
|
||||
return func(yield func(llm.StreamEvent, error) bool) {
|
||||
for _, ev := range []llm.StreamEvent{
|
||||
{Type: llm.SEMessageStart, Usage: tc.start},
|
||||
{Type: llm.SETextDelta, Text: "answer"},
|
||||
{Type: llm.SEMessageDelta, StopReason: "end_turn", Usage: tc.delta},
|
||||
{Type: llm.SEMessageStop},
|
||||
} {
|
||||
if !yield(ev, nil) {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}}
|
||||
out, err := (SideQuestionService{Provider: p}).Answer(t.Context(), llm.CompletionRequest{}, true, nil)
|
||||
if err != nil || out.Text != "answer" || out.Usage != tc.want {
|
||||
t.Fatalf("answer=%+v err=%v, want usage %+v", out, err, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user