fix: format the sql

KusionStack · Aug 29, 2024 · 2dccf66 · 2dccf66
1 parent da5116c
commit 2dccf66
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 5 deletions.
diff --git a/pkg/core/handler/search/search.go b/pkg/core/handler/search/search.go
@@ -82,6 +82,7 @@ func SearchForResource(searchMgr *search.SearchManager, aiMgr *ai.AIManager, sea
 		res, err := searchStorage.Search(ctx, searchQuery, searchPattern, &storage.Pagination{Page: searchPage, PageSize: searchPageSize})
 		if err != nil {
 			if searchPattern == storage.NLPatternType {
+				logger.Info("-------------" + err.Error() + "-------------")
 				fixedQuery, fixErr := aiMgr.FixSQL(query, searchQuery, err.Error())
 				if fixErr != nil {
 					handler.FailureRender(ctx, w, r, err)

diff --git a/pkg/core/manager/ai/util_test.go b/pkg/core/manager/ai/util_test.go
@@ -30,7 +30,7 @@ func TestExtractSelectSQL(t *testing.T) {
 			name: "NormalCase",
 			sql: "Q: 所有kind=namespace " +
 				"Schema_links: [kind, namespace] " +
-				"SQL: select * from resources where kind='namespace'",
+				"SQL: select * from resources where kind='namespace';",
 			expected: "select * from resources where kind='namespace'",
 		},
 	}

diff --git a/pkg/infra/ai/prompts.go b/pkg/infra/ai/prompts.go
@@ -21,7 +21,7 @@ const (
     You are an AI specialized in writing SQL queries.
     Please convert the text %s to sql.
     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".
+    The SQL should begin with "select * from" and end with ";".
 
     1. The database now only supports one table resources.
 
@@ -54,15 +54,17 @@ const (
 
     Q: find the kind which is not equal to pod
     Schema_links: [kind, pod]
-    SQL: select * from resources where kind!='Pod'
+    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]
-    SQL: select * from resources where kind='Deployment' and 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_'
+    SQL: select * from resources where namespace notlike 'banan_';
+
+    Please convert the text to sql.
     `
 
 	sql_fix_prompt = `