Merge pull request #103 from meysamhadeli/refactor/refactor-render-co…

…de-highlight-in-terminal

refactor: refactor highlight code changes in terminal

cmd/code.go

@@ -38,9 +38,7 @@ func handleCodeCommand(rootDependencies *RootDependencies) {
 	defer cancel()
 
 	var requestedContext string
-	var fullContext []models.FileData
-	var fullContextCodes []string
-
+	var fullContext *models.FullContextData
 	const maxTokens = 8000           // Max tokens per embedding request
 	const requestDelay = time.Second // requests per second (adjust based on rate limits)
 	var currentChunks []general_models.ChunkData
@@ -62,7 +60,7 @@ func handleCodeCommand(rootDependencies *RootDependencies) {
 	spinnerLoadContext, _ := spinner.Start("Loading Context...")
 
 	// Get all data files from the root directory
-	fullContext, fullContextCodes, err := rootDependencies.Analyzer.GetProjectFiles(rootDependencies.Cwd)
+	fullContext, err := rootDependencies.Analyzer.GetProjectFiles(rootDependencies.Cwd)
 
 	if err != nil {
 		spinnerLoadContext.Stop()
@@ -112,7 +110,7 @@ func handleCodeCommand(rootDependencies *RootDependencies) {
 
 		spinnerEmbeddingContext, _ := spinner.Start("Embedding Context...")
 
-		for _, dataFile := range fullContext {
+		for _, dataFile := range fullContext.FileData {
 			// Split file into chunks of up to 8000 tokens
 			fileChunks, err := rootDependencies.TokenManagement.SplitTokenIntoChunks(dataFile.Code, maxTokens)
 			if err != nil {
@@ -202,7 +200,7 @@ startLoop: // Label for the start loop
 			var aiResponseBuilder strings.Builder
 
 			chatRequestOperation := func() error {
-				finalPrompt, userInputPrompt := rootDependencies.Analyzer.GeneratePrompt(fullContextCodes, rootDependencies.ChatHistory.GetHistory(), userInput, requestedContext)
+				finalPrompt, userInputPrompt := rootDependencies.Analyzer.GeneratePrompt(fullContext.RawCodes, rootDependencies.ChatHistory.GetHistory(), userInput, requestedContext)
 
 				// Step 7: Send the relevant code and user input to the AI API
 				responseChan := rootDependencies.CurrentChatProvider.ChatCompletionRequest(ctx, userInputPrompt, finalPrompt)
@@ -247,7 +245,7 @@ startLoop: // Label for the start loop
 					topN := -1
 
 					// Step 6: Find relevant code chunks based on the user query embedding
-					fullContextCodes = rootDependencies.Store.FindRelevantChunks(queryEmbedding[0], topN, rootDependencies.Config.AIProviderConfig.EmbeddingsModel, rootDependencies.Config.AIProviderConfig.Threshold)
+					fullContext.RawCodes = rootDependencies.Store.FindRelevantChunks(queryEmbedding[0], topN, rootDependencies.Config.AIProviderConfig.EmbeddingsModel, rootDependencies.Config.AIProviderConfig.Threshold)
 					return nil
 				}
 

code_analyzer/analyzer.go

@@ -58,14 +58,13 @@ func NewCodeAnalyzer(cwd string, isRAG bool) contracts.ICodeAnalyzer {
 	return &CodeAnalyzer{Cwd: cwd, IsRAG: isRAG}
 }
 
-func (analyzer *CodeAnalyzer) GetProjectFiles(rootDir string) ([]models.FileData, []string, error) {
-	var result []models.FileData
-	var codes []string
+func (analyzer *CodeAnalyzer) GetProjectFiles(rootDir string) (*models.FullContextData, error) {
+	var result models.FullContextData
 
 	// Retrieve the ignore patterns from .gitignore, if it exists
 	gitIgnorePatterns, err := utils.GetGitignorePatterns(rootDir)
 	if err != nil {
-		return nil, nil, err
+		return nil, err
 	}
 
 	// Walk the directory tree and find all files
@@ -115,19 +114,19 @@ func (analyzer *CodeAnalyzer) GetProjectFiles(rootDir string) ([]models.FileData
 			codeParts := analyzer.ProcessFile(relativePath, content)
 
 			// Append the file data to the result
-			result = append(result, models.FileData{RelativePath: relativePath, Code: string(content), TreeSitterCode: strings.Join(codeParts, "\n")})
+			result.FileData = append(result.FileData, models.FileData{RelativePath: relativePath, Code: string(content), TreeSitterCode: strings.Join(codeParts, "\n")})
 
-			codes = append(codes, fmt.Sprintf("**File: %s**\n\n%s", relativePath, strings.Join(codeParts, "\n")))
+			result.RawCodes = append(result.RawCodes, fmt.Sprintf("**File: %s**\n\n%s", relativePath, strings.Join(codeParts, "\n")))
 		}
 
 		return nil
 	})
 
 	if err != nil {
-		return nil, nil, err
+		return nil, err
 	}
 
-	return result, codes, nil
+	return &result, nil
 }
 
 // ProcessFile processes a single file using Tree-sitter for syntax analysis (for .cs files).
@@ -361,7 +360,17 @@ func (analyzer *CodeAnalyzer) ApplyChanges(relativePath, diff string) error {
 	var updatedContent []string
 
 	for _, line := range diffLines {
-		updatedContent = append(updatedContent, line)
+		trimmedLine := strings.TrimSpace(line)
+		if strings.HasPrefix(trimmedLine, "-") {
+			// Ignore lines that start with "-", effectively deleting them
+			continue
+		} else if strings.HasPrefix(trimmedLine, "+") {
+			// Add lines that start with "+", but remove the "+" symbol
+			updatedContent = append(updatedContent, strings.ReplaceAll(trimmedLine, "+", " "))
+		} else {
+			// Keep all other lines as they are
+			updatedContent = append(updatedContent, line)
+		}
 	}
 
 	// Handle deletion if code is empty

code_analyzer/analyzer_test.go

@@ -134,13 +134,13 @@ func TestGetProjectFiles(t *testing.T) {
 	_ = os.WriteFile(testFilePath, []byte("package main\nfunc main() {}"), 0644)
 	_ = os.WriteFile(ignoreFilePath, []byte("node_modules\n"), 0644)
 
-	files, codes, err := analyzer.GetProjectFiles(relativePathTestDir)
+	fullContext, err := analyzer.GetProjectFiles(relativePathTestDir)
 
 	assert.NoError(t, err)
-	assert.Len(t, files, 1)
-	assert.Len(t, codes, 1)
+	assert.Len(t, fullContext.FileData, 1)
+	assert.Len(t, fullContext.RawCodes, 1)
 
-	for _, file := range files {
+	for _, file := range fullContext.FileData {
 		assert.NotEmpty(t, file.RelativePath)
 		assert.Equal(t, "test.go", filepath.Base(file.RelativePath))
 	}

code_analyzer/contracts/analyzer.go

@@ -5,7 +5,7 @@ import (
 )
 
 type ICodeAnalyzer interface {
-	GetProjectFiles(rootDir string) ([]models.FileData, []string, error)
+	GetProjectFiles(rootDir string) (*models.FullContextData, error)
 	ProcessFile(filePath string, sourceCode []byte) []string
 	GeneratePrompt(codes []string, history []string, userInput string, requestedContext string) (string, string)
 	ExtractCodeChanges(text string) []models.CodeChange

code_analyzer/models/file_data.go

@@ -6,3 +6,8 @@ type FileData struct {
 	Code           string
 	TreeSitterCode string
 }
+
+type FullContextData struct {
+	FileData []FileData
+	RawCodes []string
+}

utils/markdown_generator.go

@@ -1,16 +1,34 @@
 package utils
 
 import (
+	"fmt"
 	"github.com/alecthomas/chroma/v2/quick"
 	"os"
+	"strings"
 )
 
+var isCodeBlock = false
+
 // RenderAndPrintMarkdown handles the rendering of markdown content,
-func RenderAndPrintMarkdown(content string, language string, theme string) error {
+func RenderAndPrintMarkdown(line string, language string, theme string) error {
 
-	err := quick.Highlight(os.Stdout, content, language, "terminal256", theme)
-	if err != nil {
-		return err
+	if strings.HasPrefix(line, "```") {
+		isCodeBlock = !isCodeBlock
+	}
+	// Process the line based on its prefix
+	if strings.HasPrefix(line, "+") && isCodeBlock {
+		line = "\x1b[92m" + line + "\x1b[0m"
+		fmt.Print(line)
+	} else if strings.HasPrefix(line, "-") && isCodeBlock {
+		line = "\x1b[91m" + line + "\x1b[0m"
+		fmt.Print(line)
+	} else {
+		// Render the processed line
+		err := quick.Highlight(os.Stdout, line, language, "terminal256", theme)
+		if err != nil {
+			return err
+		}
 	}
+
 	return nil
 }