66 lines
1.7 KiB
Go
66 lines
1.7 KiB
Go
package runtime
|
|
|
|
import (
|
|
agentllm "agent-native/runtime/llm"
|
|
core "agent-system-core"
|
|
"context"
|
|
"fmt"
|
|
)
|
|
|
|
// Stop before a further request when reported total usage exhausts the run budget.
|
|
// A single in-flight request can exceed the remaining input allowance; output is capped.
|
|
func budgetedStream(stream agentllm.StreamFunc, budget core.RunBudget, outputCap int) agentllm.StreamFunc {
|
|
if budget.MaxTokens == nil || *budget.MaxTokens <= 0 {
|
|
return stream
|
|
}
|
|
limit := *budget.MaxTokens
|
|
used := 0
|
|
unknown := false
|
|
return func(ctx context.Context, model string, messages []agentllm.Message, tools []agentllm.ToolDecl, opts agentllm.StreamOptions) <-chan agentllm.Event {
|
|
left := limit - used
|
|
if left <= 0 || unknown {
|
|
out := make(chan agentllm.Event, 1)
|
|
reason := "run token budget exhausted"
|
|
if unknown {
|
|
reason = "provider usage unavailable; further requests stopped"
|
|
}
|
|
out <- agentllm.Event{Kind: agentllm.KindError, Err: fmt.Errorf("budget_exceeded: %s", reason), ErrorMessage: reason}
|
|
close(out)
|
|
return out
|
|
}
|
|
if opts.MaxTokens <= 0 || opts.MaxTokens > left {
|
|
opts.MaxTokens = left
|
|
}
|
|
if outputCap > 0 && opts.MaxTokens > outputCap {
|
|
opts.MaxTokens = outputCap
|
|
}
|
|
in := stream(ctx, model, messages, tools, opts)
|
|
out := make(chan agentllm.Event)
|
|
go func() {
|
|
defer close(out)
|
|
reported := false
|
|
for e := range in {
|
|
if e.Kind == agentllm.KindMessageEnd {
|
|
n := e.Usage.TotalTokens
|
|
if n == 0 {
|
|
n = e.Usage.PromptTokens + e.Usage.CompletionTokens
|
|
}
|
|
if n > 0 {
|
|
used += n
|
|
reported = true
|
|
}
|
|
}
|
|
select {
|
|
case out <- e:
|
|
case <-ctx.Done():
|
|
return
|
|
}
|
|
}
|
|
if !reported {
|
|
unknown = true
|
|
}
|
|
}()
|
|
return out
|
|
}
|
|
}
|