agentbudget

package module
v0.3.0 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Apr 1, 2026 License: Apache-2.0 Imports: 8 Imported by: 0

README

AgentBudget Go SDK

Real-time cost enforcement for AI agent sessions — Go edition.

Go Reference License

Install

go get github.com/AgentBudget/agentbudget/sdks/go

Go 1.21+. Zero external dependencies.

Quickstart

import agentbudget "github.com/AgentBudget/agentbudget/sdks/go"

budget, err := agentbudget.New("$5.00")
if err != nil {
    log.Fatal(err)
}

session := budget.NewSession()
defer session.Close()

// Call your LLM, then record token usage:
resp, _ := openaiClient.CreateChatCompletion(ctx, req)
if err := session.WrapUsage(
    resp.Model,
    int64(resp.Usage.PromptTokens),
    int64(resp.Usage.CompletionTokens),
); err != nil {
    // *agentbudget.BudgetExhausted → hard limit hit
    // *agentbudget.LoopDetected    → repeated calls detected
    log.Fatal(err)
}

fmt.Printf("spent: $%.4f  remaining: $%.4f\n", session.Spent(), session.Remaining())
fmt.Printf("%+v\n", session.Report())

OpenAI Example

package main

import (
    "context"
    "fmt"
    "log"
    "os"

    openai "github.com/sashabaranov/go-openai"
    agentbudget "github.com/AgentBudget/agentbudget/sdks/go"
)

func main() {
    budget, _ := agentbudget.New("$5.00",
        agentbudget.WithOnSoftLimit(func(r agentbudget.Report) {
            fmt.Printf("⚠️  90%% budget used — $%.4f spent\n", r.TotalSpent)
        }),
        agentbudget.WithOnHardLimit(func(r agentbudget.Report) {
            fmt.Printf("🛑 hard limit hit: $%.4f\n", r.TotalSpent)
        }),
    )

    client := openai.NewClient(os.Getenv("OPENAI_API_KEY"))
    session := budget.NewSession()
    defer session.Close()

    resp, err := client.CreateChatCompletion(context.Background(), openai.ChatCompletionRequest{
        Model: openai.GPT4o,
        Messages: []openai.ChatCompletionMessage{
            {Role: openai.ChatMessageRoleUser, Content: "Summarize this market report..."},
        },
    })
    if err != nil {
        log.Fatal(err)
    }

    if err := session.WrapUsage(
        resp.Model,
        int64(resp.Usage.PromptTokens),
        int64(resp.Usage.CompletionTokens),
    ); err != nil {
        log.Fatalf("budget error: %v", err)
    }

    fmt.Println(resp.Choices[0].Message.Content)
    fmt.Printf("\nspent: $%.4f  remaining: $%.4f\n", session.Spent(), session.Remaining())
}

Anthropic Example

import (
    anthropic "github.com/anthropics/anthropic-sdk-go"
    agentbudget "github.com/AgentBudget/agentbudget/sdks/go"
)

session := budget.NewSession()
defer session.Close()

resp, err := anthropicClient.Messages.New(ctx, anthropic.MessageNewParams{
    Model:     anthropic.F(anthropic.ModelClaude3Opus20240229),
    MaxTokens: anthropic.F(int64(1024)),
    Messages: anthropic.F([]anthropic.MessageParam{
        anthropic.NewUserMessage(anthropic.NewTextBlock("Research competitors...")),
    }),
})
if err != nil {
    log.Fatal(err)
}

if err := session.WrapUsage(
    string(resp.Model),
    resp.Usage.InputTokens,
    resp.Usage.OutputTokens,
); err != nil {
    log.Fatalf("budget error: %v", err)
}

Tracking Tool Costs

// Track any fixed-cost action: API calls, database queries, web scraping, etc.
result, err := callSerpAPI(query)
if trackErr := session.Track(0.01, "serp_api"); trackErr != nil {
    log.Fatal(trackErr)
}

Circuit Breaker

budget, _ := agentbudget.New("$5.00",
    agentbudget.WithSoftLimit(0.8),                    // warn at 80%
    agentbudget.WithLoopDetection(10, 60.0),           // 10 calls in 60s
    agentbudget.WithOnSoftLimit(func(r agentbudget.Report) {
        alertOpsTeam(r)
    }),
    agentbudget.WithOnHardLimit(func(r agentbudget.Report) {
        log.Printf("agent %s exceeded budget", r.SessionID)
    }),
    agentbudget.WithOnLoopDetected(func(r agentbudget.Report) {
        log.Printf("loop detected in session %s", r.SessionID)
    }),
)

Three protection levels:

  • Soft limit — fires callback once when spending exceeds threshold. Agent can finish gracefully.
  • Hard limit — WrapUsage/Track returns *BudgetExhausted. No further spend allowed.
  • Loop detection — same model or tool called N times in a window returns *LoopDetected.

Finalization Reserve

Prevent agents from being cut off mid-response. Reserve a fraction for the final step:

budget, _ := agentbudget.New("$1.00",
    agentbudget.WithFinalizationReserve(0.05), // hard limit fires at $0.95
)

session := budget.NewSession()

// Check before the final call:
if session.WouldExceed(estimatedFinalCost) {
    return "Here's what I completed so far..."
}
// Safe — won't hit the hard limit
session.WrapUsage("gpt-4o", inputTokens, outputTokens)

Child Sessions

Allocate sub-budgets to parallel tasks:

parent := budget.NewSession()
defer parent.Close()

child := parent.ChildSession(2.0) // capped at $2 or parent's remaining, whichever is lower
defer func() {
    child.Close()
    // Roll up child spend to parent
    _ = parent.Track(child.Spent(), "child:"+child.ID())
}()

_ = child.WrapUsage("gpt-4o-mini", inputTokens, outputTokens)

Custom Model Pricing

agentbudget.RegisterModel("my-model-v1", 5.00, 20.00) // per million tokens

Cost Report

report := session.Report()
// Report{
//     SessionID:    "sess_...",
//     Budget:       5.0,
//     TotalSpent:   3.421,
//     Remaining:    1.579,
//     Breakdown:    map[llm:[...] tools:[...]],
//     DurationSecs: 34.2,
//     TerminatedBy: "",     // "" | "budget_exhausted" | "loop_detected"
//     EventCount:   12,
// }

Error Handling

err := session.WrapUsage("gpt-4o", inputTokens, outputTokens)

var exhausted *agentbudget.BudgetExhausted
var loop *agentbudget.LoopDetected

switch {
case errors.As(err, &exhausted):
    fmt.Printf("hard limit: $%.4f spent of $%.2f budget\n", exhausted.Spent, exhausted.Budget)
case errors.As(err, &loop):
    fmt.Printf("loop on key: %s\n", loop.Key)
}

API Reference

agentbudget.New(maxSpend, opts...) (*Budget, error)
Option Default Description
WithSoftLimit(fraction) 0.9 Fraction of budget at which soft-limit callback fires
WithLoopDetection(n, secs) 10, 60 Calls per window before LoopDetected
WithFinalizationReserve(fraction) 0.0 Fraction held back from hard limit
WithOnSoftLimit(fn) nil Called once when soft limit triggers
WithOnHardLimit(fn) nil Called when BudgetExhausted is returned
WithOnLoopDetected(fn) nil Called when LoopDetected is returned
(*Budget).NewSession(opts...) *Session
Option Description
WithSessionID(id) Set a custom session ID
(*Session) methods
Method Description
WrapUsage(model, inputTokens, outputTokens) Record LLM call cost
Track(cost, toolName, metadata...) Record fixed-cost action
WouldExceed(cost) bool Pre-flight check without recording
Spent() float64 Total dollars spent
Remaining() float64 Dollars remaining
ChildSession(maxSpend) *Session Create a sub-budget session
Report() Report Structured cost summary
Close() Mark session ended

License

Apache 2.0 — see LICENSE

Documentation

Overview

Package agentbudget provides real-time cost enforcement for AI agent sessions.

Set a hard dollar limit on any LLM session. Wrap API responses, track tool costs, and get automatic circuit breaking — no surprise bills.

Quickstart 

budget, err := agentbudget.New("$5.00")
if err != nil {
    log.Fatal(err)
}

session := budget.NewSession()
defer session.Close()

// After calling your LLM, record usage:
err = session.WrapUsage("gpt-4o", resp.Usage.PromptTokens, resp.Usage.CompletionTokens)
if err != nil {
    // *agentbudget.BudgetExhausted or *agentbudget.LoopDetected
}

fmt.Printf("spent: $%.4f  remaining: $%.4f\n", session.Spent(), session.Remaining())

OpenAI Example 

client := openai.NewClient(os.Getenv("OPENAI_API_KEY"))
budget, _ := agentbudget.New(5.0)
session := budget.NewSession()
defer session.Close()

resp, _ := client.CreateChatCompletion(ctx, openai.ChatCompletionRequest{
    Model:    openai.GPT4o,
    Messages: messages,
})
if err := session.WrapUsage(
    resp.Model,
    int64(resp.Usage.PromptTokens),
    int64(resp.Usage.CompletionTokens),
); err != nil {
    log.Fatal(err)
}

Anthropic Example 

client, _ := anthropic.NewClient(os.Getenv("ANTHROPIC_API_KEY"))
resp, _ := client.Messages.New(ctx, anthropic.MessageNewParams{...})
if err := session.WrapUsage(
    string(resp.Model),
    resp.Usage.InputTokens,
    resp.Usage.OutputTokens,
); err != nil {
    log.Fatal(err)
}
Example (BasicSession) 
package main

import (
	"fmt"
	"log"

	agentbudget "github.com/AgentBudget/agentbudget/sdks/go"
)

func main() {
	budget, err := agentbudget.New("$5.00")
	if err != nil {
		log.Fatal(err)
	}

	session := budget.NewSession()
	defer session.Close()

	// After your LLM call, record usage:
	err = session.WrapUsage("gpt-4o", 1000, 500)
	if err != nil {
		log.Fatal(err)
	}

	fmt.Printf("spent:     $%.4f\n", session.Spent())
	fmt.Printf("remaining: $%.4f\n", session.Remaining())
}
Example (BudgetExhausted) 
package main

import (
	"errors"
	"fmt"

	agentbudget "github.com/AgentBudget/agentbudget/sdks/go"
)

func main() {
	budget, _ := agentbudget.New(0.001) // very small budget
	session := budget.NewSession()
	defer session.Close()

	err := session.WrapUsage("gpt-4o", 10000, 5000)
	if err != nil {
		var exhausted *agentbudget.BudgetExhausted
		if errors.As(err, &exhausted) {
			fmt.Printf("hard limit hit: budget=%.4f spent=%.4f\n",
				exhausted.Budget, exhausted.Spent)
		}
	}
}
Example (Callbacks) 
package main

import (
	"fmt"

	agentbudget "github.com/AgentBudget/agentbudget/sdks/go"
)

func main() {
	budget, _ := agentbudget.New("$1.00",
		agentbudget.WithSoftLimit(0.8),
		agentbudget.WithOnSoftLimit(func(r agentbudget.Report) {
			fmt.Printf("⚠️  soft limit: spent $%.4f of $%.2f\n", r.TotalSpent, r.Budget)
		}),
		agentbudget.WithOnHardLimit(func(r agentbudget.Report) {
			fmt.Printf("🛑 hard limit hit after $%.4f\n", r.TotalSpent)
		}),
	)

	session := budget.NewSession()
	defer session.Close()

	_ = session.WrapUsage("gpt-4o", 100000, 50000)
}
Example (ChildSession) 
package main

import (
	"fmt"

	agentbudget "github.com/AgentBudget/agentbudget/sdks/go"
)

func main() {
	budget, _ := agentbudget.New(10.0)
	parent := budget.NewSession()
	defer parent.Close()

	child := parent.ChildSession(2.0)
	defer func() {
		child.Close()
		// Roll up child spend to parent manually
		if child.Spent() > 0 {
			_ = parent.Track(child.Spent(), "child:"+child.ID())
		}
	}()

	_ = child.WrapUsage("gpt-4o-mini", 5000, 2000)
	fmt.Printf("child spent: $%.4f\n", child.Spent())
}
Example (CustomPricing) 
package main

import (
	"fmt"
	"log"

	agentbudget "github.com/AgentBudget/agentbudget/sdks/go"
)

func main() {
	if err := agentbudget.RegisterModel("my-model", 5.00, 20.00); err != nil {
		log.Fatal(err)
	}

	budget, _ := agentbudget.New(5.0)
	session := budget.NewSession()
	defer session.Close()

	_ = session.WrapUsage("my-model", 1000, 500)
	fmt.Printf("spent: $%.6f\n", session.Spent())
}
Example (TrackTool) 
package main

import (
	"fmt"
	"log"

	agentbudget "github.com/AgentBudget/agentbudget/sdks/go"
)

func main() {
	budget, _ := agentbudget.New(10.0)
	session := budget.NewSession()
	defer session.Close()

	// Track a tool/API call with a known cost
	if err := session.Track(0.05, "serp_api"); err != nil {
		log.Fatal(err)
	}

	fmt.Printf("spent: $%.4f\n", session.Spent())
}
Example (WouldExceed) 
package main

import (
	"fmt"

	agentbudget "github.com/AgentBudget/agentbudget/sdks/go"
)

func main() {
	budget, _ := agentbudget.New(1.0, agentbudget.WithFinalizationReserve(0.05))
	session := budget.NewSession()
	defer session.Close()

	// Do work...
	_ = session.WrapUsage("gpt-4o", 50000, 25000)

	// Before a final expensive call, check first
	estimatedFinalCost := 0.10
	if session.WouldExceed(estimatedFinalCost) {
		fmt.Println("not enough budget for final call — wrapping up")
		return
	}

	_ = session.WrapUsage("gpt-4o", 10000, 5000)
}

Index

Examples

Constants

This section is empty.

Variables

This section is empty.

Functions

func CalculateCost 

func CalculateCost(model string, inputTokens, outputTokens int64) (float64, bool)

CalculateCost returns the cost in USD for the given model and token counts. Returns (cost, true) if the model is found, (0, false) otherwise.

func RegisterModel 

func RegisterModel(model string, inputPerMillion, outputPerMillion float64) error

RegisterModel adds or overrides pricing for a model at runtime. Prices are in USD per million tokens.

Types

type Budget 

type Budget struct {
	// contains filtered or unexported fields
}

Budget is the entry point for AgentBudget. Create one per application or agent type, then call NewSession() for each agent run.

func New 

func New(maxSpend any, opts ...Option) (*Budget, error)

New creates a Budget with the given maximum spend. maxSpend may be a float64, int, or a string like "$5.00" or "5.00".

func (*Budget) MaxSpend 

func (b *Budget) MaxSpend() float64

MaxSpend returns the total configured budget including any finalization reserve.

func (*Budget) NewSession 

func (b *Budget) NewSession(opts ...SessionOption) *Session

NewSession creates a new budget-enforced session. Each call to the agent should use its own session.

type BudgetExhausted 

type BudgetExhausted struct {
	Budget float64
	Spent  float64
}

BudgetExhausted is raised when a recorded cost pushes total spend over the budget.

func (*BudgetExhausted) Error 

func (e *BudgetExhausted) Error() string

type CircuitBreaker 

type CircuitBreaker struct {
	// contains filtered or unexported fields
}

CircuitBreaker monitors budget usage (soft limit) and detects runaway loops.

func (*CircuitBreaker) CheckBudget 

func (cb *CircuitBreaker) CheckBudget(spent, budget float64) string

CheckBudget returns a warning message if the soft limit is reached (fires once). Returns empty string if not triggered.

func (*CircuitBreaker) CheckLoop 

func (cb *CircuitBreaker) CheckLoop(key string) bool

CheckLoop records a call for the given key and returns true if a loop is detected.

type CostEvent 

type CostEvent struct {
	Cost         float64
	Type         CostType
	Timestamp    time.Time
	Model        string // set for LLM events
	InputTokens  int64  // set for LLM events
	OutputTokens int64  // set for LLM events
	ToolName     string // set for tool events
	Metadata     map[string]any
}

CostEvent records a single costed action within a session.

type CostType 

type CostType string

CostType distinguishes LLM calls from manual tool/API cost entries. 

const (
	CostTypeLLM  CostType = "llm"
	CostTypeTool CostType = "tool"
)

type InvalidBudget 

type InvalidBudget struct {
	Value string
}

InvalidBudget is raised when the budget value cannot be parsed or is non-positive.

func (*InvalidBudget) Error 

func (e *InvalidBudget) Error() string

type Ledger 

type Ledger struct {
	// contains filtered or unexported fields
}

Ledger is a thread-safe running total of costs within a session.

func (*Ledger) Breakdown 

func (l *Ledger) Breakdown() map[string]any

Breakdown returns a summary map grouped by type and model/tool.

func (*Ledger) Budget 

func (l *Ledger) Budget() float64

Budget returns the configured budget ceiling.

func (*Ledger) Events 

func (l *Ledger) Events() []CostEvent

Events returns a copy of all recorded events.

func (*Ledger) Record 

func (l *Ledger) Record(event CostEvent) error

Record adds a cost event. Returns *BudgetExhausted if the budget is exceeded.

func (*Ledger) Remaining 

func (l *Ledger) Remaining() float64

Remaining returns how much budget is left.

func (*Ledger) Spent 

func (l *Ledger) Spent() float64

Spent returns the total amount spent so far.

func (*Ledger) WouldExceed 

func (l *Ledger) WouldExceed(cost float64) bool

WouldExceed returns true if adding cost would exceed the budget, without recording anything.

type LoopDetected 

type LoopDetected struct {
	Key string
}

LoopDetected is raised when the same model or tool is called too many times within the configured time window.

func (*LoopDetected) Error 

func (e *LoopDetected) Error() string

type Option 

type Option func(*budgetConfig)

Option configures a Budget.

func WithFinalizationReserve 

func WithFinalizationReserve(fraction float64) Option

WithFinalizationReserve reserves a fraction of the budget for a final step. E.g. WithFinalizationReserve(0.05) on a $1.00 budget fires the hard limit at $0.95.

func WithLoopDetection 

func WithLoopDetection(maxCalls int, windowSeconds float64) Option

WithLoopDetection configures loop detection. maxCalls is the max number of repeated calls in window seconds before LoopDetected is raised.

func WithOnHardLimit 

func WithOnHardLimit(fn func(Report)) Option

WithOnHardLimit sets a callback invoked when BudgetExhausted is raised.

func WithOnLoopDetected 

func WithOnLoopDetected(fn func(Report)) Option

WithOnLoopDetected sets a callback invoked when LoopDetected is raised.

func WithOnSoftLimit 

func WithOnSoftLimit(fn func(Report)) Option

WithOnSoftLimit sets a callback invoked once when spending exceeds the soft limit.

func WithSoftLimit 

func WithSoftLimit(fraction float64) Option

WithSoftLimit sets the fraction of budget at which the soft-limit callback fires. Default: 0.9 (90%).

type Report 

type Report struct {
	SessionID    string         `json:"session_id"`
	Budget       float64        `json:"budget"`
	TotalSpent   float64        `json:"total_spent"`
	Remaining    float64        `json:"remaining"`
	Breakdown    map[string]any `json:"breakdown"`
	DurationSecs *float64       `json:"duration_seconds"`
	TerminatedBy string         `json:"terminated_by"` // "" | "budget_exhausted" | "loop_detected"
	EventCount   int            `json:"event_count"`
}

Report is the structured cost report returned by Session.Report().

type Session 

type Session struct {
	// contains filtered or unexported fields
}

Session tracks costs for a single agent run. Create via Budget.NewSession(). It is safe to use from multiple goroutines.

func (*Session) ChildSession 

func (s *Session) ChildSession(maxSpend float64, opts ...SessionOption) *Session

ChildSession creates a sub-session with its own budget cap. When the child is closed via Close(), its total spend is NOT automatically charged to the parent — call parent.Track(child.Spent(), "child:"+child.ID()) after closing if you want rollup.

func (*Session) Close 

func (s *Session) Close()

Close marks the session as ended. Safe to call multiple times.

func (*Session) ID 

func (s *Session) ID() string

ID returns the session identifier.

func (*Session) Remaining 

func (s *Session) Remaining() float64

Remaining returns dollars left in the budget.

func (*Session) Report 

func (s *Session) Report() Report

Report returns a structured cost summary for this session.

func (*Session) Spent 

func (s *Session) Spent() float64

Spent returns total dollars spent so far.

func (*Session) Track 

func (s *Session) Track(cost float64, toolName string, metadata ...map[string]any) error

Track manually records a fixed-cost action (tool call, API call, etc.). Returns *BudgetExhausted or *LoopDetected if limits are reached.

func (*Session) WouldExceed 

func (s *Session) WouldExceed(cost float64) bool

WouldExceed returns true if the given cost would exceed the remaining budget. Does not record anything.

func (*Session) WrapUsage 

func (s *Session) WrapUsage(model string, inputTokens, outputTokens int64) error

WrapUsage records the cost of a completed LLM call given its model and token counts. Returns *BudgetExhausted if the budget is exceeded, *LoopDetected if a loop is found.

type SessionOption 

type SessionOption func(*sessionConfig)

SessionOption configures a Session.

func WithSessionID 

func WithSessionID(id string) SessionOption

WithSessionID sets a custom session identifier.

Source Files

View all Source files

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.