Add deeper reference capability for more context

Core/Sources/HostApp/Benchmark/Data/Entity/StoredSuggestionDTO.swift

@@ -1,12 +1,18 @@
-struct StoredSuggestionDTO: Codable {
+import Foundation
+
+struct StoredSuggestionDTO: Encodable {
     let fileURL: String
     let id: String
     let suggestionText: String
     let position: CursorPositionDTO
     let range: CursorRangeDTO
     let createdAt: String
     let relevantSymbols: [RelevantSymbolsDTO]
+    var relevantSymbolCount: Int {
+        relevantSymbols.count
+    }
     let model: String
+    let relevantFileScanningDurationInSeconds: Double?
 
     struct CursorPositionDTO: Codable {
         let line: Int
@@ -26,4 +32,28 @@ struct StoredSuggestionDTO: Codable {
         let endLine: Int
         let kind: String
     }
+
+    enum CodingKeys: String, CodingKey {
+        case fileURL, id, suggestionText, position, range, createdAt, relevantSymbols, model
+        case relevantFileScanningDurationInSeconds
+        case relevantSymbolCount // write, but ignore on decoding
+    }
+
+    func encode(to encoder: Encoder) throws {
+        var container = encoder.container(keyedBy: CodingKeys.self)
+        try container.encode(fileURL, forKey: .fileURL)
+        try container.encode(id, forKey: .id)
+        try container.encode(suggestionText, forKey: .suggestionText)
+        try container.encode(position, forKey: .position)
+        try container.encode(range, forKey: .range)
+        try container.encode(createdAt, forKey: .createdAt)
+        try container.encode(relevantSymbols, forKey: .relevantSymbols)
+        try container.encode(relevantSymbolCount, forKey: .relevantSymbolCount)
+        try container.encode(model, forKey: .model)
+
+        if let value = relevantFileScanningDurationInSeconds {
+            let rounded = Double(round(100 * value) / 100)
+            try container.encode(rounded, forKey: .relevantFileScanningDurationInSeconds)
+        }
+    }
 }

Core/Sources/HostApp/Benchmark/Data/Manager/MultiFileContextBenchmarkManager.swift

@@ -29,6 +29,17 @@ class MultiFileContextBenchmarkManager: BenchmarkManager {
     var selectedGenAIModel: AnyPublisher<GenAILanguageModel, Never> {
         selectedGenAIModelSubject.eraseToAnyPublisher()
     }
+
+    private let contextLevelLimitSubject = CurrentValueSubject<ContextLevelLimit, Never>(.firstLevel)
+    var contextLevelLimit: AnyPublisher<ContextLevelLimit, Never> {
+        contextLevelLimitSubject.eraseToAnyPublisher()
+    }
+
+    private var contextFileAmountLimitSubject = CurrentValueSubject<Int?, Never>(nil)
+    var contextFileAmountLimit: AnyPublisher<Int?, Never> {
+        contextFileAmountLimitSubject.eraseToAnyPublisher()
+    }
+
     private var cancellables = Set<AnyCancellable>()
 
     init(benchmarkSettingsRepository: BenchmarkSettingsRepository) {
@@ -93,10 +104,11 @@ class MultiFileContextBenchmarkManager: BenchmarkManager {
         let entrypoint = metadata.mapToEntrypoint(prefixing: benchmarkDirectory.path)
         let content: String = (try? String(contentsOf: entrypoint.fileURL, encoding: .utf8)) ?? ""
 
-        let relevantSymbols: [SymbolContent] = await retrieveRelevantSymbolsForFileContent(
+        let relevantSymbolsSummary: RelevantSymbolsSummary? = await retrieveRelevantSymbolsForFileContent(
             file: FileContent(fileURL: entrypoint.fileURL.path, content: content),
             workspace: workspace
         )
+        let relevantSymbols = relevantSymbolsSummary?.symbolsAtLevel.flatMap { $0 } ?? []
         let suggestionRequest = SuggestionProvider.SuggestionRequest(
             fileURL: entrypoint.fileURL,
             relativePath: entrypoint.fileURL.path.replacingOccurrences(of: benchmarkDirectory.path, with: ""),
@@ -122,7 +134,8 @@ class MultiFileContextBenchmarkManager: BenchmarkManager {
                 suggestion: firstSuggestion,
                 fileURL: entrypoint.fileURL,
                 relevantSymbolsFromRequest: relevantSymbols,
-                model: selectedGenAIModelSubject.value
+                model: selectedGenAIModelSubject.value,
+                relevantFileScanningDurationInSeconds: relevantSymbolsSummary?.durationInSeconds
             )
         } catch {
             return nil
@@ -196,10 +209,11 @@ class MultiFileContextBenchmarkManager: BenchmarkManager {
         let entrypoint = metadata.mapToEntrypoint(prefixing: benchmarkDirectory.path)
         let content: String = (try? String(contentsOf: entrypoint.fileURL, encoding: .utf8)) ?? ""
 
-        let relevantSymbols: [SymbolContent] = await retrieveRelevantSymbolsForFileContent(
+        let relevantSymbolsSummary: RelevantSymbolsSummary? = await retrieveRelevantSymbolsForFileContent(
             file: FileContent(fileURL: entrypoint.fileURL.path, content: content),
             workspace: workspace
         )
+        let relevantSymbols = relevantSymbolsSummary?.symbolsAtLevel.flatMap { $0 } ?? []
         let limitedRelevantSymbols = Array(relevantSymbols.prefix(10))
 
         let suggestionRequest = SuggestionRequest(
@@ -227,25 +241,32 @@ class MultiFileContextBenchmarkManager: BenchmarkManager {
                 ),
                 fileURL: entrypoint.fileURL,
                 relevantSymbolsFromRequest: relevantSymbols,
-                model: selectedGenAIModelSubject.value
+                model: selectedGenAIModelSubject.value,
+                relevantFileScanningDurationInSeconds: relevantSymbolsSummary?.durationInSeconds
             )
         } catch {
             print("CK \(error)")
             return nil
         }
     }
 
-    private func retrieveRelevantSymbolsForFileContent(file: FileContent, workspace: Workspace) async -> [SymbolContent] {
-        let multiFileContextManager = MultiFileContextManager(
+    private func retrieveRelevantSymbolsForFileContent(
+        file: FileContent,
+        workspace: Workspace
+    ) async -> RelevantSymbolsSummary? {
+        guard isMultiFileEnabledSubject.value else { return nil }
+
+        let manager = MultiFileContextManager(
             workspaceProvider: ManualWorkspaceProvider(workspace: workspace),
             parser: SwiftProgrammingLanguageSyntaxParser()
         )
-        if isMultiFileEnabledSubject.value {
-            let symbols = await multiFileContextManager.retrieveRelevantSymbolsForFileContent(file: file, ignoreWithinPaths: ["/Benchmark/"])
-            return Array(symbols.values)
-        } else {
-            return []
-        }
+
+        return await manager.expandRelevantSymbols(
+            from: file,
+            ignoreWithinPaths: ["/Benchmark/"],
+            maxDepth: contextLevelLimitSubject.value.indexLimit,   // e.g. First=0, Second=1, …, nil=no limit
+            maxFiles: contextFileAmountLimitSubject.value             // Int?, nil=no cap
+        )
     }
 
     private func applyCodeSuggestion(suggestion: SuggestionBasic.CodeSuggestion, at fileURL: URL) async {
@@ -458,6 +479,16 @@ class MultiFileContextBenchmarkManager: BenchmarkManager {
         currentStates[directory] = currentStatesInDirectory
         taskStatesSubject.send(currentStates)
     }
+
+    @MainActor
+    func saveContextLevelLimit(_ limit: ContextLevelLimit) {
+        contextLevelLimitSubject.send(limit)
+    }
+
+    @MainActor
+    func saveContextFileAmountLimit(_ limit: Int?) {
+        contextFileAmountLimitSubject.send(limit)
+    }
 }
 
 extension MetadataDTO {
@@ -534,7 +565,8 @@ extension SuggestionResponse {
             ),
             createdAt: timestamp,
             relevantSymbols: relevantSymbolsFromRequest.map { $0.toStoredDTO() },
-            model: model.id
+            model: model.id,
+            relevantFileScanningDurationInSeconds: relevantFileScanningDurationInSeconds
         )
     }
 }

Core/Sources/HostApp/Benchmark/Domain/Entity/ContextLevelLimit.swift

@@ -0,0 +1,21 @@
+enum ContextLevelLimit: CaseIterable {
+    case firstLevel
+    case secondLevel
+    case thirdLevel
+
+    var name: String {
+        switch self {
+        case .firstLevel: return "First Level"
+        case .secondLevel: return "Second Level"
+        case .thirdLevel: return "Third Level"
+        }
+    }
+
+    var indexLimit: Int? {
+        switch self {
+        case .firstLevel: return 0
+        case .secondLevel: return 1
+        case .thirdLevel: return 2
+        }
+    }
+}

Core/Sources/HostApp/Benchmark/Domain/Entity/SuggestionResponse.swift

@@ -7,4 +7,5 @@ struct SuggestionResponse {
     let fileURL: URL
     let relevantSymbolsFromRequest: [SymbolContent]
     let model: GenAILanguageModel
+    let relevantFileScanningDurationInSeconds: Double?
 }

Core/Sources/HostApp/Benchmark/Domain/Manager/BenchmarkManager.swift

@@ -10,4 +10,12 @@ protocol BenchmarkManager {
     func changeGenAIModel(to newModel: GenAILanguageModel)
     @MainActor
     func updateMultiFileContextState(_ newValue: Bool)
+
+    var contextLevelLimit: AnyPublisher<ContextLevelLimit, Never> { get }
+    @MainActor
+    func saveContextLevelLimit(_ limit: ContextLevelLimit)
+
+    var contextFileAmountLimit: AnyPublisher<Int?, Never> { get }
+    @MainActor
+    func saveContextFileAmountLimit(_ limit: Int?)
 }

Core/Sources/HostApp/Benchmark/Presentation/View/BenchmarkView.swift

@@ -28,6 +28,28 @@ struct BenchmarkView: View {
                         ConfigurationButtonView(module: module)
                     }
                     .padding(.vertical)
+
+                    HStack(spacing: 20) {
+                        Spacer()
+                        Picker("Context Level Limit: ", selection: $viewModel.contextLevelLimit) {
+                            ForEach(ContextLevelLimit.allCases, id: \.self) { limit in
+                                Text(limit.name).tag(limit)
+                            }
+                        }
+                        .fixedSize(horizontal: true, vertical: false)
+
+                        Text("Context File Amount Limit:")
+                        let numberFormatter: NumberFormatter = {
+                            let f = NumberFormatter()
+                            f.numberStyle = .none
+                            return f
+                        }()
+                        TextField("Optional", value: $viewModel.contextFileAmountLimit, formatter: numberFormatter)
+                            .textFieldStyle(RoundedBorderTextFieldStyle())
+                            .frame(width: 100)
+                    }
+                    .padding(.bottom)
+
                     ForEach(viewModel.benchmarkDirectories, id: \.self) { directory in
                         BenchmarkDirectoryEntryView(
                             directory: directory,

Core/Sources/HostApp/Benchmark/Presentation/ViewModel/BenchmarkViewModel.swift

@@ -16,6 +16,18 @@ class BenchmarkViewModel: ObservableObject {
             Task { await benchmarkManager.changeGenAIModel(to: selectedLanguageModel) }
         }
     }
+    @Published var contextLevelLimit: ContextLevelLimit = .firstLevel {
+        didSet {
+            guard oldValue != contextLevelLimit else { return }
+            Task { await benchmarkManager.saveContextLevelLimit(contextLevelLimit) }
+        }
+    }
+    @Published var contextFileAmountLimit: Int? = nil {
+        didSet {
+            guard oldValue != contextFileAmountLimit else { return }
+            Task { await benchmarkManager.saveContextFileAmountLimit(contextFileAmountLimit) }
+        }
+    }
 
     private let benchmarkSettingsRepository: BenchmarkSettingsRepository
     private let benchmarkManager: BenchmarkManager
@@ -35,6 +47,12 @@ class BenchmarkViewModel: ObservableObject {
         benchmarkManager.selectedGenAIModel
             .assign(to: \.selectedLanguageModel, on: self)
             .store(in: &cancellables)
+        benchmarkManager.contextLevelLimit
+            .assign(to: \.contextLevelLimit, on: self)
+            .store(in: &cancellables)
+        benchmarkManager.contextFileAmountLimit
+            .assign(to: \.contextFileAmountLimit, on: self)
+            .store(in: &cancellables)
     }
 
     func loadBenchmarkDirectories() {

Core/Sources/Service/MultiFileContext/Entity/RelevantSymbolsSummary.swift

@@ -0,0 +1,4 @@
+public struct RelevantSymbolsSummary {
+    public let symbolsAtLevel: [[SymbolContent]]
+    public let durationInSeconds: Double
+}

Core/Sources/Service/MultiFileContext/MultiFileContextManager.swift

@@ -62,6 +62,87 @@ public class MultiFileContextManager {
         return result
     }
 
+    private func readFile(_ urlString: String) -> FileContent? {
+        guard let url = URL(string: urlString) else { return nil }
+        guard let content = try? String(contentsOf: url, encoding: .utf8) else { return nil }
+        return FileContent(fileURL: urlString, content: content)
+    }
+
+    public func expandRelevantSymbols(
+        from root: FileContent,
+        ignoreWithinPaths: [String] = [],
+        maxDepth: Int? = nil,
+        maxFiles: Int? = nil
+    ) async -> RelevantSymbolsSummary {
+        @inline(__always)
+        func isIgnored(_ url: String) -> Bool {
+            ignoreWithinPaths.contains { url.contains($0) }
+        }
+
+        let start = Date()
+
+        let level0 = await retrieveRelevantSymbolsForFileContent(file: root, ignoreWithinPaths: ignoreWithinPaths)
+        var level0Unique: [SymbolContent] = []
+        var includedFiles = Set<String>()
+        for s in level0.values {
+            if includedFiles.insert(s.fileURL).inserted {
+                level0Unique.append(s)
+            }
+        }
+        var levels: [[SymbolContent]] = [level0Unique]
+
+        var visited = Set<String>([root.fileURL])
+        var depth = 0
+        var remainingFiles = maxFiles
+
+        while maxDepth == nil || depth < maxDepth! {
+            let prev = levels.last ?? []
+            var nextURLs = Set(prev.map(\.fileURL))
+            nextURLs.subtract(visited)
+            nextURLs = nextURLs.filter { !isIgnored($0) }
+
+            if nextURLs.isEmpty { break }
+
+            if let remaining = remainingFiles, nextURLs.count > remaining {
+                nextURLs = Set(nextURLs.prefix(remaining))
+            }
+
+            var rawNextLevel: [SymbolContent] = []
+            for url in nextURLs {
+                visited.insert(url)
+                guard let fc = readFile(url) else { continue }
+                let dict = await retrieveRelevantSymbolsForFileContent(file: fc, ignoreWithinPaths: ignoreWithinPaths)
+                rawNextLevel.append(contentsOf: dict.values)
+            }
+
+            var nextLevel: [SymbolContent] = []
+            var seenThisLevel = Set<String>()
+            for s in rawNextLevel {
+                guard !includedFiles.contains(s.fileURL) else { continue }
+                if seenThisLevel.insert(s.fileURL).inserted {
+                    nextLevel.append(s)
+                }
+            }
+
+            if nextLevel.isEmpty { break }
+
+            includedFiles.formUnion(nextLevel.map(\.fileURL))
+
+            levels.append(nextLevel)
+            depth += 1
+
+            if remainingFiles != nil {
+                remainingFiles! -= nextURLs.count
+                if remainingFiles! <= 0 { break }
+            }
+        }
+
+        return RelevantSymbolsSummary(
+            symbolsAtLevel: levels,
+            durationInSeconds: Date().timeIntervalSince(start)
+        )
+    }
+
     public func retrieveRelevantSymbolsForFileContent(
         file: FileContent,
         ignoreWithinPaths: [String] = []
@@ -171,7 +252,7 @@ extension String {
             let before = (range.lowerBound == startIndex) ? nil : self[index(before: range.lowerBound)]
             let after  = (range.upperBound == endIndex)   ? nil : self[range.upperBound]
 
-            let boundaryBefore = before.map { !isIdentChar($0) } ?? true
+            let boundaryBefore = before.map { !isIdentChar($0) && $0 != "." } ?? true
             let boundaryAfter  = after.map  { !isIdentChar($0) } ?? true
 
             if boundaryBefore && boundaryAfter {