diff --git a/internal/manager/biz/aiops/graph/budget_stop_model.go b/internal/manager/biz/aiops/graph/budget_stop_model.go index eccfc2d70..828ff306d 100644 --- a/internal/manager/biz/aiops/graph/budget_stop_model.go +++ b/internal/manager/biz/aiops/graph/budget_stop_model.go @@ -9,6 +9,8 @@ import ( einomodel "github.com/cloudwego/eino/components/model" "github.com/cloudwego/eino/schema" + + "github.com/ongridio/ongrid/internal/pkg/llm" ) type budgetStopModel struct { @@ -23,6 +25,10 @@ func wrapBudgetStopModel(inner einomodel.ToolCallingChatModel) einomodel.ToolCal } func (m *budgetStopModel) Generate(ctx context.Context, input []*schema.Message, opts ...einomodel.Option) (*schema.Message, error) { + // 回调只能标记预算拒绝;在普通推理和工具熔断总结之前终止模型调用。 + if err := llm.BudgetRejectionFromContext(ctx); err != nil { + return nil, err + } if env, ok := latestTerminalToolBudget(input); ok { return m.synthesizeAfterToolBudget(ctx, input, env.Tool, opts...) } @@ -38,6 +44,9 @@ func (m *budgetStopModel) Generate(ctx context.Context, input []*schema.Message, } func (m *budgetStopModel) Stream(ctx context.Context, input []*schema.Message, opts ...einomodel.Option) (*schema.StreamReader[*schema.Message], error) { + if err := llm.BudgetRejectionFromContext(ctx); err != nil { + return nil, err + } if env, ok := latestTerminalToolBudget(input); ok { msg, err := m.synthesizeAfterToolBudget(ctx, input, env.Tool, opts...) if err != nil { diff --git a/internal/manager/biz/aiops/graph/token_budget_test.go b/internal/manager/biz/aiops/graph/token_budget_test.go new file mode 100644 index 000000000..496b8ecd1 --- /dev/null +++ b/internal/manager/biz/aiops/graph/token_budget_test.go @@ -0,0 +1,86 @@ +package graph + +import ( + "context" + "errors" + "io" + "testing" + + "github.com/cloudwego/eino/compose" + "github.com/cloudwego/eino/schema" + + "github.com/ongridio/ongrid/internal/pkg/llm" +) + +func TestBuildReActGraph_TokenBudgetStopsModelCalls(t *testing.T) { + t.Parallel() + for _, mode := range []string{"invoke", "stream"} { + for _, tc := range []struct { + name string + limit int + toolBudget bool + wantReject bool + }{ + {name: "exhausted", limit: 1, wantReject: true}, + {name: "exhausted_before_synthesis", limit: 1, toolBudget: true, wantReject: true}, + {name: "available", limit: 100000}, + {name: "available_for_synthesis", limit: 100000, toolBudget: true}, + {name: "unlimited", limit: 0}, + } { + t.Run(mode+"/"+tc.name, func(t *testing.T) { + t.Parallel() + inner := newScriptedChatModel(makeAssistantNoTools("diagnosis complete")) + g, err := BuildReActGraph(inner, nil, Config{}) + if err != nil { + t.Fatal(err) + } + handler := llm.NewBudgetCallbackHandler(llm.NewInMemoryBudget(tc.limit), 0) + input := &Input{UserText: "Check the service health and summarize the evidence."} + if tc.toolBudget { + input = &Input{History: []*schema.Message{ + schema.UserMessage(input.UserText), + makeAssistantToolCall("", "call_1", "query_logql", `{}`), + schema.ToolMessage(`{"status":"call_budget_exceeded","tool":"query_logql","final_answer_required":true}`, "call_1", schema.WithToolName("query_logql")), + }} + } + err = runBudgetGraph(context.Background(), g, input, mode, compose.WithCallbacks(handler)) + wantCalls := int32(1) + if tc.wantReject { + wantCalls = 0 + if !errors.Is(err, llm.ErrBudgetExceeded) { + t.Errorf("error = %v, want ErrBudgetExceeded", err) + } + if handler.Stats().Rejects != 1 { + t.Errorf("budget stats = %+v, want one rejection", handler.Stats()) + } + } else if err != nil { + t.Errorf("unexpected error: %v", err) + } + if calls := inner.generateCalls(); calls != wantCalls { + t.Errorf("underlying model calls = %d, want %d", calls, wantCalls) + } + }) + } + } +} + +func runBudgetGraph(ctx context.Context, g compose.Runnable[*Input, *Output], input *Input, mode string, opts ...compose.Option) error { + if mode == "invoke" { + _, err := g.Invoke(ctx, input, opts...) + return err + } + stream, err := g.Stream(ctx, input, opts...) + if err != nil { + return err + } + defer stream.Close() + for { + _, err := stream.Recv() + if errors.Is(err, io.EOF) { + return nil + } + if err != nil { + return err + } + } +} diff --git a/internal/pkg/llm/budget_callback.go b/internal/pkg/llm/budget_callback.go index 53ab5e48c..fda2c1c70 100644 --- a/internal/pkg/llm/budget_callback.go +++ b/internal/pkg/llm/budget_callback.go @@ -7,8 +7,8 @@ // Design intent: // - OnStart estimates prompt tokens (cheap rule-of-thumb: joined content // length / 4) and asks BudgetChecker.Check. On rejection we mark the -// context so OnEnd / OnError can short-circuit reporting and PR-2 graph -// code can surface ErrBudgetExceeded. +// context so the graph's ChatModel wrapper can return ErrBudgetExceeded +// before invoking the underlying model. // - OnEnd reads schema.ResponseMeta.Usage from the model's reply and // records it via BudgetChecker.Record. // - We do NOT touch the legacy budget gate inside openaiClient.Chat — it @@ -19,9 +19,8 @@ // // Why not return an error from OnStart? eino's callback contract gives // handlers no way to short-circuit a component. We attach the rejection -// to ctx so the graph node (PR-2) can read it back out and fail-fast -// before the network call. PR-1 ships the plumbing; PR-2 wires the check -// into the node. +// to ctx so the graph's ChatModel wrapper can read it back and fail before +// either a normal model call or a tool-budget synthesis call. package llm import ( @@ -39,9 +38,9 @@ import ( type budgetRejectKey struct{} // BudgetRejectionFromContext returns the budget-rejection error attached -// to ctx by BudgetCallbackHandler.OnStart, or nil if none. Graph nodes -// (PR-2) call this immediately after invoking the model to convert the -// soft signal into a hard error. +// to ctx by BudgetCallbackHandler.OnStart, or nil if none. The graph's +// ChatModel wrapper checks it before invoking the model, including Stream, +// to convert the callback signal into a hard error. func BudgetRejectionFromContext(ctx context.Context) error { v := ctx.Value(budgetRejectKey{}) if v == nil {