opencode/yao
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add assistant info retrieval and chat mode management

Browse source 
- Introduced `GetInfo` method in the Assistant struct to return basic assistant information with optional locale support for internationalization.
- Implemented `GetInfoByIDs` function to retrieve information for multiple assistants based on their IDs, enhancing batch processing capabilities.
- Updated chat buffer management to include a mode parameter, allowing for dynamic switching between chat and task modes.
- Enhanced message handling to store and retrieve the mode associated with each message, improving context tracking during chat sessions.
- Revised tests to validate the new functionalities, ensuring accurate retrieval of assistant information and proper mode management in chat operations.
This commit is contained in:
Max 2025-12-10 12:15:55 +08:00
parent 2e19ccc914
commit 24ad563070
18 changed files with 842 additions and 283 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

48
agent/assistant/assistant.go
View file

@ -349,6 +349,54 @@ func (ast *Assistant) Clone() *Assistant {
return clone
}
// GetInfo returns the basic info of the assistant with optional locale
func (ast *Assistant) GetInfo(locale ...string) *store.AssistantInfo {
if ast == nil {
return nil
}
loc := ""
if len(locale) > 0 {
loc = locale[0]
}
info := &store.AssistantInfo{
AssistantID: ast.ID,
Avatar: ast.Avatar,
}
// Apply i18n translation if locale is provided
if loc != "" {
info.Name = ast.GetName(loc)
info.Description = ast.GetDescription(loc)
} else {
info.Name = ast.Name
info.Description = ast.Description
}
return info
}
// GetInfoByIDs retrieves basic info for multiple assistants by their IDs
// Returns a map of assistant_id -> AssistantInfo
func GetInfoByIDs(ids []string, locale ...string) map[string]*store.AssistantInfo {
result := make(map[string]*store.AssistantInfo)
if len(ids) == 0 {
return result
}
for _, id := range ids {
ast, err := Get(id)
if err != nil || ast == nil {
continue
}
result[id] = ast.GetInfo(locale...)
}
return result
}
// Update updates the assistant properties
func (ast *Assistant) Update(data map[string]interface{}) error {
if ast == nil {

25
agent/assistant/chat.go
View file

@ -243,14 +243,16 @@ func (ast *Assistant) InitBuffer(ctx *agentcontext.Context) {
requestID = uuid.New().String()
}
// Get connector from options
// Get connector and mode from options
connector := ""
mode := ""
if ctx.Stack.Options != nil {
connector = ctx.Stack.Options.Connector
mode = ctx.Stack.Options.Mode
}
ctx.Buffer = agentcontext.NewChatBuffer(ctx.ChatID, requestID, ast.ID, connector)
log.Trace("[CHAT] Buffer initialized: chatID=%s, requestID=%s, assistantID=%s, connector=%s", ctx.ChatID, requestID, ast.ID, connector)
ctx.Buffer = agentcontext.NewChatBuffer(ctx.ChatID, requestID, ast.ID, connector, mode)
log.Trace("[CHAT] Buffer initialized: chatID=%s, requestID=%s, assistantID=%s, connector=%s, mode=%s", ctx.ChatID, requestID, ast.ID, connector, mode)
}
// BufferUserInput adds user input messages to the buffer
@ -341,7 +343,7 @@ func (ast *Assistant) FlushBuffer(ctx *agentcontext.Context, finalStatus string,
}
}
// 2. Update chat last_message_at and last_connector
// 2. Update chat last_message_at, last_connector, and last_mode
if len(messages) > 0 {
now := time.Now()
updates := map[string]interface{}{
@ -351,6 +353,10 @@ func (ast *Assistant) FlushBuffer(ctx *agentcontext.Context, finalStatus string,
if connector := ctx.Buffer.Connector(); connector != "" {
updates["last_connector"] = connector
}
// Also update last_mode if available
if mode := ctx.Buffer.Mode(); mode != "" {
updates["last_mode"] = mode
}
if updateErr := chatStore.UpdateChat(ctx.ChatID, updates); updateErr != nil {
log.Trace("[CHAT] Failed to update chat: %v", updateErr)
}
@ -388,6 +394,7 @@ func (ast *Assistant) convertBufferedMessages(buffered []*agentcontext.BufferedM
ThreadID: msg.ThreadID,
AssistantID: msg.AssistantID,
Connector: msg.Connector,
Mode: msg.Mode,
Sequence: msg.Sequence,
Metadata: msg.Metadata,
CreatedAt: msg.CreatedAt,
@ -454,7 +461,6 @@ func (ast *Assistant) EnsureChat(ctx *agentcontext.Context) error {
chat := &storetypes.Chat{
ChatID: ctx.ChatID,
AssistantID: ast.ID,
Mode: "chat",
Status: "active",
Share: "private",
Sort: 0,
@ -487,6 +493,15 @@ func GetChatStore() storetypes.ChatStore {
return storage
}
// GetStore returns the full store instance (implements both ChatStore and AssistantStore)
// Returns nil if storage is not configured
func GetStore() storetypes.Store {
if storage == nil {
return nil
}
return storage
}
// =============================================================================
// Deprecated methods (kept for compatibility)
// =============================================================================

69
agent/assistant/chat_test.go
View file

@ -571,7 +571,6 @@ func TestFlushBuffer(t *testing.T) {
err := chatStore.CreateChat(&storetypes.Chat{
ChatID: chatID,
AssistantID: ast.ID,
Mode: "chat",
Status: "active",
Share: "private",
CreatedAt: time.Now(),
@ -620,7 +619,6 @@ func TestFlushBuffer(t *testing.T) {
err := chatStore.CreateChat(&storetypes.Chat{
ChatID: chatID,
AssistantID: ast.ID,
Mode: "chat",
Status: "active",
Share: "private",
CreatedAt: time.Now(),
@ -670,7 +668,6 @@ func TestFlushBuffer(t *testing.T) {
err := chatStore.CreateChat(&storetypes.Chat{
ChatID: chatID,
AssistantID: ast.ID,
Mode: "chat",
Status: "active",
Share: "private",
CreatedAt: time.Now(),
@ -696,6 +693,71 @@ func TestFlushBuffer(t *testing.T) {
chatStore.DeleteChat(chatID)
t.Logf("✓ Buffer flushed on interrupt: resume records saved with interrupted status")
})
t.Run("FlushWithModeAndConnector", func(t *testing.T) {
chatID := fmt.Sprintf("test_flush_mode_%s", uuid.New().String()[:8])
ctx := agentcontext.New(context.Background(), nil, chatID)
ctx.Space = plan.NewMemorySharedSpace()
// Enter stack with connector and mode options
opts := &agentcontext.Options{
Connector: "deepseek.v3",
Mode: "task",
}
_, _, done := agentcontext.EnterStack(ctx, ast.ID, opts)
defer done()
ast.InitBuffer(ctx)
// Verify buffer has correct connector and mode
require.NotNil(t, ctx.Buffer, "Buffer should be initialized")
assert.Equal(t, "deepseek.v3", ctx.Buffer.Connector(), "Buffer should have connector set")
assert.Equal(t, "task", ctx.Buffer.Mode(), "Buffer should have mode set")
// Ensure chat exists
err := chatStore.CreateChat(&storetypes.Chat{
ChatID: chatID,
AssistantID: ast.ID,
Status: "active",
Share: "private",
CreatedAt: time.Now(),
UpdatedAt: time.Now(),
})
require.NoError(t, err)
// Add some messages to buffer
ctx.Buffer.AddUserInput("Test question for mode", "")
ctx.Buffer.AddAssistantMessage("M1", "text", map[string]interface{}{"content": "Test answer with mode"}, "", "", ast.ID, nil)
// Flush buffer
ast.FlushBuffer(ctx, agentcontext.StepStatusCompleted, nil)
// Verify messages were saved with connector and mode
messages, err := chatStore.GetMessages(chatID, storetypes.MessageFilter{})
assert.NoError(t, err)
assert.Len(t, messages, 2, "Should have 2 messages saved")
// Assistant message should have connector and mode
var assistantMsg *storetypes.Message
for _, msg := range messages {
if msg.Role == "assistant" {
assistantMsg = msg
break
}
}
require.NotNil(t, assistantMsg, "Should find assistant message")
assert.Equal(t, "deepseek.v3", assistantMsg.Connector, "Message should have connector")
assert.Equal(t, "task", assistantMsg.Mode, "Message should have mode")
// Verify chat was updated with last_connector and last_mode
chat, err := chatStore.GetChat(chatID)
assert.NoError(t, err)
assert.Equal(t, "deepseek.v3", chat.LastConnector, "Chat should have last_connector updated")
assert.Equal(t, "task", chat.LastMode, "Chat should have last_mode updated")
// Cleanup
chatStore.DeleteChat(chatID)
t.Logf("✓ Buffer flushed with mode and connector: connector=%s, mode=%s", chat.LastConnector, chat.LastMode)
})
}
func TestEnsureChat(t *testing.T) {
@ -741,7 +803,6 @@ func TestEnsureChat(t *testing.T) {
ChatID: chatID,
AssistantID: ast.ID,
Title: "Existing Chat",
Mode: "chat",
Status: "active",
Share: "private",
CreatedAt: time.Now(),

22
agent/context/buffer.go
View file

@ -19,6 +19,7 @@ type ChatBuffer struct {
requestID string
assistantID string
connector string // Current connector ID (for data analysis)
mode string // Current chat mode (chat or task)
// Message buffer
messages []*BufferedMessage
@ -47,6 +48,7 @@ type BufferedMessage struct {
ThreadID string `json:"thread_id,omitempty"`
AssistantID string `json:"assistant_id,omitempty"`
Connector string `json:"connector,omitempty"` // Connector ID used for this message
Mode string `json:"mode,omitempty"` // Chat mode used for this message (chat or task)
Sequence int `json:"sequence"`
Metadata map[string]interface{} `json:"metadata,omitempty"`
CreatedAt time.Time `json:"created_at"`
@ -97,12 +99,13 @@ const (
)
// NewChatBuffer creates a new chat buffer
func NewChatBuffer(chatID, requestID, assistantID, connector string) *ChatBuffer {
func NewChatBuffer(chatID, requestID, assistantID, connector, mode string) *ChatBuffer {
return &ChatBuffer{
chatID: chatID,
requestID: requestID,
assistantID: assistantID,
connector: connector,
mode: mode,
messages: make([]*BufferedMessage, 0),
steps: make([]*BufferedStep, 0),
}
@ -135,6 +138,11 @@ func (b *ChatBuffer) AddMessage(msg *BufferedMessage) {
msg.CreatedAt = time.Now()
}
// Set mode from buffer if not provided
if msg.Mode == "" && b.mode != "" {
msg.Mode = b.mode
}
// Auto-increment sequence
b.msgSequence++
msg.Sequence = b.msgSequence
@ -447,6 +455,18 @@ func (b *ChatBuffer) SetConnector(connector string) {
b.connector = connector
}
// Mode returns the current chat mode
func (b *ChatBuffer) Mode() string {
return b.mode
}
// SetMode updates the chat mode (when user switches mode)
func (b *ChatBuffer) SetMode(mode string) {
b.mu.Lock()
defer b.mu.Unlock()
b.mode = mode
}
// =============================================================================
// Helper Functions
// =============================================================================

158
agent/context/buffer_test.go
View file

@ -17,7 +17,7 @@ import (
func TestBufferNewChatBuffer(t *testing.T) {
t.Run("CreateWithAllFields", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-123", "req-456", "assistant-789", "")
buffer := context.NewChatBuffer("chat-123", "req-456", "assistant-789", "", "")
assert.NotNil(t, buffer)
assert.Equal(t, "chat-123", buffer.ChatID())
@ -29,7 +29,7 @@ func TestBufferNewChatBuffer(t *testing.T) {
})
t.Run("CreateWithEmptyFields", func(t *testing.T) {
buffer := context.NewChatBuffer("", "", "", "")
buffer := context.NewChatBuffer("", "", "", "", "")
assert.NotNil(t, buffer)
assert.Empty(t, buffer.ChatID())
@ -44,7 +44,7 @@ func TestBufferNewChatBuffer(t *testing.T) {
func TestBufferAddMessage(t *testing.T) {
t.Run("AddSingleMessage", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "")
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "", "")
msg := &context.BufferedMessage{
Role: "assistant",
@ -65,7 +65,7 @@ func TestBufferAddMessage(t *testing.T) {
})
t.Run("AddMultipleMessages", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "")
buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "", "")
for i := 0; i < 5; i++ {
buffer.AddMessage(&context.BufferedMessage{
@ -85,14 +85,14 @@ func TestBufferAddMessage(t *testing.T) {
})
t.Run("AddNilMessage", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3", "")
buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3", "", "")
buffer.AddMessage(nil)
assert.Equal(t, 0, buffer.GetMessageCount())
})
t.Run("AddMessageWithExistingID", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-4", "req-4", "assistant-4", "")
buffer := context.NewChatBuffer("chat-4", "req-4", "assistant-4", "", "")
msg := &context.BufferedMessage{
MessageID: "custom-id-123",
@ -107,7 +107,7 @@ func TestBufferAddMessage(t *testing.T) {
})
t.Run("AddMessageWithExistingTimestamp", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-5", "req-5", "assistant-5", "")
buffer := context.NewChatBuffer("chat-5", "req-5", "assistant-5", "", "")
customTime := time.Date(2024, 1, 1, 12, 0, 0, 0, time.UTC)
msg := &context.BufferedMessage{
@ -125,7 +125,7 @@ func TestBufferAddMessage(t *testing.T) {
func TestBufferAddUserInput(t *testing.T) {
t.Run("AddStringContent", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "")
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "", "")
buffer.AddUserInput("What is the weather?", "")
messages := buffer.GetMessages()
@ -137,7 +137,7 @@ func TestBufferAddUserInput(t *testing.T) {
})
t.Run("AddUserInputWithName", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "")
buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "", "")
buffer.AddUserInput("Hello", "John")
messages := buffer.GetMessages()
@ -146,7 +146,7 @@ func TestBufferAddUserInput(t *testing.T) {
})
t.Run("AddComplexContent", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3", "")
buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3", "", "")
complexContent := []map[string]interface{}{
{"type": "text", "text": "Look at this image"},
{"type": "image_url", "image_url": map[string]string{"url": "https://example.com/image.jpg"}},
@ -163,7 +163,7 @@ func TestBufferAddUserInput(t *testing.T) {
func TestBufferAddAssistantMessage(t *testing.T) {
t.Run("AddTextMessage", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "")
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "", "")
buffer.AddAssistantMessage(
"M1",
"text",
@ -186,7 +186,7 @@ func TestBufferAddAssistantMessage(t *testing.T) {
})
t.Run("SkipEventMessage", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "")
buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "", "")
buffer.AddAssistantMessage(
"E1",
"event",
@ -199,7 +199,7 @@ func TestBufferAddAssistantMessage(t *testing.T) {
})
t.Run("AddRetrievalMessage", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3", "")
buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3", "", "")
buffer.AddAssistantMessage(
"M2",
"retrieval",
@ -218,7 +218,7 @@ func TestBufferAddAssistantMessage(t *testing.T) {
})
t.Run("AddToolCallMessage", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-4", "req-4", "assistant-4", "")
buffer := context.NewChatBuffer("chat-4", "req-4", "assistant-4", "", "")
buffer.AddAssistantMessage(
"M3",
"tool_call",
@ -236,7 +236,7 @@ func TestBufferAddAssistantMessage(t *testing.T) {
})
t.Run("AddCustomTypeMessage", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-5", "req-5", "assistant-5", "")
buffer := context.NewChatBuffer("chat-5", "req-5", "assistant-5", "", "")
buffer.AddAssistantMessage(
"M4",
"custom_chart",
@ -255,7 +255,7 @@ func TestBufferAddAssistantMessage(t *testing.T) {
func TestBufferGetMessages(t *testing.T) {
t.Run("GetMessagesReturnsSliceCopy", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "")
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "", "")
buffer.AddUserInput("Hello", "")
messages1 := buffer.GetMessages()
@ -268,7 +268,7 @@ func TestBufferGetMessages(t *testing.T) {
})
t.Run("GetEmptyMessages", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "")
buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "", "")
messages := buffer.GetMessages()
assert.NotNil(t, messages)
@ -277,7 +277,7 @@ func TestBufferGetMessages(t *testing.T) {
}
func TestBufferGetMessageCount(t *testing.T) {
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "")
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "", "")
assert.Equal(t, 0, buffer.GetMessageCount())
buffer.AddUserInput("Message 1", "")
@ -293,7 +293,7 @@ func TestBufferGetMessageCount(t *testing.T) {
func TestBufferBeginStep(t *testing.T) {
t.Run("BeginStepWithStack", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "")
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "", "")
stack := &context.Stack{
ID: "stack-123",
@ -319,7 +319,7 @@ func TestBufferBeginStep(t *testing.T) {
})
t.Run("BeginStepWithNilStack", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "")
buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "", "")
step := buffer.BeginStep(context.StepTypeInput, nil, nil)
@ -330,7 +330,7 @@ func TestBufferBeginStep(t *testing.T) {
})
t.Run("BeginMultipleSteps", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3", "")
buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3", "", "")
step1 := buffer.BeginStep(context.StepTypeInput, nil, nil)
step2 := buffer.BeginStep(context.StepTypeHookCreate, nil, nil)
@ -345,7 +345,7 @@ func TestBufferBeginStep(t *testing.T) {
})
t.Run("BeginStepWithSpaceSnapshot", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-4", "req-4", "assistant-4", "")
buffer := context.NewChatBuffer("chat-4", "req-4", "assistant-4", "", "")
// Set space snapshot before beginning step
buffer.SetSpaceSnapshot(map[string]interface{}{
@ -363,7 +363,7 @@ func TestBufferBeginStep(t *testing.T) {
func TestBufferCompleteStep(t *testing.T) {
t.Run("CompleteCurrentStep", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "")
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "", "")
buffer.BeginStep(context.StepTypeLLM, map[string]interface{}{"prompt": "Hello"}, nil)
buffer.CompleteStep(map[string]interface{}{"response": "Hi there!"})
@ -376,7 +376,7 @@ func TestBufferCompleteStep(t *testing.T) {
})
t.Run("CompleteWithNoCurrentStep", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "")
buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "", "")
// Should not panic
buffer.CompleteStep(map[string]interface{}{"response": "test"})
@ -384,7 +384,7 @@ func TestBufferCompleteStep(t *testing.T) {
})
t.Run("CompleteMultipleStepsSequentially", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3", "")
buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3", "", "")
buffer.BeginStep(context.StepTypeInput, nil, nil)
buffer.CompleteStep(map[string]interface{}{"done": true})
@ -405,7 +405,7 @@ func TestBufferCompleteStep(t *testing.T) {
func TestBufferFailCurrentStep(t *testing.T) {
t.Run("FailWithError", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "")
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "", "")
buffer.BeginStep(context.StepTypeLLM, nil, nil)
buffer.FailCurrentStep(context.ResumeStatusFailed, fmt.Errorf("API error: rate limit exceeded"))
@ -417,7 +417,7 @@ func TestBufferFailCurrentStep(t *testing.T) {
})
t.Run("FailWithInterrupted", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "")
buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "", "")
buffer.BeginStep(context.StepTypeLLM, nil, nil)
buffer.FailCurrentStep(context.ResumeStatusInterrupted, nil)
@ -429,7 +429,7 @@ func TestBufferFailCurrentStep(t *testing.T) {
})
t.Run("FailAlreadyCompletedStep", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3", "")
buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3", "", "")
buffer.BeginStep(context.StepTypeLLM, nil, nil)
buffer.CompleteStep(map[string]interface{}{"done": true})
@ -443,7 +443,7 @@ func TestBufferFailCurrentStep(t *testing.T) {
})
t.Run("FailWithNoCurrentStep", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-4", "req-4", "assistant-4", "")
buffer := context.NewChatBuffer("chat-4", "req-4", "assistant-4", "", "")
// Should not panic
buffer.FailCurrentStep(context.ResumeStatusFailed, fmt.Errorf("error"))
@ -452,12 +452,12 @@ func TestBufferFailCurrentStep(t *testing.T) {
func TestBufferGetCurrentStep(t *testing.T) {
t.Run("NoCurrentStep", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "")
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "", "")
assert.Nil(t, buffer.GetCurrentStep())
})
t.Run("HasCurrentStep", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "")
buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "", "")
buffer.BeginStep(context.StepTypeLLM, nil, nil)
current := buffer.GetCurrentStep()
@ -466,7 +466,7 @@ func TestBufferGetCurrentStep(t *testing.T) {
})
t.Run("CurrentStepClearedAfterComplete", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3", "")
buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3", "", "")
buffer.BeginStep(context.StepTypeLLM, nil, nil)
buffer.CompleteStep(nil)
@ -476,7 +476,7 @@ func TestBufferGetCurrentStep(t *testing.T) {
func TestBufferGetStepsForResume(t *testing.T) {
t.Run("CompletedSuccessfully", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "")
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "", "")
buffer.BeginStep(context.StepTypeInput, nil, nil)
buffer.CompleteStep(nil)
@ -489,7 +489,7 @@ func TestBufferGetStepsForResume(t *testing.T) {
})
t.Run("FailedRequest", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "")
buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "", "")
buffer.BeginStep(context.StepTypeInput, nil, nil)
buffer.CompleteStep(nil)
@ -505,7 +505,7 @@ func TestBufferGetStepsForResume(t *testing.T) {
})
t.Run("InterruptedRequest", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3", "")
buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3", "", "")
buffer.BeginStep(context.StepTypeInput, nil, nil)
buffer.CompleteStep(nil)
@ -523,7 +523,7 @@ func TestBufferGetStepsForResume(t *testing.T) {
func TestBufferGetAllSteps(t *testing.T) {
t.Run("GetStepsReturnsSliceCopy", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "")
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "", "")
buffer.BeginStep(context.StepTypeLLM, nil, nil)
steps1 := buffer.GetAllSteps()
@ -535,7 +535,7 @@ func TestBufferGetAllSteps(t *testing.T) {
})
t.Run("GetEmptySteps", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "")
buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "", "")
steps := buffer.GetAllSteps()
assert.NotNil(t, steps)
@ -549,7 +549,7 @@ func TestBufferGetAllSteps(t *testing.T) {
func TestBufferSpaceSnapshot(t *testing.T) {
t.Run("SetAndGetSnapshot", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "")
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "", "")
snapshot := map[string]interface{}{
"user_id": "user-123",
@ -566,7 +566,7 @@ func TestBufferSpaceSnapshot(t *testing.T) {
})
t.Run("SnapshotIsCopy", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "")
buffer := context.NewChatBuffer("chat-2", "req-2", "assistant-2", "", "")
original := map[string]interface{}{"key": "original"}
buffer.SetSpaceSnapshot(original)
@ -580,7 +580,7 @@ func TestBufferSpaceSnapshot(t *testing.T) {
})
t.Run("GetSnapshotReturnsCopy", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3", "")
buffer := context.NewChatBuffer("chat-3", "req-3", "assistant-3", "", "")
buffer.SetSpaceSnapshot(map[string]interface{}{"key": "value"})
retrieved1 := buffer.GetSpaceSnapshot()
@ -591,13 +591,13 @@ func TestBufferSpaceSnapshot(t *testing.T) {
})
t.Run("GetNilSnapshot", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-4", "req-4", "assistant-4", "")
buffer := context.NewChatBuffer("chat-4", "req-4", "assistant-4", "", "")
snapshot := buffer.GetSpaceSnapshot()
assert.Nil(t, snapshot)
})
t.Run("SetNilSnapshot", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-5", "req-5", "assistant-5", "")
buffer := context.NewChatBuffer("chat-5", "req-5", "assistant-5", "", "")
buffer.SetSpaceSnapshot(map[string]interface{}{"key": "value"})
buffer.SetSpaceSnapshot(nil)
@ -612,7 +612,7 @@ func TestBufferSpaceSnapshot(t *testing.T) {
func TestBufferIdentityMethods(t *testing.T) {
t.Run("SetAssistantID", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-original", "")
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-original", "", "")
assert.Equal(t, "assistant-original", buffer.AssistantID())
@ -621,22 +621,22 @@ func TestBufferIdentityMethods(t *testing.T) {
})
t.Run("ChatID", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-test", "req-test", "assistant-test", "")
buffer := context.NewChatBuffer("chat-test", "req-test", "assistant-test", "", "")
assert.Equal(t, "chat-test", buffer.ChatID())
})
t.Run("RequestID", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-test", "req-test", "assistant-test", "")
buffer := context.NewChatBuffer("chat-test", "req-test", "assistant-test", "", "")
assert.Equal(t, "req-test", buffer.RequestID())
})
t.Run("Connector", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-test", "req-test", "assistant-test", "openai")
buffer := context.NewChatBuffer("chat-test", "req-test", "assistant-test", "openai", "")
assert.Equal(t, "openai", buffer.Connector())
})
t.Run("SetConnector", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai")
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai", "")
assert.Equal(t, "openai", buffer.Connector())
// Simulate user switching connector mid-conversation
@ -645,14 +645,14 @@ func TestBufferIdentityMethods(t *testing.T) {
})
t.Run("EmptyConnector", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-test", "req-test", "assistant-test", "")
buffer := context.NewChatBuffer("chat-test", "req-test", "assistant-test", "", "")
assert.Equal(t, "", buffer.Connector())
})
}
func TestBufferConnectorInMessages(t *testing.T) {
t.Run("MessageInheritsConnector", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai")
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai", "")
// Add assistant message - should inherit connector from buffer
buffer.AddAssistantMessage(
@ -668,7 +668,7 @@ func TestBufferConnectorInMessages(t *testing.T) {
})
t.Run("MessageConnectorUpdatesWithBuffer", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai")
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai", "")
// First message with openai
buffer.AddAssistantMessage(
@ -696,7 +696,7 @@ func TestBufferConnectorInMessages(t *testing.T) {
})
t.Run("UserInputDoesNotSetConnector", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai")
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai", "")
// User input doesn't have connector (it's set by the system based on which model processes it)
buffer.AddUserInput("Hello", "")
@ -709,7 +709,7 @@ func TestBufferConnectorInMessages(t *testing.T) {
})
t.Run("MultipleConnectorSwitches", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai")
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai", "")
// Simulate a conversation with multiple connector switches
connectors := []string{"openai", "anthropic", "openai", "google"}
@ -737,7 +737,7 @@ func TestBufferConnectorInMessages(t *testing.T) {
// =============================================================================
func TestBufferConcurrentMessageOperations(t *testing.T) {
buffer := context.NewChatBuffer("chat-concurrent", "req-concurrent", "assistant-concurrent", "")
buffer := context.NewChatBuffer("chat-concurrent", "req-concurrent", "assistant-concurrent", "", "")
var wg sync.WaitGroup
numGoroutines := 100
@ -770,7 +770,7 @@ func TestBufferConcurrentMessageOperations(t *testing.T) {
}
func TestBufferConcurrentStepOperations(t *testing.T) {
buffer := context.NewChatBuffer("chat-concurrent", "req-concurrent", "assistant-concurrent", "")
buffer := context.NewChatBuffer("chat-concurrent", "req-concurrent", "assistant-concurrent", "", "")
var wg sync.WaitGroup
numGoroutines := 50
@ -794,7 +794,7 @@ func TestBufferConcurrentStepOperations(t *testing.T) {
}
func TestBufferConcurrentReadWrite(t *testing.T) {
buffer := context.NewChatBuffer("chat-rw", "req-rw", "assistant-rw", "")
buffer := context.NewChatBuffer("chat-rw", "req-rw", "assistant-rw", "", "")
var wg sync.WaitGroup
done := make(chan bool)
@ -868,7 +868,7 @@ func TestBufferStepStatusConstants(t *testing.T) {
func TestBufferEdgeCases(t *testing.T) {
t.Run("LargeNumberOfMessages", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-large", "req-large", "assistant-large", "")
buffer := context.NewChatBuffer("chat-large", "req-large", "assistant-large", "", "")
// Add 10000 messages
for i := 0; i < 10000; i++ {
@ -885,7 +885,7 @@ func TestBufferEdgeCases(t *testing.T) {
})
t.Run("MessageWithEmptyProps", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-empty", "req-empty", "assistant-empty", "")
buffer := context.NewChatBuffer("chat-empty", "req-empty", "assistant-empty", "", "")
buffer.AddMessage(&context.BufferedMessage{
Role: "assistant",
@ -899,7 +899,7 @@ func TestBufferEdgeCases(t *testing.T) {
})
t.Run("StepWithEmptyInput", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-step", "req-step", "assistant-step", "")
buffer := context.NewChatBuffer("chat-step", "req-step", "assistant-step", "", "")
step := buffer.BeginStep(context.StepTypeLLM, nil, nil)
assert.Nil(t, step.Input)
@ -910,7 +910,7 @@ func TestBufferEdgeCases(t *testing.T) {
})
t.Run("AllMessageTypes", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-types", "req-types", "assistant-types", "")
buffer := context.NewChatBuffer("chat-types", "req-types", "assistant-types", "", "")
messageTypes := []string{
"text", "image", "loading", "tool_call", "tool_result",
@ -926,7 +926,7 @@ func TestBufferEdgeCases(t *testing.T) {
})
t.Run("AllStepTypes", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-step-types", "req-step-types", "assistant-step-types", "")
buffer := context.NewChatBuffer("chat-step-types", "req-step-types", "assistant-step-types", "", "")
stepTypes := []string{
context.StepTypeInput, context.StepTypeHookCreate, context.StepTypeLLM,
@ -949,7 +949,7 @@ func TestBufferEdgeCases(t *testing.T) {
func TestBufferCompleteWorkflow(t *testing.T) {
t.Run("SuccessfulChatFlow", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-workflow", "req-workflow", "assistant-main", "")
buffer := context.NewChatBuffer("chat-workflow", "req-workflow", "assistant-main", "", "")
// 1. User input
buffer.AddUserInput("What's the weather in San Francisco?", "John")
@ -993,7 +993,7 @@ func TestBufferCompleteWorkflow(t *testing.T) {
})
t.Run("InterruptedChatFlow", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-interrupted", "req-interrupted", "assistant-main", "")
buffer := context.NewChatBuffer("chat-interrupted", "req-interrupted", "assistant-main", "", "")
// Set space snapshot
buffer.SetSpaceSnapshot(map[string]interface{}{
@ -1024,7 +1024,7 @@ func TestBufferCompleteWorkflow(t *testing.T) {
})
t.Run("A2ACallWithDelegation", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-a2a", "req-a2a", "assistant-main", "")
buffer := context.NewChatBuffer("chat-a2a", "req-a2a", "assistant-main", "", "")
mainStack := &context.Stack{ID: "stack-main", Depth: 0}
childStack := &context.Stack{ID: "stack-child", ParentID: "stack-main", Depth: 1}
@ -1060,7 +1060,7 @@ func TestBufferCompleteWorkflow(t *testing.T) {
})
t.Run("ConcurrentAgentCalls", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-concurrent-a2a", "req-concurrent-a2a", "assistant-main", "")
buffer := context.NewChatBuffer("chat-concurrent-a2a", "req-concurrent-a2a", "assistant-main", "", "")
// Main assistant spawns multiple concurrent calls
buffer.BeginStep(context.StepTypeInput, nil, nil)
@ -1105,7 +1105,7 @@ func TestBufferCompleteWorkflow(t *testing.T) {
func TestBufferMessageSequence(t *testing.T) {
t.Run("SequenceAutoIncrement", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-seq", "req-seq", "assistant-seq", "")
buffer := context.NewChatBuffer("chat-seq", "req-seq", "assistant-seq", "", "")
for i := 0; i < 10; i++ {
buffer.AddMessage(&context.BufferedMessage{
@ -1121,7 +1121,7 @@ func TestBufferMessageSequence(t *testing.T) {
})
t.Run("MixedMessageTypes", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-mixed", "req-mixed", "assistant-mixed", "")
buffer := context.NewChatBuffer("chat-mixed", "req-mixed", "assistant-mixed", "", "")
buffer.AddUserInput("Hello", "")
buffer.AddAssistantMessage("M1", "text", nil, "", "", "", nil)
@ -1142,7 +1142,7 @@ func TestBufferMessageSequence(t *testing.T) {
func TestBufferStepSequence(t *testing.T) {
t.Run("SequenceAutoIncrement", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-step-seq", "req-step-seq", "assistant-step-seq", "")
buffer := context.NewChatBuffer("chat-step-seq", "req-step-seq", "assistant-step-seq", "", "")
for i := 0; i < 5; i++ {
buffer.BeginStep(context.StepTypeLLM, nil, nil)
@ -1163,10 +1163,10 @@ func TestBufferStepSequence(t *testing.T) {
func TestBufferMultipleRequests(t *testing.T) {
t.Run("NewBufferPerRequest", func(t *testing.T) {
// Simulate multiple requests with separate buffers
buffer1 := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "")
buffer1 := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "", "")
buffer1.AddUserInput("Request 1", "")
buffer2 := context.NewChatBuffer("chat-1", "req-2", "assistant-1", "")
buffer2 := context.NewChatBuffer("chat-1", "req-2", "assistant-1", "", "")
buffer2.AddUserInput("Request 2", "")
// Buffers should be independent
@ -1187,7 +1187,7 @@ func TestBufferMultipleRequests(t *testing.T) {
func TestBufferStreamingMessage(t *testing.T) {
t.Run("AddStreamingMessage", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai")
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai", "")
buffer.AddStreamingMessage(
"msg-stream-1",
@ -1211,7 +1211,7 @@ func TestBufferStreamingMessage(t *testing.T) {
})
t.Run("AppendMessageContent", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai")
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai", "")
// Add streaming message
buffer.AddStreamingMessage(
@ -1235,7 +1235,7 @@ func TestBufferStreamingMessage(t *testing.T) {
})
t.Run("AppendToNonExistentMessage", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai")
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai", "")
// Try to append to non-existent message
ok := buffer.AppendMessageContent("non-existent", "content")
@ -1243,7 +1243,7 @@ func TestBufferStreamingMessage(t *testing.T) {
})
t.Run("AppendToCompletedMessage", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai")
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai", "")
// Add and complete streaming message
buffer.AddStreamingMessage(
@ -1260,7 +1260,7 @@ func TestBufferStreamingMessage(t *testing.T) {
})
t.Run("CompleteStreamingMessage", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai")
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai", "")
// Add streaming message
buffer.AddStreamingMessage(
@ -1289,7 +1289,7 @@ func TestBufferStreamingMessage(t *testing.T) {
})
t.Run("CompleteNonExistentMessage", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai")
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai", "")
content, ok := buffer.CompleteStreamingMessage("non-existent")
assert.False(t, ok)
@ -1297,7 +1297,7 @@ func TestBufferStreamingMessage(t *testing.T) {
})
t.Run("StreamingMessageWorkflow", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "deepseek")
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "deepseek", "")
// Simulate a typical streaming workflow:
// 1. SendStream sends initial content
@ -1332,7 +1332,7 @@ func TestBufferStreamingMessage(t *testing.T) {
})
t.Run("MixedStreamingAndRegularMessages", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai")
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai", "")
// Add user input (regular)
buffer.AddUserInput("Hello", "user1")
@ -1366,7 +1366,7 @@ func TestBufferStreamingMessage(t *testing.T) {
})
t.Run("StreamingMessageWithEmptyInitialContent", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai")
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai", "")
// Add streaming message with nil props
buffer.AddStreamingMessage(
@ -1386,7 +1386,7 @@ func TestBufferStreamingMessage(t *testing.T) {
})
t.Run("ConcurrentStreamingOperations", func(t *testing.T) {
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai")
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai", "")
// Add streaming message
buffer.AddStreamingMessage(

4
agent/context/context.go
View file

@ -399,8 +399,8 @@ func (ctx *Context) GetMessageMetadata(messageID string) *MessageMetadata {
// InitBuffer initializes the chat buffer for this context
// Should be called at the start of Stream() to begin buffering messages and steps
func (ctx *Context) InitBuffer(assistantID, connector string) *ChatBuffer {
ctx.Buffer = NewChatBuffer(ctx.ChatID, ctx.RequestID(), assistantID, connector)
func (ctx *Context) InitBuffer(assistantID, connector, mode string) *ChatBuffer {
ctx.Buffer = NewChatBuffer(ctx.ChatID, ctx.RequestID(), assistantID, connector, mode)
return ctx.Buffer
}

28
agent/context/openapi.go
View file

@ -67,6 +67,7 @@ func GetCompletionRequest(c *gin.Context, cache store.Store) (*CompletionRequest
opts := &Options{
Context: c.Request.Context(),
Skip: GetSkip(c, completionReq),
Mode: GetMode(c, completionReq),
}
// Try to extract custom connector from model field
@ -377,6 +378,33 @@ func GetRoute(c *gin.Context, req *CompletionRequest) string {
return ""
}
// GetMode extracts mode from request with priority:
// 1. Query parameter "mode"
// 2. Header "X-Yao-Mode"
// 3. CompletionRequest metadata "mode" (from payload)
func GetMode(c *gin.Context, req *CompletionRequest) string {
// Priority 1: Query parameter
if mode := c.Query("mode"); mode != "" {
return mode
}
// Priority 2: Header
if mode := c.GetHeader("X-Yao-Mode"); mode != "" {
return mode
}
// Priority 3: From CompletionRequest metadata
if req != nil && req.Metadata != nil {
if mode, ok := req.Metadata["mode"]; ok {
if modeStr, ok := mode.(string); ok && modeStr != "" {
return modeStr
}
}
}
return ""
}
// GetSkip extracts skip configuration from request with priority:
// 1. CompletionRequest.Skip (from payload body) - Priority
// 2. Individual query parameters: "skip_history", "skip_trace"

100
agent/context/openapi_test.go
View file

@ -1145,3 +1145,103 @@ func TestGetSkip_FromBodyViaParseRequest(t *testing.T) {
t.Error("Expected GetSkip to return Trace=false")
}
}
func TestGetMode_FromQuery(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
gin.SetMode(gin.TestMode)
req := httptest.NewRequest("GET", "/chat/completions?mode=task", nil)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = req
mode := GetMode(c, nil)
if mode != "task" {
t.Errorf("Expected mode 'task' from query, got '%s'", mode)
}
}
func TestGetMode_FromHeader(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
gin.SetMode(gin.TestMode)
req := httptest.NewRequest("GET", "/chat/completions", nil)
req.Header.Set("X-Yao-Mode", "chat")
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = req
mode := GetMode(c, nil)
if mode != "chat" {
t.Errorf("Expected mode 'chat' from header, got '%s'", mode)
}
}
func TestGetMode_FromMetadata(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
gin.SetMode(gin.TestMode)
req := httptest.NewRequest("GET", "/chat/completions", nil)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = req
completionReq := &CompletionRequest{
Metadata: map[string]interface{}{
"mode": "task",
},
}
mode := GetMode(c, completionReq)
if mode != "task" {
t.Errorf("Expected mode 'task' from metadata, got '%s'", mode)
}
}
func TestGetMode_Priority(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
gin.SetMode(gin.TestMode)
// Query has highest priority
req := httptest.NewRequest("GET", "/chat/completions?mode=query_mode", nil)
req.Header.Set("X-Yao-Mode", "header_mode")
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = req
completionReq := &CompletionRequest{
Metadata: map[string]interface{}{
"mode": "metadata_mode",
},
}
mode := GetMode(c, completionReq)
if mode != "query_mode" {
t.Errorf("Expected mode 'query_mode' (query has priority), got '%s'", mode)
}
}
func TestGetMode_Empty(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
gin.SetMode(gin.TestMode)
req := httptest.NewRequest("GET", "/chat/completions", nil)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = req
mode := GetMode(c, nil)
if mode != "" {
t.Errorf("Expected empty mode, got '%s'", mode)
}
}

59
agent/store/CHAT_STORAGE_DESIGN.md
View file

@ -92,7 +92,7 @@ Stores chat metadata and session information.
| `title` | string(500) | Yes | - | Chat title |
| `assistant_id` | string(200) | No | Yes | Associated assistant ID |
| `last_connector` | string(200) | Yes | Yes | Last used connector ID |
| `mode` | string(50) | No | - | Chat mode (default: "chat") |
| `last_mode` | string(50) | Yes | - | Last used chat mode (chat/task) |
| `status` | enum | No | Yes | Status: `active`, `archived` |
| `public` | boolean | No | - | Whether shared across all teams |
| `share` | enum | No | Yes | Sharing scope: `private`, `team` |
@ -141,23 +141,24 @@ Stores user-visible messages (both user input and assistant responses).
**Table Name:** `agent_message`
| Column | Type | Nullable | Index | Description |
| -------------- | ----------- | -------- | ----- | ------------------------------------------ |
| `id` | ID | No | PK | Auto-increment primary key |
| `message_id` | string(64) | No | - | Message identifier (unique within request) |
| `chat_id` | string(64) | No | Yes | Parent chat ID |
| `request_id` | string(64) | Yes | Yes | Request ID for grouping |
| `role` | enum | No | Yes | Role: `user`, `assistant` |
| `type` | string(50) | No | - | Message type (text, image, loading, etc.) |
| `props` | json | No | - | Message properties (content, url, etc.) |
| `block_id` | string(64) | Yes | Yes | Block grouping ID |
| `thread_id` | string(64) | Yes | Yes | Thread grouping ID |
| `assistant_id` | string(200) | Yes | Yes | Assistant ID (join to get name/avatar) |
| `connector` | string(200) | Yes | Yes | Connector ID used for this message |
| `sequence` | integer | No | - | Message order within chat (in composite) |
| `metadata` | json | Yes | - | Additional metadata |
| `created_at` | timestamp | No | Yes | Creation timestamp |
| `updated_at` | timestamp | No | - | Last update timestamp |
| Column | Type | Nullable | Index | Description |
| -------------- | ----------- | -------- | ----- | ------------------------------------------- |
| `id` | ID | No | PK | Auto-increment primary key |
| `message_id` | string(64) | No | - | Message identifier (unique within request) |
| `chat_id` | string(64) | No | Yes | Parent chat ID |
| `request_id` | string(64) | Yes | Yes | Request ID for grouping |
| `role` | enum | No | Yes | Role: `user`, `assistant` |
| `type` | string(50) | No | - | Message type (text, image, loading, etc.) |
| `props` | json | No | - | Message properties (content, url, etc.) |
| `block_id` | string(64) | Yes | Yes | Block grouping ID |
| `thread_id` | string(64) | Yes | Yes | Thread grouping ID |
| `assistant_id` | string(200) | Yes | Yes | Assistant ID (join to get name/avatar) |
| `connector` | string(200) | Yes | Yes | Connector ID used for this message |
| `mode` | string(50) | Yes | - | Chat mode used for this message (chat/task) |
| `sequence` | integer | No | - | Message order within chat (in composite) |
| `metadata` | json | Yes | - | Additional metadata |
| `created_at` | timestamp | No | Yes | Creation timestamp |
| `updated_at` | timestamp | No | - | Last update timestamp |
**Indexes:**
@ -171,6 +172,20 @@ Stores user-visible messages (both user input and assistant responses).
| `idx_msg_thread` | `thread_id` | index |
| `idx_msg_assistant` | `assistant_id` | index |
**Message Ordering:**
Messages are ordered by `created_at` first, then by `sequence` within the same timestamp. This ensures correct chronological order when there are multiple requests with overlapping sequence numbers:
```sql
ORDER BY created_at ASC, sequence ASC
```
**Why this ordering?**
- `sequence` is assigned per-request, so different requests may have the same sequence numbers
- `created_at` groups messages by request time, ensuring messages from earlier requests appear first
- Within the same request (same `created_at`), `sequence` preserves the internal ordering
**Message Types:**
All message types are stored, including built-in types and custom types. See `agent/output/BUILTIN_TYPES.md` for built-in Props structures.
@ -813,8 +828,8 @@ type Chat struct {
ChatID string `json:"chat_id"`
Title string `json:"title,omitempty"`
AssistantID string `json:"assistant_id"`
LastConnector string `json:"last_connector,omitempty"` // Last used connector ID
Mode string `json:"mode"`
LastConnector string `json:"last_connector,omitempty"` // Last used connector ID (updated on each message)
LastMode string `json:"last_mode,omitempty"` // Last used chat mode (updated on each message)
Status string `json:"status"` // "active" or "archived"
Public bool `json:"public"` // Whether shared across all teams
Share string `json:"share"` // "private" or "team"
@ -837,6 +852,7 @@ type Message struct {
ThreadID string `json:"thread_id,omitempty"`
AssistantID string `json:"assistant_id,omitempty"`
Connector string `json:"connector,omitempty"` // Connector ID used for this message
Mode string `json:"mode,omitempty"` // Chat mode used for this message (chat or task)
Sequence int `json:"sequence"`
Metadata map[string]interface{} `json:"metadata,omitempty"`
CreatedAt time.Time `json:"created_at"`
@ -1585,7 +1601,8 @@ GET /v1/chat/sessions/chat_123
"chat_id": "chat_123",
"title": "Weather Query",
"assistant_id": "weather_assistant",
"mode": "chat",
"last_connector": "deepseek.v3",
"last_mode": "chat",
"status": "active",
"public": false,
"share": "private",

20
agent/store/types/types.go
View file

@ -30,11 +30,11 @@ type Chat struct {
Title string `json:"title,omitempty"`
AssistantID string `json:"assistant_id"`
LastConnector string `json:"last_connector,omitempty"` // Last used connector ID (updated on each message)
Mode string `json:"mode"`
Status string `json:"status"` // "active" or "archived"
Public bool `json:"public"` // Whether shared across all teams
Share string `json:"share"` // "private" or "team"
Sort int `json:"sort"` // Sort order for display
LastMode string `json:"last_mode,omitempty"` // Last used chat mode (updated on each message)
Status string `json:"status"` // "active" or "archived"
Public bool `json:"public"` // Whether shared across all teams
Share string `json:"share"` // "private" or "team"
Sort int `json:"sort"` // Sort order for display
LastMessageAt *time.Time `json:"last_message_at,omitempty"`
Metadata map[string]interface{} `json:"metadata,omitempty"`
CreatedAt time.Time `json:"created_at"`
@ -109,6 +109,7 @@ type Message struct {
ThreadID string `json:"thread_id,omitempty"`
AssistantID string `json:"assistant_id,omitempty"`
Connector string `json:"connector,omitempty"` // Connector ID used for this message
Mode string `json:"mode,omitempty"` // Chat mode used for this message (chat or task)
Sequence int `json:"sequence"`
Metadata map[string]interface{} `json:"metadata,omitempty"`
CreatedAt time.Time `json:"created_at"`
@ -198,6 +199,15 @@ type AssistantList struct {
Total int `json:"total"` // Total number of items across all pages
}
// AssistantInfo contains basic assistant information for display
// Used in chat history to show assistant details with i18n support
type AssistantInfo struct {
AssistantID string `json:"assistant_id"`
Name string `json:"name"`
Avatar string `json:"avatar,omitempty"`
Description string `json:"description,omitempty"`
}
// Tag represents a tag
type Tag struct {
Value string `json:"value"`

14
agent/store/xun/chat.go
View file

@ -42,9 +42,6 @@ func (store *Xun) CreateChat(chat *types.Chat) error {
}
// Set defaults
if chat.Mode == "" {
chat.Mode = "chat"
}
if chat.Status == "" {
chat.Status = "active"
}
@ -56,7 +53,6 @@ func (store *Xun) CreateChat(chat *types.Chat) error {
data := map[string]interface{}{
"chat_id": chat.ChatID,
"assistant_id": chat.AssistantID,
"mode": chat.Mode,
"status": chat.Status,
"public": chat.Public,
"share": chat.Share,
@ -65,6 +61,11 @@ func (store *Xun) CreateChat(chat *types.Chat) error {
"updated_at": time.Now(),
}
// Handle last_mode (nullable)
if chat.LastMode != "" {
data["last_mode"] = chat.LastMode
}
// Handle nullable fields
if chat.Title != "" {
data["title"] = chat.Title
@ -72,6 +73,9 @@ func (store *Xun) CreateChat(chat *types.Chat) error {
if chat.LastConnector != "" {
data["last_connector"] = chat.LastConnector
}
if chat.LastMode != "" {
data["last_mode"] = chat.LastMode
}
if chat.LastMessageAt != nil {
data["last_message_at"] = *chat.LastMessageAt
}
@ -337,7 +341,7 @@ func (store *Xun) rowToChat(data map[string]interface{}) (*types.Chat, error) {
Title: getString(data, "title"),
AssistantID: getString(data, "assistant_id"),
LastConnector: getString(data, "last_connector"),
Mode: getString(data, "mode"),
LastMode: getString(data, "last_mode"),
Status: getString(data, "status"),
Public: getBool(data, "public"),
Share: getString(data, "share"),

13
agent/store/xun/chat_test.go
View file

@ -29,7 +29,6 @@ func TestCreateChat(t *testing.T) {
chat := &types.Chat{
AssistantID: "test_assistant",
Title: "Test Chat",
Mode: "chat",
Status: "active",
Share: "private",
}
@ -55,7 +54,7 @@ func TestCreateChat(t *testing.T) {
AssistantID: "test_assistant",
LastConnector: "openai",
Title: "Full Chat",
Mode: "task",
LastMode: "task",
Status: "active",
Public: true,
Share: "team",
@ -84,8 +83,8 @@ func TestCreateChat(t *testing.T) {
if retrieved.LastConnector != "openai" {
t.Errorf("Expected last_connector 'openai', got '%s'", retrieved.LastConnector)
}
if retrieved.Mode != "task" {
t.Errorf("Expected mode 'task', got '%s'", retrieved.Mode)
if retrieved.LastMode != "task" {
t.Errorf("Expected last_mode 'task', got '%s'", retrieved.LastMode)
}
if !retrieved.Public {
t.Error("Expected public to be true")
@ -183,8 +182,9 @@ func TestCreateChat(t *testing.T) {
t.Fatalf("Failed to retrieve chat: %v", err)
}
if retrieved.Mode != "chat" {
t.Errorf("Expected default mode 'chat', got '%s'", retrieved.Mode)
// last_mode is nullable, so it should be empty by default
if retrieved.LastMode != "" {
t.Errorf("Expected default last_mode to be empty, got '%s'", retrieved.LastMode)
}
if retrieved.Status != "active" {
t.Errorf("Expected default status 'active', got '%s'", retrieved.Status)
@ -1103,7 +1103,6 @@ func TestChatCompleteWorkflow(t *testing.T) {
chat := &types.Chat{
AssistantID: "workflow_assistant",
Title: "Workflow Test Chat",
Mode: "chat",
Status: "active",
}

9
agent/store/xun/message.go
View file

@ -69,6 +69,7 @@ func (store *Xun) SaveMessages(chatID string, messages []*types.Message) error {
"thread_id": nil,
"assistant_id": nil,
"connector": nil,
"mode": nil,
"metadata": nil,
"created_at": now,
"updated_at": now,
@ -90,6 +91,9 @@ func (store *Xun) SaveMessages(chatID string, messages []*types.Message) error {
if msg.Connector != "" {
row["connector"] = msg.Connector
}
if msg.Mode != "" {
row["mode"] = msg.Mode
}
if msg.Metadata != nil {
metadataJSON, err := jsoniter.MarshalToString(msg.Metadata)
if err != nil {
@ -147,8 +151,8 @@ func (store *Xun) GetMessages(chatID string, filter types.MessageFilter) ([]*typ
qb.Limit(1000000).Offset(filter.Offset)
}
// Order by sequence
qb.OrderBy("sequence", "asc")
// Order by created_at first, then by sequence within the same request
qb.OrderBy("created_at", "asc").OrderBy("sequence", "asc")
rows, err := qb.Get()
if err != nil {
@ -332,6 +336,7 @@ func (store *Xun) rowToMessage(data map[string]interface{}) (*types.Message, err
ThreadID: getString(data, "thread_id"),
AssistantID: getString(data, "assistant_id"),
Connector: getString(data, "connector"),
Mode: getString(data, "mode"),
Sequence: getInt(data, "sequence"),
}

199
agent/store/xun/message_test.go
View file

@ -282,6 +282,84 @@ func TestSaveMessages(t *testing.T) {
t.Logf("Successfully saved and retrieved messages with different connectors")
})
t.Run("SaveMessageWithMode", func(t *testing.T) {
modeChat := &types.Chat{
AssistantID: "test_assistant",
}
err := store.CreateChat(modeChat)
if err != nil {
t.Fatalf("Failed to create chat: %v", err)
}
defer store.DeleteChat(modeChat.ChatID)
// Save messages with different modes
messages := []*types.Message{
{
Role: "user",
Type: "user_input",
Props: map[string]interface{}{"content": "Hello in chat mode"},
Sequence: 1,
Mode: "chat",
Connector: "deepseek.v3",
AssistantID: "test_assistant",
},
{
Role: "assistant",
Type: "text",
Props: map[string]interface{}{"content": "Hi there in chat mode!"},
Sequence: 2,
Mode: "chat",
Connector: "deepseek.v3",
AssistantID: "test_assistant",
},
{
Role: "user",
Type: "user_input",
Props: map[string]interface{}{"content": "Now run a task"},
Sequence: 3,
Mode: "task",
Connector: "deepseek.v3",
AssistantID: "test_assistant",
},
{
Role: "assistant",
Type: "text",
Props: map[string]interface{}{"content": "Running task!"},
Sequence: 4,
Mode: "task",
Connector: "deepseek.v3",
AssistantID: "test_assistant",
},
}
err = store.SaveMessages(modeChat.ChatID, messages)
if err != nil {
t.Fatalf("Failed to save messages: %v", err)
}
// Retrieve and verify modes
retrieved, err := store.GetMessages(modeChat.ChatID, types.MessageFilter{})
if err != nil {
t.Fatalf("Failed to get messages: %v", err)
}
if len(retrieved) != 4 {
t.Fatalf("Expected 4 messages, got %d", len(retrieved))
}
// Verify each message has correct mode
for _, msg := range retrieved {
if msg.Sequence <= 2 && msg.Mode != "chat" {
t.Errorf("Expected mode 'chat' for sequence %d, got '%s'", msg.Sequence, msg.Mode)
}
if msg.Sequence > 2 && msg.Mode != "task" {
t.Errorf("Expected mode 'task' for sequence %d, got '%s'", msg.Sequence, msg.Mode)
}
}
t.Logf("Successfully saved and retrieved messages with different modes")
})
t.Run("SaveMessageWithEmptyConnector", func(t *testing.T) {
emptyConnChat := &types.Chat{
AssistantID: "test_assistant",
@ -533,6 +611,127 @@ func TestGetMessages(t *testing.T) {
t.Errorf("Expected 0 messages from non-existent chat, got %d", len(retrieved))
}
})
t.Run("OrderByCreatedAtThenSequence", func(t *testing.T) {
// This test verifies that messages are ordered by created_at first, then by sequence
// This is important when there are multiple request_ids with overlapping sequence numbers
orderChat := &types.Chat{
AssistantID: "test_assistant",
}
err := store.CreateChat(orderChat)
if err != nil {
t.Fatalf("Failed to create chat: %v", err)
}
defer store.DeleteChat(orderChat.ChatID)
// Simulate two separate requests with overlapping sequence numbers
// SaveMessages uses time.Now() for created_at, so we need to call it twice
// with a small delay to ensure different timestamps
// Request 1: sequences 1, 2
req1Messages := []*types.Message{
{
Role: "user",
Type: "user_input",
Props: map[string]interface{}{"content": "Request 1 - Message 1"},
Sequence: 1,
RequestID: "order_req_001",
},
{
Role: "assistant",
Type: "text",
Props: map[string]interface{}{"content": "Request 1 - Response 1"},
Sequence: 2,
RequestID: "order_req_001",
},
}
err = store.SaveMessages(orderChat.ChatID, req1Messages)
if err != nil {
t.Fatalf("Failed to save request 1 messages: %v", err)
}
// Delay to ensure different created_at timestamps
// Database timestamp precision may only be to second level
time.Sleep(1100 * time.Millisecond)
// Request 2: sequences 1, 2 (same as request 1, but later created_at)
req2Messages := []*types.Message{
{
Role: "user",
Type: "user_input",
Props: map[string]interface{}{"content": "Request 2 - Message 1"},
Sequence: 1, // Same sequence as req1, but later created_at
RequestID: "order_req_002",
},
{
Role: "assistant",
Type: "text",
Props: map[string]interface{}{"content": "Request 2 - Response 1"},
Sequence: 2, // Same sequence as req1, but later created_at
RequestID: "order_req_002",
},
}
err = store.SaveMessages(orderChat.ChatID, req2Messages)
if err != nil {
t.Fatalf("Failed to save request 2 messages: %v", err)
}
// Retrieve messages
retrieved, err := store.GetMessages(orderChat.ChatID, types.MessageFilter{})
if err != nil {
t.Fatalf("Failed to get messages: %v", err)
}
if len(retrieved) != 4 {
t.Fatalf("Expected 4 messages, got %d", len(retrieved))
}
// Verify order: should be chronological by created_at, then by sequence
// Messages from req_001 should come before req_002
expectedOrder := []struct {
requestID string
sequence int
content string
}{
{"order_req_001", 1, "Request 1 - Message 1"},
{"order_req_001", 2, "Request 1 - Response 1"},
{"order_req_002", 1, "Request 2 - Message 1"},
{"order_req_002", 2, "Request 2 - Response 1"},
}
for i, expected := range expectedOrder {
msg := retrieved[i]
if msg.RequestID != expected.requestID {
t.Errorf("Message %d: expected RequestID '%s', got '%s'", i, expected.requestID, msg.RequestID)
}
if msg.Sequence != expected.sequence {
t.Errorf("Message %d: expected Sequence %d, got %d", i, expected.sequence, msg.Sequence)
}
content, _ := msg.Props["content"].(string)
if content != expected.content {
t.Errorf("Message %d: expected content '%s', got '%s'", i, expected.content, content)
}
}
// Additional verification: ensure created_at is non-decreasing
for i := 1; i < len(retrieved); i++ {
if retrieved[i].CreatedAt.Before(retrieved[i-1].CreatedAt) {
t.Errorf("Message %d created_at (%v) is before message %d created_at (%v)",
i, retrieved[i].CreatedAt, i-1, retrieved[i-1].CreatedAt)
}
// If same created_at, sequence should be increasing
if retrieved[i].CreatedAt.Equal(retrieved[i-1].CreatedAt) {
if retrieved[i].Sequence < retrieved[i-1].Sequence {
t.Errorf("Messages with same created_at: message %d sequence (%d) < message %d sequence (%d)",
i, retrieved[i].Sequence, i-1, retrieved[i-1].Sequence)
}
}
}
t.Logf("Successfully verified message ordering: created_at first, then sequence")
})
}
// TestUpdateMessage tests updating messages

288
data/bindata.go
View file

File diff suppressed because it is too large Load diff

52
openapi/chat/session.go
View file

@ -363,13 +363,59 @@ func GetMessages(c *gin.Context) {
return
}
// Get locale from query parameter or Accept-Language header
locale := getLocale(c)
// Collect unique assistant IDs from messages and fetch their info
assistantIDs := collectAssistantIDs(messages)
assistants := assistant.GetInfoByIDs(assistantIDs, locale)
response.RespondWithSuccess(c, response.StatusOK, gin.H{
"chat_id": chatID,
"messages": messages,
"count": len(messages),
"chat_id": chatID,
"messages": messages,
"count": len(messages),
"assistants": assistants,
})
}
// getLocale extracts locale from request
// Priority: 1. Query param "locale", 2. Accept-Language header
func getLocale(c *gin.Context) string {
// Priority 1: Query parameter
if locale := c.Query("locale"); locale != "" {
return strings.ToLower(locale)
}
// Priority 2: Header Accept-Language
if acceptLang := c.GetHeader("Accept-Language"); acceptLang != "" {
// Parse Accept-Language header (e.g., "en-US,en;q=0.9,zh;q=0.8")
// Take the first language
parts := strings.Split(acceptLang, ",")
if len(parts) > 0 {
// Remove quality value if present
lang := strings.Split(parts[0], ";")[0]
return strings.ToLower(strings.TrimSpace(lang))
}
}
return ""
}
// collectAssistantIDs extracts unique assistant IDs from messages
func collectAssistantIDs(messages []*storetypes.Message) []string {
seen := make(map[string]bool)
var ids []string
for _, msg := range messages {
if msg.AssistantID != "" && !seen[msg.AssistantID] {
seen[msg.AssistantID] = true
ids = append(ids, msg.AssistantID)
}
}
return ids
}
// =============================================================================
// Helper Functions
// =============================================================================

9
yao/models/agent/chat.mod.yao
View file

@ -50,13 +50,12 @@
"index": true
},
{
"name": "mode",
"name": "last_mode",
"type": "string",
"label": "Mode",
"comment": "Chat mode (default: chat)",
"label": "Last Mode",
"comment": "Last used chat mode (updated on each message)",
"length": 50,
"nullable": false,
"default": "chat"
"nullable": true
},
{
"name": "status",

8
yao/models/agent/message.mod.yao
View file

@ -100,6 +100,14 @@
"nullable": true,
"index": true
},
{
"name": "mode",
"type": "string",
"label": "Mode",
"comment": "Chat mode used for this message (chat or task)",
"length": 50,
"nullable": true
},
{
"name": "sequence",
"type": "integer",