feat(completion): add subcommand parameter

jigsaw373 · Apr 21, 2023 · 296072e · 296072e
1 parent 0154ff3
commit 296072e
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/cmd/cli/completion.go b/cmd/cli/completion.go
@@ -59,7 +59,7 @@ func newOAIClients() (oaiClients, error) {
 	return clients, nil
 }
 
-func gptCompletion(ctx context.Context, client oaiClients, prompts []string, deploymentName string) (string, error) {
+func completion(ctx context.Context, client oaiClients, prompts []string, deploymentName string, subcommand string) (string, error) {
 	temp := float32(*sensitivity)
 
 	maxTokens, err := calculateMaxTokens(prompts, deploymentName)
@@ -68,7 +68,7 @@ func gptCompletion(ctx context.Context, client oaiClients, prompts []string, dep
 	}
 
 	var prompt strings.Builder
-	_, err = fmt.Fprintf(&prompt, "You are a Terraform HCL generator, only generate valid Terraform HCL templates.")
+	_, err = fmt.Fprint(&prompt, subcommand)
 
 	if err != nil {
 		return "", fmt.Errorf("error prompt string builder: %w", err)

diff --git a/cmd/cli/root.go b/cmd/cli/root.go
@@ -16,7 +16,10 @@ import (
 	"github.com/walles/env"
 )
 
-const version = "0.0.1"
+const (
+	version         = "0.0.1"
+	applySubCommand = "You are a Terraform HCL generator, only generate valid Terraform HCL without provider templates."
+)
 
 var (
 	openAIDeploymentName = flag.String("openai-deployment-name", env.GetOr("OPENAI_DEPLOYMENT_NAME", env.String, "text-davinci-003"), "The deployment name used for the model in OpenAI service.")
@@ -69,7 +72,7 @@ func run(args []string) error {
 		return fmt.Errorf("error run command: %w", err)
 	}
 
-	completion, err := gptCompletion(ctx, oaiClients, args, *openAIDeploymentName)
+	completion, err := completion(ctx, oaiClients, args, *openAIDeploymentName, applySubCommand)
 	if err != nil {
 		return err
 	}