Implement connector management in chat system
- Enhanced chat buffer and message handling to support dynamic connector switching, allowing users to change the connector during a chat session. - Updated the `ChatBuffer` and `BufferedMessage` structures to include a `connector` field, enabling tracking of the connector used for each message. - Modified the `EnsureChat` method to skip chat creation when history is disabled, improving chat session management. - Revised tests to validate the new connector functionality, ensuring accurate message retrieval and connector state management. - Updated `CHAT_STORAGE_DESIGN.md` to reflect the addition of the `last_connector` field in chat metadata and message structures.
This commit is contained in:
parent
279ae161c6
commit
a8a1103b6b
14 changed files with 616 additions and 224 deletions
|
|
@ -243,8 +243,14 @@ func (ast *Assistant) InitBuffer(ctx *agentcontext.Context) {
|
|||
requestID = uuid.New().String()
|
||||
}
|
||||
|
||||
ctx.Buffer = agentcontext.NewChatBuffer(ctx.ChatID, requestID, ast.ID)
|
||||
log.Trace("[CHAT] Buffer initialized: chatID=%s, requestID=%s, assistantID=%s", ctx.ChatID, requestID, ast.ID)
|
||||
// Get connector from options
|
||||
connector := ""
|
||||
if ctx.Stack.Options != nil {
|
||||
connector = ctx.Stack.Options.Connector
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
// BufferUserInput adds user input messages to the buffer
|
||||
|
|
@ -335,13 +341,18 @@ func (ast *Assistant) FlushBuffer(ctx *agentcontext.Context, finalStatus string,
|
|||
}
|
||||
}
|
||||
|
||||
// 2. Update chat last_message_at
|
||||
// 2. Update chat last_message_at and last_connector
|
||||
if len(messages) > 0 {
|
||||
now := time.Now()
|
||||
if updateErr := chatStore.UpdateChat(ctx.ChatID, map[string]interface{}{
|
||||
updates := map[string]interface{}{
|
||||
"last_message_at": now,
|
||||
}); updateErr != nil {
|
||||
log.Trace("[CHAT] Failed to update last_message_at: %v", updateErr)
|
||||
}
|
||||
// Also update last_connector if available
|
||||
if connector := ctx.Buffer.Connector(); connector != "" {
|
||||
updates["last_connector"] = connector
|
||||
}
|
||||
if updateErr := chatStore.UpdateChat(ctx.ChatID, updates); updateErr != nil {
|
||||
log.Trace("[CHAT] Failed to update chat: %v", updateErr)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -376,6 +387,7 @@ func (ast *Assistant) convertBufferedMessages(buffered []*agentcontext.BufferedM
|
|||
BlockID: msg.BlockID,
|
||||
ThreadID: msg.ThreadID,
|
||||
AssistantID: msg.AssistantID,
|
||||
Connector: msg.Connector,
|
||||
Sequence: msg.Sequence,
|
||||
Metadata: msg.Metadata,
|
||||
CreatedAt: msg.CreatedAt,
|
||||
|
|
@ -422,6 +434,11 @@ func (ast *Assistant) EnsureChat(ctx *agentcontext.Context) error {
|
|||
return nil // No chat ID, skip
|
||||
}
|
||||
|
||||
// Skip if history is disabled
|
||||
if ctx.Stack != nil && ctx.Stack.Options != nil && ctx.Stack.Options.Skip != nil && ctx.Stack.Options.Skip.History {
|
||||
return nil // Skip.History is true, don't create chat session
|
||||
}
|
||||
|
||||
chatStore := GetChatStore()
|
||||
if chatStore == nil {
|
||||
return nil // No store, skip
|
||||
|
|
@ -445,6 +462,11 @@ func (ast *Assistant) EnsureChat(ctx *agentcontext.Context) error {
|
|||
UpdatedAt: time.Now(),
|
||||
}
|
||||
|
||||
// Set last_connector from options (user selected connector)
|
||||
if ctx.Stack != nil && ctx.Stack.Options != nil && ctx.Stack.Options.Connector != "" {
|
||||
chat.LastConnector = ctx.Stack.Options.Connector
|
||||
}
|
||||
|
||||
// Set permission fields from authorized info
|
||||
if ctx.Authorized != nil {
|
||||
chat.CreatedBy = ctx.Authorized.UserID
|
||||
|
|
|
|||
|
|
@ -799,6 +799,34 @@ func TestEnsureChat(t *testing.T) {
|
|||
t.Logf("✓ Chat created with permission fields: user=%s, team=%s, tenant=%s",
|
||||
chat.CreatedBy, chat.TeamID, chat.TenantID)
|
||||
})
|
||||
|
||||
t.Run("SkipHistoryEnabled", func(t *testing.T) {
|
||||
chatID := fmt.Sprintf("test_ensure_skip_%s", uuid.New().String()[:8])
|
||||
|
||||
// Create context
|
||||
ctx := agentcontext.New(context.Background(), nil, chatID)
|
||||
|
||||
// Set up stack with Skip.History = true
|
||||
ctx.Stack = &agentcontext.Stack{
|
||||
ID: "test_stack",
|
||||
AssistantID: ast.ID,
|
||||
Depth: 0,
|
||||
Options: &agentcontext.Options{
|
||||
Skip: &agentcontext.Skip{
|
||||
History: true,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
// EnsureChat should NOT create chat when Skip.History is true
|
||||
err := ast.EnsureChat(ctx)
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Verify chat was NOT created
|
||||
_, err = chatStore.GetChat(chatID)
|
||||
assert.Error(t, err, "Chat should not be created when Skip.History is true")
|
||||
t.Logf("✓ Chat not created when Skip.History is true")
|
||||
})
|
||||
}
|
||||
|
||||
func TestConvertBufferedTypes(t *testing.T) {
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ type ChatBuffer struct {
|
|||
chatID string
|
||||
requestID string
|
||||
assistantID string
|
||||
connector string // Current connector ID (for data analysis)
|
||||
|
||||
// Message buffer
|
||||
messages []*BufferedMessage
|
||||
|
|
@ -45,6 +46,7 @@ type BufferedMessage struct {
|
|||
BlockID string `json:"block_id,omitempty"`
|
||||
ThreadID string `json:"thread_id,omitempty"`
|
||||
AssistantID string `json:"assistant_id,omitempty"`
|
||||
Connector string `json:"connector,omitempty"` // Connector ID used for this message
|
||||
Sequence int `json:"sequence"`
|
||||
Metadata map[string]interface{} `json:"metadata,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
|
|
@ -94,11 +96,12 @@ const (
|
|||
)
|
||||
|
||||
// NewChatBuffer creates a new chat buffer
|
||||
func NewChatBuffer(chatID, requestID, assistantID string) *ChatBuffer {
|
||||
func NewChatBuffer(chatID, requestID, assistantID, connector string) *ChatBuffer {
|
||||
return &ChatBuffer{
|
||||
chatID: chatID,
|
||||
requestID: requestID,
|
||||
assistantID: assistantID,
|
||||
connector: connector,
|
||||
messages: make([]*BufferedMessage, 0),
|
||||
steps: make([]*BufferedStep, 0),
|
||||
}
|
||||
|
|
@ -170,6 +173,7 @@ func (b *ChatBuffer) AddAssistantMessage(msgType string, props map[string]interf
|
|||
BlockID: blockID,
|
||||
ThreadID: threadID,
|
||||
AssistantID: assistantID,
|
||||
Connector: b.connector, // Use current connector
|
||||
Metadata: metadata,
|
||||
})
|
||||
}
|
||||
|
|
@ -342,6 +346,18 @@ func (b *ChatBuffer) SetAssistantID(assistantID string) {
|
|||
b.assistantID = assistantID
|
||||
}
|
||||
|
||||
// Connector returns the current connector ID
|
||||
func (b *ChatBuffer) Connector() string {
|
||||
return b.connector
|
||||
}
|
||||
|
||||
// SetConnector updates the connector ID (when user switches connector)
|
||||
func (b *ChatBuffer) SetConnector(connector string) {
|
||||
b.mu.Lock()
|
||||
defer b.mu.Unlock()
|
||||
b.connector = connector
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Helper Functions
|
||||
// =============================================================================
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
"text",
|
||||
map[string]interface{}{"content": "Hello, how can I help?"},
|
||||
|
|
@ -184,7 +184,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(
|
||||
"event",
|
||||
map[string]interface{}{"event": "message_start"},
|
||||
|
|
@ -196,7 +196,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(
|
||||
"retrieval",
|
||||
map[string]interface{}{
|
||||
|
|
@ -214,7 +214,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(
|
||||
"tool_call",
|
||||
map[string]interface{}{
|
||||
|
|
@ -231,7 +231,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(
|
||||
"custom_chart",
|
||||
map[string]interface{}{
|
||||
|
|
@ -249,7 +249,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()
|
||||
|
|
@ -262,7 +262,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)
|
||||
|
|
@ -271,7 +271,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", "")
|
||||
|
|
@ -287,7 +287,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",
|
||||
|
|
@ -313,7 +313,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)
|
||||
|
||||
|
|
@ -324,7 +324,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)
|
||||
|
|
@ -339,7 +339,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{}{
|
||||
|
|
@ -357,7 +357,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!"})
|
||||
|
|
@ -370,7 +370,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"})
|
||||
|
|
@ -378,7 +378,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})
|
||||
|
|
@ -399,7 +399,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"))
|
||||
|
|
@ -411,7 +411,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)
|
||||
|
|
@ -423,7 +423,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})
|
||||
|
|
@ -437,7 +437,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"))
|
||||
|
|
@ -446,12 +446,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()
|
||||
|
|
@ -460,7 +460,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)
|
||||
|
||||
|
|
@ -470,7 +470,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)
|
||||
|
|
@ -483,7 +483,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)
|
||||
|
|
@ -499,7 +499,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)
|
||||
|
|
@ -517,7 +517,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()
|
||||
|
|
@ -529,7 +529,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)
|
||||
|
|
@ -543,7 +543,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",
|
||||
|
|
@ -560,7 +560,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)
|
||||
|
|
@ -574,7 +574,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()
|
||||
|
|
@ -585,13 +585,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)
|
||||
|
||||
|
|
@ -606,7 +606,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())
|
||||
|
||||
|
|
@ -615,14 +615,111 @@ 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")
|
||||
assert.Equal(t, "openai", buffer.Connector())
|
||||
})
|
||||
|
||||
t.Run("SetConnector", func(t *testing.T) {
|
||||
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai")
|
||||
assert.Equal(t, "openai", buffer.Connector())
|
||||
|
||||
// Simulate user switching connector mid-conversation
|
||||
buffer.SetConnector("anthropic")
|
||||
assert.Equal(t, "anthropic", buffer.Connector())
|
||||
})
|
||||
|
||||
t.Run("EmptyConnector", func(t *testing.T) {
|
||||
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")
|
||||
|
||||
// Add assistant message - should inherit connector from buffer
|
||||
buffer.AddAssistantMessage(
|
||||
"text",
|
||||
map[string]interface{}{"content": "Hello"},
|
||||
"block-1", "thread-1", "assistant-1", nil,
|
||||
)
|
||||
|
||||
messages := buffer.GetMessages()
|
||||
require.Len(t, messages, 1)
|
||||
assert.Equal(t, "openai", messages[0].Connector)
|
||||
})
|
||||
|
||||
t.Run("MessageConnectorUpdatesWithBuffer", func(t *testing.T) {
|
||||
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai")
|
||||
|
||||
// First message with openai
|
||||
buffer.AddAssistantMessage(
|
||||
"text",
|
||||
map[string]interface{}{"content": "Using OpenAI"},
|
||||
"", "", "assistant-1", nil,
|
||||
)
|
||||
|
||||
// User switches connector
|
||||
buffer.SetConnector("anthropic")
|
||||
|
||||
// Second message with anthropic
|
||||
buffer.AddAssistantMessage(
|
||||
"text",
|
||||
map[string]interface{}{"content": "Now using Claude"},
|
||||
"", "", "assistant-1", nil,
|
||||
)
|
||||
|
||||
messages := buffer.GetMessages()
|
||||
require.Len(t, messages, 2)
|
||||
assert.Equal(t, "openai", messages[0].Connector, "First message should use openai")
|
||||
assert.Equal(t, "anthropic", messages[1].Connector, "Second message should use anthropic")
|
||||
})
|
||||
|
||||
t.Run("UserInputDoesNotSetConnector", func(t *testing.T) {
|
||||
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", "")
|
||||
|
||||
messages := buffer.GetMessages()
|
||||
require.Len(t, messages, 1)
|
||||
// User input messages don't have connector field set by AddUserInput
|
||||
// Connector is only set for assistant messages
|
||||
assert.Equal(t, "", messages[0].Connector)
|
||||
})
|
||||
|
||||
t.Run("MultipleConnectorSwitches", func(t *testing.T) {
|
||||
buffer := context.NewChatBuffer("chat-1", "req-1", "assistant-1", "openai")
|
||||
|
||||
// Simulate a conversation with multiple connector switches
|
||||
connectors := []string{"openai", "anthropic", "openai", "google"}
|
||||
for i, conn := range connectors {
|
||||
buffer.SetConnector(conn)
|
||||
buffer.AddAssistantMessage(
|
||||
"text",
|
||||
map[string]interface{}{"content": fmt.Sprintf("Message %d", i+1)},
|
||||
"", "", "assistant-1", nil,
|
||||
)
|
||||
}
|
||||
|
||||
messages := buffer.GetMessages()
|
||||
require.Len(t, messages, 4)
|
||||
|
||||
for i, msg := range messages {
|
||||
assert.Equal(t, connectors[i], msg.Connector, "Message %d should have connector %s", i+1, connectors[i])
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
|
|
@ -630,7 +727,7 @@ func TestBufferIdentityMethods(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
|
||||
|
|
@ -663,7 +760,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
|
||||
|
|
@ -687,7 +784,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)
|
||||
|
|
@ -761,7 +858,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++ {
|
||||
|
|
@ -778,7 +875,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",
|
||||
|
|
@ -792,7 +889,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)
|
||||
|
|
@ -803,7 +900,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",
|
||||
|
|
@ -819,7 +916,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,
|
||||
|
|
@ -842,7 +939,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")
|
||||
|
|
@ -886,7 +983,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{}{
|
||||
|
|
@ -917,7 +1014,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}
|
||||
|
|
@ -953,7 +1050,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)
|
||||
|
|
@ -997,7 +1094,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{
|
||||
|
|
@ -1013,7 +1110,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("text", nil, "", "", "", nil)
|
||||
|
|
@ -1034,7 +1131,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)
|
||||
|
|
@ -1055,10 +1152,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
|
||||
|
|
|
|||
|
|
@ -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 string) *ChatBuffer {
|
||||
ctx.Buffer = NewChatBuffer(ctx.ChatID, ctx.RequestID(), assistantID)
|
||||
func (ctx *Context) InitBuffer(assistantID, connector string) *ChatBuffer {
|
||||
ctx.Buffer = NewChatBuffer(ctx.ChatID, ctx.RequestID(), assistantID, connector)
|
||||
return ctx.Buffer
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -91,6 +91,7 @@ Stores chat metadata and session information.
|
|||
| `chat_id` | string(64) | No | Unique | Unique chat identifier |
|
||||
| `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") |
|
||||
| `status` | enum | No | Yes | Status: `active`, `archived` |
|
||||
| `public` | boolean | No | - | Whether shared across all teams |
|
||||
|
|
@ -129,6 +130,7 @@ These fields are automatically managed by the framework and used for access cont
|
|||
| Name | Columns | Type |
|
||||
| -------------------- | ----------------- | ----- |
|
||||
| `idx_chat_assistant` | `assistant_id` | index |
|
||||
| `idx_chat_last_conn` | `last_connector` | index |
|
||||
| `idx_chat_status` | `status` | index |
|
||||
| `idx_chat_share` | `share` | index |
|
||||
| `idx_chat_last_msg` | `last_message_at` | index |
|
||||
|
|
@ -151,6 +153,7 @@ Stores user-visible messages (both user input and assistant responses).
|
|||
| `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 |
|
||||
|
|
@ -809,6 +812,7 @@ 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"`
|
||||
Status string `json:"status"` // "active" or "archived"
|
||||
Public bool `json:"public"` // Whether shared across all teams
|
||||
|
|
@ -831,6 +835,7 @@ type Message struct {
|
|||
BlockID string `json:"block_id,omitempty"`
|
||||
ThreadID string `json:"thread_id,omitempty"`
|
||||
AssistantID string `json:"assistant_id,omitempty"`
|
||||
Connector string `json:"connector,omitempty"` // Connector ID used for this message
|
||||
Sequence int `json:"sequence"`
|
||||
Metadata map[string]interface{} `json:"metadata,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ 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 (updated on each message)
|
||||
Mode string `json:"mode"`
|
||||
Status string `json:"status"` // "active" or "archived"
|
||||
Public bool `json:"public"` // Whether shared across all teams
|
||||
|
|
@ -107,6 +108,7 @@ type Message struct {
|
|||
BlockID string `json:"block_id,omitempty"`
|
||||
ThreadID string `json:"thread_id,omitempty"`
|
||||
AssistantID string `json:"assistant_id,omitempty"`
|
||||
Connector string `json:"connector,omitempty"` // Connector ID used for this message
|
||||
Sequence int `json:"sequence"`
|
||||
Metadata map[string]interface{} `json:"metadata,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
|
|
|
|||
|
|
@ -69,6 +69,9 @@ func (store *Xun) CreateChat(chat *types.Chat) error {
|
|||
if chat.Title != "" {
|
||||
data["title"] = chat.Title
|
||||
}
|
||||
if chat.LastConnector != "" {
|
||||
data["last_connector"] = chat.LastConnector
|
||||
}
|
||||
if chat.LastMessageAt != nil {
|
||||
data["last_message_at"] = *chat.LastMessageAt
|
||||
}
|
||||
|
|
@ -330,14 +333,15 @@ func (store *Xun) ListChats(filter types.ChatFilter) (*types.ChatList, error) {
|
|||
// rowToChat converts a database row to a Chat struct
|
||||
func (store *Xun) rowToChat(data map[string]interface{}) (*types.Chat, error) {
|
||||
chat := &types.Chat{
|
||||
ChatID: getString(data, "chat_id"),
|
||||
Title: getString(data, "title"),
|
||||
AssistantID: getString(data, "assistant_id"),
|
||||
Mode: getString(data, "mode"),
|
||||
Status: getString(data, "status"),
|
||||
Public: getBool(data, "public"),
|
||||
Share: getString(data, "share"),
|
||||
Sort: getInt(data, "sort"),
|
||||
ChatID: getString(data, "chat_id"),
|
||||
Title: getString(data, "title"),
|
||||
AssistantID: getString(data, "assistant_id"),
|
||||
LastConnector: getString(data, "last_connector"),
|
||||
Mode: getString(data, "mode"),
|
||||
Status: getString(data, "status"),
|
||||
Public: getBool(data, "public"),
|
||||
Share: getString(data, "share"),
|
||||
Sort: getInt(data, "sort"),
|
||||
}
|
||||
|
||||
// Handle timestamps
|
||||
|
|
|
|||
|
|
@ -53,6 +53,7 @@ func TestCreateChat(t *testing.T) {
|
|||
now := time.Now()
|
||||
chat := &types.Chat{
|
||||
AssistantID: "test_assistant",
|
||||
LastConnector: "openai",
|
||||
Title: "Full Chat",
|
||||
Mode: "task",
|
||||
Status: "active",
|
||||
|
|
@ -80,6 +81,9 @@ func TestCreateChat(t *testing.T) {
|
|||
if retrieved.Title != "Full Chat" {
|
||||
t.Errorf("Expected title 'Full Chat', got '%s'", retrieved.Title)
|
||||
}
|
||||
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)
|
||||
}
|
||||
|
|
@ -313,6 +317,83 @@ func TestUpdateChat(t *testing.T) {
|
|||
_ = store.DeleteChat(chat.ChatID)
|
||||
})
|
||||
|
||||
t.Run("UpdateLastConnector", func(t *testing.T) {
|
||||
chat := &types.Chat{
|
||||
AssistantID: "test_assistant",
|
||||
LastConnector: "openai",
|
||||
}
|
||||
err := store.CreateChat(chat)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create chat: %v", err)
|
||||
}
|
||||
|
||||
// Verify initial connector
|
||||
retrieved, err := store.GetChat(chat.ChatID)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to retrieve chat: %v", err)
|
||||
}
|
||||
if retrieved.LastConnector != "openai" {
|
||||
t.Errorf("Expected last_connector 'openai', got '%s'", retrieved.LastConnector)
|
||||
}
|
||||
|
||||
// Update to different connector (simulating user switching connector)
|
||||
err = store.UpdateChat(chat.ChatID, map[string]interface{}{
|
||||
"last_connector": "anthropic",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to update chat: %v", err)
|
||||
}
|
||||
|
||||
// Verify updated connector
|
||||
retrieved, err = store.GetChat(chat.ChatID)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to retrieve chat: %v", err)
|
||||
}
|
||||
if retrieved.LastConnector != "anthropic" {
|
||||
t.Errorf("Expected last_connector 'anthropic', got '%s'", retrieved.LastConnector)
|
||||
}
|
||||
|
||||
// Clean up
|
||||
_ = store.DeleteChat(chat.ChatID)
|
||||
})
|
||||
|
||||
t.Run("UpdateLastConnectorAndLastMessageAt", func(t *testing.T) {
|
||||
// This simulates what FlushBuffer does
|
||||
chat := &types.Chat{
|
||||
AssistantID: "test_assistant",
|
||||
LastConnector: "openai",
|
||||
}
|
||||
err := store.CreateChat(chat)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create chat: %v", err)
|
||||
}
|
||||
|
||||
// Update both fields together (like FlushBuffer does)
|
||||
now := time.Now()
|
||||
err = store.UpdateChat(chat.ChatID, map[string]interface{}{
|
||||
"last_message_at": now,
|
||||
"last_connector": "claude",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to update chat: %v", err)
|
||||
}
|
||||
|
||||
retrieved, err := store.GetChat(chat.ChatID)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to retrieve chat: %v", err)
|
||||
}
|
||||
|
||||
if retrieved.LastConnector != "claude" {
|
||||
t.Errorf("Expected last_connector 'claude', got '%s'", retrieved.LastConnector)
|
||||
}
|
||||
if retrieved.LastMessageAt == nil {
|
||||
t.Error("Expected last_message_at to be set")
|
||||
}
|
||||
|
||||
// Clean up
|
||||
_ = store.DeleteChat(chat.ChatID)
|
||||
})
|
||||
|
||||
t.Run("UpdateMultipleFields", func(t *testing.T) {
|
||||
chat := &types.Chat{
|
||||
AssistantID: "test_assistant",
|
||||
|
|
|
|||
|
|
@ -68,6 +68,7 @@ func (store *Xun) SaveMessages(chatID string, messages []*types.Message) error {
|
|||
"block_id": nil,
|
||||
"thread_id": nil,
|
||||
"assistant_id": nil,
|
||||
"connector": nil,
|
||||
"metadata": nil,
|
||||
"created_at": now,
|
||||
"updated_at": now,
|
||||
|
|
@ -86,6 +87,9 @@ func (store *Xun) SaveMessages(chatID string, messages []*types.Message) error {
|
|||
if msg.AssistantID != "" {
|
||||
row["assistant_id"] = msg.AssistantID
|
||||
}
|
||||
if msg.Connector != "" {
|
||||
row["connector"] = msg.Connector
|
||||
}
|
||||
if msg.Metadata != nil {
|
||||
metadataJSON, err := jsoniter.MarshalToString(msg.Metadata)
|
||||
if err != nil {
|
||||
|
|
@ -327,6 +331,7 @@ func (store *Xun) rowToMessage(data map[string]interface{}) (*types.Message, err
|
|||
BlockID: getString(data, "block_id"),
|
||||
ThreadID: getString(data, "thread_id"),
|
||||
AssistantID: getString(data, "assistant_id"),
|
||||
Connector: getString(data, "connector"),
|
||||
Sequence: getInt(data, "sequence"),
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -208,6 +208,121 @@ func TestSaveMessages(t *testing.T) {
|
|||
}
|
||||
})
|
||||
|
||||
t.Run("SaveMessageWithConnector", func(t *testing.T) {
|
||||
connChat := &types.Chat{
|
||||
AssistantID: "test_assistant",
|
||||
}
|
||||
err := store.CreateChat(connChat)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create chat: %v", err)
|
||||
}
|
||||
defer store.DeleteChat(connChat.ChatID)
|
||||
|
||||
// Save messages with different connectors
|
||||
messages := []*types.Message{
|
||||
{
|
||||
Role: "user",
|
||||
Type: "user_input",
|
||||
Props: map[string]interface{}{"content": "Hello"},
|
||||
Sequence: 1,
|
||||
Connector: "openai",
|
||||
AssistantID: "test_assistant",
|
||||
},
|
||||
{
|
||||
Role: "assistant",
|
||||
Type: "text",
|
||||
Props: map[string]interface{}{"content": "Hi there!"},
|
||||
Sequence: 2,
|
||||
Connector: "openai",
|
||||
AssistantID: "test_assistant",
|
||||
},
|
||||
{
|
||||
Role: "user",
|
||||
Type: "user_input",
|
||||
Props: map[string]interface{}{"content": "Switch to Claude"},
|
||||
Sequence: 3,
|
||||
Connector: "anthropic",
|
||||
AssistantID: "test_assistant",
|
||||
},
|
||||
{
|
||||
Role: "assistant",
|
||||
Type: "text",
|
||||
Props: map[string]interface{}{"content": "Now using Claude!"},
|
||||
Sequence: 4,
|
||||
Connector: "anthropic",
|
||||
AssistantID: "test_assistant",
|
||||
},
|
||||
}
|
||||
|
||||
err = store.SaveMessages(connChat.ChatID, messages)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to save messages: %v", err)
|
||||
}
|
||||
|
||||
// Retrieve and verify connectors
|
||||
retrieved, err := store.GetMessages(connChat.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 connector
|
||||
for _, msg := range retrieved {
|
||||
if msg.Sequence <= 2 && msg.Connector != "openai" {
|
||||
t.Errorf("Expected connector 'openai' for sequence %d, got '%s'", msg.Sequence, msg.Connector)
|
||||
}
|
||||
if msg.Sequence > 2 && msg.Connector != "anthropic" {
|
||||
t.Errorf("Expected connector 'anthropic' for sequence %d, got '%s'", msg.Sequence, msg.Connector)
|
||||
}
|
||||
}
|
||||
|
||||
t.Logf("Successfully saved and retrieved messages with different connectors")
|
||||
})
|
||||
|
||||
t.Run("SaveMessageWithEmptyConnector", func(t *testing.T) {
|
||||
emptyConnChat := &types.Chat{
|
||||
AssistantID: "test_assistant",
|
||||
}
|
||||
err := store.CreateChat(emptyConnChat)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create chat: %v", err)
|
||||
}
|
||||
defer store.DeleteChat(emptyConnChat.ChatID)
|
||||
|
||||
// Save message without connector
|
||||
messages := []*types.Message{
|
||||
{
|
||||
Role: "user",
|
||||
Type: "text",
|
||||
Props: map[string]interface{}{"content": "No connector"},
|
||||
Sequence: 1,
|
||||
// Connector is empty
|
||||
},
|
||||
}
|
||||
|
||||
err = store.SaveMessages(emptyConnChat.ChatID, messages)
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to save message: %v", err)
|
||||
}
|
||||
|
||||
retrieved, err := store.GetMessages(emptyConnChat.ChatID, types.MessageFilter{})
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to get messages: %v", err)
|
||||
}
|
||||
|
||||
if len(retrieved) != 1 {
|
||||
t.Fatalf("Expected 1 message, got %d", len(retrieved))
|
||||
}
|
||||
|
||||
// Empty connector should be stored as empty string
|
||||
if retrieved[0].Connector != "" {
|
||||
t.Errorf("Expected empty connector, got '%s'", retrieved[0].Connector)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("SaveEmptyMessages", func(t *testing.T) {
|
||||
err := store.SaveMessages(chat.ChatID, []*types.Message{})
|
||||
if err != nil {
|
||||
|
|
|
|||
288
data/bindata.go
288
data/bindata.go
File diff suppressed because it is too large
Load diff
|
|
@ -40,6 +40,15 @@
|
|||
"nullable": false,
|
||||
"index": true
|
||||
},
|
||||
{
|
||||
"name": "last_connector",
|
||||
"type": "string",
|
||||
"label": "Last Connector",
|
||||
"comment": "Last used connector ID (updated on each message)",
|
||||
"length": 200,
|
||||
"nullable": true,
|
||||
"index": true
|
||||
},
|
||||
{
|
||||
"name": "mode",
|
||||
"type": "string",
|
||||
|
|
@ -122,4 +131,3 @@
|
|||
},
|
||||
"option": { "timestamps": true, "soft_deletes": true, "permission": true }
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -92,6 +92,15 @@
|
|||
"nullable": true,
|
||||
"index": true
|
||||
},
|
||||
{
|
||||
"name": "connector",
|
||||
"type": "string",
|
||||
"label": "Connector",
|
||||
"comment": "Connector ID used for this message",
|
||||
"length": 200,
|
||||
"nullable": true,
|
||||
"index": true
|
||||
},
|
||||
{
|
||||
"name": "sequence",
|
||||
"type": "integer",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue