feat(model): LLM token 预算可配置化 — models.json每模型maxTokens + LLM_FALLBACK_MAX_TOKENS回退兜底
- config.go: LLMConfig.FallbackMaxTokens(LLM_FALLBACK_MAX_TOKENS, 默认8192) - runtimeconfig/models.go: Group.MaxTokens(models.json每模型独立预算, 负值拒绝) - model/fallback.go: Chain.WithTokenBudget + applyBudget; runtime.go buildSnapshot接入(active.MaxTokens, env.FallbackMaxTokens) - 优先级: 显式Request.MaxTokens > 组maxTokens > 仅多步链末端的LLM_FALLBACK_MAX_TOKENS > 适配器默认8192; 单步链不动env兜底(避免OpenAI主模型行为回归) - gateway admin: AdminLlmProvider.MaxTokens 透传/回显/负值400 - .env.example + .env(LLM_FALLBACK_MAX_TOKENS=8192); Docs/api-client v0.27 + api-reference + design/04 - 新增测试(config/runtimeconfig/budget_test/admin_write), build/vet/test绿+3平台交叉编译
This commit is contained in:
AskaEth committed
1 parent
f347681034
commit
53002f1e11
13 files changed
+384
-15
No files matched your search
@@ -0,0 +1,185 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/config/runtimeconfig"
|
||||
)
|
||||
|
||||
// newFallbackDispatchServer serves two chain steps from one endpoint: the
|
||||
// "main" model fails so the chain moves on, every other model succeeds.
|
||||
func newFallbackDispatchServer(t *testing.T, ch chan<- capturedRequest) *httptest.Server {
|
||||
t.Helper()
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
var body map[string]any
|
||||
_ = json.NewDecoder(r.Body).Decode(&body)
|
||||
ch <- capturedRequest{method: r.Method, path: r.URL.Path, body: body}
|
||||
if body["model"] == "main" {
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
_, _ = io.WriteString(w, `{"error":"down"}`)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = io.WriteString(w, okCompletion)
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
return srv
|
||||
}
|
||||
|
||||
func chatOnce(t *testing.T, chain *Chain) {
|
||||
t.Helper()
|
||||
if _, err := chain.Chat(context.Background(),
|
||||
Request{Messages: []Message{TextMessage(RoleUser, "hi")}}, Options{}); err != nil {
|
||||
t.Fatalf("Chain.Chat error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestChain_Chat_GroupBudgetAppliesToPrimary asserts the active group's own
|
||||
// maxTokens is applied to a request that carries no explicit budget.
|
||||
func TestChain_Chat_GroupBudgetAppliesToPrimary(t *testing.T) {
|
||||
t.Parallel()
|
||||
ch := make(chan capturedRequest, 1)
|
||||
srv := newCompletionServer(t, http.StatusOK, "application/json", okCompletion, ch)
|
||||
chain := NewChain(Step{Model: "m", Adapter: NewOpenAIChat(config.LLMConfig{BaseURL: srv.URL, Model: "m"})}).
|
||||
WithTokenBudget(4096, 8192)
|
||||
|
||||
chatOnce(t, chain)
|
||||
if got := (<-ch).body["max_tokens"]; got != float64(4096) {
|
||||
t.Errorf("body.max_tokens = %v, want the group budget 4096", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestChain_Chat_ExplicitRequestWinsOverGroupBudget asserts an explicit
|
||||
// Request.MaxTokens is never overridden (the small review/judge budgets).
|
||||
func TestChain_Chat_ExplicitRequestWinsOverGroupBudget(t *testing.T) {
|
||||
t.Parallel()
|
||||
ch := make(chan capturedRequest, 1)
|
||||
srv := newCompletionServer(t, http.StatusOK, "application/json", okCompletion, ch)
|
||||
chain := NewChain(Step{Model: "m", Adapter: NewOpenAIChat(config.LLMConfig{BaseURL: srv.URL, Model: "m"})}).
|
||||
WithTokenBudget(4096, 8192)
|
||||
|
||||
if _, err := chain.Chat(context.Background(),
|
||||
Request{Messages: []Message{TextMessage(RoleUser, "hi")}, MaxTokens: 64}, Options{}); err != nil {
|
||||
t.Fatalf("Chain.Chat error: %v", err)
|
||||
}
|
||||
if got := (<-ch).body["max_tokens"]; got != float64(64) {
|
||||
t.Errorf("body.max_tokens = %v, want the explicit 64", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestChain_Chat_FallbackEndAppliesEnvBudget asserts the last-resort budget only
|
||||
// reaches the final fallback step, never the primary.
|
||||
func TestChain_Chat_FallbackEndAppliesEnvBudget(t *testing.T) {
|
||||
t.Parallel()
|
||||
primaryCh := make(chan capturedRequest, 1)
|
||||
primary := newCompletionServer(t, http.StatusInternalServerError, "application/json", `{"error":"down"}`, primaryCh)
|
||||
backupCh := make(chan capturedRequest, 1)
|
||||
backup := newCompletionServer(t, http.StatusOK, "application/json", okCompletion, backupCh)
|
||||
chain := NewChain(
|
||||
Step{Model: "main", Adapter: NewOpenAIChat(config.LLMConfig{BaseURL: primary.URL, Model: "main"})},
|
||||
Step{Model: "backup", Adapter: NewOpenAIChat(config.LLMConfig{BaseURL: backup.URL, Model: "backup"})},
|
||||
).WithTokenBudget(0, 3000)
|
||||
|
||||
chatOnce(t, chain)
|
||||
if got := (<-primaryCh).body["max_tokens"]; got != nil {
|
||||
t.Errorf("primary body.max_tokens = %v, want omitted (no budget on the primary)", got)
|
||||
}
|
||||
if got := (<-backupCh).body["max_tokens"]; got != float64(3000) {
|
||||
t.Errorf("fallback body.max_tokens = %v, want the env fallback budget 3000", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestChain_Chat_GroupBudgetWinsOverFallbackEnv asserts a group maxTokens beats
|
||||
// the last-resort env budget even on the final fallback step.
|
||||
func TestChain_Chat_GroupBudgetWinsOverFallbackEnv(t *testing.T) {
|
||||
t.Parallel()
|
||||
primaryCh := make(chan capturedRequest, 1)
|
||||
primary := newCompletionServer(t, http.StatusInternalServerError, "application/json", `{"error":"down"}`, primaryCh)
|
||||
backupCh := make(chan capturedRequest, 1)
|
||||
backup := newCompletionServer(t, http.StatusOK, "application/json", okCompletion, backupCh)
|
||||
chain := NewChain(
|
||||
Step{Model: "main", Adapter: NewOpenAIChat(config.LLMConfig{BaseURL: primary.URL, Model: "main"})},
|
||||
Step{Model: "backup", Adapter: NewOpenAIChat(config.LLMConfig{BaseURL: backup.URL, Model: "backup"})},
|
||||
).WithTokenBudget(2000, 999)
|
||||
|
||||
chatOnce(t, chain)
|
||||
if got := (<-backupCh).body["max_tokens"]; got != float64(2000) {
|
||||
t.Errorf("fallback body.max_tokens = %v, want the group budget 2000", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestChain_Chat_SingleStepIgnoresFallbackEnv asserts a chain with no fallback
|
||||
// model (one step = the primary) never gets the last-resort budget, so the
|
||||
// adapter default is preserved.
|
||||
func TestChain_Chat_SingleStepIgnoresFallbackEnv(t *testing.T) {
|
||||
t.Parallel()
|
||||
ch := make(chan capturedRequest, 1)
|
||||
srv := newCompletionServer(t, http.StatusOK, "application/json", okCompletion, ch)
|
||||
chain := NewChain(Step{Model: "m", Adapter: NewOpenAIChat(config.LLMConfig{BaseURL: srv.URL, Model: "m"})}).
|
||||
WithTokenBudget(0, 3000)
|
||||
|
||||
chatOnce(t, chain)
|
||||
if got := (<-ch).body["max_tokens"]; got != nil {
|
||||
t.Errorf("body.max_tokens = %v, want omitted (single-step chain is not a fallback)", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRuntimeLLM_GroupMaxTokensEndToEnd proves the models.json maxTokens field
|
||||
// reaches the provider request body through the real runtime stack.
|
||||
func TestRuntimeLLM_GroupMaxTokensEndToEnd(t *testing.T) {
|
||||
ch := make(chan capturedRequest, 1)
|
||||
srv := newCompletionServer(t, http.StatusOK, "application/json", okCompletion, ch)
|
||||
path := filepath.Join(t.TempDir(), "models.json")
|
||||
if err := runtimeconfig.SaveModels(path, []runtimeconfig.Group{
|
||||
{ID: "openai", Provider: ProviderOpenAI, BaseURL: srv.URL, APIKey: "sk-1", Model: "m", MaxTokens: 1234},
|
||||
}); err != nil {
|
||||
t.Fatalf("SaveModels: %v", err)
|
||||
}
|
||||
m := NewRuntimeLLM(path, rtEnv(), nil)
|
||||
if got := m.Groups()[0].MaxTokens; got != 1234 {
|
||||
t.Fatalf("Groups()[0].MaxTokens = %d, want 1234", got)
|
||||
}
|
||||
if _, err := m.Chat(context.Background(),
|
||||
Request{Messages: []Message{TextMessage(RoleUser, "hi")}}, Options{}); err != nil {
|
||||
t.Fatalf("RuntimeLLM.Chat error: %v", err)
|
||||
}
|
||||
if got := (<-ch).body["max_tokens"]; got != float64(1234) {
|
||||
t.Errorf("body.max_tokens = %v, want the group budget 1234", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRuntimeLLM_FallbackEndAppliesEnvBudget proves LLM_FALLBACK_MAX_TOKENS is
|
||||
// wired from config through buildSnapshot into the final fallback step.
|
||||
func TestRuntimeLLM_FallbackEndAppliesEnvBudget(t *testing.T) {
|
||||
ch := make(chan capturedRequest, 2)
|
||||
srv := newFallbackDispatchServer(t, ch)
|
||||
path := filepath.Join(t.TempDir(), "models.json")
|
||||
if err := runtimeconfig.SaveModels(path, []runtimeconfig.Group{
|
||||
{ID: "openai", Provider: ProviderOpenAI, BaseURL: srv.URL, APIKey: "sk-1", Model: "main", Fallback: "backup"},
|
||||
}); err != nil {
|
||||
t.Fatalf("SaveModels: %v", err)
|
||||
}
|
||||
env := rtEnv()
|
||||
env.FallbackMaxTokens = 2500
|
||||
m := NewRuntimeLLM(path, env, nil)
|
||||
if _, err := m.Chat(context.Background(),
|
||||
Request{Messages: []Message{TextMessage(RoleUser, "hi")}}, Options{}); err != nil {
|
||||
t.Fatalf("RuntimeLLM.Chat error: %v", err)
|
||||
}
|
||||
primary := <-ch
|
||||
backup := <-ch
|
||||
if got := primary.body["max_tokens"]; got != nil {
|
||||
t.Errorf("primary body.max_tokens = %v, want omitted", got)
|
||||
}
|
||||
if got := backup.body["max_tokens"]; got != float64(2500) {
|
||||
t.Errorf("fallback body.max_tokens = %v, want LLM_FALLBACK_MAX_TOKENS 2500", got)
|
||||
}
|
||||
}
|
||||
@@ -19,6 +19,13 @@ type Step struct {
|
||||
// single provider is expected.
|
||||
type Chain struct {
|
||||
steps []Step
|
||||
// groupMaxTokens is the active model group's own budget (models.json
|
||||
// maxTokens). When > 0 it applies to every step of the group's chain and
|
||||
// wins over fallbackMaxTokens.
|
||||
groupMaxTokens int
|
||||
// fallbackMaxTokens is the last-resort budget (LLM_FALLBACK_MAX_TOKENS). It
|
||||
// only applies to the final step of a chain that actually has a fallback.
|
||||
fallbackMaxTokens int
|
||||
}
|
||||
|
||||
// NewChain builds a fallback chain from the given steps.
|
||||
@@ -26,6 +33,37 @@ func NewChain(steps ...Step) *Chain {
|
||||
return &Chain{steps: steps}
|
||||
}
|
||||
|
||||
// WithTokenBudget sets the chain's token budgets: the active group's own budget
|
||||
// and the last-resort fallback budget. It must be called before the chain is
|
||||
// published to concurrent callers (the runtime builds a fresh chain per
|
||||
// snapshot and swaps it in atomically), so the fields are never mutated in
|
||||
// flight.
|
||||
func (c *Chain) WithTokenBudget(groupMaxTokens, fallbackMaxTokens int) *Chain {
|
||||
c.groupMaxTokens = groupMaxTokens
|
||||
c.fallbackMaxTokens = fallbackMaxTokens
|
||||
return c
|
||||
}
|
||||
|
||||
// applyBudget resolves the effective MaxTokens for step idx, honoring the
|
||||
// precedence: an explicit Request.MaxTokens > the group's MaxTokens > the
|
||||
// last-resort LLM_FALLBACK_MAX_TOKENS > the adapter default. The last-resort
|
||||
// budget only applies when the chain really has a fallback step (len > 1) and
|
||||
// only to that final step; a single-step chain is a primary model, not a
|
||||
// fallback, so its adapter default is left untouched.
|
||||
func (c *Chain) applyBudget(req Request, idx int) Request {
|
||||
if req.MaxTokens != 0 {
|
||||
return req
|
||||
}
|
||||
if c.groupMaxTokens > 0 {
|
||||
req.MaxTokens = c.groupMaxTokens
|
||||
return req
|
||||
}
|
||||
if c.fallbackMaxTokens > 0 && len(c.steps) > 1 && idx == len(c.steps)-1 {
|
||||
req.MaxTokens = c.fallbackMaxTokens
|
||||
}
|
||||
return req
|
||||
}
|
||||
|
||||
// NewOpenAIChatChain builds a fallback chain from a single supplier config: the
|
||||
// primary model first, then cfg.Fallback (when set and distinct), both served
|
||||
// by an OpenAIChatAdapter against the same OpenAI-compatible endpoint. It
|
||||
@@ -45,8 +83,8 @@ func (c *Chain) Chat(ctx context.Context, req Request, opts Options) (Response,
|
||||
return Response{}, ErrNoSteps
|
||||
}
|
||||
var lastErr error
|
||||
for _, s := range c.steps {
|
||||
stepReq := req
|
||||
for i, s := range c.steps {
|
||||
stepReq := c.applyBudget(req, i)
|
||||
if s.Model != "" {
|
||||
stepReq.Model = s.Model
|
||||
}
|
||||
@@ -68,8 +106,8 @@ func (c *Chain) ChatStream(ctx context.Context, req Request, onEvent func(Stream
|
||||
return Response{}, ErrNoSteps
|
||||
}
|
||||
var lastErr error
|
||||
for _, s := range c.steps {
|
||||
stepReq := req
|
||||
for i, s := range c.steps {
|
||||
stepReq := c.applyBudget(req, i)
|
||||
if s.Model != "" {
|
||||
stepReq.Model = s.Model
|
||||
}
|
||||
|
||||
@@ -169,6 +169,12 @@ func buildSnapshot(groups []runtimeconfig.Group, env config.LLMConfig) (*runtime
|
||||
if err != nil && !errors.Is(err, ErrNoModel) {
|
||||
return nil, err
|
||||
}
|
||||
if chain != nil {
|
||||
// Budgets are resolved inside the chain: the active group's own
|
||||
// maxTokens (models.json) wins, and env.FallbackMaxTokens only caps the
|
||||
// last step of a chain that really has a fallback (design/04 §3.1).
|
||||
chain = chain.WithTokenBudget(active.MaxTokens, env.FallbackMaxTokens)
|
||||
}
|
||||
router := NewRouter(DefaultModelForProvider(cfg),
|
||||
WithVisionModel(env.VisionModel),
|
||||
WithDefaultProvider(cfg.Provider),
|
||||
|
||||
Reference in new issue
Block a user