84 lines
2.3 KiB
Go
84 lines
2.3 KiB
Go
package ai
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/zeromicro/go-zero/core/stores/redis"
|
|
)
|
|
|
|
const defaultRedisNamespace = "haixun:dev:v1"
|
|
|
|
// RedisModelsCache shares model-list entries and refresh locks across processes.
|
|
type RedisModelsCache struct {
|
|
client *redis.Redis
|
|
namespace string
|
|
}
|
|
|
|
func NewRedisModelsCache(conf redis.RedisConf, namespace string) *RedisModelsCache {
|
|
return &RedisModelsCache{client: redis.MustNewRedis(conf), namespace: normalizeRedisNamespace(namespace)}
|
|
}
|
|
|
|
func (c *RedisModelsCache) Get(ctx context.Context, key ModelsCacheKey) (ModelsCacheValue, bool, error) {
|
|
raw, err := c.client.GetCtx(ctx, c.valueKey(key))
|
|
if err != nil {
|
|
return ModelsCacheValue{}, false, err
|
|
}
|
|
if raw == "" {
|
|
return ModelsCacheValue{}, false, nil
|
|
}
|
|
var value ModelsCacheValue
|
|
if err := json.Unmarshal([]byte(raw), &value); err != nil {
|
|
return ModelsCacheValue{}, false, err
|
|
}
|
|
return value, true, nil
|
|
}
|
|
|
|
func (c *RedisModelsCache) Set(ctx context.Context, key ModelsCacheKey, value ModelsCacheValue, ttl time.Duration) error {
|
|
raw, err := json.Marshal(value)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return c.client.SetexCtx(ctx, c.valueKey(key), string(raw), durationSeconds(ttl))
|
|
}
|
|
|
|
func (c *RedisModelsCache) AcquireRefresh(ctx context.Context, key ModelsCacheKey, token string, ttl time.Duration) (bool, error) {
|
|
return c.client.SetnxExCtx(ctx, c.lockKey(key), token, durationSeconds(ttl))
|
|
}
|
|
|
|
func (c *RedisModelsCache) ReleaseRefresh(ctx context.Context, key ModelsCacheKey, token string) error {
|
|
_, err := c.client.EvalCtx(ctx, `
|
|
if redis.call('GET', KEYS[1]) == ARGV[1] then
|
|
return redis.call('DEL', KEYS[1])
|
|
end
|
|
return 0
|
|
`, []string{c.lockKey(key)}, token)
|
|
return err
|
|
}
|
|
|
|
func (c *RedisModelsCache) valueKey(key ModelsCacheKey) string {
|
|
return fmt.Sprintf("%s:ai:models:{%s:%s}:value", c.namespace, key.Provider, key.CredentialFingerprint)
|
|
}
|
|
|
|
func (c *RedisModelsCache) lockKey(key ModelsCacheKey) string {
|
|
return fmt.Sprintf("%s:ai:models:{%s:%s}:lock", c.namespace, key.Provider, key.CredentialFingerprint)
|
|
}
|
|
|
|
func normalizeRedisNamespace(namespace string) string {
|
|
if namespace = strings.Trim(strings.TrimSpace(namespace), ":"); namespace != "" {
|
|
return namespace
|
|
}
|
|
return defaultRedisNamespace
|
|
}
|
|
|
|
func durationSeconds(ttl time.Duration) int {
|
|
seconds := int(ttl.Seconds())
|
|
if seconds < 1 {
|
|
return 1
|
|
}
|
|
return seconds
|
|
}
|