55 lines
1.7 KiB
Go
55 lines
1.7 KiB
Go
package ice
|
|
|
|
// BudgetConfig controls how the context window is divided among sources.
|
|
type BudgetConfig struct {
|
|
NumCtx int
|
|
SystemReserve int // tokens reserved for system prompt
|
|
RecentReserve int // tokens reserved for recent conversation
|
|
ConversationPct float64 // fraction of remaining budget for past conversations
|
|
MemoryPct float64 // fraction of remaining budget for memories
|
|
CodePct float64 // fraction of remaining budget for code context
|
|
}
|
|
|
|
// DefaultBudgetConfig returns sensible defaults for a given context window.
|
|
func DefaultBudgetConfig(numCtx int) BudgetConfig {
|
|
return BudgetConfig{
|
|
NumCtx: numCtx,
|
|
SystemReserve: 1500,
|
|
RecentReserve: 2000,
|
|
ConversationPct: 0.40,
|
|
MemoryPct: 0.20,
|
|
CodePct: 0.40,
|
|
}
|
|
}
|
|
|
|
// Calculate allocates token budgets given how many tokens the current prompt uses.
|
|
func (bc BudgetConfig) Calculate(promptTokens int) Budget {
|
|
// Use 75% of numCtx as total available.
|
|
available := int(float64(bc.NumCtx) * 0.75)
|
|
available -= bc.SystemReserve
|
|
available -= bc.RecentReserve
|
|
available -= promptTokens
|
|
|
|
if available < 0 {
|
|
available = 0
|
|
}
|
|
|
|
return Budget{
|
|
Total: available,
|
|
System: bc.SystemReserve,
|
|
Recent: bc.RecentReserve,
|
|
Conversation: int(float64(available) * bc.ConversationPct),
|
|
Memory: int(float64(available) * bc.MemoryPct),
|
|
Code: int(float64(available) * bc.CodePct),
|
|
}
|
|
}
|
|
|
|
// estimateTokens returns a rough token count for a string (chars / 4).
|
|
func estimateTokens(s string) int {
|
|
n := len(s) / 4
|
|
if n == 0 && len(s) > 0 {
|
|
n = 1
|
|
}
|
|
return n
|
|
}
|