token efficiency: fix tier classification, expand cold tier, add usage log and compaction model override
- Fix dead-code bug: ClassifyTier now parses inner tool names from call_tool/pipe arguments via ResultTierArgs - Add file_grep, file_glob, lsp_symbols, lsp_completion, lsp_diagnostics, reasoning_think to cold tier - Add append-only JSONL usage log (~/.local/share/ollie/usage.jsonl) - Add compaction model override: config > OLLIE_COMPACTION_MODEL env > per-backend default (haiku/gpt-4o-mini/flash) - Expose CompactionModel on Core interface
This commit is contained in:
parent
2b914ce150
commit
10cbded305
|
|
@ -0,0 +1,32 @@
|
|||
package agent
|
||||
|
||||
import "os"
|
||||
|
||||
// defaultCompactionModels maps backend names to cheap models suitable for compaction.
|
||||
var defaultCompactionModels = map[string]string{
|
||||
"anthropic": "claude-3-5-haiku-latest",
|
||||
"openai": "gpt-4o-mini",
|
||||
"openrouter": "gpt-4o-mini",
|
||||
"gemini": "gemini-2.0-flash",
|
||||
}
|
||||
|
||||
// backendNamer is the subset of backend.Backend needed for compaction model resolution.
|
||||
type backendNamer interface {
|
||||
Name() string
|
||||
}
|
||||
|
||||
// resolveCompactionModel returns the model to use for compaction.
|
||||
// Priority: agent config > OLLIE_COMPACTION_MODEL env > per-backend default.
|
||||
// Returns "" if no override is configured (use the session's current model).
|
||||
func resolveCompactionModel(cfgModel string, b backendNamer) string {
|
||||
if cfgModel != "" {
|
||||
return cfgModel
|
||||
}
|
||||
if env := os.Getenv("OLLIE_COMPACTION_MODEL"); env != "" {
|
||||
return env
|
||||
}
|
||||
if b != nil {
|
||||
return defaultCompactionModels[b.Name()]
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
|
@ -0,0 +1,52 @@
|
|||
package agent
|
||||
|
||||
import (
|
||||
"os"
|
||||
"testing"
|
||||
)
|
||||
|
||||
type mockBackendForCompact struct {
|
||||
name string
|
||||
model string
|
||||
}
|
||||
|
||||
func (b *mockBackendForCompact) Name() string { return b.name }
|
||||
func (b *mockBackendForCompact) Model() string { return b.model }
|
||||
|
||||
func TestResolveCompactionModel_ConfigWins(t *testing.T) {
|
||||
b := &mockBackendForCompact{name: "anthropic", model: "claude-sonnet-4-5"}
|
||||
got := resolveCompactionModel("my-custom-model", b)
|
||||
if got != "my-custom-model" {
|
||||
t.Errorf("got %q; want my-custom-model", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveCompactionModel_EnvOverridesDefault(t *testing.T) {
|
||||
t.Setenv("OLLIE_COMPACTION_MODEL", "env-model")
|
||||
b := &mockBackendForCompact{name: "anthropic", model: "claude-sonnet-4-5"}
|
||||
got := resolveCompactionModel("", b)
|
||||
if got != "env-model" {
|
||||
t.Errorf("got %q; want env-model", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveCompactionModel_BackendDefault(t *testing.T) {
|
||||
os.Unsetenv("OLLIE_COMPACTION_MODEL")
|
||||
tests := []struct {
|
||||
backend string
|
||||
want string
|
||||
}{
|
||||
{"anthropic", "claude-3-5-haiku-latest"},
|
||||
{"openai", "gpt-4o-mini"},
|
||||
{"openrouter", "gpt-4o-mini"},
|
||||
{"gemini", "gemini-2.0-flash"},
|
||||
{"ollama", ""},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
b := &mockBackendForCompact{name: tt.backend}
|
||||
got := resolveCompactionModel("", b)
|
||||
if got != tt.want {
|
||||
t.Errorf("backend=%s: got %q; want %q", tt.backend, got, tt.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -121,11 +121,11 @@ func BuildRuntime(cfg *config.Config, d tools.Dispatcher, cwd string, env []stri
|
|||
}
|
||||
}
|
||||
|
||||
var tierFn func(string) ResultTier
|
||||
var tierFn func(string, json.RawMessage) ResultTier
|
||||
if srv, ok := d.GetServer("execute"); ok {
|
||||
if tc, ok := srv.(tools.TierClassifier); ok {
|
||||
tierFn = func(name string) ResultTier {
|
||||
switch tc.ResultTier(name) {
|
||||
tierFn = func(name string, args json.RawMessage) ResultTier {
|
||||
switch tc.ResultTierArgs(name, args) {
|
||||
case "cold":
|
||||
return TierCold
|
||||
case "warm":
|
||||
|
|
@ -137,10 +137,11 @@ func BuildRuntime(cfg *config.Config, d tools.Dispatcher, cwd string, env []stri
|
|||
}
|
||||
}
|
||||
|
||||
var backendName, modelName string
|
||||
var backendName, modelName, compactionModel string
|
||||
if cfg != nil {
|
||||
backendName = cfg.Backend
|
||||
modelName = cfg.Model
|
||||
compactionModel = cfg.CompactionModel
|
||||
}
|
||||
|
||||
return &Runtime{
|
||||
|
|
@ -155,6 +156,7 @@ func BuildRuntime(cfg *config.Config, d tools.Dispatcher, cwd string, env []stri
|
|||
MaxSteps: maxSteps,
|
||||
CfgBackend: backendName,
|
||||
CfgModel: modelName,
|
||||
CompactionModel: compactionModel,
|
||||
Messages: messages,
|
||||
}
|
||||
}
|
||||
|
|
@ -653,6 +655,13 @@ func (s *agent) runCompact(ctx context.Context, trigger string) (int, error) {
|
|||
if pre.Context != "" {
|
||||
s.session.appendUserMessage(pre.Context)
|
||||
}
|
||||
// Use a cheaper model for compaction if configured.
|
||||
compactModel := resolveCompactionModel(s.runtime.CompactionModel, s.runtime.Backend)
|
||||
origModel := s.runtime.Backend.Model()
|
||||
if compactModel != "" && compactModel != origModel {
|
||||
s.runtime.Backend.SetModel(compactModel)
|
||||
defer s.runtime.Backend.SetModel(origModel)
|
||||
}
|
||||
n, _, err := s.session.compact(ctx, s.runtime.Backend)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
|
|
@ -820,6 +829,18 @@ func (s *agent) GenerationParams() backend.GenerationParams {
|
|||
return s.runtime.GenParams
|
||||
}
|
||||
|
||||
func (s *agent) CompactionModel() string {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return s.runtime.CompactionModel
|
||||
}
|
||||
|
||||
func (s *agent) SetCompactionModel(model string) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.runtime.CompactionModel = model
|
||||
}
|
||||
|
||||
func (s *agent) SetGenerationParams(params backend.GenerationParams) error {
|
||||
if s.IsRunning() {
|
||||
return fmt.Errorf("cannot change params while agent is running")
|
||||
|
|
@ -1168,6 +1189,7 @@ func (s *agent) executeTurn(ctx context.Context, input string) string {
|
|||
|
||||
if s.session != nil {
|
||||
s.session.recordTurnCost(s.cfg.Backend.Model())
|
||||
appendUsageLog(s.sessionID, s.cfg.Backend.Name(), s.cfg.Backend.Model(), s.session)
|
||||
if s.session.LastTurnCostUSD > 0 {
|
||||
s.emit(Event{Role: "info", Content: fmt.Sprintf("costLast=$%.4f\n", s.session.LastTurnCostUSD)})
|
||||
}
|
||||
|
|
|
|||
|
|
@ -47,7 +47,7 @@ type agentConfig struct {
|
|||
Tools []backend.Tool
|
||||
Exec toolExecutor
|
||||
ClassifyTool func(name string) bool // nil=treat all as serial; true=parallel-read-safe
|
||||
ClassifyTier func(name string) ResultTier // nil=TierHot for all; tools self-classify retention
|
||||
ClassifyTier func(name string, args json.RawMessage) ResultTier // nil=TierHot for all; tools self-classify retention
|
||||
ToolResultMaxBytes int // 0=unlimited; truncate tool results beyond this size
|
||||
Output EventHandler
|
||||
preamble string // compiled system+agent prompt sent as the system role
|
||||
|
|
@ -337,7 +337,7 @@ func run(ctx context.Context, cfg agentConfig, state state) error {
|
|||
emit(cfg, Event{Role: "tool", Name: tc.Name, Content: result})
|
||||
tier := TierHot
|
||||
if !isErr && cfg.ClassifyTier != nil {
|
||||
tier = cfg.ClassifyTier(tc.Name)
|
||||
tier = cfg.ClassifyTier(tc.Name, tc.Arguments)
|
||||
}
|
||||
return toolResult{ToolCallID: tc.ID, Name: tc.Name, Content: result, ContentBlocks: resultBlocks, IsError: isErr, Tier: tier}, false
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
package agent
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
|
||||
"ollie/pkg/backend"
|
||||
"ollie/pkg/tools"
|
||||
)
|
||||
|
|
@ -24,7 +26,7 @@ type Runtime struct {
|
|||
Tools []backend.Tool
|
||||
Exec toolExecutor
|
||||
ClassifyTool toolClassifier
|
||||
ClassifyTier func(string) ResultTier
|
||||
ClassifyTier func(string, json.RawMessage) ResultTier
|
||||
GenParams backend.GenerationParams
|
||||
MaxSteps int
|
||||
ToolResultMaxBytes int // 0 means use defaultToolResultMaxBytes
|
||||
|
|
@ -33,6 +35,9 @@ type Runtime struct {
|
|||
// BuildRuntime returns.
|
||||
CfgBackend string
|
||||
CfgModel string
|
||||
// CompactionModel overrides the model used for context compaction.
|
||||
// Resolved from: agent config > OLLIE_COMPACTION_MODEL env > per-backend default.
|
||||
CompactionModel string
|
||||
// Messages holds non-fatal warnings from runtime construction
|
||||
// (e.g. tool listing errors). Surfaced once at session start.
|
||||
Messages []string
|
||||
|
|
|
|||
|
|
@ -116,6 +116,12 @@ type Core interface {
|
|||
// Returns an error if the agent is currently running.
|
||||
SetGenerationParams(params backend.GenerationParams) error
|
||||
|
||||
// CompactionModel returns the model override used for context compaction.
|
||||
CompactionModel() string
|
||||
|
||||
// SetCompactionModel changes the model used for context compaction.
|
||||
SetCompactionModel(model string)
|
||||
|
||||
// SetEnv injects a session-scoped environment variable into execute_code
|
||||
// subprocesses. Does not affect the daemon process environment.
|
||||
SetEnv(key, value string)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,60 @@
|
|||
package agent
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"time"
|
||||
|
||||
"ollie/pkg/paths"
|
||||
)
|
||||
|
||||
// usageEntry is one line in the append-only usage JSONL log.
|
||||
type usageEntry struct {
|
||||
Timestamp string `json:"ts"`
|
||||
SessionID string `json:"session"`
|
||||
Backend string `json:"backend"`
|
||||
Model string `json:"model"`
|
||||
InputTokens int `json:"input"`
|
||||
OutputTokens int `json:"output"`
|
||||
CachedTokens int `json:"cached"`
|
||||
CreationTokens int `json:"creation"`
|
||||
CostUSD float64 `json:"cost_usd"`
|
||||
}
|
||||
|
||||
// usageLogPath returns the path for the usage log file.
|
||||
// Respects OLLIE_USAGE_LOG env var; defaults to ~/.local/share/ollie/usage.jsonl.
|
||||
func usageLogPath() string {
|
||||
if p := os.Getenv("OLLIE_USAGE_LOG"); p != "" {
|
||||
return p
|
||||
}
|
||||
return filepath.Join(paths.DataDir(), "usage.jsonl")
|
||||
}
|
||||
|
||||
// appendUsageLog writes a single usage entry to the append-only JSONL log.
|
||||
func appendUsageLog(sessionID, backendName, model string, sess *Session) {
|
||||
path := usageLogPath()
|
||||
if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil {
|
||||
return
|
||||
}
|
||||
f, err := os.OpenFile(path, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
entry := usageEntry{
|
||||
Timestamp: time.Now().UTC().Format(time.RFC3339),
|
||||
SessionID: sessionID,
|
||||
Backend: backendName,
|
||||
Model: model,
|
||||
InputTokens: sess.turnInputTokens,
|
||||
OutputTokens: sess.turnOutputTokens,
|
||||
CachedTokens: sess.turnCachedTokens,
|
||||
CreationTokens: sess.turnCreationTokens,
|
||||
CostUSD: sess.LastTurnCostUSD,
|
||||
}
|
||||
data, _ := json.Marshal(entry)
|
||||
data = append(data, '\n')
|
||||
f.Write(data) //nolint:errcheck
|
||||
}
|
||||
|
|
@ -84,6 +84,7 @@ type Config struct {
|
|||
ResponseFormat string `json:"responseFormat,omitempty"`
|
||||
Stop []string `json:"stop,omitempty"`
|
||||
Verbosity string `json:"verbosity,omitempty"`
|
||||
CompactionModel string `json:"compactionModel,omitempty"`
|
||||
}
|
||||
|
||||
// Load parses a Config from r.
|
||||
|
|
|
|||
|
|
@ -1,17 +1,26 @@
|
|||
package execute
|
||||
|
||||
import "strings"
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// defaultColdTools are tools whose results are consumed immediately and don't
|
||||
// need to persist verbatim in the message history.
|
||||
var defaultColdTools = map[string]string{
|
||||
"file_read": "cold",
|
||||
"file_grep": "cold",
|
||||
"file_glob": "cold",
|
||||
"memory_recall": "cold",
|
||||
"web_fetch": "cold",
|
||||
"web_search": "cold",
|
||||
"lsp_definition": "cold",
|
||||
"lsp_references": "cold",
|
||||
"lsp_hover": "cold",
|
||||
"lsp_symbols": "cold",
|
||||
"lsp_completion": "cold",
|
||||
"lsp_diagnostics": "cold",
|
||||
"reasoning_think": "cold",
|
||||
}
|
||||
|
||||
// ResultTier implements tools.TierClassifier. Checks the built-in table first,
|
||||
|
|
@ -27,6 +36,83 @@ func (e *Server) ResultTier(name string) string {
|
|||
return detectTier(code)
|
||||
}
|
||||
|
||||
// ResultTierArgs classifies the tier using both the outer tool name and its
|
||||
// arguments. For call_tool and pipe, it inspects inner tool names and returns
|
||||
// the coldest tier found. Tier ordering: cold < warm < hot.
|
||||
func (e *Server) ResultTierArgs(name string, args json.RawMessage) string {
|
||||
switch name {
|
||||
case "call_tool":
|
||||
return e.tierFromCallToolArgs(args)
|
||||
case "pipe":
|
||||
return e.tierFromPipeArgs(args)
|
||||
case "execute_code":
|
||||
return "warm"
|
||||
default:
|
||||
return e.ResultTier(name)
|
||||
}
|
||||
}
|
||||
|
||||
func (e *Server) tierFromCallToolArgs(args json.RawMessage) string {
|
||||
var parsed struct {
|
||||
Calls []struct {
|
||||
Tool string `json:"tool"`
|
||||
Parallel []struct {
|
||||
Tool string `json:"tool"`
|
||||
} `json:"parallel"`
|
||||
} `json:"calls"`
|
||||
}
|
||||
if json.Unmarshal(args, &parsed) != nil || len(parsed.Calls) == 0 {
|
||||
return "hot"
|
||||
}
|
||||
coldest := "hot"
|
||||
for _, c := range parsed.Calls {
|
||||
if c.Tool != "" {
|
||||
coldest = colderTier(coldest, e.ResultTier(c.Tool))
|
||||
}
|
||||
for _, p := range c.Parallel {
|
||||
if p.Tool != "" {
|
||||
coldest = colderTier(coldest, e.ResultTier(p.Tool))
|
||||
}
|
||||
}
|
||||
}
|
||||
return coldest
|
||||
}
|
||||
|
||||
func (e *Server) tierFromPipeArgs(args json.RawMessage) string {
|
||||
var parsed struct {
|
||||
Stages []struct {
|
||||
Tool string `json:"tool"`
|
||||
Parallel []struct {
|
||||
Tool string `json:"tool"`
|
||||
} `json:"parallel"`
|
||||
} `json:"stages"`
|
||||
}
|
||||
if json.Unmarshal(args, &parsed) != nil || len(parsed.Stages) == 0 {
|
||||
return "hot"
|
||||
}
|
||||
coldest := "hot"
|
||||
for _, s := range parsed.Stages {
|
||||
if s.Tool != "" {
|
||||
coldest = colderTier(coldest, e.ResultTier(s.Tool))
|
||||
}
|
||||
for _, p := range s.Parallel {
|
||||
if p.Tool != "" {
|
||||
coldest = colderTier(coldest, e.ResultTier(p.Tool))
|
||||
}
|
||||
}
|
||||
}
|
||||
return coldest
|
||||
}
|
||||
|
||||
// colderTier returns the colder of two tiers. cold < warm < hot.
|
||||
func colderTier(a, b string) string {
|
||||
order := map[string]int{"cold": 0, "warm": 1, "hot": 2}
|
||||
if order[b] < order[a] {
|
||||
return b
|
||||
}
|
||||
return a
|
||||
}
|
||||
|
||||
// detectTier scans the first 10 lines of a tool script for an
|
||||
// "ollie:tier" annotation. Matches "ollie:tier hot", "ollie:tier warm",
|
||||
// or "ollie:tier cold" anywhere in the line.
|
||||
|
|
|
|||
|
|
@ -0,0 +1,97 @@
|
|||
package execute
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestResultTierArgs_CallToolCold(t *testing.T) {
|
||||
s := &Server{}
|
||||
args := json.RawMessage(`{"calls":[{"tool":"file_read","args":["/tmp/x"]}]}`)
|
||||
if got := s.ResultTierArgs("call_tool", args); got != "cold" {
|
||||
t.Errorf("ResultTierArgs(call_tool, file_read) = %q; want cold", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResultTierArgs_CallToolMultipleCold(t *testing.T) {
|
||||
s := &Server{}
|
||||
args := json.RawMessage(`{"calls":[{"tool":"file_read","args":["/a"]},{"tool":"file_grep","args":["pat"]}]}`)
|
||||
if got := s.ResultTierArgs("call_tool", args); got != "cold" {
|
||||
t.Errorf("ResultTierArgs(call_tool, file_read+file_grep) = %q; want cold", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResultTierArgs_CallToolMixedHot(t *testing.T) {
|
||||
s := &Server{}
|
||||
// file_write is not in defaultColdTools and has no script, so falls back to hot
|
||||
args := json.RawMessage(`{"calls":[{"tool":"file_read","args":["/a"]},{"tool":"file_write","args":["/b","x"]}]}`)
|
||||
if got := s.ResultTierArgs("call_tool", args); got != "cold" {
|
||||
t.Errorf("ResultTierArgs(call_tool, file_read+file_write) = %q; want cold (coldest wins)", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResultTierArgs_CallToolParallel(t *testing.T) {
|
||||
s := &Server{}
|
||||
args := json.RawMessage(`{"calls":[{"parallel":[{"tool":"file_grep","args":["x"]},{"tool":"lsp_hover","args":["/f","1","1"]}]}]}`)
|
||||
if got := s.ResultTierArgs("call_tool", args); got != "cold" {
|
||||
t.Errorf("ResultTierArgs(call_tool, parallel cold tools) = %q; want cold", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResultTierArgs_PipeCold(t *testing.T) {
|
||||
s := &Server{}
|
||||
args := json.RawMessage(`{"stages":[{"tool":"file_grep","args":["TODO"]},{"code":"wc -l"}]}`)
|
||||
if got := s.ResultTierArgs("pipe", args); got != "cold" {
|
||||
t.Errorf("ResultTierArgs(pipe, file_grep+code) = %q; want cold", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResultTierArgs_ExecuteCode(t *testing.T) {
|
||||
s := &Server{}
|
||||
args := json.RawMessage(`{"steps":[{"code":"date"}]}`)
|
||||
if got := s.ResultTierArgs("execute_code", args); got != "warm" {
|
||||
t.Errorf("ResultTierArgs(execute_code) = %q; want warm", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResultTierArgs_BadJSON(t *testing.T) {
|
||||
s := &Server{}
|
||||
if got := s.ResultTierArgs("call_tool", json.RawMessage(`{invalid`)); got != "hot" {
|
||||
t.Errorf("ResultTierArgs(bad json) = %q; want hot", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResultTierArgs_EmptyCalls(t *testing.T) {
|
||||
s := &Server{}
|
||||
if got := s.ResultTierArgs("call_tool", json.RawMessage(`{"calls":[]}`)); got != "hot" {
|
||||
t.Errorf("ResultTierArgs(empty calls) = %q; want hot", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResultTier_DefaultColdTools(t *testing.T) {
|
||||
s := &Server{}
|
||||
coldTools := []string{"file_read", "file_grep", "file_glob", "memory_recall",
|
||||
"web_fetch", "web_search", "lsp_definition", "lsp_references",
|
||||
"lsp_hover", "lsp_symbols", "lsp_completion", "lsp_diagnostics", "reasoning_think"}
|
||||
for _, name := range coldTools {
|
||||
if got := s.ResultTier(name); got != "cold" {
|
||||
t.Errorf("ResultTier(%q) = %q; want cold", name, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestColderTier(t *testing.T) {
|
||||
tests := []struct{ a, b, want string }{
|
||||
{"hot", "cold", "cold"},
|
||||
{"cold", "hot", "cold"},
|
||||
{"warm", "cold", "cold"},
|
||||
{"hot", "warm", "warm"},
|
||||
{"hot", "hot", "hot"},
|
||||
{"cold", "cold", "cold"},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
if got := colderTier(tt.a, tt.b); got != tt.want {
|
||||
t.Errorf("colderTier(%q, %q) = %q; want %q", tt.a, tt.b, got, tt.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -126,4 +126,8 @@ type ParallelClassifier interface {
|
|||
// or "cold" (immediately summarized). Returns "hot" for unknown tools.
|
||||
type TierClassifier interface {
|
||||
ResultTier(name string) string
|
||||
// ResultTierArgs classifies the tier using both the outer tool name and its
|
||||
// arguments. For multiplexed tools (call_tool, pipe), this inspects the
|
||||
// inner tool names. Returns the coldest tier among all inner tools.
|
||||
ResultTierArgs(name string, args json.RawMessage) string
|
||||
}
|
||||
|
|
|
|||
Reference in New Issue