diff --git a/core/cli/agent_worker.go b/core/cli/agent_worker.go index 7ba862062..4811288f7 100644 --- a/core/cli/agent_worker.go +++ b/core/cli/agent_worker.go @@ -445,6 +445,7 @@ func handleMCPCIJob(shutdownCtx context.Context, data []byte, apiURL, apiToken s // Create LLM client pointing back to the frontend API llm := clients.NewLocalAILLM(task.Model, apiToken, apiURL) + modelCfg.ApplyAgentGenerationCap(llm) // Build cogito options ctx, cancel := context.WithTimeout(shutdownCtx, jobTimeout) diff --git a/core/cli/agent_worker_generation_cap_test.go b/core/cli/agent_worker_generation_cap_test.go new file mode 100644 index 000000000..d28bdd628 --- /dev/null +++ b/core/cli/agent_worker_generation_cap_test.go @@ -0,0 +1,98 @@ +package cli + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "sync" + "time" + + "github.com/modelcontextprotocol/go-sdk/mcp" + "gopkg.in/yaml.v3" + + "github.com/mudler/LocalAI/core/config" + mcpTools "github.com/mudler/LocalAI/core/http/endpoints/mcp" + "github.com/mudler/LocalAI/core/services/jobs" + "github.com/mudler/LocalAI/core/services/testutil" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("handleMCPCIJob generation cap", func() { + var ( + mcpFixture *httptest.Server + llmAPI *httptest.Server + mu sync.Mutex + sentCaps []int + ) + + BeforeEach(func() { + srv := mcp.NewServer(&mcp.Implementation{Name: "fixture", Version: "v0.0.1"}, nil) + srv.AddTool(&mcp.Tool{ + Name: "noop", + Description: "fixture tool", + InputSchema: json.RawMessage(`{"type":"object","properties":{}}`), + }, func(context.Context, *mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: "ok"}}}, nil + }) + mcpFixture = httptest.NewServer(mcp.NewStreamableHTTPHandler(func(*http.Request) *mcp.Server { return srv }, nil)) + + // The LLM API records the output cap of every completion the job sends + // and then fails it: what this spec pins is the request, not the run. + sentCaps = nil + llmAPI = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var body struct { + MaxTokens int `json:"max_tokens"` + } + if err := json.NewDecoder(r.Body).Decode(&body); err == nil { + mu.Lock() + sentCaps = append(sentCaps, body.MaxTokens) + mu.Unlock() + } + http.Error(w, "fixture", http.StatusBadRequest) + })) + }) + + AfterEach(func() { + mcpTools.CloseMCPSessions("capped") + mcpTools.CloseMCPSessions("uncapped") + llmAPI.Close() + mcpFixture.Close() + }) + + runJob := func(name, agentYAML string) []int { + var cfg config.ModelConfig + Expect(yaml.Unmarshal([]byte(fmt.Sprintf("name: %s\n%smcp:\n remote: |\n {\"mcpServers\": {\"fixture\": {\"url\": %q}}}\n", name, agentYAML, mcpFixture.URL)), &cfg)).To(Succeed()) + + payload, err := json.Marshal(jobs.JobEvent{ + JobID: "job-" + name, + TaskID: "task-" + name, + Job: &jobs.JobRecord{ID: "job-" + name}, + Task: &jobs.TaskRecord{ID: "task-" + name, Model: name, Prompt: "do it"}, + ModelConfig: &cfg, + }) + Expect(err).ToNot(HaveOccurred()) + + Expect(handleMCPCIJob(context.Background(), payload, llmAPI.URL, "", testutil.NewFakeBus(), 2*time.Second)).To(Succeed()) + + mu.Lock() + defer mu.Unlock() + return append([]int(nil), sentCaps...) + } + + It("sends agent.max_generation_tokens as the output cap of the job's completions", func() { + caps := runJob("capped", "agent:\n max_generation_tokens: 1234\n") + Expect(caps).ToNot(BeEmpty()) + Expect(caps).To(HaveEach(1234)) + }) + + It("leaves the client default in place when the model sets no cap", func() { + caps := runJob("uncapped", "") + Expect(caps).ToNot(BeEmpty()) + Expect(caps).ToNot(ContainElement(1234)) + Expect(caps).To(HaveEach(BeNumerically(">", 0))) + }) +}) diff --git a/core/config/agent_generation_cap_test.go b/core/config/agent_generation_cap_test.go new file mode 100644 index 000000000..744e700c8 --- /dev/null +++ b/core/config/agent_generation_cap_test.go @@ -0,0 +1,81 @@ +package config_test + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "sync" + + "github.com/mudler/cogito/clients" + "github.com/sashabaranov/go-openai" + "gopkg.in/yaml.v3" + + "github.com/mudler/LocalAI/core/config" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("agent.max_generation_tokens", func() { + var ( + srv *httptest.Server + mu sync.Mutex + caps []int + ) + + BeforeEach(func() { + caps = nil + srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var body struct { + MaxTokens int `json:"max_tokens"` + } + Expect(json.NewDecoder(r.Body).Decode(&body)).To(Succeed()) + mu.Lock() + caps = append(caps, body.MaxTokens) + mu.Unlock() + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"choices":[{"index":0,"finish_reason":"stop","message":{"role":"assistant","content":"ok"}}]}`)) + })) + DeferCleanup(srv.Close) + }) + + // agentLLM builds the client the way the agent job paths do: the model + // config comes from YAML, the client points at the LocalAI API, and the + // config's cap is applied before the loop runs. + agentLLM := func(modelYAML string) *clients.LocalAIClient { + var cfg config.ModelConfig + Expect(yaml.Unmarshal([]byte(modelYAML), &cfg)).To(Succeed()) + llm := clients.NewLocalAILLM(cfg.Name, "", srv.URL) + cfg.ApplyAgentGenerationCap(llm) + return llm + } + + sentCap := func(llm *clients.LocalAIClient, req openai.ChatCompletionRequest) int { + req.Messages = []openai.ChatCompletionMessage{{Role: "user", Content: "hi"}} + _, _, err := llm.CreateChatCompletion(context.Background(), req) + Expect(err).ToNot(HaveOccurred()) + mu.Lock() + defer mu.Unlock() + Expect(caps).ToNot(BeEmpty()) + return caps[len(caps)-1] + } + + It("sends the configured cap on every agent completion", func() { + llm := agentLLM("name: m\nagent:\n max_generation_tokens: 2048\n") + Expect(sentCap(llm, openai.ChatCompletionRequest{})).To(Equal(2048)) + Expect(sentCap(llm, openai.ChatCompletionRequest{})).To(Equal(2048)) + }) + + It("keeps the client default when the model sets no cap", func() { + bare := sentCap(clients.NewLocalAILLM("m", "", srv.URL), openai.ChatCompletionRequest{}) + Expect(bare).To(BeNumerically(">", 0)) + llm := agentLLM("name: m\n") + Expect(sentCap(llm, openai.ChatCompletionRequest{})).To(Equal(bare)) + }) + + It("lets a request that carries its own cap win, so cogito's length retry can raise it", func() { + llm := agentLLM("name: m\nagent:\n max_generation_tokens: 2048\n") + Expect(sentCap(llm, openai.ChatCompletionRequest{MaxTokens: 4096})).To(Equal(4096)) + }) +}) diff --git a/core/config/meta/registry.go b/core/config/meta/registry.go index 5900600c5..af4717ef6 100644 --- a/core/config/meta/registry.go +++ b/core/config/meta/registry.go @@ -946,6 +946,15 @@ func DefaultRegistry() map[string]FieldMetaOverride { }, // --- MCP --- + "agent.max_generation_tokens": { + Section: "agent", + Label: "Max Generation Tokens", + Description: "Output-token cap for each completion of an agent job (/api/agent/jobs). Bounds a single completion that keeps generating, which max_iterations and loop_detection do not. Raise it for reasoning models that think at length. 0 or empty keeps the client default (16384).", + Component: "number", + Min: f64(0), + Order: 120, + Advanced: true, + }, "mcp.remote": { Section: "mcp", Label: "Remote MCP Servers", diff --git a/core/config/model_config.go b/core/config/model_config.go index 82e12bca4..dff45dfe4 100644 --- a/core/config/model_config.go +++ b/core/config/model_config.go @@ -920,6 +920,22 @@ type AgentConfig struct { LoopDetection int `yaml:"loop_detection,omitempty" json:"loop_detection,omitempty"` MaxAdjustmentAttempts int `yaml:"max_adjustment_attempts,omitempty" json:"max_adjustment_attempts,omitempty"` ForceReasoningTool bool `yaml:"force_reasoning_tool,omitempty" json:"force_reasoning_tool,omitempty"` + // MaxGenerationTokens caps the tokens a single completion may generate in + // the agent loop. Neither max_iterations nor loop_detection stops one + // completion that keeps generating; without a per-model value every model + // gets the client's fixed default, which is too low for some reasoning + // models and far above what others need. 0 keeps the client default. + MaxGenerationTokens int `yaml:"max_generation_tokens,omitempty" json:"max_generation_tokens,omitempty"` +} + +// ApplyAgentGenerationCap puts agent.max_generation_tokens on the LLM client +// the agent loop uses. The client applies it to every request that carries no +// cap of its own, so a larger cap cogito sets when it retries a completion +// truncated by length still wins. Unset leaves the client default in place. +func (c *ModelConfig) ApplyAgentGenerationCap(llm interface{ SetMaxTokens(int) }) { + if c.Agent.MaxGenerationTokens > 0 { + llm.SetMaxTokens(c.Agent.MaxGenerationTokens) + } } // HasMCPServers returns true if any MCP servers (remote or stdio) are configured. diff --git a/core/services/agentpool/agent_jobs.go b/core/services/agentpool/agent_jobs.go index d14aa439f..f782ba868 100644 --- a/core/services/agentpool/agent_jobs.go +++ b/core/services/agentpool/agent_jobs.go @@ -900,6 +900,7 @@ func (s *AgentJobService) ExecuteJobInternal(job schema.Job, task schema.Task, c // Create LLM client defaultLLM := clients.NewLocalAILLM(modelConfig.Name, apiKey, "http://127.0.0.1:"+port) + modelConfig.ApplyAgentGenerationCap(defaultLLM) // Initialize traces slice job.Traces = []schema.JobTrace{} diff --git a/core/services/agentpool/agent_jobs_generation_cap_test.go b/core/services/agentpool/agent_jobs_generation_cap_test.go new file mode 100644 index 000000000..770446557 --- /dev/null +++ b/core/services/agentpool/agent_jobs_generation_cap_test.go @@ -0,0 +1,120 @@ +package agentpool_test + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "os" + "path/filepath" + "sync" + "time" + + "github.com/modelcontextprotocol/go-sdk/mcp" + + "github.com/mudler/LocalAI/core/config" + mcpTools "github.com/mudler/LocalAI/core/http/endpoints/mcp" + "github.com/mudler/LocalAI/core/schema" + "github.com/mudler/LocalAI/core/services/agentpool" + "github.com/mudler/LocalAI/core/templates" + "github.com/mudler/LocalAI/pkg/model" + "github.com/mudler/LocalAI/pkg/system" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("AgentJobService generation cap", func() { + var ( + tempDir string + mcpFixture *httptest.Server + llmAPI *httptest.Server + mu sync.Mutex + sentCaps []int + ) + + BeforeEach(func() { + var err error + tempDir, err = os.MkdirTemp("", "agent_jobs_cap_test") + Expect(err).NotTo(HaveOccurred()) + + srv := mcp.NewServer(&mcp.Implementation{Name: "fixture", Version: "v0.0.1"}, nil) + srv.AddTool(&mcp.Tool{ + Name: "noop", + Description: "fixture tool", + InputSchema: json.RawMessage(`{"type":"object","properties":{}}`), + }, func(context.Context, *mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: "ok"}}}, nil + }) + mcpFixture = httptest.NewServer(mcp.NewStreamableHTTPHandler(func(*http.Request) *mcp.Server { return srv }, nil)) + + // Stands in for the LocalAI API the job calls back into: it records the + // output cap of every completion and fails it, since this spec pins the + // request rather than the run. + sentCaps = nil + llmAPI = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var body struct { + MaxTokens int `json:"max_tokens"` + } + if err := json.NewDecoder(r.Body).Decode(&body); err == nil { + mu.Lock() + sentCaps = append(sentCaps, body.MaxTokens) + mu.Unlock() + } + http.Error(w, "fixture", http.StatusBadRequest) + })) + }) + + AfterEach(func() { + mcpTools.CloseMCPSessions("capped") + mcpTools.CloseMCPSessions("uncapped") + llmAPI.Close() + mcpFixture.Close() + _ = os.RemoveAll(tempDir) + }) + + runJob := func(name, agentYAML string) []int { + body := fmt.Sprintf("name: %s\n%smcp:\n remote: |\n {\"mcpServers\": {\"fixture\": {\"url\": %q}}}\n", name, agentYAML, mcpFixture.URL) + Expect(os.WriteFile(filepath.Join(tempDir, name+".yaml"), []byte(body), 0644)).To(Succeed()) + + systemState := &system.SystemState{} + systemState.Model.ModelsPath = tempDir + appConfig := config.NewApplicationConfig( + config.WithDynamicConfigDir(tempDir), + config.WithContext(context.Background()), + ) + appConfig.SystemState = systemState + u, err := url.Parse(llmAPI.URL) + Expect(err).NotTo(HaveOccurred()) + appConfig.APIAddress = u.Host + + configLoader := config.NewModelConfigLoader(tempDir) + Expect(configLoader.LoadModelConfigsFromPath(tempDir)).To(Succeed()) + service := agentpool.NewAgentJobService(appConfig, model.NewModelLoader(systemState), configLoader, templates.NewEvaluator(tempDir)) + + task := schema.Task{ID: "task-" + name, Name: name, Model: name, Prompt: "do it", Enabled: true} + job := schema.Job{ID: "job-" + name, TaskID: task.ID} + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + _ = service.ExecuteJobInternal(job, task, ctx) + + mu.Lock() + defer mu.Unlock() + return append([]int(nil), sentCaps...) + } + + It("sends agent.max_generation_tokens as the output cap of the job's completions", func() { + caps := runJob("capped", "agent:\n max_generation_tokens: 1234\n") + Expect(caps).ToNot(BeEmpty()) + Expect(caps).To(HaveEach(1234)) + }) + + It("leaves the client default in place when the model sets no cap", func() { + caps := runJob("uncapped", "") + Expect(caps).ToNot(BeEmpty()) + Expect(caps).ToNot(ContainElement(1234)) + Expect(caps).To(HaveEach(BeNumerically(">", 0))) + }) +}) diff --git a/docs/content/advanced/model-configuration.md b/docs/content/advanced/model-configuration.md index f014be7d0..accd57ba5 100644 --- a/docs/content/advanced/model-configuration.md +++ b/docs/content/advanced/model-configuration.md @@ -1015,6 +1015,7 @@ Agent/autonomous agent configuration: | `agent.enable_planning` | bool | Enable planning capabilities | | `agent.enable_mcp_prompts` | bool | Enable MCP prompts | | `agent.enable_plan_re_evaluator` | bool | Enable plan re-evaluation | +| `agent.max_generation_tokens` | int | Output-token cap for each completion in the agent loop (agent jobs, `/api/agent/jobs`). Bounds a single completion that keeps generating, which `max_iterations` and `loop_detection` do not. A request that already carries a cap keeps it, so a length retry can still raise it. `0` (default) keeps the client default of 16384 | ## Reasoning Configuration diff --git a/docs/content/features/mcp.md b/docs/content/features/mcp.md index 79c4be153..427140b9e 100644 --- a/docs/content/features/mcp.md +++ b/docs/content/features/mcp.md @@ -191,6 +191,7 @@ different tools share one loaded model. #### Agent Configuration (`agent`) - **`max_iterations`**: Maximum number of MCP tool execution loop iterations (default: 10). Each iteration allows the model to call tools and receive results before generating the next response. +- **`max_generation_tokens`**: Output-token cap for each completion of an MCP agent job (default: the client default, 16384). `max_iterations` bounds the number of turns, not the length of one; this bounds a single completion that keeps generating. Raise it for reasoning models that think at length, lower it for models that should only pick a tool and answer briefly. It applies to agent jobs (`/api/agent/jobs`, also when run on a distributed agent worker); an MCP tool loop on `/v1/chat/completions` uses the request's own `max_tokens`. ## Agent-scoped MCP servers