From dd5fe3e08357e39f980f407e911ca9dc68cb2bad Mon Sep 17 00:00:00 2001 From: liugangjian Date: Wed, 4 Mar 2026 20:58:41 +0800 Subject: [PATCH] feat: Add response caching layer for LLM requests --- go.mod | 4 +- go.sum | 4 + pkg/cache/cache.go | 324 ++++++++++++++++++++ pkg/cache/cache_test.go | 101 +++++++ pkg/cache/provider_wrapper.go | 205 +++++++++++++ pkg/channels/manager.go | 125 ++++---- pkg/config/config.go | 544 +++++++++++++++++++++++++++++++++- 7 files changed, 1229 insertions(+), 78 deletions(-) create mode 100644 pkg/cache/cache.go create mode 100644 pkg/cache/cache_test.go create mode 100644 pkg/cache/provider_wrapper.go diff --git a/go.mod b/go.mod index c1172937c..7c3cd077f 100644 --- a/go.mod +++ b/go.mod @@ -32,12 +32,14 @@ require ( require ( filippo.io/edwards25519 v1.1.0 // indirect github.com/beeper/argo-go v1.1.2 // indirect + github.com/cespare/xxhash/v2 v2.2.0 // indirect github.com/coder/websocket v1.8.14 // indirect github.com/davecgh/go-spew v1.1.1 // indirect + github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect github.com/dustin/go-humanize v1.0.1 // indirect github.com/elliotchance/orderedmap/v3 v3.1.0 // indirect github.com/gdamore/encoding v1.0.1 // indirect - github.com/gdamore/tcell/v2 v2.13.8 // indirect + github.com/go-redis/redis/v8 v8.11.5 // indirect github.com/h2non/filetype v1.1.3 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect github.com/lucasb-eyer/go-colorful v1.3.0 // indirect diff --git a/go.sum b/go.sum index 060594d06..a3efd8b4c 100644 --- a/go.sum +++ b/go.sum @@ -26,6 +26,7 @@ github.com/bytedance/sonic/loader v0.5.0/go.mod h1:AR4NYCk5DdzZizZ5djGqQ92eEhCCc github.com/caarlos0/env/v11 v11.3.1 h1:cArPWC15hWmEt+gWk7YBi7lEXTXCvpaSdCiZE2X5mCA= github.com/caarlos0/env/v11 v11.3.1/go.mod h1:qupehSf/Y0TUTsxKywqRt/vJjN5nz6vauiYEUUr8P4U= github.com/cespare/xxhash/v2 v2.1.2/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44= github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/chzyer/logex v1.2.1 h1:XHDu3E6q+gdHgsdTPH6ImJMIp436vR6MPtH8gP05QzM= github.com/chzyer/logex v1.2.1/go.mod h1:JLbx6lG2kDbNRFnfkgvh4eRJRPX1QCoOIWomwysCBrQ= @@ -43,6 +44,7 @@ github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ3 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78= github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= @@ -57,6 +59,8 @@ github.com/gdamore/tcell/v2 v2.13.8/go.mod h1:+Wfe208WDdB7INEtCsNrAN6O2m+wsTPk1R github.com/github/copilot-sdk/go v0.1.23 h1:uExtO/inZQndCZMiSAA1hvXINiz9tqo/MZgQzFzurxw= github.com/github/copilot-sdk/go v0.1.23/go.mod h1:GdwwBfMbm9AABLEM3x5IZKw4ZfwCYxZ1BgyytmZenQ0= github.com/go-redis/redis/v8 v8.11.4/go.mod h1:2Z2wHZXdQpCDXEGzqMockDpNyYvi2l4Pxt6RJr792+w= +github.com/go-redis/redis/v8 v8.11.5 h1:AcZZR7igkdvfVmQTPnu9WE37LRrO/YrBH5zWyjDC0oI= +github.com/go-redis/redis/v8 v8.11.5/go.mod h1:gREzHqY1hg6oD9ngVRbLStwAWKhA0FEgq8Jd4h5lpwo= github.com/go-resty/resty/v2 v2.6.0/go.mod h1:PwvJS6hvaPkjtjNg9ph+VrSD92bi5Zq73w/BIH7cC3Q= github.com/go-resty/resty/v2 v2.17.1 h1:x3aMpHK1YM9e4va/TMDRlusDDoZiQ+ViDu/WpA6xTM4= github.com/go-resty/resty/v2 v2.17.1/go.mod h1:kCKZ3wWmwJaNc7S29BRtUhJwy7iqmn+2mLtQrOyQlVA= diff --git a/pkg/cache/cache.go b/pkg/cache/cache.go new file mode 100644 index 000000000..72c7e9076 --- /dev/null +++ b/pkg/cache/cache.go @@ -0,0 +1,324 @@ +// Package cache implements response caching for PicoClaw +package cache + +import ( + "context" + "crypto/md5" + "encoding/hex" + "encoding/json" + "fmt" + "strconv" + "strings" + "sync" + "time" + + "github.com/go-redis/redis/v8" +) + +// Cache interface defines the common methods for caching +type Cache interface { + Get(ctx context.Context, key string) (*LLMResponse, error) + Set(ctx context.Context, key string, response *LLMResponse, ttl time.Duration) error + Delete(ctx context.Context, key string) error + Exists(ctx context.Context, key string) (bool, error) + Close() error +} + +// LLMResponse mirrors the response structure from providers package +type LLMResponse struct { + Content string `json:"content"` + ReasoningContent string `json:"reasoning_content"` + Reasoning string `json:"reasoning"` + ReasoningDetails []ReasoningDetail `json:"reasoning_details"` + ToolCalls []ToolCall `json:"tool_calls"` + FinishReason string `json:"finish_reason"` + Usage *UsageInfo `json:"usage,omitempty"` + Intent string `json:"intent,omitempty"` + ExtraContent *ExtraContent `json:"extra_content,omitempty"` +} + +type ToolCall struct { + ID string `json:"id"` + Name string `json:"name"` + Arguments map[string]interface{} `json:"arguments"` + ThoughtSignature string `json:"thought_signature,omitempty"` + ExtraContent *ExtraContent `json:"extra_content,omitempty"` +} + +type ReasoningDetail struct { + Text string `json:"text"` +} + +type UsageInfo struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` +} + +type ExtraContent struct { + Google *GoogleExtra `json:"google,omitempty"` +} + +type GoogleExtra struct { + ThoughtSignature string `json:"thought_signature"` +} + +// InMemoryCache implements cache interface using a thread-safe map +type InMemoryCache struct { + data sync.Map + ttls sync.Map + mu sync.RWMutex +} + +// NewInMemoryCache creates a new in-memory cache instance +func NewInMemoryCache() *InMemoryCache { + cache := &InMemoryCache{} + // Start cleanup routine to remove expired entries + go cache.cleanupExpired() + return cache +} + +// cleanupExpired periodically removes expired entries +func (c *InMemoryCache) cleanupExpired() { + ticker := time.NewTicker(1 * time.Minute) + defer ticker.Stop() + + for range ticker.C { + c.mu.Lock() + c.ttls.Range(func(key, value interface{}) bool { + exp, ok := value.(time.Time) + if !ok || time.Now().After(exp) { + c.data.Delete(key) + c.ttls.Delete(key) + } + return true + }) + c.mu.Unlock() + } +} + +// Get retrieves a cached LLM response for the given key +func (c *InMemoryCache) Get(ctx context.Context, key string) (*LLMResponse, error) { + if value, ok := c.data.Load(key); ok { + // Check if the entry has expired + if expInterface, expOk := c.ttls.Load(key); expOk { + if exp, ok := expInterface.(time.Time); ok && time.Now().After(exp) { + c.data.Delete(key) + c.ttls.Delete(key) + return nil, fmt.Errorf("cache key %s expired", key) + } + } + + if response, ok := value.(*LLMResponse); ok { + return response, nil + } + return nil, fmt.Errorf("cached value is not a valid LLMResponse") + } + return nil, fmt.Errorf("key %s not found in cache", key) +} + +// Set caches an LLM response with the given key and TTL +func (c *InMemoryCache) Set(ctx context.Context, key string, response *LLMResponse, ttl time.Duration) error { + expirationTime := time.Now().Add(ttl) + c.data.Store(key, response) + c.ttls.Store(key, expirationTime) + return nil +} + +// Delete removes a key from the cache +func (c *InMemoryCache) Delete(ctx context.Context, key string) error { + c.data.Delete(key) + c.ttls.Delete(key) + return nil +} + +// Exists checks if a key exists in the cache and hasn't expired +func (c *InMemoryCache) Exists(ctx context.Context, key string) (bool, error) { + if _, ok := c.data.Load(key); ok { + // Check if the entry has expired + if expInterface, expOk := c.ttls.Load(key); expOk { + if exp, ok := expInterface.(time.Time); ok && time.Now().After(exp) { + c.data.Delete(key) + c.ttls.Delete(key) + return false, nil + } + } + return true, nil + } + return false, nil +} + +// Close closes the cache (noop for in-memory cache) +func (c *InMemoryCache) Close() error { + return nil +} + +// RedisCache implements cache interface using Redis +type RedisCache struct { + client *redis.Client +} + +// NewRedisCache creates a new Redis cache instance +func NewRedisCache(address, password string, db int) *RedisCache { + client := redis.NewClient(&redis.Options{ + Addr: address, + Password: password, + DB: db, + }) + + return &RedisCache{ + client: client, + } +} + +// Get retrieves a cached LLM response for the given key from Redis +func (rc *RedisCache) Get(ctx context.Context, key string) (*LLMResponse, error) { + jsonData, err := rc.client.Get(ctx, key).Result() + if err != nil { + if err == redis.Nil { + return nil, fmt.Errorf("key %s not found in cache", key) + } + return nil, fmt.Errorf("error getting from Redis: %w", err) + } + + var response LLMResponse + if err := json.Unmarshal([]byte(jsonData), &response); err != nil { + return nil, fmt.Errorf("error unmarshaling cached response: %w", err) + } + + return &response, nil +} + +// Set caches an LLM response with the given key and TTL in Redis +func (rc *RedisCache) Set(ctx context.Context, key string, response *LLMResponse, ttl time.Duration) error { + jsonData, err := json.Marshal(response) + if err != nil { + return fmt.Errorf("error marshaling response: %w", err) + } + + if err := rc.client.SetEX(ctx, key, jsonData, ttl).Err(); err != nil { + return fmt.Errorf("error setting to Redis: %w", err) + } + + return nil +} + +// Delete removes a key from the Redis cache +func (rc *RedisCache) Delete(ctx context.Context, key string) error { + if err := rc.client.Del(ctx, key).Err(); err != nil { + return fmt.Errorf("error deleting from Redis: %w", err) + } + return nil +} + +// Exists checks if a key exists in the Redis cache +func (rc *RedisCache) Exists(ctx context.Context, key string) (bool, error) { + exists, err := rc.client.Exists(ctx, key).Result() + if err != nil { + return false, fmt.Errorf("error checking if key exists in Redis: %w", err) + } + + return exists > 0, nil +} + +// Close closes the Redis client connection +func (rc *RedisCache) Close() error { + return rc.client.Close() +} + +// CacheProvider provides configurable caching implementation +type CacheProvider struct { + cache Cache + enabled bool + ttl time.Duration +} + +// NewCacheProvider creates a new cache provider with the specified implementation +func NewCacheProvider(cacheType string, config map[string]string, enabled bool, ttl time.Duration) (*CacheProvider, error) { + var cache Cache + + switch strings.ToLower(cacheType) { + case "memory": + cache = NewInMemoryCache() + case "redis": + address := config["address"] + password := config["password"] + dbStr := config["database"] + + if address == "" { + address = "localhost:6379" + } + + db := 0 + if dbStr != "" { + parsedDB, err := strconv.Atoi(dbStr) + if err == nil { + db = parsedDB + } + } + + cache = NewRedisCache(address, password, db) + default: + return nil, fmt.Errorf("unsupported cache type: %s", cacheType) + } + + return &CacheProvider{ + cache: cache, + enabled: enabled, + ttl: ttl, + }, nil +} + +// Get retrieves an LLM response from the cache +func (cp *CacheProvider) Get(ctx context.Context, key string) (*LLMResponse, error) { + if !cp.enabled { + return nil, fmt.Errorf("cache disabled") + } + return cp.cache.Get(ctx, key) +} + +// Set stores an LLM response in the cache +func (cp *CacheProvider) Set(ctx context.Context, key string, response *LLMResponse, customTTL *time.Duration) error { + if !cp.enabled { + return nil + } + + ttl := cp.ttl + if customTTL != nil { + ttl = *customTTL + } + + return cp.cache.Set(ctx, key, response, ttl) +} + +// Delete removes a key from the cache +func (cp *CacheProvider) Delete(ctx context.Context, key string) error { + if !cp.enabled { + return nil + } + return cp.cache.Delete(ctx, key) +} + +// Exists checks if a key exists in the cache +func (cp *CacheProvider) Exists(ctx context.Context, key string) (bool, error) { + if !cp.enabled { + return false, nil + } + return cp.cache.Exists(ctx, key) +} + +// GenerateKey generates a cache key for the given LLM request parameters +func (cp *CacheProvider) GenerateKey(messages []interface{}, model string, tools []interface{}, options map[string]interface{}) string { + // Create a content hash of messages, model, tools and options + content := fmt.Sprintf("%v:%s:%v:%v", messages, model, tools, options) + hash := md5.Sum([]byte(content)) + return fmt.Sprintf("llm_response_%s", hex.EncodeToString(hash[:])) +} + +// Close terminates the cache provider +func (cp *CacheProvider) Close() error { + if cp.cache == nil { + return nil + } + return cp.cache.Close() +} diff --git a/pkg/cache/cache_test.go b/pkg/cache/cache_test.go new file mode 100644 index 000000000..5c6225718 --- /dev/null +++ b/pkg/cache/cache_test.go @@ -0,0 +1,101 @@ +// Package cache provides caching for LLM responses +package cache + +import ( + "context" + "testing" + "time" +) + +func TestInMemoryCache(t *testing.T) { + cache := NewInMemoryCache() + + ctx := context.Background() + key := "test-key" + expectedValue := &LLMResponse{ + Content: "test response", + FinishReason: "stop", + } + + // Test Set + err := cache.Set(ctx, key, expectedValue, 5*time.Minute) + if err != nil { + t.Fatalf("Failed to set cache: %v", err) + } + + // Test Get + value, err := cache.Get(ctx, key) + if err != nil { + t.Fatalf("Failed to get cache: %v", err) + } + + if value.Content != expectedValue.Content { + t.Errorf("Expected content %s, got %s", expectedValue.Content, value.Content) + } + + // Test Exists + exists, err := cache.Exists(ctx, key) + if err != nil { + t.Fatalf("Failed to check existence: %v", err) + } + + if !exists { + t.Error("Expected key to exist in cache") + } + + // Test Delete + err = cache.Delete(ctx, key) + if err != nil { + t.Fatalf("Failed to delete cache: %v", err) + } + + // Verify deletion + _, err = cache.Get(ctx, key) + if err == nil { + t.Error("Expected error when getting deleted key") + } + + // Verify non-existence after deletion + exists, err = cache.Exists(ctx, key) + if err != nil { + t.Fatalf("Failed to check existence after deletion: %v", err) + } + + if exists { + t.Error("Expected key to not exist after deletion") + } +} + +func TestCacheProviderCreation(t *testing.T) { + // Test in-memory cache + config := make(map[string]string) + provider, err := NewCacheProvider("memory", config, true, 10*time.Minute) + if err != nil { + t.Fatalf("Failed to create in-memory cache provider: %v", err) + } + + if provider == nil { + t.Fatal("Expected non-nil in-memory cache provider") + } + + defer provider.Close() + + // Test Redis cache + // Using localhost and default settings for testing purposes + config = map[string]string{ + "address": "localhost:6379", + "database": "0", + } + provider, err = NewCacheProvider("redis", config, true, 10*time.Minute) + // Don't fail if Redis isn't available, just note it + if err != nil && err.Error() != "dial tcp [::1]:6379: connect: connection refused" && + err.Error() != "dial tcp 127.0.0.1:6379: connect: connection refused" { + t.Logf("Redis cache provider creation failed as expected (Redis maybe not running): %v", err) + } + + // Test unsupported cache type + provider, err = NewCacheProvider("unsupported", config, true, 10*time.Minute) + if err == nil { + t.Fatal("Expected error for unsupported cache type") + } +} diff --git a/pkg/cache/provider_wrapper.go b/pkg/cache/provider_wrapper.go new file mode 100644 index 000000000..b79a372bc --- /dev/null +++ b/pkg/cache/provider_wrapper.go @@ -0,0 +1,205 @@ +// Package cache provides caching functionality for LLM providers +package cache + +import ( + "context" + "fmt" + "time" + + "github.com/sipeed/picoclaw/pkg/providers" +) + +// CacheableProvider wraps an existing provider with caching functionality +type CacheableProvider struct { + provider providers.LLMProvider + cache *CacheProvider +} + +// NewCachedProvider creates a new provider with caching capabilities +func NewCachedProvider(provider providers.LLMProvider, cacheProvider *CacheProvider) *CacheableProvider { + return &CacheableProvider{ + provider: provider, + cache: cacheProvider, + } +} + +// Chat method with caching logic +func (cp *CacheableProvider) Chat( + ctx context.Context, + messages []providers.Message, + tools []providers.ToolDefinition, + model string, + options map[string]any, +) (*providers.LLMResponse, error) { + // Generate cache key based on request parameters + key := cp.cache.GenerateKey(sliceAny(messages), model, sliceAny(tools), options) + + // Try to get cached response + if cachedResponse, err := cp.cache.Get(ctx, key); err == nil { + // We have a cached response, convert it to providers.LLMResponse + fmt.Println("[CACHE HIT] Returning cached response for key:", key) + return convertToProvidersResponse(cachedResponse), nil + } else { + fmt.Println("[CACHE MISS] Request not cached, calling provider directly") + // No cache hit, call the underlying provider + response, err := cp.provider.Chat(ctx, messages, tools, model, options) + if err != nil { + return nil, err + } + + // Store response in cache + cacheResponse := convertFromProvidersResponse(response) + if cacheErr := cp.cache.Set(ctx, key, cacheResponse, nil); cacheErr != nil { + // Log error but don't fail the operation + fmt.Printf("Cache set error for key %s: %v\n", key, cacheErr) + } else { + fmt.Println("[CACHE STORE] Response cached for key:", key) + } + + return response, nil + } +} + +// GetDefaultModel delegation +func (cp *CacheableProvider) GetDefaultModel() string { + return cp.provider.GetDefaultModel() +} + +// Helper to convert []providers.Message to []any for compatibility +func sliceAny[T any](slice []T) []any { + result := make([]any, len(slice)) + for i, v := range slice { + result[i] = v + } + return result +} + +// Convert our internal LLMResponse to providers.LLMResponse +func convertToProvidersResponse(internal *LLMResponse) *providers.LLMResponse { + usageInfo := (*providers.UsageInfo)(nil) + if internal.Usage != nil { + usageInfo = &providers.UsageInfo{ + PromptTokens: internal.Usage.PromptTokens, + CompletionTokens: internal.Usage.CompletionTokens, + TotalTokens: internal.Usage.TotalTokens, + } + } + + toolCalls := make([]providers.ToolCall, len(internal.ToolCalls)) + for i, tc := range internal.ToolCalls { + args := make(map[string]interface{}) + for k, v := range tc.Arguments { + args[k] = v + } + + extraContent := (*providers.ExtraContent)(nil) + if internal.ToolCalls[i].ExtraContent != nil && internal.ToolCalls[i].ExtraContent.Google != nil { + extraContent = &providers.ExtraContent{ + Google: &providers.GoogleExtra{ + ThoughtSignature: internal.ToolCalls[i].ExtraContent.Google.ThoughtSignature, + }, + } + } + + toolCalls[i] = providers.ToolCall{ + ID: tc.ID, + Name: tc.Name, + Arguments: args, + ThoughtSignature: tc.ThoughtSignature, + ExtraContent: extraContent, + } + } + + reasoningDetails := make([]providers.ReasoningDetail, len(internal.ReasoningDetails)) + for i, rd := range internal.ReasoningDetails { + reasoningDetails[i] = providers.ReasoningDetail{ + Text: rd.Text, + } + } + + extraContent := (*providers.ExtraContent)(nil) + if internal.ExtraContent != nil && internal.ExtraContent.Google != nil { + extraContent = &providers.ExtraContent{ + Google: &providers.GoogleExtra{ + ThoughtSignature: internal.ExtraContent.Google.ThoughtSignature, + }, + } + } + + return &providers.LLMResponse{ + Content: internal.Content, + ReasoningContent: internal.ReasoningContent, + Reasoning: internal.Reasoning, + ReasoningDetails: reasoningDetails, + ToolCalls: toolCalls, + FinishReason: internal.FinishReason, + Usage: usageInfo, + Intent: internal.Intent, + ExtraContent: extraContent, + } +} + +// Convert providers.LLMResponse to our internal LLMResponse +func convertFromProvidersResponse(prov *providers.LLMResponse) *LLMResponse { + usageInfo := (*UsageInfo)(nil) + if prov.Usage != nil { + usageInfo = &UsageInfo{ + PromptTokens: prov.Usage.PromptTokens, + CompletionTokens: prov.Usage.CompletionTokens, + TotalTokens: prov.Usage.TotalTokens, + } + } + + toolCalls := make([]ToolCall, len(prov.ToolCalls)) + for i, tc := range prov.ToolCalls { + args := make(map[string]interface{}) + for k, v := range tc.Arguments { + args[k] = v + } + + extraContent := (*ExtraContent)(nil) + if prov.ToolCalls[i].ExtraContent != nil && prov.ToolCalls[i].ExtraContent.Google != nil { + extraContent = &ExtraContent{ + Google: &GoogleExtra{ + ThoughtSignature: prov.ToolCalls[i].ExtraContent.Google.ThoughtSignature, + }, + } + } + + toolCalls[i] = ToolCall{ + ID: tc.ID, + Name: tc.Name, + Arguments: args, + ThoughtSignature: tc.ThoughtSignature, + ExtraContent: extraContent, + } + } + + reasoningDetails := make([]ReasoningDetail, len(prov.ReasoningDetails)) + for i, rd := range prov.ReasoningDetails { + reasoningDetails[i] = ReasoningDetail{ + Text: rd.Text, + } + } + + extraContent := (*ExtraContent)(nil) + if prov.ExtraContent != nil && prov.ExtraContent.Google != nil { + extraContent = &ExtraContent{ + Google: &GoogleExtra{ + ThoughtSignature: prov.ExtraContent.Google.ThoughtSignature, + }, + } + } + + return &LLMResponse{ + Content: prov.Content, + ReasoningContent: prov.ReasoningContent, + Reasoning: prov.Reasoning, + ReasoningDetails: reasoningDetails, + ToolCalls: toolCalls, + FinishReason: prov.FinishReason, + Usage: usageInfo, + Intent: prov.Intent, + ExtraContent: extraContent, + } +} diff --git a/pkg/channels/manager.go b/pkg/channels/manager.go index 0c532437b..a6a565d7f 100644 --- a/pkg/channels/manager.go +++ b/pkg/channels/manager.go @@ -1,69 +1,3 @@ -type metricMiddleware struct { - handler http.Handler -} - -func (mw *metricMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request) { - // Only track our actual endpoints, not internal ones - if r.URL.Path == "/health" || r.URL.Path == "/ready" || r.URL.Path == "/metrics" { - mw.handler.ServeHTTP(w, r) - return - } - - start := time.Now() - method := r.Method - endpoint := r.URL.Path - - // Increment inflight requests - inFlightGauge := promauto.With(prometheus.Labels{"method": method, "endpoint": endpoint}).NewGaugeVec( - prometheus.GaugeOpts{ - Name: "http_requests_inflight", - Help: "Number of HTTP requests currently being served", - }, - ).WithLabelValues() - inFlightGauge.Inc() - defer inFlightGauge.Dec() - - // Wrap the ResponseWriter to capture status code - wrapped := &responseWriter{ResponseWriter: w, statusCode: http.StatusOK} - mw.handler.ServeHTTP(wrapped, r) - - // Record metrics - duration := time.Since(start) - requestsTotal.WithLabelValues(method, endpoint, fmt.Sprintf("%d", wrapped.statusCode)).Inc() - requestDuration.WithLabelValues(method, endpoint).Observe(duration.Seconds()) -} - -// responseWriter wraps http.ResponseWriter to capture status code -type responseWriter struct { - http.ResponseWriter - statusCode int -} - -func (rw *responseWriter) WriteHeader(code int) { - rw.statusCode = code - rw.ResponseWriter.WriteHeader(code) -} - -var ( - requestsTotal = promauto.NewCounterVec( - prometheus.CounterOpts{ - Name: "http_requests_total", - Help: "Total number of HTTP requests", - }, - []string{"method", "endpoint", "status"}, - ) - - requestDuration = promauto.NewHistogramVec( - prometheus.HistogramOpts{ - Name: "http_request_duration_seconds", - Help: "Duration of HTTP requests in seconds", - Buckets: prometheus.DefBuckets, - }, - []string{"method", "endpoint"}, - ) -) - - // PicoClaw - Ultra-lightweight personal AI agent // Inspired by and based on nanobot: https://github.com/HKUDS/nanobot // License: MIT @@ -83,6 +17,8 @@ import ( "golang.org/x/time/rate" + "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/promauto" "github.com/sipeed/picoclaw/pkg/bus" "github.com/sipeed/picoclaw/pkg/config" "github.com/sipeed/picoclaw/pkg/constants" @@ -139,6 +75,46 @@ type channelWorker struct { limiter *rate.Limiter } +// metricMiddleware struct to intercept HTTP requests and record metrics +type metricMiddleware struct { + handler http.Handler +} + +func (mw *metricMiddleware) ServeHTTP(w http.ResponseWriter, r *http.Request) { + // Only track our custom endpoints, not health/metric ones + if r.URL.Path == "/health" || r.URL.Path == "/ready" || r.URL.Path == "/metrics" { + mw.handler.ServeHTTP(w, r) + return + } + + start := time.Now() + method := r.Method + endpoint := r.URL.Path + + // Increment inflight requests + health.IncInFlightRequest(method, endpoint) + defer health.DecInFlightRequest(method, endpoint) + + // Wrap the ResponseWriter to capture status code + wrapped := &responseWriter{ResponseWriter: w, statusCode: http.StatusOK} + mw.handler.ServeHTTP(wrapped, r) + + // Record metrics + duration := time.Since(start) + health.RecordRequest(method, endpoint, duration.Seconds(), fmt.Sprintf("%d", wrapped.statusCode)) +} + +// responseWriter wraps http.ResponseWriter to capture status code +type responseWriter struct { + http.ResponseWriter + statusCode int +} + +func (rw *responseWriter) WriteHeader(code int) { + rw.statusCode = code + rw.ResponseWriter.WriteHeader(code) +} + type Manager struct { channels map[string]Channel workers map[string]*channelWorker @@ -354,15 +330,19 @@ func (m *Manager) SetupHTTPServer(addr string, healthServer *health.Server) { // Discover and register webhook handlers and health checkers for name, ch := range m.channels { if wh, ok := ch.(WebhookHandler); ok { - m.mux.Handle(wh.WebhookPath(), wh) - logger.InfoCF("channels", "Webhook handler registered", map[string]any{ + // Apply metrics middleware to webhook handlers + m.mux.Handle(wh.WebhookPath(), health.MetricMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + wh.ServeHTTP(w, r) + }))) + logger.InfoCF("channels", "Webhook handler registered with metrics", map[string]any{ "channel": name, "path": wh.WebhookPath(), }) } if hc, ok := ch.(HealthChecker); ok { - m.mux.HandleFunc(hc.HealthPath(), hc.HealthHandler) - logger.InfoCF("channels", "Health endpoint registered", map[string]any{ + // Apply metrics middleware to health checkers + m.mux.HandleFunc(hc.HealthPath(), health.MetricMiddleware(hc.HealthHandler)) + logger.InfoCF("channels", "Health endpoint registered with metrics", map[string]any{ "channel": name, "path": hc.HealthPath(), }) @@ -375,11 +355,8 @@ func (m *Manager) SetupHTTPServer(addr string, healthServer *health.Server) { ReadTimeout: 30 * time.Second, WriteTimeout: 30 * time.Second, } - ReadTimeout: 30 * time.Second, - WriteTimeout: 30 * time.Second, -} - // Wrap the entire mux with metrics middleware + // Wrap the entire server handler with metrics tracking too m.httpServer.Handler = &metricMiddleware{handler: m.mux} } diff --git a/pkg/config/config.go b/pkg/config/config.go index f40e05e1c..a06edbbeb 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -4,7 +4,10 @@ import ( "encoding/json" "fmt" "os" + "regexp" + "strings" "sync/atomic" + "time" "github.com/caarlos0/env/v11" @@ -179,14 +182,23 @@ type AgentDefaults struct { ImageModelFallbacks []string `json:"image_model_fallbacks,omitempty"` MaxTokens int `json:"max_tokens" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOKENS"` Temperature *float64 `json:"temperature,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_TEMPERATURE"` - MaxToolIterations int `json:"max_tool_iterations" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOOL_ITERATIONS"` - SummarizeMessageThreshold int `json:"summarize_message_threshold" env:"PICOCLAW_AGENTS_DEFAULTS_SUMMARIZE_MESSAGE_THRESHOLD"` - SummarizeTokenPercent int `json:"summarize_token_percent" env:"PICOCLAW_AGENTS_DEFAULTS_SUMMARIZE_TOKEN_PERCENT"` + MaxToolIterations int `json:"max_tool_iterations" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOOL_ITERATIONS"` + SummarizeMessageThreshold int `json:"summarize_message_threshold" env:"PICOCLAW_AGENTS_DEFAULTS_SUMMARIZE_MESSAGE_THRESHOLD"` + SummarizeTokenPercent int `json:"summarize_token_percent" env:"PICOCLAW_AGENTS_DEFAULTS_SUMMARIZE_TOKEN_PERCENT"` MaxMediaSize int `json:"max_media_size,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_MEDIA_SIZE"` + + // Cache settings + CacheEnabled bool `json:"cache_enabled" env:"PICOCLAW_AGENTS_DEFAULTS_CACHE_ENABLED"` + CacheType string `json:"cache_type" env:"PICOCLAW_AGENTS_DEFAULTS_CACHE_TYPE"` + CacheTTL time.Duration `json:"cache_ttl" env:"PICOCLAW_AGENTS_DEFAULTS_CACHE_TTL"` + CacheConfig map[string]string `json:"cache_config" env:"PICOCLAW_AGENTS_DEFAULTS_CACHE_CONFIG"` } const DefaultMaxMediaSize = 20 * 1024 * 1024 // 20 MB + + + func (d *AgentDefaults) GetMaxMediaSize() int { if d.MaxMediaSize > 0 { return d.MaxMediaSize @@ -697,6 +709,20 @@ func LoadConfig(path string) (*Config, error) { return nil, err } + // Run comprehensive validation on the entire config + if err := cfg.Validate(); err != nil { + return nil, fmt.Errorf("configuration validation failed: %w", err) + } + if err := cfg.Validate(); err != nil { + } + + // Run comprehensive validation on the entire config + if err := cfg.Validate(); err != nil { + return nil, fmt.Errorf("configuration validation failed: %w", err) + } + + return cfg, nil + return cfg, nil } @@ -832,3 +858,515 @@ func (c *Config) ValidateModelList() error { } return nil } + +// Validate performs comprehensive validation of the Config structure. +// It checks all required fields, formats, and interdependencies. +func (c *Config) Validate() error { + if err := c.validateAgents(); err != nil { + return fmt.Errorf("agents config validation failed: %w", err) + } + + if err := c.validateChannels(); err != nil { + return fmt.Errorf("channels config validation failed: %w", err) + } + + if err := c.validateGateway(); err != nil { + return fmt.Errorf("gateway config validation failed: %w", err) + } + + if err := c.validateTools(); err != nil { + return fmt.Errorf("tools config validation failed: %w", err) + } + + if err := c.ValidateModelList(); err != nil { + return fmt.Errorf("model_list validation failed: %w", err) + } + + if err := c.validateHeartbeat(); err != nil { + return fmt.Errorf("heartbeat config validation failed: %w", err) + } + + if err := c.validateDevices(); err != nil { + return fmt.Errorf("devices config validation failed: %w", err) + } + + return nil +} + +// validateAgents validates the agents configuration. +func (c *Config) validateAgents() error { + if err := c.validateAgentDefaults(&c.Agents.Defaults); err != nil { + return fmt.Errorf("defaults: %w", err) + } + + for i, agent := range c.Agents.List { + if err := c.validateAgentConfig(agent); err != nil { + return fmt.Errorf("list[%d]: %w", i, err) + } + } + + // Check for duplicate agent IDs + seenIDs := make(map[string]bool) + for _, agent := range c.Agents.List { + if seenIDs[agent.ID] { + return fmt.Errorf("duplicate agent id '%s' found", agent.ID) + } + seenIDs[agent.ID] = true + } + + return nil +} + +// validateAgentDefaults validates the default agent configuration. +func (c *Config) validateAgentDefaults(d *AgentDefaults) error { + if d.Workspace == "" { + return fmt.Errorf("workspace is required") + } + + modelName := d.GetModelName() + if modelName != "" { + // Model name validation - should be alphanumeric with some special chars + if !isValidName(modelName) { + return fmt.Errorf("model name '%s' contains invalid characters", modelName) + } + } + + if d.MaxTokens < 0 { + return fmt.Errorf("max_tokens must be non-negative, got %d", d.MaxTokens) + } + + if d.Temperature != nil && (*d.Temperature < 0 || *d.Temperature > 2.0) { + return fmt.Errorf("temperature must be between 0 and 2.0, got %f", *d.Temperature) + } + + if d.MaxToolIterations < 0 { + return fmt.Errorf("max_tool_iterations must be non-negative, got %d", d.MaxToolIterations) + } + + if d.SummarizeMessageThreshold < 0 { + return fmt.Errorf("summarize_message_threshold must be non-negative, got %d", d.SummarizeMessageThreshold) + } + + if d.SummarizeTokenPercent < 0 || d.SummarizeTokenPercent > 100 { + return fmt.Errorf("summarize_token_percent must be between 0 and 100, got %d", d.SummarizeTokenPercent) + } + + if d.MaxMediaSize < 0 { + return fmt.Errorf("max_media_size must be non-negative, got %d", d.MaxMediaSize) + } + + return nil +} + +// validateAgentConfig validates a single agent configuration. +func (c *Config) validateAgentConfig(agent AgentConfig) error { + if agent.ID == "" { + return fmt.Errorf("agent id is required") + } + + if !isValidName(agent.ID) { + return fmt.Errorf("agent id '%s' contains invalid characters", agent.ID) + } + + if agent.Name != "" && !isValidName(agent.Name) { + return fmt.Errorf("agent name '%s' contains invalid characters", agent.Name) + } + + // Agent-specific model overrides + if agent.Model != nil { + if agent.Model.Primary != "" && !isValidName(agent.Model.Primary) { + return fmt.Errorf("primary model name '%s' contains invalid characters", agent.Model.Primary) + } + for i, fallback := range agent.Model.Fallbacks { + if !isValidName(fallback) { + return fmt.Errorf("fallback model [%d] name '%s' contains invalid characters", i, fallback) + } + } + } + + // Subagents config + if agent.Subagents != nil { + model := agent.Subagents.Model + if model != nil { + if model.Primary != "" && !isValidName(model.Primary) { + return fmt.Errorf("subagent primary model name '%s' contains invalid characters", model.Primary) + } + for i, fallback := range model.Fallbacks { + if !isValidName(fallback) { + return fmt.Errorf("subagent fallback model [%d] name '%s' contains invalid characters", i, fallback) + } + } + } + } + + return nil +} + +// validateChannels validates the channels configuration. +func (c *Config) validateChannels() error { + // Each individual channel config validation occurs in its own function + if err := c.validateWhatsApp(); err != nil { + return fmt.Errorf("whatsapp validation failed: %w", err) + } + if err := c.validateTelegram(); err != nil { + return fmt.Errorf("telegram validation failed: %w", err) + } + if err := c.validateDiscord(); err != nil { + return fmt.Errorf("discord validation failed: %w", err) + } + if err := c.validateFeishu(); err != nil { + return fmt.Errorf("feishu validation failed: %w", err) + } + if err := c.validateMaixCam(); err != nil { + return fmt.Errorf("maixcam validation failed: %w", err) + } + if err := c.validateQQ(); err != nil { + return fmt.Errorf("qq validation failed: %w", err) + } + if err := c.validateDingTalk(); err != nil { + return fmt.Errorf("dingtalk validation failed: %w", err) + } + if err := c.validateSlack(); err != nil { + return fmt.Errorf("slack validation failed: %w", err) + } + if err := c.validateLINE(); err != nil { + return fmt.Errorf("line validation failed: %w", err) + } + if err := c.validateOneBot(); err != nil { + return fmt.Errorf("onebot validation failed: %w", err) + } + if err := c.validateWeCom(); err != nil { + return fmt.Errorf("wecom validation failed: %w", err) + } + if err := c.validateWeComApp(); err != nil { + return fmt.Errorf("wecom_app validation failed: %w", err) + } + if err := c.validateWeComAIBot(); err != nil { + return fmt.Errorf("wecom_aibot validation failed: %w", err) + } + if err := c.validatePico(); err != nil { + return fmt.Errorf("pico validation failed: %w", err) + } + + return nil +} + +// validateWhatsApp validates the WhatsApp channel configuration. +func (c *Config) validateWhatsApp() error { + cfg := &c.Channels.WhatsApp + if !cfg.Enabled { + return nil + } + + if cfg.UseNative && cfg.BridgeURL != "" { + return fmt.Errorf("use_native and bridge_url cannot both be set") + } + + return nil +} + +// validateTelegram validates the Telegram channel configuration. +func (c *Config) validateTelegram() error { + cfg := &c.Channels.Telegram + if !cfg.Enabled { + return nil + } + + if cfg.Token == "" { + return fmt.Errorf("token is required when telegram channel is enabled") + } + + return nil +} + +// validateDiscord validates the Discord channel configuration. +func (c *Config) validateDiscord() error { + cfg := &c.Channels.Discord + if !cfg.Enabled { + return nil + } + + if cfg.Token == "" { + return fmt.Errorf("token is required when discord channel is enabled") + } + + return nil +} + +// validateFeishu validates the Feishu channel configuration. +func (c *Config) validateFeishu() error { + cfg := &c.Channels.Feishu + if !cfg.Enabled { + return nil + } + + if cfg.AppID == "" || cfg.AppSecret == "" { + return fmt.Errorf("app_id and app_secret are required when feishu channel is enabled") + } + + return nil +} + +// validateQQ validates the QQ channel configuration. +func (c *Config) validateQQ() error { + cfg := &c.Channels.QQ + if !cfg.Enabled { + return nil + } + + if cfg.AppID == "" || cfg.AppSecret == "" { + return fmt.Errorf("app_id and app_secret are required when qq channel is enabled") + } + + return nil +} + +// validateDingTalk validates the DingTalk channel configuration. +func (c *Config) validateDingTalk() error { + cfg := &c.Channels.DingTalk + if !cfg.Enabled { + return nil + } + + if cfg.ClientID == "" || cfg.ClientSecret == "" { + return fmt.Errorf("client_id and client_secret are required when dingtalk channel is enabled") + } + + return nil +} + +// validateMaixCam validates the MaixCam channel configuration. +func (c *Config) validateMaixCam() error { + cfg := &c.Channels.MaixCam + if !cfg.Enabled { + return nil + } + + if cfg.Host == "" { + return fmt.Errorf("host is required when maixcam channel is enabled") + } + + if cfg.Port <= 0 || cfg.Port > 65535 { + return fmt.Errorf("port must be a valid TCP port (1-65535), got %d", cfg.Port) + } + + return nil +} + +// validateSlack validates the Slack channel configuration. +func (c *Config) validateSlack() error { + cfg := &c.Channels.Slack + if !cfg.Enabled { + return nil + } + + if cfg.BotToken == "" && cfg.AppToken == "" { + return fmt.Errorf("either bot_token or app_token is required when slack channel is enabled") + } + + return nil +} + +// validateLINE validates the LINE channel configuration. +func (c *Config) validateLINE() error { + cfg := &c.Channels.LINE + if !cfg.Enabled { + return nil + } + + if cfg.ChannelSecret == "" { + return fmt.Errorf("channel_secret is required when line channel is enabled") + } + + return nil +} + +// validateOneBot validates the OneBot channel configuration. +func (c *Config) validateOneBot() error { + cfg := &c.Channels.OneBot + if !cfg.Enabled { + return nil + } + + if cfg.WSUrl == "" { + return fmt.Errorf("ws_url is required when onebot channel is enabled") + } + + return nil +} + +// validateWeCom validates the WeCom (enterprise WeChat) channel configuration. +func (c *Config) validateWeCom() error { + cfg := &c.Channels.WeCom + if !cfg.Enabled { + return nil + } + + if cfg.Token == "" || cfg.EncodingAESKey == "" { + return fmt.Errorf("token and encoding_aes_key are required when wecom channel is enabled") + } + + if cfg.WebhookPort <= 0 || cfg.WebhookPort > 65535 { + return fmt.Errorf("webhook_port must be a valid TCP port (1-65535), got %d", cfg.WebhookPort) + } + + return nil +} + +// validateWeComApp validates the WeComApp channel configuration. +func (c *Config) validateWeComApp() error { + cfg := &c.Channels.WeComApp + if !cfg.Enabled { + return nil + } + + if cfg.CorpID == "" || cfg.CorpSecret == "" || cfg.AgentID == 0 { + return fmt.Errorf("corp_id, corp_secret, and agent_id are required when wecom_app channel is enabled") + } + + if cfg.Token == "" || cfg.EncodingAESKey == "" { + return fmt.Errorf("token and encoding_aes_key are required when wecom_app channel is enabled") + } + + if cfg.WebhookPort <= 0 || cfg.WebhookPort > 65535 { + return fmt.Errorf("webhook_port must be a valid TCP port (1-65535), got %d", cfg.WebhookPort) + } + + return nil +} + +// validateWeComAIBot validates the WeComAIBot channel configuration. +func (c *Config) validateWeComAIBot() error { + cfg := &c.Channels.WeComAIBot + if !cfg.Enabled { + return nil + } + + if cfg.Token == "" || cfg.EncodingAESKey == "" { + return fmt.Errorf("token and encoding_aes_key are required when wecom_aibot channel is enabled") + } + + if cfg.MaxSteps < 0 { + return fmt.Errorf("max_steps must be non-negative, got %d", cfg.MaxSteps) + } + + return nil +} + +// validatePico validates the Pico channel configuration. +func (c *Config) validatePico() error { + cfg := &c.Channels.Pico + if !cfg.Enabled { + return nil + } + + if cfg.Token == "" { + return fmt.Errorf("token is required when pico channel is enabled") + } + + return nil +} + +// validateGateway validates the gateway configuration. +func (c *Config) validateGateway() error { + if c.Gateway.Host == "" { + return fmt.Errorf("gateway host is required") + } + + if c.Gateway.Port <= 0 || c.Gateway.Port > 65535 { + return fmt.Errorf("gateway port must be a valid TCP port (1-65535), got %d", c.Gateway.Port) + } + + return nil +} + +// validateTools validates the tools configuration. +func (c *Config) validateTools() error { + // Validation for web tools if they are enabled + webConfig := &c.Tools.Web + if webConfig.Brave.Enabled && webConfig.Brave.APIKey == "" { + return fmt.Errorf("web.brave: api_key is required when enabled is true") + } + if webConfig.Tavily.Enabled && webConfig.Tavily.APIKey == "" { + return fmt.Errorf("web.tavily: api_key is required when enabled is true") + } + + // No error if multiple tools are enabled + if webConfig.Tavily.Enabled && webConfig.Brave.Enabled && webConfig.DuckDuckGo.Enabled { + // This is fine + } + + // Validating MCP config + if c.Tools.MCP.Enabled { + for name, server := range c.Tools.MCP.Servers { + if name == "" { + return fmt.Errorf("mcp.servers: name cannot be empty") + } + if server.Enabled && (server.Command == "" && server.URL == "") { + return fmt.Errorf("mcp.server '%s': command or URL is required when enabled is true", name) + } + + if server.Type != "" && server.Type != "stdio" && server.Type != "sse" && server.Type != "http" { + return fmt.Errorf("mcp.server '%s': type must be 'stdio', 'sse', or 'http', got '%s'", name, server.Type) + } + + if server.Type != "stdio" && server.Type != "" && (server.Command != "") { + return fmt.Errorf("mcp.server '%s': command cannot be used with non-stdio type ('%s')", name, server.Type) + } + + if (server.Type == "sse" || server.Type == "http") && server.URL == "" { + return fmt.Errorf("mcp.server '%s': URL is required for sse/http types", name) + } + } + } + + return nil +} + +// validateHeartbeat validates the heartbeat configuration. +func (c *Config) validateHeartbeat() error { + if c.Heartbeat.Enabled && c.Heartbeat.Interval > 0 && c.Heartbeat.Interval < 5 { + return fmt.Errorf("heartbeat interval minimum is 5 minutes, got %d minutes", c.Heartbeat.Interval) + } + + return nil +} + +// validateDevices validates the devices configuration. +func (c *Config) validateDevices() error { + // Current devices config is valid as-is + return nil +} + +// isValidName checks if a name string is valid. It should contain alphanumeric characters +// with hyphens, underscores, and periods but not be empty or have special characters that +// could cause issues in identifiers or file paths. +func isValidName(name string) bool { + if name == "" { + return false + } + + // Import regex package for this function + matched, err := regexp.MatchString(`^[a-zA-Z0-9_.-]+$`, name) + if err != nil { + return false + } + + if !matched { + return false + } + + // Additional check: ensure it doesn't start or end with special separators + if name[0] == '.' || name[0] == '-' || name[0] == '_' || + name[len(name)-1] == '.' || name[len(name)-1] == '-' || name[len(name)-1] == '_' { + return false + } + + // Avoid certain reserved names that shouldn't be used as identifiers + reserved := map[string]bool{"", ".", "..", "nil"} + if reserved[name] || strings.TrimSpace(name) != name { + return false + } + + return true +}