diff --git a/pkg/agent/compaction.go b/pkg/agent/compaction.go new file mode 100644 index 0000000..45bf4b4 --- /dev/null +++ b/pkg/agent/compaction.go @@ -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 "" +} diff --git a/pkg/agent/compaction_test.go b/pkg/agent/compaction_test.go new file mode 100644 index 0000000..324f450 --- /dev/null +++ b/pkg/agent/compaction_test.go @@ -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) + } + } +} diff --git a/pkg/agent/core.go b/pkg/agent/core.go index 305ccfe..fe9d4cc 100644 --- a/pkg/agent/core.go +++ b/pkg/agent/core.go @@ -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,25 +137,27 @@ 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{ - Dispatcher: d, - Tools: allTools, - Exec: exec, - ClassifyTool: classify, - ClassifyTier: tierFn, - Hooks: hooks, - Preamble: preamble, - GenParams: genParams, - MaxSteps: maxSteps, - CfgBackend: backendName, - CfgModel: modelName, - Messages: messages, + Dispatcher: d, + Tools: allTools, + Exec: exec, + ClassifyTool: classify, + ClassifyTier: tierFn, + Hooks: hooks, + Preamble: preamble, + GenParams: genParams, + 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)}) } diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 2a0cc43..48790b8 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -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 @@ -335,10 +335,10 @@ func run(ctx context.Context, cfg agentConfig, state state) error { resultCache.Store(tc.Name+"\x00"+string(tc.Arguments), result) } emit(cfg, Event{Role: "tool", Name: tc.Name, Content: result}) - tier := TierHot - if !isErr && cfg.ClassifyTier != nil { - tier = cfg.ClassifyTier(tc.Name) - } + tier := TierHot + if !isErr && cfg.ClassifyTier != nil { + tier = cfg.ClassifyTier(tc.Name, tc.Arguments) + } return toolResult{ToolCallID: tc.ID, Name: tc.Name, Content: result, ContentBlocks: resultBlocks, IsError: isErr, Tier: tier}, false } diff --git a/pkg/agent/runtime.go b/pkg/agent/runtime.go index 190c1de..7e64629 100644 --- a/pkg/agent/runtime.go +++ b/pkg/agent/runtime.go @@ -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 diff --git a/pkg/agent/types.go b/pkg/agent/types.go index a4d27da..60e6bc1 100644 --- a/pkg/agent/types.go +++ b/pkg/agent/types.go @@ -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) diff --git a/pkg/agent/usage_log.go b/pkg/agent/usage_log.go new file mode 100644 index 0000000..8c5597f --- /dev/null +++ b/pkg/agent/usage_log.go @@ -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 +} diff --git a/pkg/config/config.go b/pkg/config/config.go index 69a28d9..9464e3e 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -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. diff --git a/pkg/tools/execute/tier.go b/pkg/tools/execute/tier.go index a12ced2..97e1ac2 100644 --- a/pkg/tools/execute/tier.go +++ b/pkg/tools/execute/tier.go @@ -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", - "memory_recall": "cold", - "web_fetch": "cold", - "web_search": "cold", - "lsp_definition": "cold", - "lsp_references": "cold", - "lsp_hover": "cold", + "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. diff --git a/pkg/tools/execute/tier_test.go b/pkg/tools/execute/tier_test.go new file mode 100644 index 0000000..f6989c2 --- /dev/null +++ b/pkg/tools/execute/tier_test.go @@ -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) + } + } +} diff --git a/pkg/tools/tools.go b/pkg/tools/tools.go index 7b58ec2..11cf0fa 100644 --- a/pkg/tools/tools.go +++ b/pkg/tools/tools.go @@ -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 }