wizardchen
266b9a5943
- Updated agentChat API to accept optional knowledge_base_ids and agent_enabled parameters. - Introduced getSession API to retrieve session details. - Modified useStream to validate knowledge_base_ids for agent-chat requests and prepare POST body accordingly. - Enhanced createKnowledgeBase API to include additional configuration options. - Refactored chat routing to remove knowledge base ID from the chat path, allowing for dynamic knowledge base selection. - Improved session management by preventing duplicate session entries in the menu store. - Updated UI components for better session handling and display. - Added context management for agent conversations to improve state persistence across sessions.
66 lines
1.8 KiB
Go
66 lines
1.8 KiB
Go
package llmcontext
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
|
|
"github.com/Tencent/WeKnora/internal/logger"
|
|
"github.com/Tencent/WeKnora/internal/models/chat"
|
|
)
|
|
|
|
// memoryStorage implements ContextStorage using in-memory storage
|
|
type memoryStorage struct {
|
|
sessions map[string][]chat.Message
|
|
mu sync.RWMutex
|
|
}
|
|
|
|
// NewMemoryStorage creates a new memory-based storage
|
|
func NewMemoryStorage() ContextStorage {
|
|
return &memoryStorage{
|
|
sessions: make(map[string][]chat.Message),
|
|
}
|
|
}
|
|
|
|
// Save saves messages for a session to memory
|
|
func (ms *memoryStorage) Save(ctx context.Context, sessionID string, messages []chat.Message) error {
|
|
ms.mu.Lock()
|
|
defer ms.mu.Unlock()
|
|
|
|
// Make a copy to avoid external modifications
|
|
messageCopy := make([]chat.Message, len(messages))
|
|
copy(messageCopy, messages)
|
|
|
|
ms.sessions[sessionID] = messageCopy
|
|
logger.Debugf(ctx, "[MemoryStorage][Session-%s] Saved %d messages to memory", sessionID, len(messages))
|
|
return nil
|
|
}
|
|
|
|
// Load loads messages for a session from memory
|
|
func (ms *memoryStorage) Load(ctx context.Context, sessionID string) ([]chat.Message, error) {
|
|
ms.mu.RLock()
|
|
defer ms.mu.RUnlock()
|
|
|
|
messages, exists := ms.sessions[sessionID]
|
|
if !exists {
|
|
logger.Debugf(ctx, "[MemoryStorage][Session-%s] No context found in memory", sessionID)
|
|
return []chat.Message{}, nil
|
|
}
|
|
|
|
// Return a copy to avoid external modifications
|
|
messageCopy := make([]chat.Message, len(messages))
|
|
copy(messageCopy, messages)
|
|
|
|
logger.Debugf(ctx, "[MemoryStorage][Session-%s] Loaded %d messages from memory", sessionID, len(messages))
|
|
return messageCopy, nil
|
|
}
|
|
|
|
// Delete deletes all messages for a session from memory
|
|
func (ms *memoryStorage) Delete(ctx context.Context, sessionID string) error {
|
|
ms.mu.Lock()
|
|
defer ms.mu.Unlock()
|
|
|
|
delete(ms.sessions, sessionID)
|
|
logger.Debugf(ctx, "[MemoryStorage][Session-%s] Deleted context from memory", sessionID)
|
|
return nil
|
|
}
|