refactor: ♻️ enhance ask more context logic (#22)

#15

src/AIAssist/Commands/InternalCommands/RunCommand.cs

@@ -40,14 +40,15 @@ public async Task<bool> ExecuteAsync(IServiceScope scope, string? input)
         // Check if more context is needed
         if (codeAssistantManager.CheckExtraContextForResponse(responseContent, out var requiredFiles))
         {
+            var style = spectreUtilities.Theme.ConsoleStyle.Highlight.CombineStyle(s => s.Underline = true);
             var confirmation = spectreUtilities.ConfirmationPrompt(
-                $"Do you want to add ${string.Join(", ", requiredFiles.Select(file => $"'{file}'"))} to the context?"
+                $"Do you want to add [{spectreUtilities.CreateStringStyle(style)}]{string.Join(", ", requiredFiles.Select(file => $"'{file}'"))}[/] to the context?"
             );
 
             if (confirmation)
             {
                 await codeAssistantManager.AddOrUpdateCodeFiles(requiredFiles);
-                var fullFilesContentForContext = await codeAssistantManager.GetCodeTreeContentsFromCache(requiredFiles);
+                var fullFilesContentForContext = await codeAssistantManager.GetCodeTreeContents(requiredFiles);
 
                 var newQueryWithAddedFiles = promptManager.FilesAddedToChat(fullFilesContentForContext);
                 spectreUtilities.SuccessTextLine(

src/AIAssist/Contracts/CodeAssist/ICodeAssist.cs

@@ -2,7 +2,10 @@ namespace AIAssist.Contracts.CodeAssist;
 
 public interface ICodeAssist
 {
+    // AddOrUpdate folders, sub-folders, files in summary
     Task LoadInitCodeFiles(string contextWorkingDirectory, IList<string>? codeFiles);
+
+    // AddOrUpdate files with full definition
     Task AddOrUpdateCodeFiles(IList<string>? codeFiles);
     Task<IEnumerable<string>> GetCodeTreeContents(IList<string>? codeFiles);
     IAsyncEnumerable<string?> QueryChatCompletionAsync(string userQuery);

src/AIAssist/Contracts/CodeAssist/ICodeAssistantManager.cs

@@ -7,7 +7,7 @@ public interface ICodeAssistantManager
     Task LoadCodeFiles(string contextWorkingDirectory, IList<string>? codeFiles);
     IAsyncEnumerable<string?> QueryAsync(string userQuery);
     Task AddOrUpdateCodeFiles(IList<string>? codeFiles);
-    Task<IEnumerable<string>> GetCodeTreeContentsFromCache(IList<string>? codeFiles);
+    Task<IEnumerable<string>> GetCodeTreeContents(IList<string>? codeFiles);
     bool CheckExtraContextForResponse(string response, out IList<string> requiredFiles);
     IList<DiffResult> ParseDiffResults(string diffContent, string contextWorkingDirectory);
     void ApplyChanges(IList<DiffResult> diffResults, string contextWorkingDirectory);

src/AIAssist/Contracts/IContextService.cs

@@ -8,8 +8,25 @@ public interface IContextService
     IList<FileItemContext> GetAllFiles();
     IList<FileItemContext> GetFiles(IList<string>? filesRelativePath);
     void ValidateLoadedFilesLimit();
-    void AddContextFolder(string contextFolder);
-    void AddOrUpdateFolder(IList<string>? foldersRelativePath);
+
+    /// <summary>
+    /// AddOrUpdate folders, sub-folders, files in root level with summary mode.
+    /// </summary>
+    /// <param name="rootContextFolder"></param>
+    void AddContextFolder(string rootContextFolder);
+
+    /// <summary>
+    /// AddOrUpdate folders with full files in root level with definition mode.
+    /// </summary>
+    /// <param name="rootFoldersRelativePath"></param>
+    void AddOrUpdateFolder(IList<string>? rootFoldersRelativePath);
+    void RemoveFolder(IList<string>? rootFoldersRelativePath);
+
+    /// <summary>
+    /// AddOrUpdate files with full definition in all levels
+    /// </summary>
+    /// <param name="filesRelativePath"></param>
     void AddOrUpdateFiles(IList<string>? filesRelativePath);
+    void RemoveFiles(IList<string>? filesRelativePath);
     void AddOrUpdateUrls(IList<string>? urls);
 }

src/AIAssist/Extensions/DependencyInjectionExtensions.cs

@@ -98,7 +98,7 @@ private static void AddChatDependencies(HostApplicationBuilder builder)
 
     private static void AddSpectreConsoleDependencies(HostApplicationBuilder builder)
     {
-        builder.Services.AddSingleton(AnsiConsole.Console);
+        builder.Services.AddSingleton<IAnsiConsole>(AnsiConsole.Console);
 
         builder.Services.AddSingleton<ISpectreUtilities>(sp =>
         {

src/AIAssist/Prompts/Templates/ask-more-context.template

@@ -1,10 +1,7 @@
-If for generating  your response, you need full code implementation not just summary of code in the context code blocks, consider following rules:
-- Don't **MODIFY** code blocks in the existing context when code blocks are not complete and you have summary of code in the content of code block and you need full code block content for generating the response.
-- If you need full code content and file for a code block in the context and currently context has summary of code, **Ask** the user to add the full code content for code block in the context to generate a response and **stop** the generating the response and wait for user files. 
-- Include the required file paths in a `markdown code block` using the format below:
+If for generating your response and changing in a file you need full code definition and currently in code context there is no `definition` section for corresponding file, consider following rules:
+- Don't **MODIFY** code blocks in the existing context and `Stop` generating response.
+- Show the list of `required files` in the bellow format, don't forget to show this it is very IMPORTANT and if you miss this your response is INCORRECT:
 
-Required Files for Context:
-```
-- relative_file_path1
-- relative_file_path2
-```
+Required files for Context:
+- test/testFile1.cs 
+- test/testFile2.cs 

src/AIAssist/Prompts/Templates/code-assist-search-replace-diff.template

@@ -34,6 +34,8 @@ For creating a new file show original path `--- /dev/null` and modified path to
 For deleting a file, show the original path as --- [relativeFilePath of deleted file] and the modified path as `+++ /dev/null`.
 For file moves or renames, show a deletion with `--- [original relativeFilePath]` and `+++ /dev/null`, and a creation at the new path with `--- /dev/null` and `+++ [new relativeFilePath]`, including any modified content.
 Keep original `code style` and `formating` during apply unified diff format, Indentation is very important, Ensure that indentation is accurate in the diffs.
+Don't Skip codes logic with adding some comments for summarize your response, it should be complete.
+Don't respond your code blocks suggestion as tree structure.
 
 For better understanding the response format based on user request here you can see a sample:
 

src/AIAssist/Prompts/Templates/code-assistant-code-block-diff.template

@@ -27,6 +27,8 @@ We have 3 `action types` for each code block:
 For each change you make, you MUST return the `FULL CONTENT` of the modified or new file. for delete file we don't need the content.
 For changes for each file you should consider a seperated `CodeBlock` section.
 Keep original `code style` and `formating` and `indentation` during apply unified diff format, `Indentation` is very IMPORTANT, Ensure to KEEP INDENTATION is accurate in the diffs.
+Don't Skip codes logic with adding some comments for summarize your response, it should be complete.
+Don't respond your code blocks suggestion as tree structure.
 
 For better understanding the response format based on user request here you can see a sample:
 

src/AIAssist/Prompts/Templates/code-assistant-unified-diff.template

@@ -34,6 +34,8 @@ For creating a new file show original path `--- /dev/null` and modified path to
 For deleting a file, show the original path as --- [relativeFilePath of deleted file] and the modified path as `+++ /dev/null`.
 For file moves or renames, show a deletion with `--- [original relativeFilePath]` and `+++ /dev/null`, and a creation at the new path with `--- /dev/null` and `+++ [new relativeFilePath]`, including any modified content.
 Keep original `code style` and `formating` during apply unified diff format, Indentation is very important, Ensure that indentation is accurate in the diffs.
+Don't Skip codes logic with adding some comments for summarize your response, it should be complete.
+Don't respond your code blocks suggestion as tree structure.
 
 For better understanding the response format based on user request here you can see a sample:
 

src/AIAssist/Services/CodeAssistStrategies/TreeSitterCodeAssistSummary.cs

@@ -13,6 +13,7 @@ ILLMClientManager llmClientManager
     public Task LoadInitCodeFiles(string contextWorkingDirectory, IList<string>? codeFiles)
     {
         contextService.AddContextFolder(contextWorkingDirectory);
+        // fully load passed files definition
         contextService.AddOrUpdateFiles(codeFiles);
 
         return Task.CompletedTask;
@@ -23,6 +24,7 @@ public Task AddOrUpdateCodeFiles(IList<string>? codeFiles)
         if (codeFiles is null || codeFiles.Count == 0)
             return Task.CompletedTask;
 
+        // fully load files definition
         contextService.AddOrUpdateFiles(codeFiles);
 
         return Task.CompletedTask;

src/AIAssist/Services/CodeAssistantManager.cs

@@ -22,39 +22,31 @@ public Task AddOrUpdateCodeFiles(IList<string>? codeFiles)
         return codeAssist.AddOrUpdateCodeFiles(codeFiles);
     }
 
-    public Task<IEnumerable<string>> GetCodeTreeContentsFromCache(IList<string>? codeFiles)
+    public Task<IEnumerable<string>> GetCodeTreeContents(IList<string>? codeFiles)
     {
         return codeAssist.GetCodeTreeContents(codeFiles);
     }
 
     public bool CheckExtraContextForResponse(string response, out IList<string> requiredFiles)
     {
         requiredFiles = new List<string>();
-        var pattern = @"Required Files for Context:\s*(?:```[\w]*\s*([\s\S]*?)\s*```|((?:- .*\r?\n?)+))";
+
+        var pattern = @"Required files for Context:\s*(- .*(?:\r?\n- .*)*)";
 
         var match = Regex.Match(response, pattern);
         if (match.Success)
         {
-            // Check for code block content first
-            var codeBlockContent = match.Groups[1].Value;
-            if (!string.IsNullOrEmpty(codeBlockContent))
-            {
-                var lines = codeBlockContent.Split(new[] { '\n', '\r' }, StringSplitOptions.RemoveEmptyEntries);
-                foreach (var line in lines)
-                {
-                    requiredFiles.Add(line.TrimStart('-').Trim());
-                }
-                return true;
-            }
-
-            // If no code block content, check for inline list
-            var inlineListContent = match.Groups[2].Value;
+            var inlineListContent = match.Groups[1].Value;
             if (!string.IsNullOrEmpty(inlineListContent))
             {
-                var lines = inlineListContent.Split(new[] { '\n', '\r' }, StringSplitOptions.RemoveEmptyEntries);
+                var lines = inlineListContent.Split(['\n', '\r'], StringSplitOptions.RemoveEmptyEntries);
                 foreach (var line in lines)
                 {
-                    requiredFiles.Add(line.TrimStart('-').Trim());
+                    if (line.StartsWith('-'))
+                    {
+                        // Remove the leading '-' and trim whitespace
+                        requiredFiles.Add(line.TrimStart('-').Trim());
+                    }
                 }
                 return true;
             }

src/AIAssist/Services/CodeFilesTreeGeneratorService.cs

@@ -81,13 +81,13 @@ private CodeFileMap AddOrUpdateCodeFileMap(CodeFile codeFile, bool useShortSumma
         var codeFileMap = new CodeFileMap
         {
             RelativePath = codeFile.RelativePath.NormalizePath(),
+            Path = codeFile.Path,
             OriginalCode = GetOriginalCode(fileCapturesGroup),
             TreeSitterFullCode = treeStructureGeneratorService.GenerateTreeSitter(fileCapturesGroup, true),
             TreeOriginalCode = originalCode,
             TreeSitterSummarizeCode = useShortSummary
                 ? treeStructureGeneratorService.GenerateTreeSitter(fileCapturesGroup, false)
                 : originalCode,
-            ReferencedCodesMap = GenerateRelatedCodeFilesMap(fileCapturesGroup),
         };
 
         if (_codeFilesMap.Any(x => x.Key.NormalizePath() == codeFile.RelativePath.NormalizePath()))
@@ -103,24 +103,6 @@ private CodeFileMap AddOrUpdateCodeFileMap(CodeFile codeFile, bool useShortSumma
         return codeFileMap;
     }
 
-    private static IEnumerable<ReferencedCodeMap> GenerateRelatedCodeFilesMap(
-        List<DefinitionCaptureItem> definitionItems
-    )
-    {
-        // Collect related references for each DefinitionCaptureItem
-        return definitionItems
-            .SelectMany(item =>
-                item.DefinitionCaptureReferences.Select(reference => new ReferencedCodeMap
-                {
-                    RelativePath = reference.RelativePath.NormalizePath(),
-                    ReferencedValue = reference.ReferencedValue,
-                    ReferencedUsage = reference.ReferencedUsage,
-                })
-            )
-            .Distinct()
-            .ToList();
-    }
-
     private IList<CodeFile> ReadCodeFiles(IList<string>? files)
     {
         ArgumentException.ThrowIfNullOrEmpty(_appOptions.ContextWorkingDirectory);
@@ -135,13 +117,13 @@ private IList<CodeFile> ReadCodeFiles(IList<string>? files)
             var relativePath = Path.GetRelativePath(_appOptions.ContextWorkingDirectory, file);
             var fileContent = File.ReadAllText(file);
 
-            applicationCodes.Add(new CodeFile(fileContent, relativePath.NormalizePath()));
+            applicationCodes.Add(new CodeFile(fileContent, relativePath.NormalizePath(), file));
         }
 
         return applicationCodes;
     }
 
-    private static string GetOriginalCode(List<DefinitionCaptureItem> definitionItems)
+    private static string GetOriginalCode(List<DefinitionCapture> definitionItems)
     {
         return definitionItems.First().OriginalCode;
     }