From fc9e6a5fcee55cdb2b40b962ff0b02b065172f38 Mon Sep 17 00:00:00 2001
From: Weetile <weetile@hotmail.com>
Date: Fri, 23 May 2025 13:14:41 +0100
Subject: [PATCH] add rules to AI

---
 main.go | 58 +++++++++++++++++++++------------------------------------
 1 file changed, 21 insertions(+), 37 deletions(-)

diff --git a/main.go b/main.go
index ff074c3..adf4d27 100644
--- a/main.go
+++ b/main.go
@@ -9,11 +9,11 @@ import (
 	"os"
 	"os/exec"
 	"strings"
-	// "time" // Removed: import and not used
 
 	"github.com/charmbracelet/huh"
 	"github.com/google/generative-ai-go/genai"
 	"github.com/sashabaranov/go-openai"
+	"google.golang.org/api/iterator" // Added for iterator.Done
 	"google.golang.org/api/option"
 )
 
@@ -21,7 +21,7 @@ import (
 const maxRedirects = 10
 
 // Constants for AI model names
-const geminiModelName = "gemini-1.5-flash-latest"
+const geminiModelName = "gemini-2.5-flash-preview-05-20"
 const openAIModelName = "gpt-4o-mini"
 
 // isValidURLAndFetchContent (no changes)
@@ -127,7 +127,9 @@ func buildAnalysisPrompt(scriptContent string) string {
 %s
 --- SCRIPT END ---
 
-Provide your analysis as a stream of text.`, scriptContent)
+
+Do not over explain yourself or include very basic information, assume the user already understands the basics of a shell script.
+Do not use Markdown symbols such as asterisks or backticks.`, scriptContent)
 }
 
 func analyzeWithGemini(scriptContent string, apiKey string) error {
@@ -139,24 +141,11 @@ func analyzeWithGemini(scriptContent string, apiKey string) error {
 	defer client.Close()
 
 	model := client.GenerativeModel(geminiModelName)
-	// Corrected: Assign directly to the SafetySettings field
 	model.SafetySettings = []*genai.SafetySetting{
-		{
-			Category:  genai.HarmCategoryDangerousContent,
-			Threshold: genai.HarmBlockNone,
-		},
-		{
-			Category:  genai.HarmCategoryHarassment,
-			Threshold: genai.HarmBlockNone,
-		},
-		{
-			Category:  genai.HarmCategorySexuallyExplicit,
-			Threshold: genai.HarmBlockNone,
-		},
-		{
-			Category:  genai.HarmCategoryHateSpeech,
-			Threshold: genai.HarmBlockNone,
-		},
+		{Category: genai.HarmCategoryDangerousContent, Threshold: genai.HarmBlockNone},
+		{Category: genai.HarmCategoryHarassment, Threshold: genai.HarmBlockNone},
+		{Category: genai.HarmCategorySexuallyExplicit, Threshold: genai.HarmBlockNone},
+		{Category: genai.HarmCategoryHateSpeech, Threshold: genai.HarmBlockNone},
 	}
 	prompt := buildAnalysisPrompt(scriptContent)
 
@@ -164,21 +153,21 @@ func analyzeWithGemini(scriptContent string, apiKey string) error {
 	iter := model.GenerateContentStream(ctx, genai.Text(prompt))
 	for {
 		resp, err := iter.Next()
-		if err == io.EOF {
+		// Corrected check: Use errors.Is with iterator.Done
+		if errors.Is(err, iterator.Done) {
 			break
 		}
 		if err != nil {
+			// This will catch other errors that are not iterator.Done
 			return fmt.Errorf("Gemini stream error: %w", err)
 		}
+
 		if resp == nil || len(resp.Candidates) == 0 || resp.Candidates[0].Content == nil || len(resp.Candidates[0].Content.Parts) == 0 {
 			continue
 		}
-		// Ensure part is genai.Text before trying to print
-		if len(resp.Candidates[0].Content.Parts) > 0 {
-			if textPart, ok := resp.Candidates[0].Content.Parts[0].(genai.Text); ok {
-				fmt.Print(textPart)
-				os.Stdout.Sync() // Ensure immediate output
-			}
+		if textPart, ok := resp.Candidates[0].Content.Parts[0].(genai.Text); ok {
+			fmt.Print(textPart)
+			os.Stdout.Sync()
 		}
 	}
 	fmt.Println("\n--- GEMINI ANALYSIS FINISHED ---")
@@ -191,14 +180,8 @@ func analyzeWithOpenAI(scriptContent string, apiKey string) error {
 
 	prompt := buildAnalysisPrompt(scriptContent)
 	messages := []openai.ChatCompletionMessage{
-		{
-			Role:    openai.ChatMessageRoleSystem,
-			Content: "You are an expert shell script analyzer.",
-		},
-		{
-			Role:    openai.ChatMessageRoleUser,
-			Content: prompt,
-		},
+		{Role: openai.ChatMessageRoleSystem, Content: "You are an expert shell script analyzer."},
+		{Role: openai.ChatMessageRoleUser, Content: prompt},
 	}
 
 	request := openai.ChatCompletionRequest{
@@ -216,7 +199,7 @@ func analyzeWithOpenAI(scriptContent string, apiKey string) error {
 
 	for {
 		response, err := stream.Recv()
-		if errors.Is(err, io.EOF) {
+		if errors.Is(err, io.EOF) { // OpenAI SDK uses io.EOF for stream end
 			break
 		}
 		if err != nil {
@@ -224,7 +207,7 @@ func analyzeWithOpenAI(scriptContent string, apiKey string) error {
 		}
 		if len(response.Choices) > 0 {
 			fmt.Print(response.Choices[0].Delta.Content)
-			os.Stdout.Sync() // Ensure immediate output
+			os.Stdout.Sync()
 		}
 	}
 	fmt.Println("\n--- OPENAI ANALYSIS FINISHED ---")
@@ -333,6 +316,7 @@ func main() {
 		}
 		err := analyzeWithGemini(string(content), geminiAPIKey)
 		if err != nil {
+			// This will now only print if there's an actual error other than iterator.Done
 			fmt.Fprintf(os.Stderr, "\nError analyzing with Gemini: %v\n", err)
 		}
 	case "analyze_openai":