64 lines
1.4 KiB
Go
64 lines
1.4 KiB
Go
package cache
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"fmt"
|
|
"time"
|
|
|
|
"github.com/redis/go-redis/v9"
|
|
)
|
|
|
|
type Cache struct {
|
|
client *redis.Client
|
|
ttl time.Duration
|
|
}
|
|
|
|
func New(addr string, ttlSeconds int) (*Cache, error) {
|
|
client := redis.NewClient(&redis.Options{
|
|
Addr: addr,
|
|
})
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
|
|
if err := client.Ping(ctx).Err(); err != nil {
|
|
return nil, fmt.Errorf("connecting to Valkey: %w", err)
|
|
}
|
|
|
|
ttl := time.Duration(ttlSeconds) * time.Second
|
|
if ttl == 0 {
|
|
ttl = 1 * time.Hour
|
|
}
|
|
|
|
return &Cache{client: client, ttl: ttl}, nil
|
|
}
|
|
|
|
func (c *Cache) Get(ctx context.Context, model string, requestBody []byte) ([]byte, error) {
|
|
key := c.cacheKey(model, requestBody)
|
|
data, err := c.client.Get(ctx, key).Bytes()
|
|
if err == redis.Nil {
|
|
return nil, nil
|
|
}
|
|
return data, err
|
|
}
|
|
|
|
func (c *Cache) Set(ctx context.Context, model string, requestBody, responseBody []byte) error {
|
|
key := c.cacheKey(model, requestBody)
|
|
return c.client.Set(ctx, key, responseBody, c.ttl).Err()
|
|
}
|
|
|
|
func (c *Cache) Ping(ctx context.Context) error {
|
|
return c.client.Ping(ctx).Err()
|
|
}
|
|
|
|
func (c *Cache) Close() error {
|
|
return c.client.Close()
|
|
}
|
|
|
|
func (c *Cache) cacheKey(model string, requestBody []byte) string {
|
|
h := sha256.New()
|
|
h.Write([]byte(model))
|
|
h.Write(requestBody)
|
|
return fmt.Sprintf("llm-gw:%x", h.Sum(nil))
|
|
}
|