120 lines
3.4 KiB
Go
120 lines
3.4 KiB
Go
package main
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/mattermost/mattermost/server/public/model"
|
|
)
|
|
|
|
// ConversationStore manages conversation context in the Mattermost KV store.
|
|
type ConversationStore struct {
|
|
api pluginKV
|
|
mu sync.RWMutex
|
|
}
|
|
|
|
// pluginKV is the subset of the plugin API needed for KV operations.
|
|
type pluginKV interface {
|
|
KVGet(key string) ([]byte, *model.AppError)
|
|
KVSet(key string, value []byte) *model.AppError
|
|
KVDelete(key string) *model.AppError
|
|
}
|
|
|
|
// NewConversationStore creates a new store backed by the plugin API.
|
|
func NewConversationStore(api pluginKV) *ConversationStore {
|
|
return &ConversationStore{api: api}
|
|
}
|
|
|
|
// conversationSession holds the messages for a single user+channel session.
|
|
type conversationSession struct {
|
|
UserID string `json:"user_id"`
|
|
ChannelID string `json:"channel_id"`
|
|
Messages []chatMsg `json:"messages"`
|
|
UpdatedAt time.Time `json:"updated_at"`
|
|
}
|
|
|
|
type chatMsg struct {
|
|
Role string `json:"role"`
|
|
Content string `json:"content"`
|
|
}
|
|
|
|
const (
|
|
maxTurns = 10
|
|
sessionTTL = 1 * time.Hour
|
|
)
|
|
|
|
func sessionKey(userID, channelID string) string {
|
|
return fmt.Sprintf("ctx_%s_%s", userID, channelID)
|
|
}
|
|
|
|
// BuildMessages retrieves session history and appends the user prompt,
|
|
// returning the full message array for the Ollama API.
|
|
func (s *ConversationStore) BuildMessages(userID, channelID, prompt string) ([]ChatMessage, error) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
session := s.loadSession(userID, channelID)
|
|
if session == nil {
|
|
session = &conversationSession{
|
|
UserID: userID,
|
|
ChannelID: channelID,
|
|
Messages: []chatMsg{},
|
|
}
|
|
}
|
|
|
|
session.Messages = append(session.Messages, chatMsg{Role: "user", Content: prompt})
|
|
|
|
// Trim to max turns (keep newest).
|
|
if len(session.Messages) > maxTurns {
|
|
session.Messages = session.Messages[len(session.Messages)-maxTurns:]
|
|
}
|
|
|
|
session.UpdatedAt = time.Now()
|
|
_ = s.saveSession(session) // non-fatal
|
|
|
|
result := make([]ChatMessage, 0, len(session.Messages))
|
|
for _, m := range session.Messages {
|
|
result = append(result, ChatMessage{Role: m.Role, Content: m.Content})
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
// Reset clears the session for a user+channel.
|
|
func (s *ConversationStore) Reset(userID, channelID string) error {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
return s.api.KVDelete(sessionKey(userID, channelID))
|
|
}
|
|
|
|
func (s *ConversationStore) loadSession(userID, channelID string) *conversationSession {
|
|
data, err := s.api.KVGet(sessionKey(userID, channelID))
|
|
if err != nil || data == nil {
|
|
return nil
|
|
}
|
|
var session conversationSession
|
|
if err := json.Unmarshal(data, &session); err != nil {
|
|
return nil
|
|
}
|
|
return &session
|
|
}
|
|
|
|
func (s *ConversationStore) saveSession(session *conversationSession) error {
|
|
data, err := json.Marshal(session)
|
|
if err != nil {
|
|
return fmt.Errorf("marshal session: %w", err)
|
|
}
|
|
err = s.api.KVSet(sessionKey(session.UserID, session.ChannelID), data)
|
|
return err
|
|
}
|
|
|
|
// stopWordMatch checks if the trimmed, lowercased prompt matches a stop word exactly.
|
|
func stopWordMatch(prompt string, stopWords map[string]bool) bool {
|
|
if stopWords == nil {
|
|
return false
|
|
}
|
|
return stopWords[strings.TrimSpace(strings.ToLower(prompt))]
|
|
}
|