129 lines
2.6 KiB
Go
129 lines
2.6 KiB
Go
package conversation
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"time"
|
|
|
|
"github.com/ajac-zero/latticelm/internal/api"
|
|
"github.com/redis/go-redis/v9"
|
|
)
|
|
|
|
// RedisStore manages conversation history in Redis with automatic expiration.
|
|
type RedisStore struct {
|
|
client *redis.Client
|
|
ttl time.Duration
|
|
}
|
|
|
|
// NewRedisStore creates a Redis-backed conversation store.
|
|
func NewRedisStore(client *redis.Client, ttl time.Duration) *RedisStore {
|
|
return &RedisStore{
|
|
client: client,
|
|
ttl: ttl,
|
|
}
|
|
}
|
|
|
|
// key returns the Redis key for a conversation ID.
|
|
func (s *RedisStore) key(id string) string {
|
|
return "conv:" + id
|
|
}
|
|
|
|
// Get retrieves a conversation by ID from Redis.
|
|
func (s *RedisStore) Get(ctx context.Context, id string) (*Conversation, error) {
|
|
data, err := s.client.Get(ctx, s.key(id)).Bytes()
|
|
if err == redis.Nil {
|
|
return nil, nil
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
var conv Conversation
|
|
if err := json.Unmarshal(data, &conv); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &conv, nil
|
|
}
|
|
|
|
// Create creates a new conversation with the given messages.
|
|
func (s *RedisStore) Create(ctx context.Context, id string, model string, messages []api.Message) (*Conversation, error) {
|
|
now := time.Now()
|
|
conv := &Conversation{
|
|
ID: id,
|
|
Messages: messages,
|
|
Model: model,
|
|
CreatedAt: now,
|
|
UpdatedAt: now,
|
|
}
|
|
|
|
data, err := json.Marshal(conv)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if err := s.client.Set(ctx, s.key(id), data, s.ttl).Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return conv, nil
|
|
}
|
|
|
|
// Append adds new messages to an existing conversation.
|
|
func (s *RedisStore) Append(ctx context.Context, id string, messages ...api.Message) (*Conversation, error) {
|
|
conv, err := s.Get(ctx, id)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if conv == nil {
|
|
return nil, nil
|
|
}
|
|
|
|
conv.Messages = append(conv.Messages, messages...)
|
|
conv.UpdatedAt = time.Now()
|
|
|
|
data, err := json.Marshal(conv)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if err := s.client.Set(ctx, s.key(id), data, s.ttl).Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return conv, nil
|
|
}
|
|
|
|
// Delete removes a conversation from Redis.
|
|
func (s *RedisStore) Delete(ctx context.Context, id string) error {
|
|
return s.client.Del(ctx, s.key(id)).Err()
|
|
}
|
|
|
|
// Size returns the number of active conversations in Redis.
|
|
func (s *RedisStore) Size() int {
|
|
var count int
|
|
var cursor uint64
|
|
ctx := context.Background()
|
|
|
|
for {
|
|
keys, nextCursor, err := s.client.Scan(ctx, cursor, "conv:*", 100).Result()
|
|
if err != nil {
|
|
return 0
|
|
}
|
|
|
|
count += len(keys)
|
|
cursor = nextCursor
|
|
|
|
if cursor == 0 {
|
|
break
|
|
}
|
|
}
|
|
|
|
return count
|
|
}
|
|
|
|
// Close closes the Redis client connection.
|
|
func (s *RedisStore) Close() error {
|
|
return s.client.Close()
|
|
}
|