mirror of
https://github.com/mudler/cogito.git
synced 2026-07-25 03:15:22 -04:00
LocalAI [bot]
fe7fd5de11
* feat(agent): expose AgentState.Background for spawned background agents Add an exported Background flag on AgentState, set true for background spawns (spawn_agent background=true) and false for foreground ones. This lets embedders tell unattended background work apart from a foreground sub-agent whose result is consumed inline — e.g. to auto-notify on completion only for background agents. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * feat: add per-run token usage accumulator (CumulativeUsage) Adds a counting LLM decorator and a CumulativeUsage field on Status so a full ExecuteTools run's token usage can be summed and exposed. Preserves StreamingLLM so wrapping does not disable streaming. * refactor: harden counting stream forwarder and document usage limits Buffer the forwarded stream channel and make the send context-aware to match the client convention and prevent goroutine leaks. Document that streaming-path usage is unpopulated by the bundled clients today. * feat: expose cumulative token usage on the ExecuteTools result Wraps the run LLM in a counting decorator and stamps the summed usage onto the returned fragment's Status.CumulativeUsage via a deferred named return, covering all exit paths. Each sub-agent run reports its own total. --------- Co-authored-by: Ettore Di Giacinto <mudler@localai.io> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
51 lines
1.8 KiB
Go
51 lines
1.8 KiB
Go
package cogito_test
|
|
|
|
import (
|
|
. "github.com/mudler/cogito"
|
|
"github.com/mudler/cogito/tests/mock"
|
|
. "github.com/onsi/ginkgo/v2"
|
|
. "github.com/onsi/gomega"
|
|
"github.com/sashabaranov/go-openai"
|
|
)
|
|
|
|
var _ = Describe("ExecuteTools cumulative usage", func() {
|
|
It("sums token usage across every LLM call in the run", func() {
|
|
mockLLM := mock.NewMockOpenAIClient()
|
|
|
|
// One tool round then a final text answer => >= 2 CreateChatCompletion
|
|
// calls plus one Ask. Each configured call reports 100 total tokens.
|
|
mockLLM.AddCreateChatCompletionFunction("search", `{"query": "test"}`)
|
|
mockTool := mock.NewMockTool("search", "Search for information")
|
|
mock.SetRunResult(mockTool, "Result")
|
|
mockLLM.SetAskResponse("Final answer")
|
|
mockLLM.SetCreateChatCompletionResponse(openai.ChatCompletionResponse{
|
|
Choices: []openai.ChatCompletionChoice{
|
|
{Message: openai.ChatCompletionMessage{Role: "assistant", Content: "No more tools needed."}},
|
|
},
|
|
})
|
|
mockLLM.SetUsage(40, 60, 100)
|
|
mockLLM.SetUsage(40, 60, 100)
|
|
mockLLM.SetUsage(40, 60, 100)
|
|
|
|
fragment := NewEmptyFragment().AddMessage(UserMessageRole, "Task")
|
|
result, err := ExecuteTools(mockLLM, fragment, WithTools(mockTool))
|
|
Expect(err).ToNot(HaveOccurred())
|
|
|
|
// Expected = the total tokens of every usage entry the mock dispensed.
|
|
expected := 0
|
|
for i := 0; i < mockLLM.CreateChatCompletionUsageIndex; i++ {
|
|
expected += mockLLM.CreateChatCompletionUsage[i].TotalTokens
|
|
}
|
|
for i := 0; i < mockLLM.AskUsageIndex; i++ {
|
|
expected += mockLLM.AskUsage[i].TotalTokens
|
|
}
|
|
|
|
Expect(expected).To(BeNumerically(">", 100), "test must drive at least two billed calls")
|
|
Expect(result.Status.CumulativeUsage.TotalTokens).To(Equal(expected))
|
|
Expect(result.Status.CumulativeUsage.TotalTokens).To(
|
|
BeNumerically(">", result.Status.LastUsage.TotalTokens),
|
|
"cumulative must exceed the last single call",
|
|
)
|
|
})
|
|
})
|