Merge pull request k8sgpt-ai#324 from patrickpichler/feature/323/stor…

…e-cache feat: use OS conform path for storing cached results
ruanxin · Apr 25, 2023 · 5d6f3c3 · 5d6f3c3
2 parents c3cc413 + 7eddb8f
commit 5d6f3c3
Show file tree

Hide file tree

Showing 7 changed files with 133 additions and 35 deletions.
diff --git a/pkg/ai/iai.go b/pkg/ai/iai.go
@@ -15,12 +15,14 @@ package ai
 
 import (
 	"context"
+
+	"github.com/k8sgpt-ai/k8sgpt/pkg/cache"
 )
 
 type IAI interface {
 	Configure(config IAIConfig, language string) error
 	GetCompletion(ctx context.Context, prompt string) (string, error)
-	Parse(ctx context.Context, prompt []string, nocache bool) (string, error)
+	Parse(ctx context.Context, prompt []string, cache cache.ICache) (string, error)
 	GetName() string
 }
 

diff --git a/pkg/ai/noopai.go b/pkg/ai/noopai.go
@@ -20,8 +20,8 @@ import (
 	"strings"
 
 	"github.com/fatih/color"
+	"github.com/k8sgpt-ai/k8sgpt/pkg/cache"
 	"github.com/k8sgpt-ai/k8sgpt/pkg/util"
-	"github.com/spf13/viper"
 )
 
 type NoOpAIClient struct {
@@ -44,7 +44,7 @@ func (c *NoOpAIClient) GetCompletion(ctx context.Context, prompt string) (string
 	return response, nil
 }
 
-func (a *NoOpAIClient) Parse(ctx context.Context, prompt []string, nocache bool) (string, error) {
+func (a *NoOpAIClient) Parse(ctx context.Context, prompt []string, cache cache.ICache) (string, error) {
 	// parse the text with the AI backend
 	inputKey := strings.Join(prompt, " ")
 	// Check for cached data
@@ -57,13 +57,13 @@ func (a *NoOpAIClient) Parse(ctx context.Context, prompt []string, nocache bool)
 		return "", err
 	}
 
-	if !viper.IsSet(cacheKey) {
-		viper.Set(cacheKey, base64.StdEncoding.EncodeToString([]byte(response)))
-		if err := viper.WriteConfig(); err != nil {
-			color.Red("error writing config: %v", err)
-			return "", nil
-		}
+  err = cache.Store(cacheKey, base64.StdEncoding.EncodeToString([]byte(response)))
+
+	if err != nil {
+		color.Red("error storing value to cache: %v", err)
+		return "", nil
 	}
+
 	return response, nil
 }
 

diff --git a/pkg/ai/openai.go b/pkg/ai/openai.go
@@ -20,12 +20,12 @@ import (
 	"fmt"
 	"strings"
 
+	"github.com/k8sgpt-ai/k8sgpt/pkg/cache"
 	"github.com/k8sgpt-ai/k8sgpt/pkg/util"
 
 	"github.com/sashabaranov/go-openai"
 
 	"github.com/fatih/color"
-	"github.com/spf13/viper"
 )
 
 type OpenAIClient struct {
@@ -70,15 +70,20 @@ func (c *OpenAIClient) GetCompletion(ctx context.Context, prompt string) (string
 	return resp.Choices[0].Message.Content, nil
 }
 
-func (a *OpenAIClient) Parse(ctx context.Context, prompt []string, nocache bool) (string, error) {
+func (a *OpenAIClient) Parse(ctx context.Context, prompt []string, cache cache.ICache) (string, error) {
 	inputKey := strings.Join(prompt, " ")
 	// Check for cached data
 	sEnc := base64.StdEncoding.EncodeToString([]byte(inputKey))
 	cacheKey := util.GetCacheKey(a.GetName(), a.language, sEnc)
 	// find in viper cache
-	if viper.IsSet(cacheKey) && !nocache {
+	if cache.Exists(cacheKey) {
 		// retrieve data from cache
-		response := viper.GetString(cacheKey)
+		response, err := cache.Load(cacheKey)
+
+		if err != nil {
+			return "", err
+		}
+
 		if response == "" {
 			color.Red("error retrieving cached data")
 			return "", nil
@@ -96,13 +101,13 @@ func (a *OpenAIClient) Parse(ctx context.Context, prompt []string, nocache bool)
 		return "", err
 	}
 
-	if !viper.IsSet(cacheKey) || nocache {
-		viper.Set(cacheKey, base64.StdEncoding.EncodeToString([]byte(response)))
-		if err := viper.WriteConfig(); err != nil {
-			color.Red("error writing config: %v", err)
-			return "", nil
-		}
+	err = cache.Store(cacheKey, base64.StdEncoding.EncodeToString([]byte(response)))
+
+	if err != nil {
+		color.Red("error storing value to cache: %v", err)
+		return "", nil
 	}
+
 	return response, nil
 }
 

diff --git a/pkg/analysis/analysis.go b/pkg/analysis/analysis.go
@@ -25,6 +25,7 @@ import (
 	"github.com/fatih/color"
 	"github.com/k8sgpt-ai/k8sgpt/pkg/ai"
 	"github.com/k8sgpt-ai/k8sgpt/pkg/analyzer"
+	"github.com/k8sgpt-ai/k8sgpt/pkg/cache"
 	"github.com/k8sgpt-ai/k8sgpt/pkg/common"
 	"github.com/k8sgpt-ai/k8sgpt/pkg/kubernetes"
 	"github.com/k8sgpt-ai/k8sgpt/pkg/util"
@@ -33,14 +34,14 @@ import (
 )
 
 type Analysis struct {
-	Context        context.Context
-	Filters        []string
-	Client         *kubernetes.Client
-	AIClient       ai.IAI
-	Results        []common.Result
-	Namespace      string
-	NoCache        bool
-	Explain        bool
+	Context   context.Context
+	Filters   []string
+	Client    *kubernetes.Client
+	AIClient  ai.IAI
+	Results   []common.Result
+	Namespace string
+	Cache     cache.ICache
+	Explain   bool
 	MaxConcurrency int
 }
 
@@ -101,13 +102,13 @@ func NewAnalysis(backend string, language string, filters []string, namespace st
 	}
 
 	return &Analysis{
-		Context:        ctx,
-		Filters:        filters,
-		Client:         client,
-		AIClient:       aiClient,
-		Namespace:      namespace,
-		NoCache:        noCache,
-		Explain:        explain,
+		Context:   ctx,
+		Filters:   filters,
+		Client:    client,
+		AIClient:  aiClient,
+		Namespace: namespace,
+		Cache:     cache.New(noCache),
+		Explain:   explain,
 		MaxConcurrency: maxConcurrency,
 	}, nil
 }
@@ -229,7 +230,7 @@ func (a *Analysis) GetAIResults(output string, anonymize bool) error {
 			}
 			texts = append(texts, failure.Text)
 		}
-		parsedText, err := a.AIClient.Parse(a.Context, texts, a.NoCache)
+		parsedText, err := a.AIClient.Parse(a.Context, texts, a.Cache)
 		if err != nil {
 			// FIXME: can we avoid checking if output is json multiple times?
 			//   maybe implement the progress bar better?

diff --git a/pkg/cache/cache.go b/pkg/cache/cache.go
@@ -0,0 +1,15 @@
+package cache
+
+type ICache interface {
+	Store(key string, data string) error
+	Load(key string) (string, error)
+	Exists(key string) bool
+}
+
+func New(noCache bool) ICache {
+	if noCache {
+		return &NoopCache{}
+	}
+
+	return &FileBasedCache{}
+}
diff --git a/pkg/cache/file_based.go b/pkg/cache/file_based.go
@@ -0,0 +1,58 @@
+package cache
+
+import (
+	"fmt"
+	"os"
+	"path/filepath"
+
+	"github.com/adrg/xdg"
+	"github.com/k8sgpt-ai/k8sgpt/pkg/util"
+)
+
+var _ (ICache) = (*FileBasedCache)(nil)
+
+type FileBasedCache struct{}
+
+func (*FileBasedCache) Exists(key string) bool {
+	path, err := xdg.CacheFile(filepath.Join("k8sgpt", key))
+
+	if err != nil {
+		fmt.Fprintln(os.Stderr, "warning: error while testing if cache key exists:", err)
+		return false
+	}
+
+	exists, err := util.FileExists(path)
+
+	if err != nil {
+		fmt.Fprintln(os.Stderr, "warning: error while testing if cache key exists:", err)
+		return false
+	}
+
+	return exists
+}
+
+func (*FileBasedCache) Load(key string) (string, error) {
+	path, err := xdg.CacheFile(filepath.Join("k8sgpt", key))
+
+	if err != nil {
+		return "", err
+	}
+
+	data, err := os.ReadFile(path)
+
+	if err != nil {
+		return "", err
+	}
+
+	return string(data), nil
+}
+
+func (*FileBasedCache) Store(key string, data string) error {
+	path, err := xdg.CacheFile(filepath.Join("k8sgpt", key))
+
+	if err != nil {
+		return err
+	}
+
+	return os.WriteFile(path, []byte(data), 0600)
+}
diff --git a/pkg/cache/noop.go b/pkg/cache/noop.go
@@ -0,0 +1,17 @@
+package cache
+
+var _ (ICache) = (*NoopCache)(nil)
+
+type NoopCache struct{}
+
+func (c *NoopCache) Store(key string, data string) error {
+	return nil
+}
+
+func (c *NoopCache) Load(key string) (string, error) {
+	return "", nil
+}
+
+func (c *NoopCache) Exists(key string) bool {
+	return false
+}