diff --git a/cmd/llm/llm.go b/cmd/llm/llm.go new file mode 100644 index 0000000..fc60b7c --- /dev/null +++ b/cmd/llm/llm.go @@ -0,0 +1,400 @@ +package custom + +import ( + "bufio" + "context" + "fmt" + "io" + "net/http" + "os" + "strconv" + "strings" + + "connectrpc.com/connect" + "github.com/bitbomdev/minefield/cmd/helpers" + apiv1 "github.com/bitbomdev/minefield/gen/api/v1" + "github.com/bitbomdev/minefield/gen/api/v1/apiv1connect" + "github.com/olekukonko/tablewriter" + chromadb "github.com/philippgille/chromem-go" + "github.com/sashabaranov/go-openai" + "github.com/spf13/cobra" +) + +const PROMPT_TEMPLATE = `You are an AI assistant that helps users understand and work with a DSL (Domain Specific Language) for querying a graph database of supply chain security artifacts. You have access to documentation and examples about this DSL through the provided context. + +If the user asks for a DSL query, convert their natural language into the appropriate DSL script. The DSL uses keywords like: dependencies, dependents, library, vuln, xor, or, and. + +If the user asks general questions about the DSL or how it works, provide helpful explanations based on the context. + +YOU CAN ONLY OUTPUT THE DSL QUERY. NO REGULAR LANGUAGE. + +You cannot output periods, commas, or other punctuation. + +If an '@' is used in a package name, even if a version is not included, leave it in and do not remove it. + +If a user asks for vulnerablities for a package they mean dependencies of type vuln, for a query, and if we want to know what a vuln affects, dependents of type library. + +Globsearch queries can be used to find anything that might exist in a node, not only names of nodes, so for types of packages, if it is inside of a purl you can find it, versions, ecosystem, etc. For vulns you can do globsearches like '*GHSA*', '*CVE*', etc, depending on what the user asks. + +Try to surrond a globsearch pattern with as much glob as you can, be as general as possible. + + +If this is a leaderboard query, you should prefix your answer with 'leaderboard:'. +If this is a regular query, you should prefix your answer with 'query:'. +If this is a globsearch query, you should prefix your answer with 'globsearch:'. + + +Context information: + +%s + +--- + +User question: %s + +Please provide a helpful response. If the question requires a DSL query, format it clearly as a DSL script.` + +// options holds the command-line options. +type options struct { + maxOutput int + showInfo bool + saveQuery string + addr string + output string + queryServiceClient apiv1connect.QueryServiceClient + leaderboardServiceClient apiv1connect.LeaderboardServiceClient + graphServiceClient apiv1connect.GraphServiceClient + vectorDBPath string +} + +// AddFlags adds command-line flags to the provided cobra command. +func (o *options) AddFlags(cmd *cobra.Command) { + cmd.Flags().IntVar(&o.maxOutput, "max-output", 10, "maximum number of results to display") + cmd.Flags().BoolVar(&o.showInfo, "show-info", true, "display the info column") + cmd.Flags().StringVar(&o.addr, "addr", "http://localhost:8089", "address of the minefield server") + cmd.Flags().StringVar(&o.vectorDBPath, "vector-db-path", "./db", "Path to the vector database") + cmd.Flags().StringVar(&o.output, "output", "table", "output format (table or json)") +} + +// Run executes the custom command with the provided arguments. +func (o *options) Run(cmd *cobra.Command, args []string) error { + + if os.Getenv("OPENAI_API_KEY") == "" { + return fmt.Errorf("OPENAI_API_KEY environment variable is not set") + } + db, err := chromadb.NewPersistentDB(o.vectorDBPath, false) + if err != nil { + return fmt.Errorf("failed to initialize ChromaDB: %w", err) + } + + c := db.GetCollection("knowledge-base", nil) + if err != nil { + return fmt.Errorf("failed to get collection from ChromaDB: %w", err) + } + + // Initialize chat messages + messages := []openai.ChatCompletionMessage{ + { + Role: openai.ChatMessageRoleSystem, + Content: PROMPT_TEMPLATE, + }, + } + + client := openai.NewClient(os.Getenv("OPENAI_API_KEY")) + + // Initialize client if not injected (for testing) + if o.queryServiceClient == nil { + o.queryServiceClient = apiv1connect.NewQueryServiceClient( + http.DefaultClient, + o.addr, + connect.WithGRPC(), + connect.WithSendGzip(), + ) + } + + if o.leaderboardServiceClient == nil { + o.leaderboardServiceClient = apiv1connect.NewLeaderboardServiceClient( + http.DefaultClient, + o.addr, + ) + } + + if o.graphServiceClient == nil { + o.graphServiceClient = apiv1connect.NewGraphServiceClient( + http.DefaultClient, + o.addr, + ) + } + + fmt.Println("Starting chat session. Type 'exit' to end.") + + reader := bufio.NewReader(os.Stdin) + for { + fmt.Print("\nYou: ") + input, err := reader.ReadString('\n') + if err != nil { + return fmt.Errorf("error reading input: %w", err) + } + + input = strings.TrimSpace(input) + if strings.ToLower(input) == "exit" { + fmt.Println("Ending chat session. Goodbye!") + return nil + } + + // Get context from ChromaDB query + resultEmbeddings, err := c.Query(context.Background(), input, 13, nil, nil) + if err != nil { + return fmt.Errorf("failed to query ChromaDB: %w", err) + } + + // Build context text from results + var contextText string + for i := 0; i < 13 && i < len(resultEmbeddings); i++ { + contextText += fmt.Sprintf("%s\n\n", resultEmbeddings[i].Content) + } + + // Add user's message + messages = append(messages, openai.ChatCompletionMessage{ + Role: openai.ChatMessageRoleUser, + Content: fmt.Sprintf(PROMPT_TEMPLATE, contextText, input), + }) + + resp, err := client.CreateChatCompletion( + context.Background(), + openai.ChatCompletionRequest{ + Model: openai.GPT4, + Messages: messages, + }, + ) + if err != nil { + return fmt.Errorf("failed to create chat completion: %w", err) + } + + script := resp.Choices[0].Message.Content + + // Execute the query and capture output + var queryResult string + if strings.TrimSpace(script) != "" { + var buf strings.Builder + + // Check if this is a leaderboard query + if strings.HasPrefix(strings.TrimSpace(script), "leaderboard:") { + + + // Remove the "leaderboard:" prefix + cleanScript := strings.TrimPrefix(strings.TrimSpace(script), "leaderboard:") + fmt.Printf("\nAssistant: I'll help you with that. I'm going to use this leaderboard query:\n\"%s\"\n", cleanScript) + + req := connect.NewRequest(&apiv1.CustomLeaderboardRequest{Script: cleanScript}) + res, err := o.leaderboardServiceClient.CustomLeaderboard(cmd.Context(), req) + + if err != nil { + queryResult = fmt.Sprintf("Leaderboard query failed: %v", err) + } else if len(res.Msg.Queries) == 0 { + queryResult = "No results found" + } else { + switch o.output { + case "json": + jsonOutput, err := helpers.FormatCustomQueriesJSON(res.Msg.Queries) + if err != nil { + queryResult = fmt.Sprintf("Failed to format JSON: %v", err) + } else { + queryResult = string(jsonOutput) + } + case "table": + err = formatLeaderboardTable(&buf, res.Msg.Queries, o.maxOutput, o.showInfo) + if err != nil { + queryResult = fmt.Sprintf("Failed to format table: %v", err) + } else { + queryResult = buf.String() + } + } + } + } else if strings.HasPrefix(strings.TrimSpace(script), "query:") { + // Remove the "query:" prefix + cleanScript := strings.TrimPrefix(strings.TrimSpace(script), "query:") + fmt.Printf("\nAssistant: I'll help you with that. I'm going to use this query:\n\"%s\"\n", cleanScript) + + req := connect.NewRequest(&apiv1.QueryRequest{Script: cleanScript}) + res, err := o.queryServiceClient.Query(cmd.Context(), req) + + if err != nil { + queryResult = fmt.Sprintf("Query failed: %v", err) + } else if len(res.Msg.Nodes) == 0 { + queryResult = "No results found" + } else { + switch o.output { + case "json": + jsonOutput, err := helpers.FormatNodeJSON(res.Msg.Nodes) + if err != nil { + queryResult = fmt.Sprintf("Failed to format JSON: %v", err) + } else { + queryResult = string(jsonOutput) + } + case "table": + err = formatTable(&buf, res.Msg.Nodes, o.maxOutput, o.showInfo) + if err != nil { + queryResult = fmt.Sprintf("Failed to format table: %v", err) + } else { + queryResult = buf.String() + } + } + } + } else if strings.HasPrefix(strings.TrimSpace(script), "globsearch:") { + // Remove the "globsearch:" prefix + pattern := strings.TrimPrefix(strings.TrimSpace(script), "globsearch:") + fmt.Printf("\nAssistant: I'll help you with that. I'm going to use this pattern:\n\"%s\"\n", pattern) + + req := connect.NewRequest(&apiv1.GetNodesByGlobRequest{Pattern: pattern}) + res, err := o.graphServiceClient.GetNodesByGlob(cmd.Context(), req) + + if err != nil { + queryResult = fmt.Sprintf("Query failed: %v", err) + } else if len(res.Msg.Nodes) == 0 { + queryResult = "No results found" + } else { + switch o.output { + case "json": + jsonOutput, err := helpers.FormatNodeJSON(res.Msg.Nodes) + if err != nil { + queryResult = fmt.Sprintf("Failed to format JSON: %v", err) + } else { + queryResult = string(jsonOutput) + } + case "table": + err = formatTableGlobSearch(&buf, res.Msg.Nodes, o.maxOutput, o.showInfo) + if err != nil { + queryResult = fmt.Sprintf("Failed to format table: %v", err) + } else { + queryResult = buf.String() + } + } + } + } else { + queryResult = "Sorry the query failed, please try again." + } + } + + // Add results to chat history + feedbackMsg := fmt.Sprintf("Here are the results of the query:\n%s\nWhat else would you like to know?", queryResult) + messages = append(messages, openai.ChatCompletionMessage{ + Role: openai.ChatMessageRoleAssistant, + Content: feedbackMsg, + }) + + fmt.Println(feedbackMsg) + } +} + +// formatTable formats the nodes into a table and writes it to the provided writer. +func formatTable(w io.Writer, nodes []*apiv1.Node, maxOutput int, showInfo bool) error { + table := tablewriter.NewWriter(w) + headers := []string{"Name", "Type", "ID"} + if showInfo { + headers = append(headers, "Info") + } + table.SetHeader(headers) + table.SetAutoWrapText(false) + table.SetRowLine(true) + + count := 0 + for _, node := range nodes { + if count >= maxOutput { + break + } + + row := []string{ + node.Name, + node.Type, + strconv.FormatUint(uint64(node.Id), 10), + } + + if showInfo { + additionalInfo := helpers.ComputeAdditionalInfo(node) + row = append(row, additionalInfo) + } + + table.Append(row) + count++ + } + + table.Render() + return nil +} + +// Add new function to format leaderboard table +func formatLeaderboardTable(w io.Writer, queries []*apiv1.Query, maxOutput int, showInfo bool) error { + table := tablewriter.NewWriter(w) + headers := []string{"Name", "Type", "ID", "Output"} + if showInfo { + headers = append(headers, "Info") + } + table.SetHeader(headers) + table.SetAutoWrapText(false) + table.SetRowLine(true) + + count := 0 + for _, query := range queries { + if count >= maxOutput { + break + } + + row := []string{ + query.Node.Name, + query.Node.Type, + strconv.FormatUint(uint64(query.Node.Id), 10), + fmt.Sprint(len(query.Output)), + } + + if showInfo { + additionalInfo := helpers.ComputeAdditionalInfo(query.Node) + row = append(row, additionalInfo) + } + + table.Append(row) + count++ + } + + table.Render() + return nil +} + +func formatTableGlobSearch(w io.Writer, nodes []*apiv1.Node, maxOutput int, showInfo bool) error { + table := tablewriter.NewWriter(w) + table.SetHeader([]string{"Name", "Type", "ID"}) + table.SetAutoWrapText(false) + table.SetAutoFormatHeaders(true) + + for i, node := range nodes { + if i >= maxOutput { + break + } + table.Append([]string{ + node.Name, + node.Type, + strconv.FormatUint(uint64(node.Id), 10), + }) + } + + table.Render() + return nil +} + +// New creates and returns a new Cobra command for executing custom query scripts. +func New() *cobra.Command { + o := &options{} + + cmd := &cobra.Command{ + Use: "llm [query]", + Short: "Create a chat session with an LLM to query the graph for leaderboards, queries, and globsearches", + Long: "Creates an LLM chat session to query the graph for leaderboards, queries, and globsearches, to end the chat session use the type 'exit'. This does use OpenAI, so you need to have the OPENAI_API_KEY environment variable set.", + Args: cobra.NoArgs, + RunE: o.Run, + DisableAutoGenTag: true, + } + + o.AddFlags(cmd) + + return cmd +} diff --git a/cmd/root/root.go b/cmd/root/root.go index 7261813..98fa7c3 100644 --- a/cmd/root/root.go +++ b/cmd/root/root.go @@ -10,6 +10,7 @@ import ( "github.com/bitbomdev/minefield/cmd/leaderboard" "github.com/bitbomdev/minefield/cmd/query" "github.com/bitbomdev/minefield/cmd/server" + llm "github.com/bitbomdev/minefield/cmd/llm" "github.com/spf13/cobra" ) @@ -54,6 +55,6 @@ func New() *cobra.Command { rootCmd.AddCommand(cache.New()) rootCmd.AddCommand(leaderboard.New()) rootCmd.AddCommand(server.New()) - + rootCmd.AddCommand(llm.New()) return rootCmd } diff --git a/cmd/server/server.go b/cmd/server/server.go index 19333ac..53da8f7 100644 --- a/cmd/server/server.go +++ b/cmd/server/server.go @@ -8,6 +8,7 @@ import ( "net/http" "os" "os/signal" + "runtime" "syscall" "time" @@ -16,6 +17,7 @@ import ( "github.com/bitbomdev/minefield/gen/api/v1/apiv1connect" "github.com/bitbomdev/minefield/pkg/graph" "github.com/bitbomdev/minefield/pkg/storages" + chromadb "github.com/philippgille/chromem-go" "github.com/rs/cors" "github.com/spf13/cobra" "golang.org/x/net/http2" @@ -23,14 +25,16 @@ import ( ) type options struct { - storage graph.Storage - concurrency int32 - addr string - StorageType string - StorageAddr string - StoragePath string - UseInMemory bool - CORS []string + storage graph.Storage + concurrency int32 + addr string + StorageType string + StorageAddr string + StoragePath string + UseInMemory bool + CORS []string + UseOpenAILLM bool + VectorDBPath string } const ( @@ -53,6 +57,8 @@ func (o *options) AddFlags(cmd *cobra.Command) { []string{"http://localhost:8089"}, "Allowed origins for CORS (e.g., 'https://app.bitbom.dev')", ) + cmd.Flags().BoolVar(&o.UseOpenAILLM, "use-openai-llm", false, "Use OpenAI LLM for graph analysis") + cmd.Flags().StringVar(&o.VectorDBPath, "vector-db-path", "./db", "Path to the vector database") } func (o *options) ProvideStorage() (graph.Storage, error) { @@ -136,6 +142,97 @@ func (o *options) startServer(server *http.Server) error { stop := make(chan os.Signal, 1) signal.Notify(stop, os.Interrupt, syscall.SIGTERM) + if o.UseOpenAILLM { + db, err := chromadb.NewPersistentDB(o.VectorDBPath, false) + if err != nil { + return fmt.Errorf("failed to initialize ChromaDB: %w", err) + } + + c, err := db.CreateCollection("knowledge-base", nil, nil) + if err != nil { + return fmt.Errorf("failed to create collection in ChromaDB: %w", err) + } + + // Initialize ChromaDB documents + err = c.AddDocuments(context.Background(), []chromadb.Document{ + { + ID: "1", + Content: "To query dependencies of a package, use the format, and only output the query: dependencies library pkg:. All three are necessary.", + }, + { + ID: "2", + Content: "To find dependents of a package, use the format, and only output the query: dependents library pkg:. All three are necessary.", + }, + { + ID: "3", + Content: "For vulnerabilities related to a package, use the format, and only output the query: dependencies vuln pkg:. All three are necessary.", + }, + { + ID: "4", + Content: "Combine queries using logical operators like and, or, and xor. All three are necessary.", + }, + { + ID: "5", + Content: "When using 'and', both conditions must be true. For example, and only output the query: dependencies library pkg:A and dependencies library pkg:B.", + }, + { + ID: "6", + Content: "When using 'or', at least one of the conditions must be true. For example, and only output the query: dependencies library pkg:A or dependencies library pkg:B.", + }, + { + ID: "7", + Content: "Use 'xor' to indicate that only one of the conditions can be true. For example, and only output the query: dependencies library pkg:A xor dependencies library pkg:B.", + }, + { + ID: "8", + Content: "You can chain multiple queries together. For example, and only output the query: dependencies library pkg:A and dependents library pkg:B or vulnerabilities vuln pkg:C.", + }, + { + ID: "9", + Content: "Package names can include versioning information. For example, and only output the query: dependencies library pkg:example-lib@1.0.0.", + }, + { + ID: "10", + Content: "Ensure that all keywords are used correctly. The keywords are: dependencies, dependents, library, vuln, xor, or, and.", + }, + { + ID: "11", + Content: "If a query does not specify a package name, it cannot be processed. Always include a package name in your queries.", + }, + { + ID: "12", + Content: "To check for multiple vulnerabilities across different packages, you can use, and only output the query: dependencies vuln pkg:A or dependencies vuln pkg:B.", + }, + { + ID: "13", + Content: "When using 'or', 'and', or 'xor', to take the answer of multiple queries and use an binary operator on the whole result with another query, or another set of queries, you can wrap a set of queries in brackets (), these can also be nested. Example: ((dependencies library pkg:A and dependencies library pkg:B) or (dependencies library pkg:C and dependencies library pkg:D)) and (dependents library pkg:E).", + }, + { + ID: "14", + Content: "Leaderboard queries are a different type of query, they run queries for every single node in the graph and return a sorted list based of length of the result.", + }, + { + ID: "15", + Content: "To run a leaderboard query, it is quite similar to a regular query, but instead of runing somthing like depedencies library pkg:A, you would run dependencies library. This would create a leaderboard which is sorted by each node's dependencies of type library.", + }, + { + ID: "16", + Content: "Leaderboards format are basicaly the same as a query, just if you do not include the node name for the last part of the query, it fills it with the node we are using for the leaderboard, which is every single node in the leaderboard. This means that to make a proper leaderboard we should have at least one part that does not include a node name, we can still combine this with another query. For example: (dependencies library) and (dependents library pkg:github.com/bitbomdev/minefield) would create a leaderboard sorted by the number of dependencies a project has that is shared with minefield. We can also repeat this 2 part query multiple times, for example : dependencies library and dependents library would work as well.", + }, + { + ID: "17", + Content: "If the user, or you are not sure about what the node's name is, you can use the glob pattern to search for nodes. For example, if you want to search for all nodes that start with 'github.com/bitbomdev', you can use the pattern 'github.com/bitbomdev*'. Try to lean get as many nodes as possible, so if they tell you the name is mineifield, and maybe the org is bitbomdev, you can use '*minefield*', since it will match all nodes that contain minefield, since they are not sure about the org.", + }, + { + ID: "18", + Content: "When glob seaching never assume the position of anything, so wrap everything can in ** on both sides.", + }, + }, runtime.NumCPU()) + if err != nil { + return fmt.Errorf("failed to add documents to ChromaDB: %w", err) + } + } + go func() { log.Printf("Server is starting on %s\n", server.Addr) if err := server.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) { diff --git a/go.mod b/go.mod index fe1dbf0..3506925 100644 --- a/go.mod +++ b/go.mod @@ -49,7 +49,9 @@ require ( github.com/mattn/go-sqlite3 v1.14.24 // indirect github.com/mschoch/smat v0.2.0 // indirect github.com/onsi/gomega v1.30.0 // indirect + github.com/philippgille/chromem-go v0.7.0 github.com/rs/cors v1.11.1 + github.com/sashabaranov/go-openai v1.36.0 github.com/sirupsen/logrus v1.9.3 // indirect github.com/spdx/tools-golang v0.5.5 // indirect github.com/stretchr/objx v0.5.2 // indirect diff --git a/go.sum b/go.sum index 7afc0d8..5ab8012 100644 --- a/go.sum +++ b/go.sum @@ -79,6 +79,8 @@ github.com/onsi/gomega v1.30.0 h1:hvMK7xYz4D3HapigLTeGdId/NcfQx1VHMJc60ew99+8= github.com/onsi/gomega v1.30.0/go.mod h1:9sxs+SwGrKI0+PWe4Fxa9tFQQBG5xSsSbMXOI8PPpoQ= github.com/package-url/packageurl-go v0.1.3 h1:4juMED3hHiz0set3Vq3KeQ75KD1avthoXLtmE3I0PLs= github.com/package-url/packageurl-go v0.1.3/go.mod h1:nKAWB8E6uk1MHqiS/lQb9pYBGH2+mdJ2PJc2s50dQY0= +github.com/philippgille/chromem-go v0.7.0 h1:4jfvfyKymjKNfGxBUhHUcj1kp7B17NL/I1P+vGh1RvY= +github.com/philippgille/chromem-go v0.7.0/go.mod h1:hTd+wGEm/fFPQl7ilfCwQXkgEUxceYh86iIdoKMolPo= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/protobom/protobom v0.5.0 h1:jJYqGpdHq99zwh0/n1SOPl1aickCBZdA8pHS9V/f+XQ= @@ -91,6 +93,8 @@ github.com/rogpeppe/go-internal v1.12.0/go.mod h1:E+RYuTGaKKdloAfM02xzb0FW3Paa99 github.com/rs/cors v1.11.1 h1:eU3gRzXLRK57F5rKMGMZURNdIG4EoAmX8k94r9wXWHA= github.com/rs/cors v1.11.1/go.mod h1:XyqrcTp5zjWr1wsJ8PIRZssZ8b/WMcMf71DJnit4EMU= github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= +github.com/sashabaranov/go-openai v1.36.0 h1:fcSrn8uGuorzPWCBp8L0aCR95Zjb/Dd+ZSML0YZy9EI= +github.com/sashabaranov/go-openai v1.36.0/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg= github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ= github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ= github.com/spdx/gordf v0.0.0-20201111095634-7098f93598fb/go.mod h1:uKWaldnbMnjsSAXRurWqqrdyZen1R7kxl8TkmWk2OyM= diff --git a/pkg/storages/sql.go b/pkg/storages/sql.go index a1fe22a..c90d197 100644 --- a/pkg/storages/sql.go +++ b/pkg/storages/sql.go @@ -279,21 +279,34 @@ func (s *SQLStorage) SaveCache(cache *graph.NodeCache) error { // SaveCaches saves multiple node caches. func (s *SQLStorage) SaveCaches(caches []*graph.NodeCache) error { - kvCaches := make([]KVStore, len(caches)) - for i, cache := range caches { - cacheKey := fmt.Sprintf("%s%d", CacheKeyPrefix, cache.ID) - data, err := cache.MarshalJSON() - if err != nil { - return fmt.Errorf("failed to marshal cache: %w", err) + const batchSize = 500 // Safe batch size considering SQLite's limits + + // Process caches in batches + for i := 0; i < len(caches); i += batchSize { + end := i + batchSize + if end > len(caches) { + end = len(caches) } - kvCaches[i] = KVStore{ - Key: cacheKey, - Value: string(data), + + batch := caches[i:end] + kvCaches := make([]KVStore, len(batch)) + for j, cache := range batch { + cacheKey := fmt.Sprintf("%s%d", CacheKeyPrefix, cache.ID) + data, err := cache.MarshalJSON() + if err != nil { + return fmt.Errorf("failed to marshal cache: %w", err) + } + kvCaches[j] = KVStore{ + Key: cacheKey, + Value: string(data), + } + } + + if err := s.DB.Save(&kvCaches).Error; err != nil { + return fmt.Errorf("failed to save caches batch: %w", err) } } - if err := s.DB.Save(&kvCaches).Error; err != nil { - return fmt.Errorf("failed to save caches: %w", err) - } + return nil }