sundynix-agentix/sundynix-dispatcher/internal/llm/cache.go

package llm

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"os"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/cloudwego/eino/components/model"
	"github.com/cloudwego/eino/schema"
)

// cacheTTLSeconds 是输出缓存的过期时长（env LLM_CACHE_TTL_S，单位秒；0=关闭缓存）。默认 60s：
// 只覆盖"短时间内的重试/重复点击/双发"——这类几乎一定是同一意图，命中省成本+提速；
// 又短到不会让助手对同一问题长期"复读"同一答案（创作型输出仍随时间变化）。
func cacheTTLSeconds() time.Duration {
	// envInt 允许 0（关闭）；默认 60 秒。
	if v := os.Getenv("LLM_CACHE_TTL_S"); v != "" {
		if n, err := strconv.Atoi(v); err == nil && n >= 0 {
			return time.Duration(n) * time.Second
		}
	}
	return 60 * time.Second
}

// cacheMax 是缓存条目上限（env LLM_CACHE_MAX，默认 512）。满了随机淘汰一条。
func cacheMax() int {
	if v := os.Getenv("LLM_CACHE_MAX"); v != "" {
		if n, err := strconv.Atoi(v); err == nil && n > 0 {
			return n
		}
	}
	return 512
}

// respCache 是带 TTL + 容量上限的 Generate 结果缓存（并发安全）。
type respCache struct {
	mu     sync.Mutex
	m      map[string]cacheEntry
	ttl    time.Duration
	max    int
	hits   int64
	misses int64
}

type cacheEntry struct {
	msg *schema.Message
	exp time.Time
}

func newRespCache(ttl time.Duration, max int) *respCache {
	return &respCache{m: make(map[string]cacheEntry), ttl: ttl, max: max}
}

func (c *respCache) get(key string) (*schema.Message, bool) {
	c.mu.Lock()
	defer c.mu.Unlock()
	e, ok := c.m[key]
	if !ok || time.Now().After(e.exp) {
		if ok {
			delete(c.m, key) // 过期清理
		}
		atomic.AddInt64(&c.misses, 1)
		return nil, false
	}
	atomic.AddInt64(&c.hits, 1)
	return e.msg, true
}

func (c *respCache) put(key string, msg *schema.Message) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if len(c.m) >= c.max {
		for k := range c.m { // 满 → 随机淘汰一条（map 迭代序随机）
			delete(c.m, k)
			break
		}
	}
	c.m[key] = cacheEntry{msg: msg, exp: time.Now().Add(c.ttl)}
}

// cachingModel 缓存 Generate 结果：同输入命中则跳过 LLM 调用（省成本+提速）。
// 只缓存 Generate（非流式）；Stream 是用户可见的创作型输出，重复率低且回放复杂，直接透传不缓存。
// 键 = keyExtra（模型名 + 绑定工具哈希）+ 消息内容哈希——确保不同模型/工具集/输入互不串味。
type cachingModel struct {
	inner    model.ToolCallingChatModel
	cache    *respCache
	keyExtra string
}

// withCache 把模型包成带缓存的（ttl<=0 则不包，直接返回原模型）。modelID 用于隔离不同模型的缓存。
func withCache(inner model.ToolCallingChatModel, modelID string) model.ToolCallingChatModel {
	ttl := cacheTTLSeconds()
	if ttl <= 0 {
		return inner
	}
	return &cachingModel{inner: inner, cache: newRespCache(ttl, cacheMax()), keyExtra: "m:" + modelID}
}

func (c *cachingModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {
	key := c.keyExtra + "|" + hashMessages(input)
	if msg, ok := c.cache.get(key); ok {
		fmt.Printf("[llm] 输出缓存命中（跳过 LLM 调用，命中/未命中=%d/%d）\n",
			atomic.LoadInt64(&c.cache.hits), atomic.LoadInt64(&c.cache.misses))
		return msg, nil
	}
	out, err := c.inner.Generate(ctx, input, opts...)
	if err == nil && out != nil {
		c.cache.put(key, out)
	}
	return out, err
}

// Stream 不缓存（用户可见的创作型输出），直接透传。
func (c *cachingModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {
	return c.inner.Stream(ctx, input, opts...)
}

// WithTools 给内层绑工具，返回共享同一缓存但键加了工具哈希的新包装（不同工具集不串味）。
func (c *cachingModel) WithTools(tools []*schema.ToolInfo) (model.ToolCallingChatModel, error) {
	inner, err := c.inner.WithTools(tools)
	if err != nil {
		return nil, err
	}
	return &cachingModel{inner: inner, cache: c.cache, keyExtra: c.keyExtra + "|tools:" + hashTools(tools)}, nil
}

// hashMessages 把消息序列（角色+内容）哈希为稳定 key。
func hashMessages(msgs []*schema.Message) string {
	h := sha256.New()
	for _, m := range msgs {
		fmt.Fprintf(h, "%s\x00%s\x01", m.Role, m.Content)
	}
	return hex.EncodeToString(h.Sum(nil))
}

// hashTools 把绑定工具的名字+描述哈希（工具集不同 → 缓存键不同）。
func hashTools(tools []*schema.ToolInfo) string {
	if len(tools) == 0 {
		return "none"
	}
	h := sha256.New()
	for _, t := range tools {
		fmt.Fprintf(h, "%s\x00%s\x01", t.Name, t.Desc)
	}
	return hex.EncodeToString(h.Sum(nil))[:16]
}