From 76801150082c02645e7c9f2ece1976ffa626d4a8 Mon Sep 17 00:00:00 2001
From: bjwswang <bjwswang@gmail.com>
Date: Fri, 25 Aug 2023 13:38:41 +0800
Subject: [PATCH] fix: optimize rbac inqury promt

Signed-off-by: bjwswang <bjwswang@gmail.com>
---
 examples/rbac/README.md  |  1 +
 examples/rbac/inquiry.go | 40 ++++++++++++++++++++++++++++++----------
 2 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/examples/rbac/README.md b/examples/rbac/README.md
index f72f2748b..6cae22fc9 100644
--- a/examples/rbac/README.md
+++ b/examples/rbac/README.md
@@ -35,6 +35,7 @@ Perform RBAC inquiry and evaluate security issues.
 
 - `-f, --file`**(Must required)**: RBAC file to be queried.
 - `-l, --language`: Language for the response (default: English). Can be any language! Chinese, Japanese, Korean, etc.
+- `-o, --output-format`: OutputFormat for AI response.Can be default,json.
 
 Note: The `file` argument is required.
 
diff --git a/examples/rbac/inquiry.go b/examples/rbac/inquiry.go
index e9f53a232..b3eb878d7 100644
--- a/examples/rbac/inquiry.go
+++ b/examples/rbac/inquiry.go
@@ -21,9 +21,6 @@ Below is a snippet of RBAC configuration:
 
   {{.Context}}
 
-
-**Answer in {{.Language}}**
-
 **RBAC Assessment**:
    - Could you please review this RBAC configuration and inform me of any security risks, permissions, or potential vulnerabilities that might be associated with it?
 
@@ -31,17 +28,36 @@ Below is a snippet of RBAC configuration:
    - On a scale of 1 to 10, with 10 being the most secure, how would you rate the security of this RBAC configuration?
 
 I appreciate your assistance in evaluating this RBAC configuration for security purposes. 
+
+Your answer should follow the below format:
+{{.OutputFormat}}
+`
+	defaultOutputFormat = `Rating: "xxx"
+Assesement: 
+1. xxx
+2. xxx
 `
+	jsonOutputFormat = `
+{
+	"rating": "xxx",
+	"assesement": [
+		"1. xxx",
+		"2. xxx",
+	]
+}
+	`
 )
 
 type InquiryData struct {
-	Context  string
-	Language string
+	Context      string
+	Language     string
+	OutputFormat string
 }
 
 var (
-	rbacFile string
-	language string
+	rbacFile     string
+	language     string
+	outputFormat string
 )
 
 func Inquiry() *cobra.Command {
@@ -58,10 +74,13 @@ func Inquiry() *cobra.Command {
 			}
 			// Prepare data for template
 			data := InquiryData{
-				Context:  strings.TrimSpace(string(rbacContent)),
-				Language: language,
+				Context:      strings.TrimSpace(string(rbacContent)),
+				Language:     language,
+				OutputFormat: defaultOutputFormat,
+			}
+			if outputFormat == "json" {
+				data.OutputFormat = jsonOutputFormat
 			}
-
 			// Create a new template and parse the RBAC inquiry template
 			tmpl := template.Must(template.New("RBACInquiry").Parse(RBACInquiryTemplate))
 
@@ -101,6 +120,7 @@ func Inquiry() *cobra.Command {
 
 	cmd.Flags().StringVarP(&rbacFile, "file", "f", "", "rbac file to be inquired")
 	cmd.Flags().StringVarP(&language, "language", "l", "English", "Language in the response")
+	cmd.Flags().StringVarP(&outputFormat, "output-format", "o", "", "OutputFormat for AI response.Can be default,json")
 
 	cmd.MarkFlagRequired("file")