Update MoveToCodeAroundLineFunction to get the code directly instead

Author: intitni

Tool/Sources/ChatContextCollectors/ActiveDocumentChatContextCollector/ActiveDocumentChatContextCollector.swift

@@ -45,19 +45,7 @@ public final class ActiveDocumentChatContextCollector: ChatContextCollector {
         var functions = [any ChatGPTFunction]()
 
         if !isSensitive {
-            // When there is a line annotation not in the focused area, the bot can move the focus
-            // area
-            // to the code covering the line of the annotation.
-
-            if let focusedContext = context.focusedContext,
-               !focusedContext.otherLineAnnotations.isEmpty
-            {
-                functions.append(MoveToCodeAroundLineFunction(contextCollector: self))
-            }
-
-            if context.focusedContext == nil, !context.lineAnnotations.isEmpty {
-                functions.append(MoveToCodeAroundLineFunction(contextCollector: self))
-            }
+            functions.append(GetCodeCodeAroundLineFunction(contextCollector: self))
         }
 
         return .init(
@@ -89,9 +77,8 @@ public final class ActiveDocumentChatContextCollector: ChatContextCollector {
         let start = """
         ## File and Code Scope
 
-        You can use the following context to answer my questions about the editing document \
-        or code. The context shows only a part of the code in the editing document, and will \
-        change during the conversation, so it may not match our conversation.
+        You can use the following context to answer my questions about the editing document.\
+        The context shows only a part of the code in the editing document.
 
         \(
             context.focusedContext == nil
@@ -122,7 +109,9 @@ public final class ActiveDocumentChatContextCollector: ChatContextCollector {
                 Ask the user to select the code in the editor to get help. Also tell them the file is in gitignore.
                 """
                 : """
-                Focused Code (start from line \(focusedContext.codeRange.start.line + 1)):
+                Focused Code (from line \(
+                    focusedContext.codeRange.start.line + 1
+                ) to line \(focusedContext.codeRange.end.line + 1)):
                 ```\(context.language.rawValue)
                 \(focusedContext.code)
                 ```
@@ -131,8 +120,8 @@ public final class ActiveDocumentChatContextCollector: ChatContextCollector {
             let fileAnnotations = focusedContext.otherLineAnnotations.isEmpty || isSensitive
                 ? ""
                 : """
-                Other Annotations:\"""
-                (They are not inside the focused code. You don't known how to handle them until you get the code at the line)
+                Out-of-scope Annotations:\"""
+                (They are not inside the focused code. You can get the code at the line for details)
                 \(
                     focusedContext.otherLineAnnotations
                         .map(convertAnnotationToText)

Tool/Sources/ChatContextCollectors/ActiveDocumentChatContextCollector/Functions/ExpandFocusRangeFunction.swift

@@ -0,0 +1,93 @@
+import ASTParser
+import Foundation
+import OpenAIService
+import SuggestionModel
+
+struct GetCodeCodeAroundLineFunction: ChatGPTFunction {
+    struct Arguments: Codable {
+        var line: Int
+    }
+
+    struct Result: ChatGPTFunctionResult {
+        var range: CursorRange
+        var content: String
+        var language: CodeLanguage
+
+        var botReadableContent: String {
+            """
+            Code in range \(range)
+            ```\(language.rawValue)
+            \(content)
+            ```
+            """
+        }
+    }
+
+    struct E: Error, LocalizedError {
+        var errorDescription: String?
+    }
+
+    var name: String {
+        "getCodeAtLine"
+    }
+
+    var description: String {
+        "Get the code at the given line. You must ONLY call it when the user give you a specific line or the user ask about an out of scope annotation."
+    }
+
+    var argumentSchema: JSONSchemaValue { [
+        .type: "object",
+        .properties: [
+            "line": [
+                .type: "number",
+                .description: "The line number in the file",
+            ],
+        ],
+        .required: ["line"],
+    ] }
+
+    weak var contextCollector: ActiveDocumentChatContextCollector?
+
+    init(contextCollector: ActiveDocumentChatContextCollector) {
+        self.contextCollector = contextCollector
+    }
+
+    func prepare(reportProgress: @escaping (String) async -> Void) async {
+        await reportProgress("Finding code around..")
+    }
+
+    func call(
+        arguments: Arguments,
+        reportProgress: @escaping (String) async -> Void
+    ) async throws -> Result {
+        guard var activeDocumentContext = contextCollector?.activeDocumentContext else {
+            throw E(errorDescription: "No active document found.")
+        }
+        await reportProgress("Reading code around line \(arguments.line)..")
+        activeDocumentContext.moveToCodeAroundLine(max(arguments.line - 1, 0))
+        guard let newContext = activeDocumentContext.focusedContext else {
+            let progress = "Failed to read code around line \(arguments.line)..)"
+            await reportProgress(progress)
+            throw E(errorDescription: progress)
+        }
+        let progress = "Finish reading code at \(newContext.codeRange)"
+        await reportProgress(progress)
+        return .init(
+            range: newContext.codeRange,
+            content: newContext.code
+                .split(omittingEmptySubsequences: false, whereSeparator: \.isNewline)
+                .enumerated()
+                .map {
+                    let (index, content) = $0
+                    if index + newContext.codeRange.start.line == arguments.line - 1 {
+                        return content + " // <--- line \(arguments.line)"
+                    } else {
+                        return content
+                    }
+                }
+                .joined(separator: "\n"),
+            language: activeDocumentContext.language
+        )
+    }
+}
+