Let hosts register tools and observe token usage
This commit was merged in pull request #4.
This commit is contained in:
@@ -50,6 +50,10 @@ type Options struct {
|
|||||||
Compact *CompactConfig
|
Compact *CompactConfig
|
||||||
// Subagent 可选:提供时注册 delegate_task(委派由宿主实现,内核只开端口)。
|
// Subagent 可选:提供时注册 delegate_task(委派由宿主实现,内核只开端口)。
|
||||||
Subagent subagent.Factory
|
Subagent subagent.Factory
|
||||||
|
// HostTools 是宿主产品工具(登记工件、卡片、知识库)。与内核工具同名时失败,不覆盖。
|
||||||
|
HostTools []tools.Tool
|
||||||
|
// MaxToolCalls 是单步工具调用上限。0 表示不限制。请求自带预算时以请求为准。
|
||||||
|
MaxToolCalls int
|
||||||
// Log 可选。nil 时丢弃日志。
|
// Log 可选。nil 时丢弃日志。
|
||||||
Log *logsys.Logger
|
Log *logsys.Logger
|
||||||
// ExecutionID 写入日志,不参与执行控制。
|
// ExecutionID 写入日志,不参与执行控制。
|
||||||
@@ -75,6 +79,19 @@ type Backend struct {
|
|||||||
lastFailure map[string]string
|
lastFailure map[string]string
|
||||||
}
|
}
|
||||||
|
|
||||||
|
type workDirKey struct{}
|
||||||
|
|
||||||
|
// WithWorkDir 把本步工作目录放进上下文,供宿主工具读取内核刚写出的文件。
|
||||||
|
func WithWorkDir(ctx context.Context, dir string) context.Context {
|
||||||
|
return context.WithValue(ctx, workDirKey{}, dir)
|
||||||
|
}
|
||||||
|
|
||||||
|
// WorkDir 返回 WithWorkDir 写入的目录。没有时返回空串。
|
||||||
|
func WorkDir(ctx context.Context) string {
|
||||||
|
dir, _ := ctx.Value(workDirKey{}).(string)
|
||||||
|
return dir
|
||||||
|
}
|
||||||
|
|
||||||
// New 创建 native 后端。
|
// New 创建 native 后端。
|
||||||
func New(opts Options) *Backend {
|
func New(opts Options) *Backend {
|
||||||
return &Backend{opts: opts, lastFailure: map[string]string{}}
|
return &Backend{opts: opts, lastFailure: map[string]string{}}
|
||||||
@@ -168,6 +185,12 @@ func (b *Backend) Execute(ctx context.Context, req exec.Request) (*exec.Result,
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
subagent.Register(reg, b.opts.Subagent)
|
subagent.Register(reg, b.opts.Subagent)
|
||||||
|
for _, tool := range b.opts.HostTools {
|
||||||
|
if reg.Has(tool.Name()) {
|
||||||
|
return &exec.Result{Success: false, ErrorReason: "host tool name conflicts with a builtin: " + tool.Name()}, nil
|
||||||
|
}
|
||||||
|
reg.Register(tool)
|
||||||
|
}
|
||||||
if err := reg.Restrict(req.ToolAllowlist); err != nil {
|
if err := reg.Restrict(req.ToolAllowlist); err != nil {
|
||||||
return &exec.Result{Success: false, ErrorReason: err.Error()}, nil
|
return &exec.Result{Success: false, ErrorReason: err.Error()}, nil
|
||||||
}
|
}
|
||||||
@@ -197,6 +220,9 @@ func (b *Backend) Execute(ctx context.Context, req exec.Request) (*exec.Result,
|
|||||||
timeout = b.opts.Timeout
|
timeout = b.opts.Timeout
|
||||||
}
|
}
|
||||||
var runBudget *runtime.Budget
|
var runBudget *runtime.Budget
|
||||||
|
if b.opts.MaxToolCalls > 0 {
|
||||||
|
runBudget = &runtime.Budget{MaxToolCalls: b.opts.MaxToolCalls}
|
||||||
|
}
|
||||||
if req.Budget != nil {
|
if req.Budget != nil {
|
||||||
runBudget = &runtime.Budget{
|
runBudget = &runtime.Budget{
|
||||||
MaxPromptTokens: req.Budget.MaxPromptTokens,
|
MaxPromptTokens: req.Budget.MaxPromptTokens,
|
||||||
@@ -204,6 +230,9 @@ func (b *Backend) Execute(ctx context.Context, req exec.Request) (*exec.Result,
|
|||||||
MaxToolCalls: req.Budget.MaxToolCalls,
|
MaxToolCalls: req.Budget.MaxToolCalls,
|
||||||
Deadline: req.Budget.Deadline,
|
Deadline: req.Budget.Deadline,
|
||||||
}
|
}
|
||||||
|
if runBudget.MaxToolCalls == 0 {
|
||||||
|
runBudget.MaxToolCalls = b.opts.MaxToolCalls
|
||||||
|
}
|
||||||
if !req.Budget.Deadline.IsZero() {
|
if !req.Budget.Deadline.IsZero() {
|
||||||
// Deadline 与 Timeout 先到者生效:把更短者作为 ctx 超时。
|
// Deadline 与 Timeout 先到者生效:把更短者作为 ctx 超时。
|
||||||
if until := time.Until(req.Budget.Deadline); timeout <= 0 || until < timeout {
|
if until := time.Until(req.Budget.Deadline); timeout <= 0 || until < timeout {
|
||||||
@@ -211,10 +240,10 @@ func (b *Backend) Execute(ctx context.Context, req exec.Request) (*exec.Result,
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
runCtx := ctx
|
runCtx := WithWorkDir(ctx, workDir)
|
||||||
cancel := func() {}
|
cancel := func() {}
|
||||||
if timeout > 0 {
|
if timeout > 0 {
|
||||||
runCtx, cancel = context.WithTimeout(ctx, timeout)
|
runCtx, cancel = context.WithTimeout(runCtx, timeout)
|
||||||
}
|
}
|
||||||
defer cancel()
|
defer cancel()
|
||||||
|
|
||||||
@@ -313,10 +342,12 @@ func (b *Backend) run(ctx context.Context, model string, budget *runtime.Budget,
|
|||||||
logs.WriteByte('\n')
|
logs.WriteByte('\n')
|
||||||
if onProcess != nil {
|
if onProcess != nil {
|
||||||
onProcess(exec.ProcessEvent{
|
onProcess(exec.ProcessEvent{
|
||||||
Type: ev.Type,
|
Type: ev.Type,
|
||||||
Text: ev.Text,
|
Text: ev.Text,
|
||||||
Tool: ev.Tool,
|
Tool: ev.Tool,
|
||||||
Detail: ev.Detail,
|
Detail: ev.Detail,
|
||||||
|
PromptTokens: ev.PromptTokens,
|
||||||
|
CompletionTokens: ev.CompletionTokens,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
|
|||||||
@@ -0,0 +1,146 @@
|
|||||||
|
package native
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"sync"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"agent-native/exec"
|
||||||
|
"agent-native/runtime/llm"
|
||||||
|
"agent-native/runtime/tools"
|
||||||
|
)
|
||||||
|
|
||||||
|
type hostTool struct{ name string }
|
||||||
|
|
||||||
|
func (h hostTool) Name() string { return h.name }
|
||||||
|
func (h hostTool) Description() string { return "host tool" }
|
||||||
|
func (h hostTool) Parameters() json.RawMessage {
|
||||||
|
return json.RawMessage(`{"type":"object","properties":{"path":{"type":"string"}}}`)
|
||||||
|
}
|
||||||
|
func (h hostTool) Execute(ctx context.Context, args json.RawMessage) (tools.Result, error) {
|
||||||
|
return tools.Result{Content: "registered " + h.name}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHostToolAndUsageEvent(t *testing.T) {
|
||||||
|
var events []exec.ProcessEvent
|
||||||
|
b := New(Options{
|
||||||
|
Model: "fake",
|
||||||
|
HostTools: []tools.Tool{hostTool{name: "register_artifact"}},
|
||||||
|
Stream: usageScript(),
|
||||||
|
})
|
||||||
|
res, err := b.Execute(context.Background(), exec.Request{
|
||||||
|
WorkDir: t.TempDir(), Prompt: "登记",
|
||||||
|
OnEvent: func(ev exec.ProcessEvent) { events = append(events, ev) },
|
||||||
|
})
|
||||||
|
if err != nil || !res.Success {
|
||||||
|
t.Fatalf("execute: err=%v res=%+v logs=%s", err, res, logsOf(res))
|
||||||
|
}
|
||||||
|
sawText, sawStart, sawEnd, sawUsage := false, false, false, false
|
||||||
|
for _, ev := range events {
|
||||||
|
switch ev.Type {
|
||||||
|
case "text_delta":
|
||||||
|
if ev.Text != "" {
|
||||||
|
sawText = true
|
||||||
|
}
|
||||||
|
case "tool_execution_start":
|
||||||
|
if ev.Tool == "register_artifact" {
|
||||||
|
sawStart = true
|
||||||
|
}
|
||||||
|
case "tool_execution_end":
|
||||||
|
if ev.Tool == "register_artifact" {
|
||||||
|
sawEnd = true
|
||||||
|
}
|
||||||
|
case "usage":
|
||||||
|
if ev.PromptTokens == 11 && ev.CompletionTokens == 7 {
|
||||||
|
sawUsage = true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !sawText || !sawStart || !sawEnd || !sawUsage {
|
||||||
|
t.Fatalf("events missing text/tool/usage: %+v", events)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHostToolSeesStepWorkDir(t *testing.T) {
|
||||||
|
var seen string
|
||||||
|
b := New(Options{
|
||||||
|
Model: "fake",
|
||||||
|
Stream: script(llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{ID: "c1", Name: "peek", Arguments: `{}`}}}),
|
||||||
|
HostTools: []tools.Tool{peekTool{got: &seen}},
|
||||||
|
})
|
||||||
|
dir := t.TempDir()
|
||||||
|
_, err := b.Execute(context.Background(), exec.Request{WorkDir: dir, Prompt: "x", Targets: []string{"output.md"}})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if seen != dir {
|
||||||
|
t.Fatalf("workdir = %q, want %s", seen, dir)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
type peekTool struct{ got *string }
|
||||||
|
|
||||||
|
func (p peekTool) Name() string { return "peek" }
|
||||||
|
func (p peekTool) Description() string { return "peek" }
|
||||||
|
func (p peekTool) Parameters() json.RawMessage { return json.RawMessage(`{"type":"object"}`) }
|
||||||
|
func (p peekTool) Execute(ctx context.Context, args json.RawMessage) (tools.Result, error) {
|
||||||
|
*p.got = WorkDir(ctx)
|
||||||
|
return tools.Result{Content: "ok"}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHostToolNameConflictFailsClosed(t *testing.T) {
|
||||||
|
b := New(Options{
|
||||||
|
Model: "fake",
|
||||||
|
Stream: script(llm.Message{Role: llm.RoleAssistant, Content: "x", StopReason: llm.StopEnd}),
|
||||||
|
HostTools: []tools.Tool{hostTool{name: "read_file"}},
|
||||||
|
})
|
||||||
|
res, err := b.Execute(context.Background(), exec.Request{
|
||||||
|
WorkDir: t.TempDir(), Prompt: "x", Targets: []string{"output.md"},
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if res.Success || res.ErrorReason == "" {
|
||||||
|
t.Fatalf("conflict must fail closed, got %+v", res)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func usageScript() llm.StreamFunc {
|
||||||
|
var mu sync.Mutex
|
||||||
|
step := 0
|
||||||
|
return func(ctx context.Context, model string, messages []llm.Message, td []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
|
||||||
|
ch := make(chan llm.Event, 4)
|
||||||
|
go func() {
|
||||||
|
defer close(ch)
|
||||||
|
mu.Lock()
|
||||||
|
defer mu.Unlock()
|
||||||
|
step++
|
||||||
|
switch step {
|
||||||
|
case 1:
|
||||||
|
ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{
|
||||||
|
Role: llm.RoleAssistant, StopReason: llm.StopToolCalls,
|
||||||
|
ToolCalls: []llm.ToolCall{{ID: "c1", Name: "register_artifact", Arguments: `{"path":"a.docx"}`}},
|
||||||
|
}, Usage: llm.Usage{PromptTokens: 4, CompletionTokens: 2}}
|
||||||
|
case 2:
|
||||||
|
ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{
|
||||||
|
Role: llm.RoleAssistant, StopReason: llm.StopToolCalls,
|
||||||
|
ToolCalls: []llm.ToolCall{{ID: "c2", Name: "write_file", Arguments: `{"path":"output.md","content":"done"}`}},
|
||||||
|
}, Usage: llm.Usage{PromptTokens: 4, CompletionTokens: 2}}
|
||||||
|
default:
|
||||||
|
ch <- llm.Event{Kind: llm.KindTextDelta, Text: "已登记"}
|
||||||
|
ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{
|
||||||
|
Role: llm.RoleAssistant, Content: "已登记", StopReason: llm.StopEnd,
|
||||||
|
}, Usage: llm.Usage{PromptTokens: 3, CompletionTokens: 3}}
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
return ch
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func logsOf(res *exec.Result) string {
|
||||||
|
if res == nil {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
return res.Logs
|
||||||
|
}
|
||||||
+6
-3
@@ -40,9 +40,9 @@ type Request struct {
|
|||||||
// 共享总预算与截止时间,避免传输、步骤与主循环次数相乘)。
|
// 共享总预算与截止时间,避免传输、步骤与主循环次数相乘)。
|
||||||
// 全部字段可选;零值字段不设限。
|
// 全部字段可选;零值字段不设限。
|
||||||
type Budget struct {
|
type Budget struct {
|
||||||
MaxPromptTokens int `json:"max_prompt_tokens,omitempty"`
|
MaxPromptTokens int `json:"max_prompt_tokens,omitempty"`
|
||||||
MaxCompletionTokens int `json:"max_completion_tokens,omitempty"`
|
MaxCompletionTokens int `json:"max_completion_tokens,omitempty"`
|
||||||
MaxToolCalls int `json:"max_tool_calls,omitempty"`
|
MaxToolCalls int `json:"max_tool_calls,omitempty"`
|
||||||
// Deadline 绝对截止时刻;与 Request.Timeout 先到者生效。
|
// Deadline 绝对截止时刻;与 Request.Timeout 先到者生效。
|
||||||
Deadline time.Time `json:"deadline,omitempty"`
|
Deadline time.Time `json:"deadline,omitempty"`
|
||||||
}
|
}
|
||||||
@@ -54,6 +54,9 @@ type ProcessEvent struct {
|
|||||||
Text string // delta 文本(Type 为 *_delta 时)
|
Text string // delta 文本(Type 为 *_delta 时)
|
||||||
Tool string // 工具名(tool_* 时)
|
Tool string // 工具名(tool_* 时)
|
||||||
Detail string // 工具参数/结果摘要(tool_execution_* 时)
|
Detail string // 工具参数/结果摘要(tool_execution_* 时)
|
||||||
|
// PromptTokens / CompletionTokens 在 Type=usage 时填写。
|
||||||
|
PromptTokens int
|
||||||
|
CompletionTokens int
|
||||||
}
|
}
|
||||||
|
|
||||||
const (
|
const (
|
||||||
|
|||||||
@@ -29,6 +29,9 @@ type Event struct {
|
|||||||
// Detail carries tool arguments/results or reasoning deltas for
|
// Detail carries tool arguments/results or reasoning deltas for
|
||||||
// process-visualization consumers; log writers may ignore it.
|
// process-visualization consumers; log writers may ignore it.
|
||||||
Detail string
|
Detail string
|
||||||
|
// PromptTokens / CompletionTokens 只在 type=usage 时有意义。
|
||||||
|
PromptTokens int
|
||||||
|
CompletionTokens int
|
||||||
}
|
}
|
||||||
|
|
||||||
// Config drives one agent session.
|
// Config drives one agent session.
|
||||||
@@ -164,6 +167,7 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) {
|
|||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
emit(cfg, Event{Type: "agent_end", Message: last})
|
emit(cfg, Event{Type: "agent_end", Message: last})
|
||||||
|
emitUsage(cfg, usage)
|
||||||
return Result{Text: last.Content, Messages: messages, StopReason: last.StopReason, Usage: usage}, nil
|
return Result{Text: last.Content, Messages: messages, StopReason: last.StopReason, Usage: usage}, nil
|
||||||
}
|
}
|
||||||
errCap := errString("max turns exceeded")
|
errCap := errString("max turns exceeded")
|
||||||
@@ -296,6 +300,10 @@ func drain(fn func() []llm.Message) []llm.Message {
|
|||||||
return fn()
|
return fn()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func emitUsage(cfg Config, usage llm.Usage) {
|
||||||
|
emit(cfg, Event{Type: "usage", PromptTokens: usage.PromptTokens, CompletionTokens: usage.CompletionTokens})
|
||||||
|
}
|
||||||
|
|
||||||
func emit(cfg Config, ev Event) {
|
func emit(cfg Config, ev Event) {
|
||||||
if cfg.OnEvent != nil {
|
if cfg.OnEvent != nil {
|
||||||
cfg.OnEvent(ev)
|
cfg.OnEvent(ev)
|
||||||
|
|||||||
@@ -45,6 +45,11 @@ func NewRegistry(root string) *Registry {
|
|||||||
|
|
||||||
func (r *Registry) Root() string { return r.root }
|
func (r *Registry) Root() string { return r.root }
|
||||||
|
|
||||||
|
func (r *Registry) Has(name string) bool {
|
||||||
|
_, ok := r.tools[name]
|
||||||
|
return ok
|
||||||
|
}
|
||||||
|
|
||||||
func (r *Registry) Register(t Tool) {
|
func (r *Registry) Register(t Tool) {
|
||||||
if _, ok := r.tools[t.Name()]; !ok {
|
if _, ok := r.tools[t.Name()]; !ok {
|
||||||
r.order = append(r.order, t.Name())
|
r.order = append(r.order, t.Name())
|
||||||
|
|||||||
Reference in New Issue
Block a user