diff --git a/pkg/agent/core.go b/pkg/agent/core.go index a8be6b7..7a5728f 100644 --- a/pkg/agent/core.go +++ b/pkg/agent/core.go @@ -349,7 +349,14 @@ func (s *agentCore) Submit(ctx context.Context, input string, handler EventHandl handle := &actionHandle{cancel: actCancel} s.currentAction.Store(handle) - s.loopcfg.Output = handler + s.loopcfg.Output = func(ev Event) { + if ev.Role == "usage" && s.session != nil { + var in, out int + fmt.Sscanf(ev.Content, "%d %d", &in, &out) + s.session.addUsage(backend.Usage{InputTokens: in, OutputTokens: out}) + } + handler(ev) + } s.loopcfg.PopInject = func() string { if p := s.pendingInject.Swap(nil); p != nil { return *p diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 92da0a1..371c52b 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -80,6 +80,12 @@ func run(ctx context.Context, cfg loopConfig, state state) error { if ev.Done { stopReason = ev.StopReason done = true + if ev.Usage.InputTokens > 0 || ev.Usage.OutputTokens > 0 { + emit(cfg, Event{ + Role: "usage", + Content: fmt.Sprintf("%d %d", ev.Usage.InputTokens, ev.Usage.OutputTokens), + }) + } break } } diff --git a/pkg/agent/session.go b/pkg/agent/session.go index ac80bf3..6781cc0 100644 --- a/pkg/agent/session.go +++ b/pkg/agent/session.go @@ -69,6 +69,10 @@ type Session struct { goal string messages []backend.Message complete bool + // Cumulative usage tracking. + TotalInputTokens int + TotalOutputTokens int + TotalRequests int } // newSession creates a new Session with an initial user message. @@ -86,6 +90,12 @@ func (s *Session) history() []backend.Message { func (s *Session) isComplete() bool { return s.complete } +func (s *Session) addUsage(u backend.Usage) { + s.TotalInputTokens += u.InputTokens + s.TotalOutputTokens += u.OutputTokens + s.TotalRequests++ +} + func (s *Session) update(assistant backend.Message, results []toolResult) error { s.messages = append(s.messages, assistant) for _, r := range results { diff --git a/pkg/tools/builtin.go b/pkg/tools/builtin.go index 1926b4e..844f8a8 100644 --- a/pkg/tools/builtin.go +++ b/pkg/tools/builtin.go @@ -90,3 +90,20 @@ func FileDefs() []ToolInfo { }, } } + +// ReasoningDefs returns the ToolInfo definitions for the reasoning_* built-in tools. +func ReasoningDefs() []ToolInfo { + return []ToolInfo{ + { + Name: "reasoning_think", + Description: "Internal reasoning tool for breaking down problems, planning approaches, analyzing constraints, and reasoning through multi-step solutions before acting. The thought content is not shown to the user.", + InputSchema: json.RawMessage(`{ + "type": "object", + "required": ["thought"], + "properties": { + "thought": {"type": "string", "description": "A reflective note or intermediate reasoning step."} + } + }`), + }, + } +} diff --git a/pkg/tools/reasoning/server.go b/pkg/tools/reasoning/server.go new file mode 100644 index 0000000..69dd5a3 --- /dev/null +++ b/pkg/tools/reasoning/server.go @@ -0,0 +1,42 @@ +package reasoning + +import ( + "context" + "encoding/json" + + "ollie/pkg/tools" +) + +// Decl returns a factory for a reasoning Server. +func Decl() func() tools.Server { return func() tools.Server { return &Server{} } } + +// Server implements tools.Server for reasoning tools. +type Server struct{} + +func (s *Server) ListTools() ([]tools.ToolInfo, error) { + return tools.ReasoningDefs(), nil +} + +func (s *Server) CallTool(_ context.Context, tool string, args json.RawMessage) (json.RawMessage, error) { + switch tool { + case "reasoning_think": + var a struct { + Thought string `json:"thought"` + } + if json.Unmarshal(args, &a) != nil || a.Thought == "" { + return json.Marshal(map[string]any{ + "content": []map[string]string{{"type": "text", "text": "error: missing required field 'thought'"}}, + }) + } + // No-op: the thought is recorded in conversation history by the loop. + return json.Marshal(map[string]any{ + "content": []map[string]string{{"type": "text", "text": ""}}, + }) + default: + return json.Marshal(map[string]string{"error": "unknown tool: " + tool}) + } +} + +func (s *Server) Close() {} + +var _ tools.Server = (*Server)(nil)