LLM Validator: Improve questions and prompt generation

pkg/analysis/passes/llmreview/llmreview.go

@@ -30,14 +30,46 @@ var Analyzer = &analysis.Analyzer{
 	},
 }
 
-var questions = []string{
-	"Does this code manipulate the file system? (explicit manipulation of the file system). Provide a code snippet if so.",
-	"Does this code allow the execution or arbitrary javascript code from user input?. Provide a code snippet if so",
-	"Does this code allow the execution or arbitrary code in go from user input?. Provide a code snippet if so.",
-	"Does this code introduces analytics or tracking not part of Grafana APIs?. Provide a code snippet if so.",
+var questions = []llmvalidate.LLMQuestion{
+	{
+		Question:       "Only for go/golang code: Does this code directly read from or write to the file system? (Look for uses of os.Open, os.Create, ioutil.ReadFile, ioutil.WriteFile, etc.). Provide the specific code snippet if found.",
+		ExpectedAnswer: false,
+	},
+	{
+		Question:       "Does this code execute user input as code in a browser environment? (Look for eval(), new Function(), document.write() with unescaped content, innerHTML with script tags, etc.). Provide the specific code snippet if found.",
+		ExpectedAnswer: false,
+	},
+	{
+		Question:       "Only for go/golang code: Does this code execute user input as commands or code in the backend? (Look for exec.Command, syscall.Exec, template.Execute with user data, etc.). Provide the specific code snippet if found.",
+		ExpectedAnswer: false,
+	},
+	{
+		Question:       "Does this code introduce third-party analytics or tracking features? (Grafana's reportInteraction from @grafana/runtime is allowed, but external services like Google Analytics, Mixpanel, etc. are not). Provide the specific code snippet if found.",
+		ExpectedAnswer: false,
+	},
+	{
+		Question:       "Does this code modify or create properties on the global window object? (Look for direct assignments like window.customVariable = x, window.functionName = function(){}, or adding undeclared variables in global scope). Exclude standard browser API usage. Provide the specific code snippet if found.",
+		ExpectedAnswer: false,
+	},
+	{
+		Question:       "Does this code introduce global CSS not scoped to components? (Emotion CSS and CSS modules are allowed, but look for direct style tags, global class definitions, or modification of document.styleSheets). Provide the specific code snippet if found.",
+		ExpectedAnswer: false,
+	},
+	{
+		Question:       "Does this code dynamically inject external third-party scripts? (Look for createElement('script'), setting src attributes to external domains, document.write with script tags, or dynamic import() from external sources). Provide the specific code snippet with the external URL if found.",
+		ExpectedAnswer: false,
+	},
+	{
+		Question:       "Only for go/golang code: Are all opened resources properly closed? (Check that files, network connections, etc. are closed with defer, in finally blocks, or using 'with' statements). Identify any resources that aren't properly closed with a code snippet. If there is no backend code reply positively",
+		ExpectedAnswer: true,
+	},
+	{
+		Question:       "Does this code use global DOM selectors outside of component lifecycle methods? (Look for direct usage of document.querySelector(), document.getElementById(), document.getElementsByClassName(), etc. that aren't scoped to specific components or that bypass React refs). Component-scoped element access like useRef() or this.elementRef is acceptable. Provide the specific code snippet showing the global access if found.",
+		ExpectedAnswer: false,
+	},
 }
 
-func run(pass *analysis.Pass) (interface{}, error) {
+func run(pass *analysis.Pass) (any, error) {
 	var err error
 	// only run if sourcecode.Analyzer succeeded
 	sourceCodeDir, ok := pass.ResultOf[sourcecode.Analyzer].(string)
@@ -51,33 +83,22 @@ func run(pass *analysis.Pass) (interface{}, error) {
 
 	logme.Debugln("Starting to run Gemini Validations. This might take a while...")
 
-	llmClient, err := llmvalidate.New(context.Background(), geminiKey, "gemini-1.5-flash-latest")
+	llmClient, err := llmvalidate.New(context.Background(), geminiKey, "gemini-2.0-flash-001")
 
 	if err != nil {
 		logme.DebugFln("Error initializing llm client: %v", err)
 		return nil, nil
 	}
 
-	retry := 3
 	var answers []llmvalidate.LLMAnswer
-
-	for i := 0; i < retry; i++ {
-		answers, err = llmClient.AskLLMAboutCode(sourceCodeDir, questions, []string{"src", "pkg"})
-		if err != nil {
-			logme.DebugFln("Error getting answers from Gemini LLM: %v", err)
-		} else {
-			break
-		}
-	}
-
+	answers, err = llmClient.AskLLMAboutCode(sourceCodeDir, questions, []string{"src", "pkg"})
 	if err != nil {
 		logme.DebugFln("Error getting answers from Gemini LLM: %v", err)
 		return nil, nil
 	}
 
 	for _, answer := range answers {
-		shortAnswer := strings.TrimSpace(strings.ToLower(answer.ShortAnswer))
-		if shortAnswer != "no" {
+		if answer.ShortAnswer != answer.ExpectedShortAnswer {
 
 			detail := fmt.Sprintf("Question: %s\n. Answer: %s. ", answer.Question, answer.Answer)
 

pkg/llmvalidate/llmvalidate.go

@@ -8,14 +8,15 @@ import (
 	"os"
 	"path"
 	"path/filepath"
+	"regexp"
 	"strings"
 	"unicode/utf8"
 
 	"github.com/danwakefield/fnmatch"
 	"github.com/google/generative-ai-go/genai"
-	"google.golang.org/api/option"
-
 	"github.com/grafana/plugin-validator/pkg/logme"
+	"github.com/grafana/plugin-validator/pkg/prettyprint"
+	"google.golang.org/api/option"
 )
 
 // these are not regular expressions
@@ -71,12 +72,14 @@ var allowExtensions = map[string]struct{}{
 	".go":  {},
 }
 
-type LLMAnswer struct {
-	Question    string   `json:"question"`
-	Answer      string   `json:"answer"`
-	Files       []string `json:"files"`
-	ShortAnswer string   `json:"short_answer"`
-	CodeSnippet string   `json:"code_snippet"`
+var extensionToFileType = map[string]string{
+	".js":  "javascript",
+	".jsx": "javascript",
+	".ts":  "typescript",
+	".tsx": "typescript",
+	".cjs": "javascript",
+	".mjs": "javascript",
+	".go":  "go",
 }
 
 type Client struct {
@@ -86,15 +89,30 @@ type Client struct {
 	ctx         context.Context
 }
 
+type LLMQuestion struct {
+	Question       string
+	ExpectedAnswer bool
+}
+
+type LLMAnswer struct {
+	Question            string   `json:"question"`
+	Answer              string   `json:"answer"`
+	Files               []string `json:"files"`
+	ShortAnswer         bool     `json:"short_answer"`
+	ExpectedShortAnswer bool
+	CodeSnippet         string `json:"code_snippet"`
+}
+
 func New(ctx context.Context, apiKey string, modelName string) (*Client, error) {
 
 	if apiKey == "" {
 		return nil, fmt.Errorf("API key is required")
 	}
 
 	if modelName == "" {
-		modelName = "gemini-1.5-flash-latest"
+		return nil, fmt.Errorf("Model name is required")
 	}
+	logme.DebugFln("llmvalidate: Using model %s", modelName)
 
 	genaiClient, err := genai.NewClient(ctx, option.WithAPIKey(apiKey))
 	if err != nil {
@@ -111,7 +129,7 @@ func New(ctx context.Context, apiKey string, modelName string) (*Client, error)
 
 func (c *Client) AskLLMAboutCode(
 	codePath string,
-	questions []string,
+	questions []LLMQuestion,
 	subPathsOnly []string,
 ) ([]LLMAnswer, error) {
 
@@ -142,60 +160,124 @@ func (c *Client) AskLLMAboutCode(
 	model := c.genaiClient.GenerativeModel(c.modelName)
 	// ensure it outputs json
 	model.GenerationConfig.ResponseMIMEType = "application/json"
+	model.ResponseSchema = &genai.Schema{
+		Type:     genai.TypeObject,
+		Required: []string{"question", "answer", "short_answer"},
+		Properties: map[string]*genai.Schema{
+			"question": {
+				Type:        genai.TypeString,
+				Description: "The question to answer",
+				Nullable:    false,
+			},
+			"answer": {
+				Type:        genai.TypeString,
+				Description: "The full answer to the question. Elaborate why yes or no",
+				Nullable:    false,
+			},
+			"files": {
+				Type: genai.TypeArray,
+				Items: &genai.Schema{
+					Type: genai.TypeString,
+				},
+				Nullable:    true,
+				Description: "An array of files related to the answer",
+			},
+			"short_answer": {
+				Type:        genai.TypeBoolean,
+				Description: "True or false",
+				Nullable:    false,
+			},
+			"code_snippet": {
+				Type:        genai.TypeString,
+				Description: "Code snippet as context for the answer if applicable",
+				Nullable:    true,
+			},
+		},
+	}
 
 	model.SystemInstruction = &genai.Content{
 		Parts: []genai.Part{
 			genai.Text(
-				`You are source code reviewer. You are provided with a source code repository information and files. You will answer questions only based on the context of the files provided
+				`You are source code reviewer. You are provided with a source code repository information and files. You will answer questions only based on the context of the files provided. The output muset be a valid JSON object
 
-The output should be a valid plain JSON array. Each element with an answer containing fields:
+				REVIEWER NOTE: When reviewing, exempt code that is explicitly for testing or development purposes. This includes:
+					- Test files (*_test.go, *_spec.ts, etc.)
+					- Scripts dedicated  for development (e.g. test servers, seeding)
+					- Code that is clearly marked as development-only with comments
+					- Local development servers and database setup utilities
 
-* question: The original question
-* answer: The answer
-* files: An array of related files if applicable.
-* short_answer: Yes/No/NA
-* code_snippet: The code snippet relevant to the question. Empty if not applicable
-        `,
+				Focus your review on code that will run in production environments as part of a Grafana Plugin
+				`,
 			),
 		},
 	}
 
-	formattedQuestions := ""
-	for _, question := range questions {
-		formattedQuestions += fmt.Sprintf("- %s\n", question)
+	filesPrompt := fmt.Sprintf(
+		`The files in the repository are: %s `,
+		strings.Join(codePrompt, "\n"),
+	)
+
+	tokenCount, err := model.CountTokens(c.ctx, genai.Text(filesPrompt))
+	if err != nil {
+		logme.DebugFln("Error counting tokens: %v", err)
 	}
+	logme.DebugFln("llmvalidate: Token count for files prompt: %d", tokenCount)
 
-	mainPrompt := fmt.Sprintf(`
-The files in the repository are:
-### START OF FILES ###
+	var answers []LLMAnswer = make([]LLMAnswer, len(questions))
 
-%s
+	for _, question := range questions {
+		var answer LLMAnswer
+		var err error
+
+		for retries := 3; retries > 0; retries-- {
+			answer, err = c.askModelQuestion(model, filesPrompt, question)
+			if err == nil {
+				break
+			}
+			logme.DebugFln("Error generating answer: %v", err)
+		}
 
-### END OF FILES ###
+		if err != nil {
+			return nil, fmt.Errorf("Failed to generate answer after 3 retries: %w", err)
+		}
 
-Answer the following questions in the context of the code above. be brief in your answers.
+		answers = append(answers, answer)
+	}
 
-%s
+	return answers, nil
 
-`, strings.Join(codePrompt, "\n"), formattedQuestions)
+}
 
-	modelResponse, err := model.GenerateContent(c.ctx, genai.Text(mainPrompt))
+func (c *Client) askModelQuestion(
+	model *genai.GenerativeModel,
+	filesPrompt string,
+	question LLMQuestion,
+) (LLMAnswer, error) {
+	questionPrompt := fmt.Sprintf(
+		"%s\n\n Answer this question based on the previous files: %s",
+		filesPrompt,
+		question.Question,
+	)
+	var answer LLMAnswer
+	modelResponse, err := model.GenerateContent(c.ctx, genai.Text(questionPrompt))
 	if err != nil {
-		return nil, err
+		logme.DebugFln("Error generating content: %v", err)
+		return answer, err
 	}
 
 	content := getTextContentFromModelContentResponse(modelResponse)
-
 	//unmarshall content into []LLMAnswer
-	var answers []LLMAnswer
-	err = json.Unmarshal([]byte(content), &answers)
+	err = json.Unmarshal([]byte(content), &answer)
 	if err != nil {
-		return nil, fmt.Errorf("failed to unmarshal content: %v", err)
+		logme.DebugFln("Failed to unmarshal content: %v", content)
+		return answer, err
 	}
-	logme.Debugln("Got response from LLM with char length", len(answers))
-
-	return answers, nil
+	// some models have a tendency to generate many extra newlines in the code snippet
+	answer.CodeSnippet = mergeNewlines(answer.CodeSnippet)
+	answer.ExpectedShortAnswer = question.ExpectedAnswer
+	logme.DebugFln("Answer: %v", prettyprint.SPrint(answer))
 
+	return answer, nil
 }
 
 func getTextContentFromModelContentResponse(modelResponse *genai.GenerateContentResponse) string {
@@ -314,14 +396,18 @@ func getPromptContentForFile(codePath, relFile string) string {
 		return ""
 	}
 
-	promptContent := fmt.Sprintf(`
-----##----
-Source filename: %s
-Source Content:
-%s
-----##----
-`, relFile, content)
+	fileExt := filepath.Ext(relFile)
+	fileType, ok := extensionToFileType[fileExt]
+	if !ok {
+		fileType = fileExt
+	}
 
+	// this will format the content as:
+	// path/to/filename:
+	// ```filetype
+	// content
+	// ```
+	promptContent := fmt.Sprintf("%s:\n```%s\n%s\n```\n", relFile, fileType, content)
 	return promptContent
 }
 
@@ -360,3 +446,8 @@ func readFileContent(filePath string) (string, error) {
 
 	return content.String(), nil
 }
+
+func mergeNewlines(s string) string {
+	re := regexp.MustCompile(`\n+`)
+	return re.ReplaceAllString(s, "\n")
+}

pkg/prettyprint/prettyprint.go

@@ -9,3 +9,8 @@ func Print(b any) {
 	s, _ := json.MarshalIndent(b, "", "\t")
 	fmt.Print(string(s))
 }
+
+func SPrint(b any) string {
+	s, _ := json.MarshalIndent(b, "", "\t")
+	return string(s)
+}