From a1b39f9b63aeb0fd66f0f8a2f9f52df9bb8381b0 Mon Sep 17 00:00:00 2001 From: jueli12 Date: Wed, 18 Sep 2024 21:58:55 +0800 Subject: [PATCH 1/5] feat: change the AI command line argument to optional --- pkg/core/handler/search/search.go | 152 +++++++++++++++--------------- 1 file changed, 76 insertions(+), 76 deletions(-) diff --git a/pkg/core/handler/search/search.go b/pkg/core/handler/search/search.go index db94a495..bc8a930f 100644 --- a/pkg/core/handler/search/search.go +++ b/pkg/core/handler/search/search.go @@ -15,15 +15,16 @@ package search import ( - "github.com/KusionStack/karpor/pkg/core/manager/ai" - "net/http" - "strconv" + "errors" + "github.com/KusionStack/karpor/pkg/core/manager/ai" + "net/http" + "strconv" - "github.com/KusionStack/karpor/pkg/core/handler" - "github.com/KusionStack/karpor/pkg/core/manager/search" - "github.com/KusionStack/karpor/pkg/infra/search/storage" - "github.com/KusionStack/karpor/pkg/util/ctxutil" - "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" + "github.com/KusionStack/karpor/pkg/core/handler" + "github.com/KusionStack/karpor/pkg/core/manager/search" + "github.com/KusionStack/karpor/pkg/infra/search/storage" + "github.com/KusionStack/karpor/pkg/util/ctxutil" + "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" ) // SearchForResource returns an HTTP handler function that returns an @@ -47,79 +48,78 @@ import ( // @Failure 500 {string} string "Internal Server Error" // @Router /rest-api/v1/search [get] func SearchForResource(searchMgr *search.SearchManager, aiMgr *ai.AIManager, searchStorage storage.SearchStorage) http.HandlerFunc { - return func(w http.ResponseWriter, r *http.Request) { - // Extract the context and logger from the request. - ctx := r.Context() - logger := ctxutil.GetLogger(ctx) + return func(w http.ResponseWriter, r *http.Request) { + // Extract the context and logger from the request. + ctx := r.Context() + logger := ctxutil.GetLogger(ctx) - // Extract URL query parameters with default value - searchQuery := r.URL.Query().Get("query") - searchPattern := r.URL.Query().Get("pattern") - searchPageSize, _ := strconv.Atoi(r.URL.Query().Get("pageSize")) - searchPage, _ := strconv.Atoi(r.URL.Query().Get("page")) - if searchPageSize <= 1 { - searchPageSize = 10 - } - if searchPage <= 1 { - searchPage = 1 - } + // Extract URL query parameters with default value + searchQuery := r.URL.Query().Get("query") + searchPattern := r.URL.Query().Get("pattern") + searchPageSize, _ := strconv.Atoi(r.URL.Query().Get("pageSize")) + searchPage, _ := strconv.Atoi(r.URL.Query().Get("page")) + if searchPageSize <= 1 { + searchPageSize = 10 + } + if searchPage <= 1 { + searchPage = 1 + } - query := searchQuery + query := searchQuery - if searchPattern == storage.NLPatternType { - if err := ai.CheckAIManager(aiMgr); err != nil { - handler.FailureRender(ctx, w, r, err) - return - } + if searchPattern == storage.NLPatternType { + if err := ai.CheckAIManager(aiMgr); err != nil { + handler.FailureRender(ctx, w, r, err) + return + } + res, err := aiMgr.ConvertTextToSQL(searchQuery) + if err != nil { + handler.FailureRender(ctx, w, r, err) + return + } + searchQuery = res + } - res, err := aiMgr.ConvertTextToSQL(searchQuery) - if err != nil { - handler.FailureRender(ctx, w, r, err) - return - } - searchQuery = res - } + logger.Info("Searching for resources...", "page", searchPage, "pageSize", searchPageSize) - logger.Info("Searching for resources...", "page", searchPage, "pageSize", searchPageSize) + res, err := searchStorage.Search(ctx, searchQuery, searchPattern, &storage.Pagination{Page: searchPage, PageSize: searchPageSize}) + if err != nil { + if searchPattern == storage.NLPatternType { + fixedQuery, fixErr := aiMgr.FixSQL(query, searchQuery, err.Error()) + if fixErr != nil { + handler.FailureRender(ctx, w, r, err) + return + } + searchQuery = fixedQuery + res, err = searchStorage.Search(ctx, searchQuery, searchPattern, &storage.Pagination{Page: searchPage, PageSize: searchPageSize}) + if err != nil { + handler.FailureRender(ctx, w, r, err) + return + } + } else { + handler.FailureRender(ctx, w, r, err) + return + } + } - res, err := searchStorage.Search(ctx, searchQuery, searchPattern, &storage.Pagination{Page: searchPage, PageSize: searchPageSize}) - if err != nil { - if searchPattern == storage.NLPatternType { - fixedQuery, fixErr := aiMgr.FixSQL(query, searchQuery, err.Error()) - if fixErr != nil { - handler.FailureRender(ctx, w, r, err) - return - } - searchQuery = fixedQuery - res, err = searchStorage.Search(ctx, searchQuery, searchPattern, &storage.Pagination{Page: searchPage, PageSize: searchPageSize}) - if err != nil { - handler.FailureRender(ctx, w, r, err) - return - } - } else { - handler.FailureRender(ctx, w, r, err) - return - } - } + if err != nil { + handler.FailureRender(ctx, w, r, err) + return + } - if err != nil { - handler.FailureRender(ctx, w, r, err) - return - } - - rt := &search.UniResourceList{} - for _, r := range res.Resources { - unObj := &unstructured.Unstructured{} - unObj.SetUnstructuredContent(r.Object) - rt.Items = append(rt.Items, search.UniResource{ - Cluster: r.Cluster, - Object: unObj, - }) - } - rt.SQLQuery = searchQuery - rt.Total = res.Total - rt.CurrentPage = searchPage - rt.PageSize = searchPageSize - handler.SuccessRender(ctx, w, r, rt) - } + rt := &search.UniResourceList{} + for _, r := range res.Resources { + unObj := &unstructured.Unstructured{} + unObj.SetUnstructuredContent(r.Object) + rt.Items = append(rt.Items, search.UniResource{ + Cluster: r.Cluster, + Object: unObj, + }) + } + rt.SQLQuery = searchQuery + rt.Total = res.Total + rt.CurrentPage = searchPage + rt.PageSize = searchPageSize + handler.SuccessRender(ctx, w, r, rt) + } } From 7ae394616fa399b7f1839e6aa3b36d645abcff1c Mon Sep 17 00:00:00 2001 From: jueli12 Date: Wed, 18 Sep 2024 22:19:39 +0800 Subject: [PATCH 2/5] validate ai manager --- pkg/core/handler/search/search.go | 135 +++++++++++++++--------------- 1 file changed, 67 insertions(+), 68 deletions(-) diff --git a/pkg/core/handler/search/search.go b/pkg/core/handler/search/search.go index bc8a930f..8a99bd51 100644 --- a/pkg/core/handler/search/search.go +++ b/pkg/core/handler/search/search.go @@ -15,7 +15,6 @@ package search import ( - "errors" "github.com/KusionStack/karpor/pkg/core/manager/ai" "net/http" "strconv" @@ -48,78 +47,78 @@ import ( // @Failure 500 {string} string "Internal Server Error" // @Router /rest-api/v1/search [get] func SearchForResource(searchMgr *search.SearchManager, aiMgr *ai.AIManager, searchStorage storage.SearchStorage) http.HandlerFunc { - return func(w http.ResponseWriter, r *http.Request) { - // Extract the context and logger from the request. - ctx := r.Context() - logger := ctxutil.GetLogger(ctx) + return func(w http.ResponseWriter, r *http.Request) { + // Extract the context and logger from the request. + ctx := r.Context() + logger := ctxutil.GetLogger(ctx) - // Extract URL query parameters with default value - searchQuery := r.URL.Query().Get("query") - searchPattern := r.URL.Query().Get("pattern") - searchPageSize, _ := strconv.Atoi(r.URL.Query().Get("pageSize")) - searchPage, _ := strconv.Atoi(r.URL.Query().Get("page")) - if searchPageSize <= 1 { - searchPageSize = 10 - } - if searchPage <= 1 { - searchPage = 1 - } + // Extract URL query parameters with default value + searchQuery := r.URL.Query().Get("query") + searchPattern := r.URL.Query().Get("pattern") + searchPageSize, _ := strconv.Atoi(r.URL.Query().Get("pageSize")) + searchPage, _ := strconv.Atoi(r.URL.Query().Get("page")) + if searchPageSize <= 1 { + searchPageSize = 10 + } + if searchPage <= 1 { + searchPage = 1 + } - query := searchQuery + query := searchQuery - if searchPattern == storage.NLPatternType { - if err := ai.CheckAIManager(aiMgr); err != nil { - handler.FailureRender(ctx, w, r, err) - return - } - res, err := aiMgr.ConvertTextToSQL(searchQuery) - if err != nil { - handler.FailureRender(ctx, w, r, err) - return - } - searchQuery = res - } + if searchPattern == storage.NLPatternType { + if err := ai.CheckAIManager(aiMgr); err != nil { + handler.FailureRender(ctx, w, r, err) + return + } + res, err := aiMgr.ConvertTextToSQL(searchQuery) + if err != nil { + handler.FailureRender(ctx, w, r, err) + return + } + searchQuery = res + } - logger.Info("Searching for resources...", "page", searchPage, "pageSize", searchPageSize) + logger.Info("Searching for resources...", "page", searchPage, "pageSize", searchPageSize) - res, err := searchStorage.Search(ctx, searchQuery, searchPattern, &storage.Pagination{Page: searchPage, PageSize: searchPageSize}) - if err != nil { - if searchPattern == storage.NLPatternType { - fixedQuery, fixErr := aiMgr.FixSQL(query, searchQuery, err.Error()) - if fixErr != nil { - handler.FailureRender(ctx, w, r, err) - return - } - searchQuery = fixedQuery - res, err = searchStorage.Search(ctx, searchQuery, searchPattern, &storage.Pagination{Page: searchPage, PageSize: searchPageSize}) - if err != nil { - handler.FailureRender(ctx, w, r, err) - return - } - } else { - handler.FailureRender(ctx, w, r, err) - return - } - } + res, err := searchStorage.Search(ctx, searchQuery, searchPattern, &storage.Pagination{Page: searchPage, PageSize: searchPageSize}) + if err != nil { + if searchPattern == storage.NLPatternType { + fixedQuery, fixErr := aiMgr.FixSQL(query, searchQuery, err.Error()) + if fixErr != nil { + handler.FailureRender(ctx, w, r, err) + return + } + searchQuery = fixedQuery + res, err = searchStorage.Search(ctx, searchQuery, searchPattern, &storage.Pagination{Page: searchPage, PageSize: searchPageSize}) + if err != nil { + handler.FailureRender(ctx, w, r, err) + return + } + } else { + handler.FailureRender(ctx, w, r, err) + return + } + } - if err != nil { - handler.FailureRender(ctx, w, r, err) - return - } + if err != nil { + handler.FailureRender(ctx, w, r, err) + return + } - rt := &search.UniResourceList{} - for _, r := range res.Resources { - unObj := &unstructured.Unstructured{} - unObj.SetUnstructuredContent(r.Object) - rt.Items = append(rt.Items, search.UniResource{ - Cluster: r.Cluster, - Object: unObj, - }) - } - rt.SQLQuery = searchQuery - rt.Total = res.Total - rt.CurrentPage = searchPage - rt.PageSize = searchPageSize - handler.SuccessRender(ctx, w, r, rt) - } + rt := &search.UniResourceList{} + for _, r := range res.Resources { + unObj := &unstructured.Unstructured{} + unObj.SetUnstructuredContent(r.Object) + rt.Items = append(rt.Items, search.UniResource{ + Cluster: r.Cluster, + Object: unObj, + }) + } + rt.SQLQuery = searchQuery + rt.Total = res.Total + rt.CurrentPage = searchPage + rt.PageSize = searchPageSize + handler.SuccessRender(ctx, w, r, rt) + } } From cb4b3ecabdd992b4521413517815c0d919a37680 Mon Sep 17 00:00:00 2001 From: jueli12 Date: Thu, 19 Sep 2024 17:54:14 +0800 Subject: [PATCH 3/5] feat: filter invalid queries for natural language search --- pkg/core/handler/search/search.go | 134 +++++++++++++++--------------- pkg/core/manager/ai/search.go | 3 + pkg/core/manager/ai/types.go | 1 + pkg/core/manager/ai/util.go | 13 ++- pkg/core/manager/ai/util_test.go | 27 ++++++ pkg/infra/ai/prompts.go | 26 ++++-- 6 files changed, 131 insertions(+), 73 deletions(-) diff --git a/pkg/core/handler/search/search.go b/pkg/core/handler/search/search.go index 8a99bd51..63a40094 100644 --- a/pkg/core/handler/search/search.go +++ b/pkg/core/handler/search/search.go @@ -47,78 +47,78 @@ import ( // @Failure 500 {string} string "Internal Server Error" // @Router /rest-api/v1/search [get] func SearchForResource(searchMgr *search.SearchManager, aiMgr *ai.AIManager, searchStorage storage.SearchStorage) http.HandlerFunc { - return func(w http.ResponseWriter, r *http.Request) { - // Extract the context and logger from the request. - ctx := r.Context() - logger := ctxutil.GetLogger(ctx) + return func(w http.ResponseWriter, r *http.Request) { + // Extract the context and logger from the request. + ctx := r.Context() + logger := ctxutil.GetLogger(ctx) - // Extract URL query parameters with default value - searchQuery := r.URL.Query().Get("query") - searchPattern := r.URL.Query().Get("pattern") - searchPageSize, _ := strconv.Atoi(r.URL.Query().Get("pageSize")) - searchPage, _ := strconv.Atoi(r.URL.Query().Get("page")) - if searchPageSize <= 1 { - searchPageSize = 10 - } - if searchPage <= 1 { - searchPage = 1 - } + // Extract URL query parameters with default value + searchQuery := r.URL.Query().Get("query") + searchPattern := r.URL.Query().Get("pattern") + searchPageSize, _ := strconv.Atoi(r.URL.Query().Get("pageSize")) + searchPage, _ := strconv.Atoi(r.URL.Query().Get("page")) + if searchPageSize <= 1 { + searchPageSize = 10 + } + if searchPage <= 1 { + searchPage = 1 + } - query := searchQuery + query := searchQuery - if searchPattern == storage.NLPatternType { - if err := ai.CheckAIManager(aiMgr); err != nil { - handler.FailureRender(ctx, w, r, err) - return - } - res, err := aiMgr.ConvertTextToSQL(searchQuery) - if err != nil { - handler.FailureRender(ctx, w, r, err) - return - } - searchQuery = res - } + if searchPattern == storage.NLPatternType { + if err := ai.CheckAIManager(aiMgr); err != nil { + handler.FailureRender(ctx, w, r, err) + return + } + res, err := aiMgr.ConvertTextToSQL(searchQuery) + if err != nil { + handler.FailureRender(ctx, w, r, err) + return + } + searchQuery = res + } - logger.Info("Searching for resources...", "page", searchPage, "pageSize", searchPageSize) + logger.Info("Searching for resources...", "page", searchPage, "pageSize", searchPageSize) - res, err := searchStorage.Search(ctx, searchQuery, searchPattern, &storage.Pagination{Page: searchPage, PageSize: searchPageSize}) - if err != nil { - if searchPattern == storage.NLPatternType { - fixedQuery, fixErr := aiMgr.FixSQL(query, searchQuery, err.Error()) - if fixErr != nil { - handler.FailureRender(ctx, w, r, err) - return - } - searchQuery = fixedQuery - res, err = searchStorage.Search(ctx, searchQuery, searchPattern, &storage.Pagination{Page: searchPage, PageSize: searchPageSize}) - if err != nil { - handler.FailureRender(ctx, w, r, err) - return - } - } else { - handler.FailureRender(ctx, w, r, err) - return - } - } + res, err := searchStorage.Search(ctx, searchQuery, searchPattern, &storage.Pagination{Page: searchPage, PageSize: searchPageSize}) + if err != nil { + if searchPattern == storage.NLPatternType { + fixedQuery, fixErr := aiMgr.FixSQL(query, searchQuery, err.Error()) + if fixErr != nil { + handler.FailureRender(ctx, w, r, err) + return + } + searchQuery = fixedQuery + res, err = searchStorage.Search(ctx, searchQuery, searchPattern, &storage.Pagination{Page: searchPage, PageSize: searchPageSize}) + if err != nil { + handler.FailureRender(ctx, w, r, err) + return + } + } else { + handler.FailureRender(ctx, w, r, err) + return + } + } - if err != nil { - handler.FailureRender(ctx, w, r, err) - return - } + if err != nil { + handler.FailureRender(ctx, w, r, err) + return + } - rt := &search.UniResourceList{} - for _, r := range res.Resources { - unObj := &unstructured.Unstructured{} - unObj.SetUnstructuredContent(r.Object) - rt.Items = append(rt.Items, search.UniResource{ - Cluster: r.Cluster, - Object: unObj, - }) - } - rt.SQLQuery = searchQuery - rt.Total = res.Total - rt.CurrentPage = searchPage - rt.PageSize = searchPageSize - handler.SuccessRender(ctx, w, r, rt) - } + rt := &search.UniResourceList{} + for _, r := range res.Resources { + unObj := &unstructured.Unstructured{} + unObj.SetUnstructuredContent(r.Object) + rt.Items = append(rt.Items, search.UniResource{ + Cluster: r.Cluster, + Object: unObj, + }) + } + rt.SQLQuery = searchQuery + rt.Total = res.Total + rt.CurrentPage = searchPage + rt.PageSize = searchPageSize + handler.SuccessRender(ctx, w, r, rt) + } } diff --git a/pkg/core/manager/ai/search.go b/pkg/core/manager/ai/search.go index 4c814518..30ea67fb 100644 --- a/pkg/core/manager/ai/search.go +++ b/pkg/core/manager/ai/search.go @@ -28,6 +28,9 @@ func (a *AIManager) ConvertTextToSQL(query string) (string, error) { if err != nil { return "", err } + if IfQueryInvalid(res) { + return "", ErrInvalidQuery + } return ExtractSelectSQL(res), nil } diff --git a/pkg/core/manager/ai/types.go b/pkg/core/manager/ai/types.go index 39e8b501..eedd6285 100644 --- a/pkg/core/manager/ai/types.go +++ b/pkg/core/manager/ai/types.go @@ -18,4 +18,5 @@ import "errors" var ( ErrMissingAuthToken = errors.New("auth token is required") + ErrInvalidQuery = errors.New("query is invalid") ) diff --git a/pkg/core/manager/ai/util.go b/pkg/core/manager/ai/util.go index effe2be3..d9d3cd1e 100644 --- a/pkg/core/manager/ai/util.go +++ b/pkg/core/manager/ai/util.go @@ -14,7 +14,18 @@ package ai -import "regexp" +import ( + "regexp" + "strings" +) + +// IfQueryInvalid check if the query is invalid +func IfQueryInvalid(sql string) bool { + if strings.Contains(strings.ToLower(sql), "error") { + return true + } + return false +} // ExtractSelectSQL extracts SQL statements that start with "SELECT * FROM" func ExtractSelectSQL(sql string) string { diff --git a/pkg/core/manager/ai/util_test.go b/pkg/core/manager/ai/util_test.go index d356d0b4..4f9264ce 100644 --- a/pkg/core/manager/ai/util_test.go +++ b/pkg/core/manager/ai/util_test.go @@ -42,3 +42,30 @@ func TestExtractSelectSQL(t *testing.T) { }) } } + +// TestIfQueryInvalid tests the IfQueryInvalid function. +func TestIfQueryInvalid(t *testing.T) { + testCases := []struct { + name string + sql string + expected bool + }{ + { + name: "ValidQueryWithoutError", + sql: "select * from resources where kind='namespace';", + expected: false, + }, + { + name: "InvalidQuery", + sql: "Error", + expected: true, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + actual := IfQueryInvalid(tc.sql) + require.Equal(t, tc.expected, actual) + }) + } +} diff --git a/pkg/infra/ai/prompts.go b/pkg/infra/ai/prompts.go index 06feceae..1ef63cb8 100644 --- a/pkg/infra/ai/prompts.go +++ b/pkg/infra/ai/prompts.go @@ -20,6 +20,7 @@ const ( text2sql_prompt = ` You are an AI specialized in writing SQL queries. Please convert the text %s to sql. + If the text is not accurate enough, please output "Error". The output tokens only need to give the SQL first, the other thought process please do not give. The SQL should begin with "select * from" and end with ";". @@ -30,23 +31,33 @@ const ( resourceVersion, labels.[key], annotations.[key], content] 2. find the schema_links for generating SQL queries for each question based on the database schema. + If there are Chinese expressions, please translate them into English. Follow are some examples. Q: find the kind which is not equal to pod A: Let’s think step by step. In the question "find the kind column which is not equal to pod", we are asked: - "find the kind" so we need column = [kind] + "find the kind" so we need column = [kind]. Based on the columns, the set of possible cell values are = [pod]. So the Schema_links are: Schema_links: [kind, pod] Q: find the kind Deployment which created before January 1, 2024, at 18:00:00 A: Let’s think step by step. In the question "find the kind Deployment which created before January 1, 2024, at 18:00:00", we are asked: - "find the kind Deployment" so we need column = [kind] - "created before" so we need column = [creationTimestamp] + "find the kind Deployment" so we need column = [kind]. + "created before" so we need column = [creationTimestamp]. Based on the columns, the set of possible cell values are = [Deployment, 2024-01-01T18:00:00Z]. So the Schema_links are: - Schema_links: [kind, creationTimestamp, Deployment, 2024-01-01T18:00:00Z] + Schema_links: [[kind, Deployment], [creationTimestamp, 2024-01-01T18:00:00Z]] + + Q: find the kind Namespace which which created + A: Let’s think step by step. In the question "find the kind", we are asked: + "find the kind Namespace " so we need column = [kind] + "created before" so we need column = [creationTimestamp] + Based on the columns, the set of possible cell values are = [kind, creationTimestamp]. + There is no creationTimestamp corresponding cell values, so the text is not accurate enough. + So the Schema_links are: + Schema_links: error 3. Use the the schema links to generate the SQL queries for each of the questions. @@ -57,14 +68,19 @@ const ( SQL: select * from resources where kind!='Pod'; Q: find the kind Deployment which created before January 1, 2024, at 18:00:00 - Schema_links: [kind, creationTimestamp, Deployment, 2024-01-01T18:00:00Z] + Schema_links: [[kind, Deployment], [creationTimestamp, 2024-01-01T18:00:00Z]] SQL: select * from resources where kind='Deployment' and creationTimestamp < '2024-01-01T18:00:00Z'; Q: find the namespace which does not contain banan Schema_links: [namespace, banan] SQL: select * from resources where namespace notlike 'banan_'; + Q: find the kind Namespace which which created + Schema_links: error + Error; + Please convert the text to sql. + If the text is not accurate enough, please output "Error". ` sql_fix_prompt = ` From a1e8ac6840bd2b535d4dead3d2f77e872116f1c6 Mon Sep 17 00:00:00 2001 From: jueli12 Date: Tue, 24 Sep 2024 14:35:09 +0800 Subject: [PATCH 4/5] fix: conflict --- pkg/core/handler/search/search.go | 150 +++++++++++++++--------------- 1 file changed, 75 insertions(+), 75 deletions(-) diff --git a/pkg/core/handler/search/search.go b/pkg/core/handler/search/search.go index 63a40094..0bc58d17 100644 --- a/pkg/core/handler/search/search.go +++ b/pkg/core/handler/search/search.go @@ -15,15 +15,15 @@ package search import ( - "github.com/KusionStack/karpor/pkg/core/manager/ai" - "net/http" - "strconv" + "github.com/KusionStack/karpor/pkg/core/manager/ai" + "net/http" + "strconv" - "github.com/KusionStack/karpor/pkg/core/handler" - "github.com/KusionStack/karpor/pkg/core/manager/search" - "github.com/KusionStack/karpor/pkg/infra/search/storage" - "github.com/KusionStack/karpor/pkg/util/ctxutil" - "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" + "github.com/KusionStack/karpor/pkg/core/handler" + "github.com/KusionStack/karpor/pkg/core/manager/search" + "github.com/KusionStack/karpor/pkg/infra/search/storage" + "github.com/KusionStack/karpor/pkg/util/ctxutil" + "k8s.io/apimachinery/pkg/apis/meta/v1/unstructured" ) // SearchForResource returns an HTTP handler function that returns an @@ -47,78 +47,78 @@ import ( // @Failure 500 {string} string "Internal Server Error" // @Router /rest-api/v1/search [get] func SearchForResource(searchMgr *search.SearchManager, aiMgr *ai.AIManager, searchStorage storage.SearchStorage) http.HandlerFunc { - return func(w http.ResponseWriter, r *http.Request) { - // Extract the context and logger from the request. - ctx := r.Context() - logger := ctxutil.GetLogger(ctx) + return func(w http.ResponseWriter, r *http.Request) { + // Extract the context and logger from the request. + ctx := r.Context() + logger := ctxutil.GetLogger(ctx) - // Extract URL query parameters with default value - searchQuery := r.URL.Query().Get("query") - searchPattern := r.URL.Query().Get("pattern") - searchPageSize, _ := strconv.Atoi(r.URL.Query().Get("pageSize")) - searchPage, _ := strconv.Atoi(r.URL.Query().Get("page")) - if searchPageSize <= 1 { - searchPageSize = 10 - } - if searchPage <= 1 { - searchPage = 1 - } + // Extract URL query parameters with default value + searchQuery := r.URL.Query().Get("query") + searchPattern := r.URL.Query().Get("pattern") + searchPageSize, _ := strconv.Atoi(r.URL.Query().Get("pageSize")) + searchPage, _ := strconv.Atoi(r.URL.Query().Get("page")) + if searchPageSize <= 1 { + searchPageSize = 10 + } + if searchPage <= 1 { + searchPage = 1 + } - query := searchQuery + query := searchQuery - if searchPattern == storage.NLPatternType { - if err := ai.CheckAIManager(aiMgr); err != nil { - handler.FailureRender(ctx, w, r, err) - return - } - res, err := aiMgr.ConvertTextToSQL(searchQuery) - if err != nil { - handler.FailureRender(ctx, w, r, err) - return - } - searchQuery = res - } + if searchPattern == storage.NLPatternType { + if err := ai.CheckAIManager(aiMgr); err != nil { + handler.FailureRender(ctx, w, r, err) + return + } + res, err := aiMgr.ConvertTextToSQL(searchQuery) + if err != nil { + handler.FailureRender(ctx, w, r, err) + return + } + searchQuery = res + } - logger.Info("Searching for resources...", "page", searchPage, "pageSize", searchPageSize) + logger.Info("Searching for resources...", "page", searchPage, "pageSize", searchPageSize) - res, err := searchStorage.Search(ctx, searchQuery, searchPattern, &storage.Pagination{Page: searchPage, PageSize: searchPageSize}) - if err != nil { - if searchPattern == storage.NLPatternType { - fixedQuery, fixErr := aiMgr.FixSQL(query, searchQuery, err.Error()) - if fixErr != nil { - handler.FailureRender(ctx, w, r, err) - return - } - searchQuery = fixedQuery - res, err = searchStorage.Search(ctx, searchQuery, searchPattern, &storage.Pagination{Page: searchPage, PageSize: searchPageSize}) - if err != nil { - handler.FailureRender(ctx, w, r, err) - return - } - } else { - handler.FailureRender(ctx, w, r, err) - return - } - } + res, err := searchStorage.Search(ctx, searchQuery, searchPattern, &storage.Pagination{Page: searchPage, PageSize: searchPageSize}) + if err != nil { + if searchPattern == storage.NLPatternType { + fixedQuery, fixErr := aiMgr.FixSQL(query, searchQuery, err.Error()) + if fixErr != nil { + handler.FailureRender(ctx, w, r, err) + return + } + searchQuery = fixedQuery + res, err = searchStorage.Search(ctx, searchQuery, searchPattern, &storage.Pagination{Page: searchPage, PageSize: searchPageSize}) + if err != nil { + handler.FailureRender(ctx, w, r, err) + return + } + } else { + handler.FailureRender(ctx, w, r, err) + return + } + } - if err != nil { - handler.FailureRender(ctx, w, r, err) - return - } + if err != nil { + handler.FailureRender(ctx, w, r, err) + return + } - rt := &search.UniResourceList{} - for _, r := range res.Resources { - unObj := &unstructured.Unstructured{} - unObj.SetUnstructuredContent(r.Object) - rt.Items = append(rt.Items, search.UniResource{ - Cluster: r.Cluster, - Object: unObj, - }) - } - rt.SQLQuery = searchQuery - rt.Total = res.Total - rt.CurrentPage = searchPage - rt.PageSize = searchPageSize - handler.SuccessRender(ctx, w, r, rt) - } + rt := &search.UniResourceList{} + for _, r := range res.Resources { + unObj := &unstructured.Unstructured{} + unObj.SetUnstructuredContent(r.Object) + rt.Items = append(rt.Items, search.UniResource{ + Cluster: r.Cluster, + Object: unObj, + }) + } + rt.SQLQuery = searchQuery + rt.Total = res.Total + rt.CurrentPage = searchPage + rt.PageSize = searchPageSize + handler.SuccessRender(ctx, w, r, rt) + } } From bcd5536efd5dbb49e5f0523e78049da985fa0f91 Mon Sep 17 00:00:00 2001 From: jueli12 Date: Tue, 24 Sep 2024 15:29:08 +0800 Subject: [PATCH 5/5] fix: duplicate if --- pkg/core/manager/ai/search.go | 2 +- pkg/core/manager/ai/util.go | 9 +++------ pkg/core/manager/ai/util_test.go | 6 +++--- 3 files changed, 7 insertions(+), 10 deletions(-) diff --git a/pkg/core/manager/ai/search.go b/pkg/core/manager/ai/search.go index 30ea67fb..62589f6b 100644 --- a/pkg/core/manager/ai/search.go +++ b/pkg/core/manager/ai/search.go @@ -28,7 +28,7 @@ func (a *AIManager) ConvertTextToSQL(query string) (string, error) { if err != nil { return "", err } - if IfQueryInvalid(res) { + if IsInvalidQuery(res) { return "", ErrInvalidQuery } return ExtractSelectSQL(res), nil diff --git a/pkg/core/manager/ai/util.go b/pkg/core/manager/ai/util.go index d9d3cd1e..275849b0 100644 --- a/pkg/core/manager/ai/util.go +++ b/pkg/core/manager/ai/util.go @@ -19,12 +19,9 @@ import ( "strings" ) -// IfQueryInvalid check if the query is invalid -func IfQueryInvalid(sql string) bool { - if strings.Contains(strings.ToLower(sql), "error") { - return true - } - return false +// IsInvalidQuery check if the query is invalid +func IsInvalidQuery(sql string) bool { + return strings.Contains(strings.ToLower(sql), "error") } // ExtractSelectSQL extracts SQL statements that start with "SELECT * FROM" diff --git a/pkg/core/manager/ai/util_test.go b/pkg/core/manager/ai/util_test.go index 4f9264ce..fb5aa047 100644 --- a/pkg/core/manager/ai/util_test.go +++ b/pkg/core/manager/ai/util_test.go @@ -43,8 +43,8 @@ func TestExtractSelectSQL(t *testing.T) { } } -// TestIfQueryInvalid tests the IfQueryInvalid function. -func TestIfQueryInvalid(t *testing.T) { +// TestIsInvalidQuery tests the IsInvalidQuery function. +func TestIsInvalidQuery(t *testing.T) { testCases := []struct { name string sql string @@ -64,7 +64,7 @@ func TestIfQueryInvalid(t *testing.T) { for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { - actual := IfQueryInvalid(tc.sql) + actual := IsInvalidQuery(tc.sql) require.Equal(t, tc.expected, actual) }) }