pkg/tools/cache.go

package tools

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"sync"
	"time"
)

// CacheEntry represents a cached tool result
type CacheEntry struct {
	Result    *ToolResult
	Timestamp time.Time
	TTL       time.Duration
}

// ToolCache provides a simple in-memory cache for tool results
type ToolCache struct {
	entries sync.Map
	ttl     time.Duration
}

// NewToolCache creates a new tool cache with default TTL
func NewToolCache(defaultTTL time.Duration) *ToolCache {
	c := &ToolCache{
		ttl: defaultTTL,
	}
	
	// Start cleanup goroutine
	go c.cleanupLoop()
	
	return c
}

// Get retrieves a cached result if available and not expired
func (c *ToolCache) Get(key string) (*ToolResult, bool) {
	val, ok := c.entries.Load(key)
	if !ok {
		return nil, false
	}

	entry := val.(CacheEntry)
	if time.Since(entry.Timestamp) > entry.TTL {
		c.entries.Delete(key)
		return nil, false
	}

	return entry.Result, true
}

// Set stores a tool result in the cache
func (c *ToolCache) Set(key string, result *ToolResult, ttl time.Duration) {
	if ttl == 0 {
		ttl = c.ttl
	}
	
	c.entries.Store(key, CacheEntry{
		Result:    result,
		Timestamp: time.Now(),
		TTL:       ttl,
	})
}

// GenerateKey generates a unique cache key for a tool execution
func (c *ToolCache) GenerateKey(toolName string, args map[string]interface{}) string {
	// Sort keys to ensure consistent hashing? 
	// JSON marshaling of maps in Go is sorted by key order since Go 1.
	// But let's rely on json.Marshal for simplicity, assuming consistent output.
	argsJSON, err := json.Marshal(args)
	if err != nil {
		// Fallback for unmarshalable args (should happen rarely for tool args)
		return fmt.Sprintf("%s:%v", toolName, args)
	}
	
	hash := sha256.Sum256(argsJSON)
	return fmt.Sprintf("%s:%s", toolName, hex.EncodeToString(hash[:]))
}

func (c *ToolCache) cleanupLoop() {
	ticker := time.NewTicker(5 * time.Minute)
	for range ticker.C {
		c.entries.Range(func(key, value interface{}) bool {
			entry := value.(CacheEntry)
			if time.Since(entry.Timestamp) > entry.TTL {
				c.entries.Delete(key)
			}
			return true
		})
	}
}

// IsCacheable checks if a tool is suitable for caching based on its name
func IsCacheable(toolName string) bool {
	switch toolName {
	case "read_file", "list_dir", "web_fetch", "web_search":
		return true
	default:
		return false
	}
}