feat(adk): add agent_with_summarization based on agent_middleware (#135)
IPender
6 months ago
committed by
GitHub
parent
d41b497bdc
commit
85592f893f
@ -0,0 +1,118 @@
|
|||||||
|
/*
|
||||||
|
* Copyright 2025 CloudWeGo Authors
|
||||||
|
*
|
||||||
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
* you may not use this file except in compliance with the License.
|
||||||
|
* You may obtain a copy of the License at
|
||||||
|
*
|
||||||
|
* http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
*
|
||||||
|
* Unless required by applicable law or agreed to in writing, software
|
||||||
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
* See the License for the specific language governing permissions and
|
||||||
|
* limitations under the License.
|
||||||
|
*/
|
||||||
|
|
||||||
|
package main
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"log"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/cloudwego/eino/adk"
|
||||||
|
"github.com/cloudwego/eino/components/tool"
|
||||||
|
"github.com/cloudwego/eino/compose"
|
||||||
|
|
||||||
|
"github.com/cloudwego/eino-examples/adk/common/model"
|
||||||
|
"github.com/cloudwego/eino-examples/adk/common/prints"
|
||||||
|
"github.com/cloudwego/eino-examples/adk/common/trace"
|
||||||
|
"github.com/cloudwego/eino-examples/adk/intro/agent_with_summarization/summarization"
|
||||||
|
"github.com/cloudwego/eino-examples/internal/logs"
|
||||||
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
summaryMaxTokensBefore = 10 * 1024
|
||||||
|
summaryMaxTokensRecent = 2 * 1024
|
||||||
|
agentMaxIterations = 30
|
||||||
|
)
|
||||||
|
|
||||||
|
func main() {
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
a, err := newAgent(ctx)
|
||||||
|
if err != nil {
|
||||||
|
logs.Fatalf("create agent failed, err=%v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
traceCloseFn, startSpanFn := trace.AppendCozeLoopCallbackIfConfigured(ctx)
|
||||||
|
defer traceCloseFn(ctx)
|
||||||
|
|
||||||
|
runner := adk.NewRunner(ctx, adk.RunnerConfig{
|
||||||
|
EnableStreaming: true, // you can disable streaming here
|
||||||
|
Agent: a,
|
||||||
|
})
|
||||||
|
|
||||||
|
query := `Write a very long report on the history of artificial intelligence.`
|
||||||
|
ctx, endSpanFn := startSpanFn(ctx, "Agent", query)
|
||||||
|
|
||||||
|
iter := runner.Query(ctx, query)
|
||||||
|
|
||||||
|
var lastMessage adk.Message
|
||||||
|
for {
|
||||||
|
event, ok := iter.Next()
|
||||||
|
if !ok {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
if event.Err != nil {
|
||||||
|
log.Fatal(event.Err)
|
||||||
|
}
|
||||||
|
|
||||||
|
prints.Event(event)
|
||||||
|
|
||||||
|
if event.Output != nil {
|
||||||
|
lastMessage, _, err = adk.GetMessage(event)
|
||||||
|
}
|
||||||
|
|
||||||
|
}
|
||||||
|
|
||||||
|
endSpanFn(ctx, lastMessage)
|
||||||
|
|
||||||
|
// wait for all span to be ended
|
||||||
|
time.Sleep(10 * time.Second)
|
||||||
|
}
|
||||||
|
|
||||||
|
func newAgent(ctx context.Context) (adk.Agent, error) {
|
||||||
|
sumMW, err := summarization.New(ctx, &summarization.Config{
|
||||||
|
Model: model.NewChatModel(),
|
||||||
|
MaxTokensBeforeSummary: summaryMaxTokensBefore,
|
||||||
|
MaxTokensForRecentMessages: summaryMaxTokensRecent,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
a, err := adk.NewChatModelAgent(ctx, &adk.ChatModelAgentConfig{
|
||||||
|
Name: "main_agent",
|
||||||
|
Description: "A long-form report assistant",
|
||||||
|
Instruction: `You are a long-form report writer working in ReAct mode.
|
||||||
|
Think step by step, call tools to expand content by repeating paragraphs, then synthesize a cohesive response.
|
||||||
|
one time call one tool, do not call multiple tools in one turn.
|
||||||
|
Each tool call should indicate the call number. After 20 tool calls, produce a final summary.`,
|
||||||
|
Model: model.NewChatModel(),
|
||||||
|
Middlewares: []adk.AgentMiddleware{sumMW},
|
||||||
|
ToolsConfig: adk.ToolsConfig{
|
||||||
|
ToolsNodeConfig: compose.ToolsNodeConfig{
|
||||||
|
Tools: []tool.BaseTool{
|
||||||
|
NewRepeatSectionsTool(),
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
MaxIterations: agentMaxIterations,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
return a, nil
|
||||||
|
}
|
||||||
@ -0,0 +1,403 @@
|
|||||||
|
/*
|
||||||
|
* Copyright 2025 CloudWeGo Authors
|
||||||
|
*
|
||||||
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
* you may not use this file except in compliance with the License.
|
||||||
|
* You may obtain a copy of the License at
|
||||||
|
*
|
||||||
|
* http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
*
|
||||||
|
* Unless required by applicable law or agreed to in writing, software
|
||||||
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
* See the License for the specific language governing permissions and
|
||||||
|
* limitations under the License.
|
||||||
|
*/
|
||||||
|
|
||||||
|
package summarization
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"github.com/pkoukk/tiktoken-go"
|
||||||
|
|
||||||
|
"github.com/cloudwego/eino/adk"
|
||||||
|
"github.com/cloudwego/eino/components/model"
|
||||||
|
"github.com/cloudwego/eino/components/prompt"
|
||||||
|
"github.com/cloudwego/eino/compose"
|
||||||
|
"github.com/cloudwego/eino/schema"
|
||||||
|
)
|
||||||
|
|
||||||
|
type TokenCounter func(ctx context.Context, msgs []adk.Message) (tokenNum []int64, err error)
|
||||||
|
|
||||||
|
// Config defines parameters for the conversation summarization middleware.
|
||||||
|
// It controls when summarization is triggered and how much recent context is retained.
|
||||||
|
// Required: Model. Optional: SystemPrompt, Counter, and token budgets.
|
||||||
|
type Config struct {
|
||||||
|
// MaxTokensBeforeSummary is the max token threshold to trigger summarization based on total context
|
||||||
|
// (system prompt + history). Uses DefaultMaxTokensBeforeSummary when <= 0.
|
||||||
|
MaxTokensBeforeSummary int
|
||||||
|
|
||||||
|
// MaxTokensForRecentMessages is the max token budget reserved for recent messages after summarization.
|
||||||
|
// Uses DefaultMaxTokensForRecentMessages when <= 0.
|
||||||
|
MaxTokensForRecentMessages int
|
||||||
|
|
||||||
|
// Counter custom token counter.
|
||||||
|
// Optional
|
||||||
|
Counter TokenCounter
|
||||||
|
|
||||||
|
// Model used to generate the summary. Must be provided.
|
||||||
|
// Required.
|
||||||
|
Model model.BaseChatModel
|
||||||
|
|
||||||
|
// SystemPrompt is the system prompt for the summarizer.
|
||||||
|
// Optional. If empty, PromptOfSummary is used.
|
||||||
|
SystemPrompt string
|
||||||
|
}
|
||||||
|
|
||||||
|
// New creates an AgentMiddleware that compacts long conversation history
|
||||||
|
// into a single summary message when the token threshold is exceeded.
|
||||||
|
// The summarizer chain is: ChatTemplate(SystemPrompt) -> ChatModel(Model).
|
||||||
|
// It applies defaults for token budgets and allows a custom Counter.
|
||||||
|
func New(ctx context.Context, cfg *Config) (adk.AgentMiddleware, error) {
|
||||||
|
if cfg == nil {
|
||||||
|
return adk.AgentMiddleware{}, fmt.Errorf("config is nil")
|
||||||
|
}
|
||||||
|
|
||||||
|
systemPrompt := cfg.SystemPrompt
|
||||||
|
if systemPrompt == "" {
|
||||||
|
systemPrompt = PromptOfSummary
|
||||||
|
}
|
||||||
|
maxBefore := DefaultMaxTokensBeforeSummary
|
||||||
|
if cfg.MaxTokensBeforeSummary > 0 {
|
||||||
|
maxBefore = cfg.MaxTokensBeforeSummary
|
||||||
|
}
|
||||||
|
maxRecent := DefaultMaxTokensForRecentMessages
|
||||||
|
if cfg.MaxTokensForRecentMessages > 0 {
|
||||||
|
maxRecent = cfg.MaxTokensForRecentMessages
|
||||||
|
}
|
||||||
|
|
||||||
|
tpl := prompt.FromMessages(schema.FString,
|
||||||
|
schema.SystemMessage(systemPrompt),
|
||||||
|
schema.UserMessage("summarize 'older_messages': "))
|
||||||
|
|
||||||
|
summarizer, err := compose.NewChain[map[string]any, *schema.Message]().
|
||||||
|
AppendChatTemplate(tpl).
|
||||||
|
AppendChatModel(cfg.Model).
|
||||||
|
Compile(ctx, compose.WithGraphName("Summarizer"))
|
||||||
|
if err != nil {
|
||||||
|
return adk.AgentMiddleware{}, fmt.Errorf("compile summarizer failed, err=%w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
sm := &summaryMiddleware{
|
||||||
|
counter: defaultCounterToken,
|
||||||
|
maxBefore: maxBefore,
|
||||||
|
maxRecent: maxRecent,
|
||||||
|
summarizer: summarizer,
|
||||||
|
}
|
||||||
|
if cfg.Counter != nil {
|
||||||
|
sm.counter = cfg.Counter
|
||||||
|
}
|
||||||
|
return adk.AgentMiddleware{BeforeChatModel: sm.BeforeModel}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
const summaryMessageFlag = "_agent_middleware_summary_message"
|
||||||
|
|
||||||
|
type summaryMiddleware struct {
|
||||||
|
counter TokenCounter
|
||||||
|
maxBefore int
|
||||||
|
maxRecent int
|
||||||
|
|
||||||
|
summarizer compose.Runnable[map[string]any, *schema.Message]
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *summaryMiddleware) BeforeModel(ctx context.Context, state *adk.ChatModelAgentState) (err error) {
|
||||||
|
if state == nil || len(state.Messages) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
messages := state.Messages
|
||||||
|
msgsToken, err := s.counter(ctx, messages)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("count token failed, err=%w", err)
|
||||||
|
}
|
||||||
|
if len(messages) != len(msgsToken) {
|
||||||
|
return fmt.Errorf("token count mismatch, msgNum=%d, tokenCountNum=%d", len(messages), len(msgsToken))
|
||||||
|
}
|
||||||
|
|
||||||
|
var total int64
|
||||||
|
for _, t := range msgsToken {
|
||||||
|
total += t
|
||||||
|
}
|
||||||
|
// Trigger summarization only when exceeding threshold
|
||||||
|
if total <= int64(s.maxBefore) {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Build blocks with user-messages, summary-message, tool-call pairings
|
||||||
|
type block struct {
|
||||||
|
msgs []*schema.Message
|
||||||
|
tokens int64
|
||||||
|
}
|
||||||
|
idx := 0
|
||||||
|
|
||||||
|
systemBlock := block{}
|
||||||
|
if idx < len(messages) {
|
||||||
|
m := messages[idx]
|
||||||
|
if m != nil && m.Role == schema.System {
|
||||||
|
systemBlock.msgs = append(systemBlock.msgs, m)
|
||||||
|
systemBlock.tokens += msgsToken[idx]
|
||||||
|
idx++
|
||||||
|
}
|
||||||
|
}
|
||||||
|
userBlock := block{}
|
||||||
|
for idx < len(messages) {
|
||||||
|
m := messages[idx]
|
||||||
|
if m == nil {
|
||||||
|
idx++
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if m.Role != schema.User {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
userBlock.msgs = append(userBlock.msgs, m)
|
||||||
|
userBlock.tokens += msgsToken[idx]
|
||||||
|
idx++
|
||||||
|
}
|
||||||
|
summaryBlock := block{}
|
||||||
|
if idx < len(messages) {
|
||||||
|
m := messages[idx]
|
||||||
|
if m != nil && m.Role == schema.Assistant {
|
||||||
|
if _, ok := m.Extra[summaryMessageFlag]; ok {
|
||||||
|
summaryBlock.msgs = append(summaryBlock.msgs, m)
|
||||||
|
summaryBlock.tokens += msgsToken[idx]
|
||||||
|
idx++
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
toolBlocks := make([]block, 0)
|
||||||
|
for i := idx; i < len(messages); i++ {
|
||||||
|
m := messages[i]
|
||||||
|
if m == nil {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
if m.Role == schema.Assistant && len(m.ToolCalls) > 0 {
|
||||||
|
b := block{msgs: []*schema.Message{m}, tokens: msgsToken[i]}
|
||||||
|
// Collect subsequent tool messages matching any tool call id
|
||||||
|
callIDs := make(map[string]struct{}, len(m.ToolCalls))
|
||||||
|
for _, tc := range m.ToolCalls {
|
||||||
|
callIDs[tc.ID] = struct{}{}
|
||||||
|
}
|
||||||
|
j := i + 1
|
||||||
|
for j < len(messages) {
|
||||||
|
nm := messages[j]
|
||||||
|
if nm == nil || nm.Role != schema.Tool {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
// Match by ToolCallID when available; if empty, include but keep boundary
|
||||||
|
if nm.ToolCallID == "" {
|
||||||
|
b.msgs = append(b.msgs, nm)
|
||||||
|
b.tokens += msgsToken[j]
|
||||||
|
} else {
|
||||||
|
if _, ok := callIDs[nm.ToolCallID]; !ok {
|
||||||
|
// Tool message not belonging to this assistant call -> end pairing
|
||||||
|
break
|
||||||
|
}
|
||||||
|
b.msgs = append(b.msgs, nm)
|
||||||
|
b.tokens += msgsToken[j]
|
||||||
|
}
|
||||||
|
j++
|
||||||
|
}
|
||||||
|
toolBlocks = append(toolBlocks, b)
|
||||||
|
i = j - 1
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
toolBlocks = append(toolBlocks, block{msgs: []*schema.Message{m}, tokens: msgsToken[i]})
|
||||||
|
}
|
||||||
|
|
||||||
|
// Split into recent and older within token budget, from newest to oldest
|
||||||
|
var recentBlocks []block
|
||||||
|
var olderBlocks []block
|
||||||
|
var recentTokens int64
|
||||||
|
for i := len(toolBlocks) - 1; i >= 0; i-- {
|
||||||
|
b := toolBlocks[i]
|
||||||
|
if recentTokens+b.tokens > int64(s.maxRecent) {
|
||||||
|
olderBlocks = append([]block{b}, olderBlocks...)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
recentBlocks = append([]block{b}, recentBlocks...)
|
||||||
|
recentTokens += b.tokens
|
||||||
|
}
|
||||||
|
|
||||||
|
joinBlocks := func(bs []block) string {
|
||||||
|
var sb strings.Builder
|
||||||
|
for _, b := range bs {
|
||||||
|
for _, m := range b.msgs {
|
||||||
|
sb.WriteString(renderMsg(m))
|
||||||
|
sb.WriteString("\n")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return sb.String()
|
||||||
|
}
|
||||||
|
|
||||||
|
olderText := joinBlocks(olderBlocks)
|
||||||
|
recentText := joinBlocks(recentBlocks)
|
||||||
|
|
||||||
|
msg, err := s.summarizer.Invoke(ctx, map[string]any{
|
||||||
|
"system_prompt": joinBlocks([]block{systemBlock}),
|
||||||
|
"user_messages": joinBlocks([]block{userBlock}),
|
||||||
|
"previous_summary": joinBlocks([]block{summaryBlock}),
|
||||||
|
"older_messages": olderText,
|
||||||
|
"recent_messages": recentText,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("summarize failed, err=%w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
summaryMsg := schema.AssistantMessage(msg.Content, nil)
|
||||||
|
msg.Name = "summary"
|
||||||
|
summaryMsg.Extra = map[string]any{
|
||||||
|
summaryMessageFlag: true,
|
||||||
|
}
|
||||||
|
|
||||||
|
// Build new state: prepend summary message, keep recent messages
|
||||||
|
newMessages := make([]*schema.Message, 0, len(messages))
|
||||||
|
newMessages = append(newMessages, systemBlock.msgs...)
|
||||||
|
newMessages = append(newMessages, userBlock.msgs...)
|
||||||
|
newMessages = append(newMessages, summaryMsg)
|
||||||
|
for _, b := range recentBlocks {
|
||||||
|
newMessages = append(newMessages, b.msgs...)
|
||||||
|
}
|
||||||
|
|
||||||
|
state.Messages = newMessages
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Render messages into strings
|
||||||
|
func renderMsg(m *schema.Message) string {
|
||||||
|
if m == nil {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
var sb strings.Builder
|
||||||
|
if m.Role == schema.Tool {
|
||||||
|
if m.ToolName != "" {
|
||||||
|
sb.WriteString("[tool:")
|
||||||
|
sb.WriteString(m.ToolName)
|
||||||
|
sb.WriteString("]\n")
|
||||||
|
} else {
|
||||||
|
sb.WriteString("[tool]\n")
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
sb.WriteString("[")
|
||||||
|
sb.WriteString(string(m.Role))
|
||||||
|
sb.WriteString("]\n")
|
||||||
|
}
|
||||||
|
if m.Content != "" {
|
||||||
|
sb.WriteString(m.Content)
|
||||||
|
sb.WriteString("\n")
|
||||||
|
}
|
||||||
|
if m.Role == schema.Assistant && len(m.ToolCalls) > 0 {
|
||||||
|
for _, tc := range m.ToolCalls {
|
||||||
|
if tc.Function.Name != "" {
|
||||||
|
sb.WriteString("tool_call: ")
|
||||||
|
sb.WriteString(tc.Function.Name)
|
||||||
|
sb.WriteString("\n")
|
||||||
|
}
|
||||||
|
if tc.Function.Arguments != "" {
|
||||||
|
sb.WriteString("args: ")
|
||||||
|
sb.WriteString(tc.Function.Arguments)
|
||||||
|
sb.WriteString("\n")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, part := range m.UserInputMultiContent {
|
||||||
|
if part.Type == schema.ChatMessagePartTypeText && part.Text != "" {
|
||||||
|
sb.WriteString(part.Text)
|
||||||
|
sb.WriteString("\n")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, part := range m.AssistantGenMultiContent {
|
||||||
|
if part.Type == schema.ChatMessagePartTypeText && part.Text != "" {
|
||||||
|
sb.WriteString(part.Text)
|
||||||
|
sb.WriteString("\n")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return sb.String()
|
||||||
|
}
|
||||||
|
|
||||||
|
func defaultCounterToken(ctx context.Context, msgs []adk.Message) (tokenNum []int64, err error) {
|
||||||
|
encoding := "cl100k_base"
|
||||||
|
tkt, err := tiktoken.GetEncoding(encoding)
|
||||||
|
if err != nil {
|
||||||
|
return nil, fmt.Errorf("get encoding failed, encoding=%v, err=%w", encoding, err)
|
||||||
|
}
|
||||||
|
tokenNum = make([]int64, len(msgs))
|
||||||
|
|
||||||
|
for i, m := range msgs {
|
||||||
|
if m == nil {
|
||||||
|
tokenNum[i] = 0
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
var sb strings.Builder
|
||||||
|
|
||||||
|
// Message role contributes to chat tokenization overhead; include it as text.
|
||||||
|
if m.Role != "" {
|
||||||
|
sb.WriteString(string(m.Role))
|
||||||
|
sb.WriteString("\n")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Core text content
|
||||||
|
if m.Content != "" {
|
||||||
|
sb.WriteString(m.Content)
|
||||||
|
sb.WriteString("\n")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Reasoning content if present
|
||||||
|
// Reasoning Content is not used by model
|
||||||
|
// if m.ReasoningContent != "" {
|
||||||
|
// sb.WriteString(m.ReasoningContent)
|
||||||
|
// sb.WriteString("\n")
|
||||||
|
// }
|
||||||
|
|
||||||
|
// Multi modal input/output text parts
|
||||||
|
for _, part := range m.UserInputMultiContent {
|
||||||
|
if part.Type == schema.ChatMessagePartTypeText && part.Text != "" {
|
||||||
|
sb.WriteString(part.Text)
|
||||||
|
sb.WriteString("\n")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, part := range m.AssistantGenMultiContent {
|
||||||
|
if part.Type == schema.ChatMessagePartTypeText && part.Text != "" {
|
||||||
|
sb.WriteString(part.Text)
|
||||||
|
sb.WriteString("\n")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Tool call textual context (name + arguments)
|
||||||
|
for _, tc := range m.ToolCalls {
|
||||||
|
if tc.Function.Name != "" {
|
||||||
|
sb.WriteString(tc.Function.Name)
|
||||||
|
sb.WriteString("\n")
|
||||||
|
}
|
||||||
|
if tc.Function.Arguments != "" {
|
||||||
|
sb.WriteString(tc.Function.Arguments)
|
||||||
|
sb.WriteString("\n")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
text := sb.String()
|
||||||
|
if text == "" {
|
||||||
|
tokenNum[i] = 0
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
tokens := tkt.Encode(text, nil, nil)
|
||||||
|
tokenNum[i] = int64(len(tokens))
|
||||||
|
}
|
||||||
|
|
||||||
|
return tokenNum, nil
|
||||||
|
}
|
||||||
@ -0,0 +1,77 @@
|
|||||||
|
/*
|
||||||
|
* Copyright 2025 CloudWeGo Authors
|
||||||
|
*
|
||||||
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
* you may not use this file except in compliance with the License.
|
||||||
|
* You may obtain a copy of the License at
|
||||||
|
*
|
||||||
|
* http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
*
|
||||||
|
* Unless required by applicable law or agreed to in writing, software
|
||||||
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
* See the License for the specific language governing permissions and
|
||||||
|
* limitations under the License.
|
||||||
|
*/
|
||||||
|
|
||||||
|
package main
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"log"
|
||||||
|
"strconv"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"github.com/cloudwego/eino/adk"
|
||||||
|
"github.com/cloudwego/eino/components/tool"
|
||||||
|
"github.com/cloudwego/eino/components/tool/utils"
|
||||||
|
)
|
||||||
|
|
||||||
|
type RepeatSectionsInput struct {
|
||||||
|
Title string `json:"title" jsonschema_description:"Section title"`
|
||||||
|
Paragraphs []string `json:"paragraphs" jsonschema_description:"Paragraphs to be repeated"`
|
||||||
|
RepeatCount int `json:"repeat_count" jsonschema_description:"Times to repeat paragraphs, default 2"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewRepeatSectionsTool() tool.InvokableTool {
|
||||||
|
t, err := utils.InferTool(
|
||||||
|
"repeat_sections",
|
||||||
|
"Repeat given paragraphs to quickly accumulate context",
|
||||||
|
func(ctx context.Context, in *RepeatSectionsInput) (string, error) {
|
||||||
|
if in.RepeatCount <= 0 {
|
||||||
|
in.RepeatCount = 2
|
||||||
|
}
|
||||||
|
var b strings.Builder
|
||||||
|
b.WriteString("## ")
|
||||||
|
b.WriteString(in.Title)
|
||||||
|
b.WriteString("\n")
|
||||||
|
idx := 0
|
||||||
|
for i := 0; i < in.RepeatCount; i++ {
|
||||||
|
for _, sec := range in.Paragraphs {
|
||||||
|
b.WriteString(strconv.Itoa(idx + 1))
|
||||||
|
b.WriteString(". ")
|
||||||
|
b.WriteString(sec)
|
||||||
|
b.WriteString("\n")
|
||||||
|
idx++
|
||||||
|
}
|
||||||
|
}
|
||||||
|
callCount := 0
|
||||||
|
times, ok := adk.GetSessionValue(ctx, "_tool_call_count")
|
||||||
|
if !ok {
|
||||||
|
times = 0
|
||||||
|
} else {
|
||||||
|
callCount = times.(int)
|
||||||
|
}
|
||||||
|
callCount++
|
||||||
|
adk.AddSessionValue(ctx, "_tool_call_count", callCount)
|
||||||
|
|
||||||
|
b.WriteString(fmt.Sprintf("Tool calls so far: %d", callCount))
|
||||||
|
return b.String(), nil
|
||||||
|
},
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
log.Fatalf("failed to create repeat_sections tool: %v", err)
|
||||||
|
}
|
||||||
|
return t
|
||||||
|
}
|
||||||
Loading…
Reference in New Issue