mirror of
https://github.com/Tencent/WeKnora.git
synced 2026-08-29 02:04:30 +08:00
feat(agent): report and attribute per-turn LLM token usage
Accumulate token usage across every LLM round of an agent turn and expose the aggregate everywhere a turn is observable: - AgentState.TurnUsage accumulates each round's reported usage, including the final-answer synthesis call; TokenUsage.Accumulate keeps the cache counters as prompt subsets and recomputes the combined cache status (unsupported survives only while every call was itself unsupported). - The agent completion event carries the aggregate, and the complete stream event exposes it both as data.usage and as a typed StreamEvent.Usage that buildStreamResponse promotes into the previously-unused StreamResponse.Usage wire field. - The assistant message persists it as a nullable usage column on both migration targets (postgres 000085 jsonb, sqlite 000012 text), for web and IM channels alike, so history reads still carry the turn's cost after the live stream is gone. - "[LLM Usage]" log lines gain ", session_id=..., principal=type:id" when the context carries them; calls outside sessions stay byte-identical. Usage from failed retry attempts is not accumulated: the per-call usage log remains the ground truth, while the turn aggregate serves as the operational signal.
This commit is contained in:
@@ -559,6 +559,7 @@ func (e *AgentEngine) runReActIteration(
|
||||
response = resp
|
||||
if response.Usage.TotalTokens > 0 {
|
||||
e.lastUsage = response.Usage
|
||||
state.TurnUsage.Accumulate(response.Usage)
|
||||
logger.Debugf(ctx, "[Agent][Round-%d] Usage: prompt=%d, completion=%d, total=%d",
|
||||
round, response.Usage.PromptTokens,
|
||||
response.Usage.CompletionTokens, response.Usage.TotalTokens)
|
||||
|
||||
@@ -144,6 +144,12 @@ Now generate the final answer:`, query, imageRequirement)
|
||||
})
|
||||
}
|
||||
|
||||
// The synthesis call is often the largest of the turn — fold its usage
|
||||
// into the turn aggregate like every ReAct round.
|
||||
if llmResult.Usage != nil {
|
||||
state.TurnUsage.Accumulate(*llmResult.Usage)
|
||||
}
|
||||
|
||||
// Safety net: strip any residual <think> blocks that may have leaked through
|
||||
fullAnswer := agenttools.StripThinkBlocks(llmResult.Content)
|
||||
logger.Infof(ctx, "[Agent][FinalAnswer] Final answer generated: %d characters", len(fullAnswer))
|
||||
@@ -195,6 +201,7 @@ func (e *AgentEngine) emitCompletionEvent(
|
||||
FinalAnswer: state.FinalAnswer,
|
||||
KnowledgeRefs: knowledgeRefsInterface,
|
||||
AgentSteps: state.RoundSteps, // Include detailed execution steps for message storage
|
||||
Usage: turnUsage(state),
|
||||
TotalSteps: len(state.RoundSteps),
|
||||
TotalDurationMs: time.Since(startTime).Milliseconds(),
|
||||
MessageID: messageID, // Include message ID for proper message update
|
||||
@@ -203,3 +210,14 @@ func (e *AgentEngine) emitCompletionEvent(
|
||||
|
||||
logger.Infof(ctx, "Agent execution completed in %d rounds", state.CurrentRound)
|
||||
}
|
||||
|
||||
// turnUsage returns the turn's aggregated LLM usage, or nil when no round
|
||||
// reported usage so the field stays absent from the completion event and the
|
||||
// persisted message alike.
|
||||
func turnUsage(state *types.AgentState) *types.TokenUsage {
|
||||
if state == nil || state.TurnUsage.TotalTokens == 0 {
|
||||
return nil
|
||||
}
|
||||
usage := state.TurnUsage
|
||||
return &usage
|
||||
}
|
||||
|
||||
@@ -0,0 +1,33 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
func TestTurnUsageNilWhenNothingReported(t *testing.T) {
|
||||
if turnUsage(nil) != nil {
|
||||
t.Fatal("nil state must yield no usage")
|
||||
}
|
||||
if turnUsage(&types.AgentState{}) != nil {
|
||||
t.Fatal("a turn whose rounds reported no usage must omit the field entirely")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTurnUsageCopiesTheAggregate(t *testing.T) {
|
||||
state := &types.AgentState{}
|
||||
state.TurnUsage.Accumulate(types.TokenUsage{PromptTokens: 100, CompletionTokens: 20, TotalTokens: 120})
|
||||
|
||||
usage := turnUsage(state)
|
||||
if usage == nil || usage.TotalTokens != 120 {
|
||||
t.Fatalf("aggregate not propagated: %+v", usage)
|
||||
}
|
||||
|
||||
// The returned pointer must be a copy: later state mutation must not
|
||||
// reach an event that has already been emitted.
|
||||
state.TurnUsage.Accumulate(types.TokenUsage{PromptTokens: 1, TotalTokens: 1})
|
||||
if usage.TotalTokens != 120 {
|
||||
t.Fatalf("emitted usage must be detached from state: %+v", usage)
|
||||
}
|
||||
}
|
||||
@@ -27,13 +27,13 @@ var versionedSQLiteColumns = map[string][]string{
|
||||
"tenants": {"api_principal_config"}, // 000064
|
||||
"users": {"is_system_admin"}, // 000053
|
||||
"knowledges": {"pending_subtasks_count"}, // 000056
|
||||
"messages": {"attachments"}, // 000034
|
||||
"messages": {"attachments", "usage"}, // 000034, 000085
|
||||
"tenant_invitations": {"token", "accepted_count"}, // 000054
|
||||
"embed_channels": {"allow_memory"}, // 000060
|
||||
"mcp_oauth_tokens": {"principal_type", "principal_id"}, // 000064
|
||||
}
|
||||
|
||||
const expectedSQLiteMigrationVersion = 11
|
||||
const expectedSQLiteMigrationVersion = 12
|
||||
|
||||
func TestSQLiteMigrationsCreateVersionedSchema(t *testing.T) {
|
||||
repoRoot := sqliteRepoRoot(t)
|
||||
|
||||
@@ -139,6 +139,7 @@ type AgentCompleteData struct {
|
||||
FinalAnswer string `json:"final_answer"`
|
||||
KnowledgeRefs []interface{} `json:"knowledge_refs,omitempty"` // []*types.SearchResult
|
||||
AgentSteps interface{} `json:"agent_steps,omitempty"` // []types.AgentStep - detailed execution steps
|
||||
Usage interface{} `json:"usage,omitempty"` // *types.TokenUsage - LLM token usage aggregated over the turn
|
||||
TotalDurationMs int64 `json:"total_duration_ms"`
|
||||
MessageID string `json:"message_id,omitempty"` // Assistant message ID
|
||||
RequestID string `json:"request_id,omitempty"`
|
||||
|
||||
@@ -646,6 +646,12 @@ func (h *AgentStreamHandler) handleComplete(ctx context.Context, evt event.Event
|
||||
}
|
||||
}
|
||||
|
||||
// Persist the turn's aggregated LLM usage with the message so history
|
||||
// reads still carry it after the live stream is gone.
|
||||
if usage, ok := data.Usage.(*types.TokenUsage); ok && usage != nil {
|
||||
h.assistantMessage.Usage = usage
|
||||
}
|
||||
|
||||
// Drain skill-generated files from the sandbox into persistent
|
||||
// storage. Best-effort: any failure is logged and the turn is
|
||||
// persisted without artifacts. Collect is a no-op when either the
|
||||
@@ -728,6 +734,13 @@ func (h *AgentStreamHandler) handleComplete(ctx context.Context, evt event.Event
|
||||
if len(h.assistantMessage.Artifacts) > 0 {
|
||||
completeData["artifacts"] = publicArtifactViews(h.assistantMessage.Artifacts)
|
||||
}
|
||||
// Carry the turn's aggregated LLM usage both inside data (map consumers)
|
||||
// and on the typed event field, which buildStreamResponse promotes to the
|
||||
// response's top-level usage.
|
||||
turnUsage, _ := data.Usage.(*types.TokenUsage)
|
||||
if turnUsage != nil {
|
||||
completeData["usage"] = turnUsage
|
||||
}
|
||||
if err := h.streamManager.AppendEvent(h.ctx, h.sessionID, h.assistantMessageID, interfaces.StreamEvent{
|
||||
ID: evt.ID,
|
||||
Type: types.ResponseTypeComplete,
|
||||
@@ -735,6 +748,7 @@ func (h *AgentStreamHandler) handleComplete(ctx context.Context, evt event.Event
|
||||
Done: true,
|
||||
Timestamp: time.Now(),
|
||||
Data: completeData,
|
||||
Usage: turnUsage,
|
||||
}); err != nil {
|
||||
logger.GetLogger(h.ctx).Errorf("Append complete event to stream failed: %v", err)
|
||||
}
|
||||
|
||||
@@ -194,6 +194,7 @@ func buildStreamResponse(evt interfaces.StreamEvent, requestID string) *types.St
|
||||
Content: evt.Content,
|
||||
Done: evt.Done,
|
||||
Data: evt.Data,
|
||||
Usage: evt.Usage,
|
||||
}
|
||||
|
||||
// Extract session_id and assistant_message_id for agent_query events
|
||||
|
||||
@@ -0,0 +1,50 @@
|
||||
package session
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestBuildStreamResponsePromotesUsageOnCompleteEvents(t *testing.T) {
|
||||
usage := &types.TokenUsage{PromptTokens: 100, CompletionTokens: 20, TotalTokens: 120}
|
||||
|
||||
response := buildStreamResponse(interfaces.StreamEvent{
|
||||
Type: types.ResponseTypeComplete,
|
||||
Done: true,
|
||||
Data: map[string]interface{}{"total_steps": 3, "usage": usage},
|
||||
Usage: usage,
|
||||
}, "req-1")
|
||||
|
||||
require.NotNil(t, response.Usage)
|
||||
assert.Equal(t, 120, response.Usage.TotalTokens)
|
||||
assert.Equal(t, usage, response.Data["usage"])
|
||||
}
|
||||
|
||||
func TestBuildStreamResponseLeavesUsageNilWhenAbsent(t *testing.T) {
|
||||
response := buildStreamResponse(interfaces.StreamEvent{
|
||||
Type: types.ResponseTypeAnswer,
|
||||
Data: map[string]interface{}{"event_id": "e-1"},
|
||||
}, "req-1")
|
||||
|
||||
assert.Nil(t, response.Usage)
|
||||
}
|
||||
|
||||
func TestStreamEventUsageSurvivesJSONRoundTrip(t *testing.T) {
|
||||
// The stream manager persists events as JSON (Redis) before the SSE loop
|
||||
// reads them back — the typed usage must survive that round trip.
|
||||
usage := &types.TokenUsage{PromptTokens: 100, CompletionTokens: 20, TotalTokens: 120}
|
||||
usage.SetPromptCacheUsage(80, 0, 20, true)
|
||||
|
||||
raw, err := json.Marshal(interfaces.StreamEvent{Type: types.ResponseTypeComplete, Usage: usage})
|
||||
require.NoError(t, err)
|
||||
|
||||
var restored interfaces.StreamEvent
|
||||
require.NoError(t, json.Unmarshal(raw, &restored))
|
||||
require.NotNil(t, restored.Usage)
|
||||
assert.Equal(t, *usage, *restored.Usage)
|
||||
}
|
||||
@@ -701,6 +701,9 @@ func applyIMCompleteDataToMessage(msg *types.Message, data event.AgentCompleteDa
|
||||
}
|
||||
msg.IsCompleted = true
|
||||
msg.AgentDurationMs = data.TotalDurationMs
|
||||
if usage, ok := data.Usage.(*types.TokenUsage); ok && usage != nil {
|
||||
msg.Usage = usage
|
||||
}
|
||||
if len(data.KnowledgeRefs) > 0 {
|
||||
refs := make([]*types.SearchResult, 0, len(data.KnowledgeRefs))
|
||||
collectIMKnowledgeReferences(&refs, data.KnowledgeRefs)
|
||||
|
||||
@@ -2,6 +2,7 @@ package chat
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/logger"
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
@@ -16,8 +17,30 @@ func logUsage(ctx context.Context, model string, u *types.TokenUsage) {
|
||||
}
|
||||
purpose, prefixFingerprint := types.LLMCallMetadataFromContext(ctx)
|
||||
logger.Infof(ctx,
|
||||
"[LLM Usage] model=%s, purpose=%s, prompt_prefix=%s, prompt_tokens=%d, completion_tokens=%d, total_tokens=%d, cached_tokens=%d, cache_read_tokens=%d, cache_write_tokens=%d, cache_miss_tokens=%d, cache_reported=%t, cache_status=%s",
|
||||
"[LLM Usage] model=%s, purpose=%s, prompt_prefix=%s, prompt_tokens=%d, completion_tokens=%d, "+
|
||||
"total_tokens=%d, cached_tokens=%d, cache_read_tokens=%d, cache_write_tokens=%d, "+
|
||||
"cache_miss_tokens=%d, cache_reported=%t, cache_status=%s%s",
|
||||
model, purpose, prefixFingerprint, u.PromptTokens, u.CompletionTokens, u.TotalTokens,
|
||||
u.CachedTokens, u.CacheReadTokens, u.CacheWriteTokens, u.CacheMissTokens,
|
||||
u.CacheReported, u.CacheStatus)
|
||||
u.CacheReported, u.CacheStatus, usageAttribution(ctx))
|
||||
}
|
||||
|
||||
// usageAttribution renders the ", session_id=…, principal=…" suffix that
|
||||
// attributes a usage line to the session and terminal principal that
|
||||
// triggered the call. Calls that run outside a session or without a resolved
|
||||
// principal (document parsing, title generation, background jobs) render an
|
||||
// empty suffix, keeping their lines byte-identical to before.
|
||||
func usageAttribution(ctx context.Context) string {
|
||||
var b strings.Builder
|
||||
if sessionID, ok := types.SessionIDFromContext(ctx); ok && sessionID != "" {
|
||||
b.WriteString(", session_id=")
|
||||
b.WriteString(sessionID)
|
||||
}
|
||||
if principal, ok := types.PrincipalFromContext(ctx); ok {
|
||||
b.WriteString(", principal=")
|
||||
b.WriteString(principal.Type)
|
||||
b.WriteString(":")
|
||||
b.WriteString(principal.ID)
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
package chat
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/Tencent/WeKnora/internal/types"
|
||||
)
|
||||
|
||||
func TestUsageAttributionEmptyOutsideSessions(t *testing.T) {
|
||||
// Calls without session or principal context (document parsing, title
|
||||
// generation, background jobs) must keep the usage line byte-identical.
|
||||
if got := usageAttribution(context.Background()); got != "" {
|
||||
t.Fatalf("expected empty suffix, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUsageAttributionCarriesSessionAndPrincipal(t *testing.T) {
|
||||
ctx := types.WithSessionID(context.Background(), "sess-123")
|
||||
ctx = types.WithPrincipal(ctx, types.Principal{Type: types.PrincipalAPIExternalUser, ID: "10000:42"})
|
||||
|
||||
got := usageAttribution(ctx)
|
||||
want := ", session_id=sess-123, principal=api_external_user:10000:42"
|
||||
if got != want {
|
||||
t.Fatalf("got %q want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUsageAttributionSessionOnly(t *testing.T) {
|
||||
ctx := types.WithSessionID(context.Background(), "sess-9")
|
||||
if got := usageAttribution(ctx); got != ", session_id=sess-9" {
|
||||
t.Fatalf("got %q", got)
|
||||
}
|
||||
}
|
||||
@@ -256,6 +256,7 @@ type AgentState struct {
|
||||
IsComplete bool `json:"is_complete"` // Whether agent has finished
|
||||
FinalAnswer string `json:"final_answer"` // The final answer to the query
|
||||
KnowledgeRefs []*SearchResult `json:"knowledge_refs"` // Collected knowledge references
|
||||
TurnUsage TokenUsage `json:"turn_usage"` // LLM token usage accumulated across every round of this turn
|
||||
}
|
||||
|
||||
// FunctionDefinition represents a function definition for LLM function calling
|
||||
|
||||
@@ -73,6 +73,79 @@ func (u *TokenUsage) MarkPromptCacheUnsupported() {
|
||||
u.CacheStatus = PromptCacheStatusUnsupported
|
||||
}
|
||||
|
||||
// Accumulate adds another call's usage into u, preserving the subset
|
||||
// semantics: prompt/completion/total and every cache counter sum
|
||||
// independently (cache counters stay subsets of the prompt count and are
|
||||
// never folded into it). CacheReported ORs, and the cache status is
|
||||
// recomputed from the combined counters so a single cache hit anywhere in
|
||||
// the accumulated calls reads as a hit.
|
||||
func (u *TokenUsage) Accumulate(other TokenUsage) {
|
||||
if u == nil {
|
||||
return
|
||||
}
|
||||
u.PromptTokens += other.PromptTokens
|
||||
u.CompletionTokens += other.CompletionTokens
|
||||
u.TotalTokens += other.TotalTokens
|
||||
u.CachedTokens += other.CachedTokens
|
||||
u.CacheReadTokens += other.CacheReadTokens
|
||||
u.CacheWriteTokens += other.CacheWriteTokens
|
||||
u.CacheMissTokens += other.CacheMissTokens
|
||||
u.CacheReported = u.CacheReported || other.CacheReported
|
||||
switch {
|
||||
case !u.CacheReported:
|
||||
u.CacheStatus = mergeUnreportedCacheStatus(u.CacheStatus, other.CacheStatus)
|
||||
case u.CacheReadTokens > 0:
|
||||
u.CacheStatus = PromptCacheStatusHit
|
||||
default:
|
||||
u.CacheStatus = PromptCacheStatusMiss
|
||||
}
|
||||
}
|
||||
|
||||
// mergeUnreportedCacheStatus folds the statuses of never-reported usage.
|
||||
// "unsupported" survives only while every accumulated call was itself
|
||||
// classified unsupported — the first accumulation adopts the incoming
|
||||
// classification, and any later call that is not known-unsupported (an
|
||||
// unreported or unclassified one) degrades the aggregate to unreported.
|
||||
func mergeUnreportedCacheStatus(accumulated, incoming PromptCacheStatus) PromptCacheStatus {
|
||||
if accumulated == "" {
|
||||
if incoming == "" {
|
||||
return PromptCacheStatusUnreported
|
||||
}
|
||||
return incoming
|
||||
}
|
||||
if accumulated == PromptCacheStatusUnsupported && incoming == PromptCacheStatusUnsupported {
|
||||
return PromptCacheStatusUnsupported
|
||||
}
|
||||
return PromptCacheStatusUnreported
|
||||
}
|
||||
|
||||
// Value persists the usage as a jsonb column (assistant messages carry the
|
||||
// turn's aggregate); nil writes SQL NULL. Mirrors the nullable-pointer
|
||||
// pattern APIPrincipalConfig uses.
|
||||
func (u *TokenUsage) Value() (driver.Value, error) {
|
||||
if u == nil {
|
||||
return nil, nil
|
||||
}
|
||||
return json.Marshal(u)
|
||||
}
|
||||
|
||||
// Scan restores a jsonb usage column; NULL leaves the receiver zero-valued.
|
||||
func (u *TokenUsage) Scan(value interface{}) error {
|
||||
if value == nil {
|
||||
return nil
|
||||
}
|
||||
var b []byte
|
||||
switch v := value.(type) {
|
||||
case []byte:
|
||||
b = v
|
||||
case string:
|
||||
b = []byte(v)
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
return json.Unmarshal(b, u)
|
||||
}
|
||||
|
||||
// LLMToolCall represents a function/tool call from the LLM
|
||||
type LLMToolCall struct {
|
||||
ID string `json:"id"`
|
||||
|
||||
@@ -9,12 +9,13 @@ import (
|
||||
|
||||
// StreamEvent represents a single event in the stream
|
||||
type StreamEvent struct {
|
||||
ID string `json:"id"` // Unique event ID
|
||||
Type types.ResponseType `json:"type"` // Event type (thinking, tool_call, tool_result, references, complete, etc.)
|
||||
Content string `json:"content"` // Event content (chunk for streaming events)
|
||||
Done bool `json:"done"` // Whether this event is done
|
||||
Timestamp time.Time `json:"timestamp"` // When this event occurred
|
||||
Data map[string]interface{} `json:"data,omitempty"` // Additional event data (references, metadata, etc.)
|
||||
ID string `json:"id"` // Unique event ID
|
||||
Type types.ResponseType `json:"type"` // Event type (thinking, tool_call, complete, etc.)
|
||||
Content string `json:"content"` // Event content (chunk for streaming events)
|
||||
Done bool `json:"done"` // Whether this event is done
|
||||
Timestamp time.Time `json:"timestamp"` // When this event occurred
|
||||
Data map[string]interface{} `json:"data,omitempty"` // Additional event data (references, metadata, etc.)
|
||||
Usage *types.TokenUsage `json:"usage,omitempty"` // LLM token usage aggregated over the turn (complete events)
|
||||
}
|
||||
|
||||
// StreamManager stream manager interface - minimal append-only design
|
||||
|
||||
@@ -285,6 +285,10 @@ type Message struct {
|
||||
IsFallback bool `json:"is_fallback,omitempty"`
|
||||
// Agent total execution duration in milliseconds (from query start to answer start)
|
||||
AgentDurationMs int64 `json:"agent_duration_ms,omitempty" gorm:"column:agent_duration_ms;default:0"`
|
||||
// LLM token usage aggregated across every round of the turn that produced this
|
||||
// assistant message. Persisted so history reads can attribute cost after the
|
||||
// live stream is gone; NULL (nil) for user messages and pre-feature rows.
|
||||
Usage *TokenUsage `json:"usage,omitempty" gorm:"type:jsonb;column:usage"`
|
||||
// RenderedContent stores the full RAG-augmented user message (with retrieved context)
|
||||
// sent to the LLM. Used to preserve retrieval context across conversation turns.
|
||||
// Empty for non-retrieval intents or assistant messages.
|
||||
|
||||
@@ -0,0 +1,119 @@
|
||||
package types
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestTokenUsageAccumulateSumsEveryCounter(t *testing.T) {
|
||||
var turn TokenUsage
|
||||
|
||||
first := TokenUsage{PromptTokens: 1000, CompletionTokens: 200, TotalTokens: 1200}
|
||||
first.SetPromptCacheUsage(800, 0, 200, true)
|
||||
second := TokenUsage{PromptTokens: 1500, CompletionTokens: 300, TotalTokens: 1800}
|
||||
second.SetPromptCacheUsage(0, 1500, 1500, true)
|
||||
|
||||
turn.Accumulate(first)
|
||||
turn.Accumulate(second)
|
||||
|
||||
if turn.PromptTokens != 2500 || turn.CompletionTokens != 500 || turn.TotalTokens != 3000 {
|
||||
t.Fatalf("token sums wrong: %+v", turn)
|
||||
}
|
||||
if turn.CacheReadTokens != 800 || turn.CacheWriteTokens != 1500 || turn.CacheMissTokens != 1700 {
|
||||
t.Fatalf("cache sums wrong: %+v", turn)
|
||||
}
|
||||
if turn.CachedTokens != turn.CacheReadTokens {
|
||||
t.Fatalf("legacy alias diverged from cache reads: %+v", turn)
|
||||
}
|
||||
if !turn.CacheReported || turn.CacheStatus != PromptCacheStatusHit {
|
||||
t.Fatalf("a hit anywhere in the turn must read as a hit: %+v", turn)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTokenUsageAccumulateStatusFollowsCombinedCounters(t *testing.T) {
|
||||
var unreported TokenUsage
|
||||
unreported.Accumulate(TokenUsage{PromptTokens: 10, TotalTokens: 10})
|
||||
if unreported.CacheReported || unreported.CacheStatus != PromptCacheStatusUnreported {
|
||||
t.Fatalf("all-unreported accumulation must stay unreported: %+v", unreported)
|
||||
}
|
||||
|
||||
var missOnly TokenUsage
|
||||
reportedMiss := TokenUsage{PromptTokens: 10, TotalTokens: 10}
|
||||
reportedMiss.SetPromptCacheUsage(0, 0, 10, true)
|
||||
missOnly.Accumulate(reportedMiss)
|
||||
if missOnly.CacheStatus != PromptCacheStatusMiss {
|
||||
t.Fatalf("reported without reads must read as miss: %+v", missOnly)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTokenUsageAccumulatePreservesUnsupported(t *testing.T) {
|
||||
var unsupported TokenUsage
|
||||
call := TokenUsage{PromptTokens: 10, TotalTokens: 10}
|
||||
call.MarkPromptCacheUnsupported()
|
||||
|
||||
unsupported.Accumulate(call)
|
||||
unsupported.Accumulate(call)
|
||||
if unsupported.CacheStatus != PromptCacheStatusUnsupported {
|
||||
t.Fatalf("all-unsupported accumulation must stay unsupported: %+v", unsupported)
|
||||
}
|
||||
|
||||
// Mixing in a merely-unreported call degrades the aggregate: the turn no
|
||||
// longer proves every provider path was incapable of reporting.
|
||||
unsupported.Accumulate(TokenUsage{PromptTokens: 5, TotalTokens: 5})
|
||||
if unsupported.CacheStatus != PromptCacheStatusUnreported {
|
||||
t.Fatalf("unsupported+unreported mix must read unreported: %+v", unsupported)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTokenUsageAccumulateOnNilReceiverIsNoOp(t *testing.T) {
|
||||
var u *TokenUsage
|
||||
u.Accumulate(TokenUsage{PromptTokens: 1, TotalTokens: 1}) // must not panic
|
||||
}
|
||||
|
||||
func TestTokenUsageValueScanRoundTrip(t *testing.T) {
|
||||
original := &TokenUsage{PromptTokens: 42, CompletionTokens: 7, TotalTokens: 49}
|
||||
original.SetPromptCacheUsage(30, 0, 12, true)
|
||||
|
||||
value, err := original.Value()
|
||||
if err != nil {
|
||||
t.Fatalf("Value failed: %v", err)
|
||||
}
|
||||
raw, ok := value.([]byte)
|
||||
if !ok {
|
||||
t.Fatalf("Value must produce bytes, got %T", value)
|
||||
}
|
||||
|
||||
var restored TokenUsage
|
||||
if err := restored.Scan(raw); err != nil {
|
||||
t.Fatalf("Scan failed: %v", err)
|
||||
}
|
||||
if restored != *original {
|
||||
t.Fatalf("round trip diverged: got %+v want %+v", restored, *original)
|
||||
}
|
||||
|
||||
// json.Marshal must agree with Value so API responses and the persisted
|
||||
// column carry the same shape.
|
||||
direct, err := json.Marshal(original)
|
||||
if err != nil {
|
||||
t.Fatalf("json.Marshal failed: %v", err)
|
||||
}
|
||||
if string(direct) != string(raw) {
|
||||
t.Fatalf("Value diverged from json.Marshal: %s vs %s", raw, direct)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTokenUsageValueNilAndScanNull(t *testing.T) {
|
||||
var u *TokenUsage
|
||||
value, err := u.Value()
|
||||
if err != nil || value != nil {
|
||||
t.Fatalf("nil usage must persist as SQL NULL, got (%v, %v)", value, err)
|
||||
}
|
||||
|
||||
restored := TokenUsage{PromptTokens: 5}
|
||||
if err := restored.Scan(nil); err != nil {
|
||||
t.Fatalf("Scan(NULL) failed: %v", err)
|
||||
}
|
||||
if restored.PromptTokens != 5 {
|
||||
t.Fatalf("Scan(NULL) must leave the receiver untouched: %+v", restored)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
ALTER TABLE messages DROP COLUMN usage;
|
||||
@@ -0,0 +1,4 @@
|
||||
-- Mirrors versioned migration 000085_message_usage:
|
||||
-- per-turn LLM token usage persisted with the assistant message.
|
||||
|
||||
ALTER TABLE messages ADD COLUMN usage TEXT;
|
||||
@@ -0,0 +1 @@
|
||||
ALTER TABLE messages DROP COLUMN IF EXISTS usage;
|
||||
@@ -0,0 +1 @@
|
||||
ALTER TABLE messages ADD COLUMN IF NOT EXISTS usage JSONB;
|
||||
Reference in New Issue
Block a user