Enhance history management in assistant with improved message filtering and conversion
- Introduce a new method to control the maximum number of history messages loaded, prioritizing user-defined options over store settings. - Update the history loading logic to filter out non-semantic message types and convert tool call and action messages into historical summaries for better context. - Refactor the message conversion functions to handle different formats and ensure clarity in the historical context provided to the LLM. - Add comprehensive tests to validate the new behavior and ensure accurate message handling in various scenarios.
This commit is contained in:
Max
parent
a6d91c866d
commit
4d7e8e8df6
3 changed files with 381 additions and 26 deletions
|
|
@ -4,6 +4,7 @@ import (
|
|||
"fmt"
|
||||
"reflect"
|
||||
|
||||
jsoniter "github.com/json-iterator/go"
|
||||
agentcontext "github.com/yaoapp/yao/agent/context"
|
||||
storetypes "github.com/yaoapp/yao/agent/store/types"
|
||||
"github.com/yaoapp/yao/trace/types"
|
||||
|
|
@ -19,6 +20,18 @@ type HistoryResult struct {
|
|||
FullMessages []agentcontext.Message // Full messages (history + clean input)
|
||||
}
|
||||
|
||||
// getHistorySize returns the history size with priority: opts.HistorySize > storeSetting.MaxSize > default (20)
|
||||
func getHistorySize(opts *agentcontext.Options) int {
|
||||
const defaultHistorySize = 20
|
||||
if opts != nil && opts.HistorySize > 0 {
|
||||
return opts.HistorySize
|
||||
}
|
||||
if setting := GetStoreSetting(); setting != nil && setting.MaxSize > 0 {
|
||||
return setting.MaxSize
|
||||
}
|
||||
return defaultHistorySize
|
||||
}
|
||||
|
||||
// WithHistory merges the input messages with chat history and traces it
|
||||
// Returns HistoryResult containing:
|
||||
// - InputMessages: cleaned input (overlap removed)
|
||||
|
|
@ -41,14 +54,11 @@ func (ast *Assistant) WithHistory(ctx *agentcontext.Context, input []agentcontex
|
|||
return result, nil
|
||||
}
|
||||
|
||||
// Get MaxSize from store setting
|
||||
maxSize := 20 // default
|
||||
if storeSetting := GetStoreSetting(); storeSetting != nil && storeSetting.MaxSize > 0 {
|
||||
maxSize = storeSetting.MaxSize
|
||||
}
|
||||
// Resolve history size: opts.HistorySize > storeSetting.MaxSize > default (20)
|
||||
maxSize := getHistorySize(opts)
|
||||
|
||||
// Load history from store
|
||||
historyMessages, err := ast.loadHistory(ctx)
|
||||
historyMessages, err := ast.loadHistory(ctx, maxSize)
|
||||
if err != nil {
|
||||
// Log warning but continue without history
|
||||
ctx.Logger.Warn("Failed to load history for chat=%s: %v", ctx.ChatID, err)
|
||||
|
|
@ -102,8 +112,8 @@ func (ast *Assistant) WithHistory(ctx *agentcontext.Context, input []agentcontex
|
|||
}
|
||||
|
||||
// loadHistory loads chat history from the store
|
||||
// Returns the most recent MaxSize messages, ordered by time (oldest first)
|
||||
func (ast *Assistant) loadHistory(ctx *agentcontext.Context) ([]agentcontext.Message, error) {
|
||||
// Returns the most recent maxSize messages, ordered by time (oldest first)
|
||||
func (ast *Assistant) loadHistory(ctx *agentcontext.Context, maxSize int) ([]agentcontext.Message, error) {
|
||||
// Check if chat ID is available
|
||||
if ctx.ChatID == "" {
|
||||
return nil, nil
|
||||
|
|
@ -115,13 +125,6 @@ func (ast *Assistant) loadHistory(ctx *agentcontext.Context) ([]agentcontext.Mes
|
|||
return nil, nil
|
||||
}
|
||||
|
||||
// Get store setting for MaxSize
|
||||
setting := GetStoreSetting()
|
||||
maxSize := 20 // default
|
||||
if setting != nil && setting.MaxSize > 0 {
|
||||
maxSize = setting.MaxSize
|
||||
}
|
||||
|
||||
// Load messages from store with limit
|
||||
filter := storetypes.MessageFilter{
|
||||
Limit: maxSize,
|
||||
|
|
@ -161,11 +164,16 @@ func (ast *Assistant) convertStoreMessageToContext(msg *storetypes.Message) *age
|
|||
return nil
|
||||
}
|
||||
|
||||
// Skip internal message types that should not be included in LLM context
|
||||
// These types are for UI/internal use only and can confuse the LLM
|
||||
// Note: "error" is kept so LLM can help troubleshoot issues
|
||||
// Handle special message types:
|
||||
// - tool_call/action: convert to historical summary text for LLM context
|
||||
// - loading/event: skip (pure UI/lifecycle signals, no semantic value)
|
||||
// - error: kept as-is so LLM can help troubleshoot issues
|
||||
switch msg.Type {
|
||||
case "tool_call", "loading", "action", "event":
|
||||
case "tool_call":
|
||||
return ast.convertToolCallToContext(msg)
|
||||
case "action":
|
||||
return ast.convertActionToContext(msg)
|
||||
case "loading", "event":
|
||||
return nil
|
||||
}
|
||||
|
||||
|
|
@ -224,6 +232,137 @@ func (ast *Assistant) extractContentFromProps(props map[string]interface{}, msgT
|
|||
return nil
|
||||
}
|
||||
|
||||
// convertToolCallToContext converts a tool_call store message to a historical summary text message.
|
||||
// This allows the LLM to understand what tools were previously called without re-invoking them.
|
||||
//
|
||||
// Supports two Props formats:
|
||||
// - Standard ToolCallProps: {"name": "tool_name", "arguments": "{...}"}
|
||||
// - Raw stream chunks: {"content": "[{\"index\":0,\"id\":\"call_...\",\"function\":{\"name\":\"tool\"}}][...]"}
|
||||
func (ast *Assistant) convertToolCallToContext(msg *storetypes.Message) *agentcontext.Message {
|
||||
if msg.Props == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Try standard ToolCallProps format first
|
||||
if name, ok := msg.Props["name"].(string); ok && name != "" {
|
||||
args, _ := msg.Props["arguments"].(string)
|
||||
const maxArgsLen = 500
|
||||
if len(args) > maxArgsLen {
|
||||
args = args[:maxArgsLen] + "..."
|
||||
}
|
||||
return &agentcontext.Message{
|
||||
Role: agentcontext.RoleAssistant,
|
||||
Content: fmt.Sprintf("[Historical Tool Call Summary] Called tool \"%s\" with arguments: %s", name, args),
|
||||
}
|
||||
}
|
||||
|
||||
// Try raw stream chunk format: {"content": "[...][...]..."}
|
||||
// Each chunk is a JSON array like [{"index":0,"id":"call_...","function":{"name":"echo__ping"}}]
|
||||
// Subsequent chunks append arguments: [{"index":0,"function":{"arguments":"..."}}]
|
||||
if raw, ok := msg.Props["content"].(string); ok && raw != "" {
|
||||
name, args := parseToolCallRawChunks(raw)
|
||||
if name == "" {
|
||||
return nil
|
||||
}
|
||||
const maxArgsLen = 500
|
||||
if len(args) > maxArgsLen {
|
||||
args = args[:maxArgsLen] + "..."
|
||||
}
|
||||
return &agentcontext.Message{
|
||||
Role: agentcontext.RoleAssistant,
|
||||
Content: fmt.Sprintf("[Historical Tool Call Summary] Called tool \"%s\" with arguments: %s", name, args),
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// parseToolCallRawChunks parses concatenated raw stream chunks to extract tool name and arguments.
|
||||
// Input format: "[{...}][{...}][{...}]" — multiple JSON arrays concatenated without separator.
|
||||
func parseToolCallRawChunks(raw string) (name, args string) {
|
||||
// Split concatenated JSON arrays: "][" is the boundary
|
||||
// e.g. "[{...}][{...}]" → ["[{...}]", "[{...}]"]
|
||||
chunks := splitJSONArrays(raw)
|
||||
|
||||
var argParts []string
|
||||
for _, chunk := range chunks {
|
||||
var items []map[string]interface{}
|
||||
if err := jsoniter.UnmarshalFromString(chunk, &items); err != nil || len(items) == 0 {
|
||||
continue
|
||||
}
|
||||
item := items[0]
|
||||
if fn, ok := item["function"].(map[string]interface{}); ok {
|
||||
if n, ok := fn["name"].(string); ok && n != "" && name == "" {
|
||||
name = n
|
||||
}
|
||||
if a, ok := fn["arguments"].(string); ok && a != "" {
|
||||
argParts = append(argParts, a)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
args = ""
|
||||
for _, part := range argParts {
|
||||
args += part
|
||||
}
|
||||
return name, args
|
||||
}
|
||||
|
||||
// splitJSONArrays splits a string of concatenated JSON arrays "[...][...][...]" into individual arrays.
|
||||
func splitJSONArrays(s string) []string {
|
||||
var result []string
|
||||
depth := 0
|
||||
start := -1
|
||||
for i, ch := range s {
|
||||
switch ch {
|
||||
case '[':
|
||||
if depth == 0 {
|
||||
start = i
|
||||
}
|
||||
depth++
|
||||
case ']':
|
||||
depth--
|
||||
if depth == 0 && start >= 0 {
|
||||
result = append(result, s[start:i+1])
|
||||
start = -1
|
||||
}
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// convertActionToContext converts an action store message to a historical summary text message.
|
||||
// This allows the LLM to understand what system actions were previously executed.
|
||||
func (ast *Assistant) convertActionToContext(msg *storetypes.Message) *agentcontext.Message {
|
||||
if msg.Props == nil {
|
||||
return nil
|
||||
}
|
||||
name, _ := msg.Props["name"].(string)
|
||||
if name == "" {
|
||||
return nil
|
||||
}
|
||||
payload := ""
|
||||
if msg.Props["payload"] != nil {
|
||||
if payloadStr, err := jsoniter.MarshalToString(msg.Props["payload"]); err == nil {
|
||||
const maxPayloadLen = 500
|
||||
if len(payloadStr) > maxPayloadLen {
|
||||
payloadStr = payloadStr[:maxPayloadLen] + "..."
|
||||
}
|
||||
payload = payloadStr
|
||||
}
|
||||
}
|
||||
if payload != "" {
|
||||
return &agentcontext.Message{
|
||||
Role: agentcontext.RoleAssistant,
|
||||
Content: fmt.Sprintf("[Historical Action Summary] Executed action \"%s\" with payload: %s", name, payload),
|
||||
}
|
||||
}
|
||||
return &agentcontext.Message{
|
||||
Role: agentcontext.RoleAssistant,
|
||||
Content: fmt.Sprintf("[Historical Action Summary] Executed action \"%s\"", name),
|
||||
}
|
||||
}
|
||||
|
||||
// findOverlapIndex finds the index in input where history messages end
|
||||
// Returns the number of input messages that overlap with history
|
||||
func (ast *Assistant) findOverlapIndex(history, input []agentcontext.Message) int {
|
||||
|
|
|
|||
|
|
@ -415,7 +415,10 @@ func TestHistoryLoading(t *testing.T) {
|
|||
}()
|
||||
|
||||
// Add various message types (only user/assistant roles allowed by DB constraint)
|
||||
// loadHistory filters by role (user/assistant only) and converts based on type
|
||||
// loadHistory filters by role (user/assistant only) and converts based on type:
|
||||
// - loading/event: skipped (no semantic value)
|
||||
// - tool_call/action: converted to historical summary text
|
||||
// - text/user_input/error: kept as-is
|
||||
err = chatStore.SaveMessages(chatID, []*storetypes.Message{
|
||||
{
|
||||
MessageID: fmt.Sprintf("filter_1_%s", reqID),
|
||||
|
|
@ -461,11 +464,9 @@ func TestHistoryLoading(t *testing.T) {
|
|||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
|
||||
// User and assistant roles are included
|
||||
// loading type messages with role=assistant are included (role-based filtering)
|
||||
// History contains: 1 user + 2 assistant (loading + text) = 3 messages
|
||||
// Plus 1 new input = 4 total
|
||||
assert.GreaterOrEqual(t, len(result.FullMessages), 3) // At least 1 user + 1 assistant from history + 1 new
|
||||
// loading type is skipped (no semantic value)
|
||||
// History: user_input + text = 2 messages; plus 1 new input = 3 total
|
||||
assert.Len(t, result.FullMessages, 3)
|
||||
|
||||
// Verify only user and assistant roles
|
||||
for _, msg := range result.FullMessages {
|
||||
|
|
@ -473,7 +474,217 @@ func TestHistoryLoading(t *testing.T) {
|
|||
"Expected user or assistant role, got: %s", msg.Role)
|
||||
}
|
||||
|
||||
t.Log("✓ Only user and assistant roles included in history")
|
||||
t.Log("✓ Loading type filtered, user/assistant roles kept")
|
||||
})
|
||||
|
||||
t.Run("ToolCallConvertedToSummary", func(t *testing.T) {
|
||||
chatID := fmt.Sprintf("test_toolcall_%s", uuid.New().String()[:8])
|
||||
ctx := newHistoryTestContext(chatID)
|
||||
reqID := uuid.New().String()[:8]
|
||||
|
||||
err := chatStore.CreateChat(&storetypes.Chat{
|
||||
ChatID: chatID,
|
||||
AssistantID: ast.ID,
|
||||
Status: "active",
|
||||
Share: "private",
|
||||
CreatedAt: time.Now(),
|
||||
UpdatedAt: time.Now(),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
defer func() {
|
||||
chatStore.DeleteMessages(chatID, nil)
|
||||
chatStore.DeleteChat(chatID)
|
||||
}()
|
||||
|
||||
// Add tool_call messages in both formats
|
||||
err = chatStore.SaveMessages(chatID, []*storetypes.Message{
|
||||
{
|
||||
MessageID: fmt.Sprintf("tc_1_%s", reqID),
|
||||
ChatID: chatID,
|
||||
RequestID: fmt.Sprintf("req_tc_%s", reqID),
|
||||
Role: "user",
|
||||
Type: "user_input",
|
||||
Props: map[string]interface{}{"content": "echo 3 ping 4"},
|
||||
Sequence: 1,
|
||||
AssistantID: ast.ID,
|
||||
CreatedAt: time.Now().Add(-3 * time.Minute),
|
||||
},
|
||||
// Raw stream chunk format (actual DB format)
|
||||
{
|
||||
MessageID: fmt.Sprintf("tc_2_%s", reqID),
|
||||
ChatID: chatID,
|
||||
RequestID: fmt.Sprintf("req_tc_%s", reqID),
|
||||
Role: "assistant",
|
||||
Type: "tool_call",
|
||||
Props: map[string]interface{}{"content": `[{"index":0,"id":"call_abc","type":"function","function":{"name":"echo__ping"}}][{"index":0,"function":{"arguments":"{\"count\":3}"}}]`},
|
||||
Sequence: 2,
|
||||
AssistantID: ast.ID,
|
||||
CreatedAt: time.Now().Add(-2 * time.Minute),
|
||||
},
|
||||
// Standard ToolCallProps format
|
||||
{
|
||||
MessageID: fmt.Sprintf("tc_3_%s", reqID),
|
||||
ChatID: chatID,
|
||||
RequestID: fmt.Sprintf("req_tc_%s", reqID),
|
||||
Role: "assistant",
|
||||
Type: "tool_call",
|
||||
Props: map[string]interface{}{"name": "echo__echo", "arguments": `{"message":"hello"}`},
|
||||
Sequence: 3,
|
||||
AssistantID: ast.ID,
|
||||
CreatedAt: time.Now().Add(-1 * time.Minute),
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
input := []agentcontext.Message{
|
||||
{Role: agentcontext.RoleUser, Content: "echo 5 ping 6"},
|
||||
}
|
||||
|
||||
result, err := ast.WithHistory(ctx, input, nil)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
|
||||
// 1 user_input + 2 tool_call summaries + 1 new input = 4
|
||||
assert.Len(t, result.FullMessages, 4)
|
||||
|
||||
// Verify tool_call messages are converted to summary text
|
||||
tcMsg1 := result.FullMessages[1]
|
||||
assert.Equal(t, agentcontext.RoleAssistant, tcMsg1.Role)
|
||||
assert.Contains(t, tcMsg1.Content, "[Historical Tool Call Summary]")
|
||||
assert.Contains(t, tcMsg1.Content, "echo__ping")
|
||||
assert.Contains(t, tcMsg1.Content, `{"count":3}`)
|
||||
|
||||
tcMsg2 := result.FullMessages[2]
|
||||
assert.Equal(t, agentcontext.RoleAssistant, tcMsg2.Role)
|
||||
assert.Contains(t, tcMsg2.Content, "[Historical Tool Call Summary]")
|
||||
assert.Contains(t, tcMsg2.Content, "echo__echo")
|
||||
assert.Contains(t, tcMsg2.Content, `{"message":"hello"}`)
|
||||
|
||||
t.Log("✓ Tool call messages converted to historical summaries")
|
||||
})
|
||||
|
||||
t.Run("ActionConvertedToSummary", func(t *testing.T) {
|
||||
chatID := fmt.Sprintf("test_action_%s", uuid.New().String()[:8])
|
||||
ctx := newHistoryTestContext(chatID)
|
||||
reqID := uuid.New().String()[:8]
|
||||
|
||||
err := chatStore.CreateChat(&storetypes.Chat{
|
||||
ChatID: chatID,
|
||||
AssistantID: ast.ID,
|
||||
Status: "active",
|
||||
Share: "private",
|
||||
CreatedAt: time.Now(),
|
||||
UpdatedAt: time.Now(),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
defer func() {
|
||||
chatStore.DeleteMessages(chatID, nil)
|
||||
chatStore.DeleteChat(chatID)
|
||||
}()
|
||||
|
||||
err = chatStore.SaveMessages(chatID, []*storetypes.Message{
|
||||
{
|
||||
MessageID: fmt.Sprintf("act_1_%s", reqID),
|
||||
ChatID: chatID,
|
||||
RequestID: fmt.Sprintf("req_act_%s", reqID),
|
||||
Role: "user",
|
||||
Type: "user_input",
|
||||
Props: map[string]interface{}{"content": "Do something"},
|
||||
Sequence: 1,
|
||||
AssistantID: ast.ID,
|
||||
CreatedAt: time.Now().Add(-2 * time.Minute),
|
||||
},
|
||||
// Action with payload
|
||||
{
|
||||
MessageID: fmt.Sprintf("act_2_%s", reqID),
|
||||
ChatID: chatID,
|
||||
RequestID: fmt.Sprintf("req_act_%s", reqID),
|
||||
Role: "assistant",
|
||||
Type: "action",
|
||||
Props: map[string]interface{}{
|
||||
"name": "robot.execute",
|
||||
"payload": map[string]interface{}{"goals": "test goal", "robot_id": "12345"},
|
||||
},
|
||||
Sequence: 2,
|
||||
AssistantID: ast.ID,
|
||||
CreatedAt: time.Now().Add(-1 * time.Minute),
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
input := []agentcontext.Message{
|
||||
{Role: agentcontext.RoleUser, Content: "What happened?"},
|
||||
}
|
||||
|
||||
result, err := ast.WithHistory(ctx, input, nil)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
|
||||
// 1 user_input + 1 action summary + 1 new input = 3
|
||||
assert.Len(t, result.FullMessages, 3)
|
||||
|
||||
actMsg := result.FullMessages[1]
|
||||
assert.Equal(t, agentcontext.RoleAssistant, actMsg.Role)
|
||||
assert.Contains(t, actMsg.Content, "[Historical Action Summary]")
|
||||
assert.Contains(t, actMsg.Content, "robot.execute")
|
||||
assert.Contains(t, actMsg.Content, "test goal")
|
||||
assert.Contains(t, actMsg.Content, "12345")
|
||||
|
||||
t.Log("✓ Action messages converted to historical summaries with payload")
|
||||
})
|
||||
|
||||
t.Run("ActionWithoutPayload", func(t *testing.T) {
|
||||
chatID := fmt.Sprintf("test_action_nopay_%s", uuid.New().String()[:8])
|
||||
ctx := newHistoryTestContext(chatID)
|
||||
reqID := uuid.New().String()[:8]
|
||||
|
||||
err := chatStore.CreateChat(&storetypes.Chat{
|
||||
ChatID: chatID,
|
||||
AssistantID: ast.ID,
|
||||
Status: "active",
|
||||
Share: "private",
|
||||
CreatedAt: time.Now(),
|
||||
UpdatedAt: time.Now(),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
defer func() {
|
||||
chatStore.DeleteMessages(chatID, nil)
|
||||
chatStore.DeleteChat(chatID)
|
||||
}()
|
||||
|
||||
err = chatStore.SaveMessages(chatID, []*storetypes.Message{
|
||||
{
|
||||
MessageID: fmt.Sprintf("actnp_1_%s", reqID),
|
||||
ChatID: chatID,
|
||||
RequestID: fmt.Sprintf("req_actnp_%s", reqID),
|
||||
Role: "assistant",
|
||||
Type: "action",
|
||||
Props: map[string]interface{}{"name": "navigate"},
|
||||
Sequence: 1,
|
||||
AssistantID: ast.ID,
|
||||
CreatedAt: time.Now(),
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
input := []agentcontext.Message{
|
||||
{Role: agentcontext.RoleUser, Content: "What happened?"},
|
||||
}
|
||||
|
||||
result, err := ast.WithHistory(ctx, input, nil)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
|
||||
// 1 action summary + 1 new input = 2
|
||||
assert.Len(t, result.FullMessages, 2)
|
||||
|
||||
actMsg := result.FullMessages[0]
|
||||
assert.Equal(t, agentcontext.RoleAssistant, actMsg.Role)
|
||||
assert.Contains(t, actMsg.Content, "[Historical Action Summary]")
|
||||
assert.Contains(t, actMsg.Content, "navigate")
|
||||
assert.NotContains(t, actMsg.Content, "payload")
|
||||
|
||||
t.Log("✓ Action without payload handled correctly")
|
||||
})
|
||||
|
||||
t.Run("ContentExtraction", func(t *testing.T) {
|
||||
|
|
|
|||
|
|
@ -319,6 +319,11 @@ type Options struct {
|
|||
// Metadata for passing custom data to hooks (e.g., scenario selection)
|
||||
Metadata map[string]any `json:"metadata,omitempty"` // Custom metadata passed to Create/Next hooks
|
||||
|
||||
// HistorySize controls the max number of history messages loaded for LLM context.
|
||||
// Priority: HistorySize > StoreSetting.MaxSize > default (20)
|
||||
// 0 means use StoreSetting or default.
|
||||
HistorySize int `json:"history_size,omitempty"`
|
||||
|
||||
// OnMessage is called for each message sent via ctx.Send()
|
||||
// Used by ctx.agent.Call with onChunk callback to receive SSE messages
|
||||
// Returns: 0 = continue, non-zero = stop
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue