Node Search (#8)

Author: codefiesta

Package.swift

@@ -17,7 +17,7 @@ let package = Package(
         ),
     ],
     dependencies: [
-        .package(url: "https://github.com/codefiesta/VimKit", from: .init(0, 4, 8))
+        .package(url: "https://github.com/codefiesta/VimKit", from: .init(0, 4, 9))
     ],
     targets: [
         .target(

Sources/VimAssistant/Extensions/Array+Extensions.swift

@@ -20,3 +20,9 @@ extension Array where Element: Comparable {
         return indices
     }
 }
+
+extension Array where Element == String {
+    func containsIgnoringCase(_ element: Element) -> Bool {
+        contains { $0.caseInsensitiveCompare(element) == .orderedSame }
+    }
+}

Sources/VimAssistant/Model/VimAssistant+Handler.swift

@@ -111,4 +111,61 @@ public class VimAssistant: ObservableObject, @unchecked Sendable {
 
         return try? JSONDecoder().decode(VimPrediction.self, from: data)
     }
+
+    /// Handles the specified prediction.
+    /// - Parameters:
+    ///   - vim: the vim object to update based on the given prediction.
+    ///   - prediction: the prediction to handle
+    func handle(vim: Vim, prediction: VimPrediction?) async {
+        guard let prediction, let bestPrediction = prediction.bestPrediction else { return }
+        let action = bestPrediction.action
+        switch action {
+        case .isolate:
+            let ids = search(vim: vim, in: prediction)
+            guard ids.isNotEmpty else { return }
+            await vim.isolate(ids: ids)
+        case .hide:
+            let ids = search(vim: vim, in: prediction)
+            guard ids.isNotEmpty else { return }
+            await vim.hide(ids: ids)
+        case .quantify:
+            break
+        case .zoomIn:
+            await vim.zoom()
+        case .zoomOut:
+            await vim.zoom(out: true)
+        case .lookLeft:
+            await vim.look(.left)
+        case .lookRight:
+            await vim.look(.right)
+        case .lookUp:
+            await vim.look(.up)
+        case .lookDown:
+            await vim.look(.down)
+        case .panLeft:
+            await vim.pan(.left)
+        case .panRight:
+            await vim.pan(.right)
+        case .panUp:
+            await vim.pan(.up)
+        case .panDown:
+            await vim.pan(.down)
+        }
+    }
+
+    /// Performs a fuzzy search using the levenshtein distance algorithm across the node tree to find the best results.
+    /// - Parameters:
+    ///   - vim: the vim object to search
+    ///   - prediction: the prediction to extract entity names from
+    /// - Returns: a set of node ids that match the prediction
+    private func search(vim: Vim, in prediction: VimPrediction) -> Set<Int> {
+        guard let tree = vim.tree else { return [] }
+        var ids: Set<Int> = .init()
+        for entity in prediction.entities {
+            let searchResults = tree.search(entity.value)
+            guard let bestResult = searchResults.first else { continue }
+            ids.formUnion(bestResult.item.ids)
+        }
+        return ids
+    }
 }

Sources/VimAssistant/Model/VimAssistant.swift

@@ -15,8 +15,8 @@ public struct VimAssistantView: View {
     @State
     var assistant: VimAssistant = .init()
 
-    /// The handler to pass prediction information to.
-    var handler: VimAssistant.Handler = .init()
+//    /// The handler to pass prediction information to.
+//    var handler: VimAssistant.Handler = .init()
 
     @State
     var inputText: String = .empty
@@ -48,7 +48,10 @@ public struct VimAssistantView: View {
             predictionView
         }
         .onChange(of: assistant.prediction) { _, prediction in
-            handler.handle(vim: vim, prediction: prediction)
+            Task {
+                await assistant.handle(vim: vim, prediction: prediction)
+            }
+//            handler.handle(vim: vim, prediction: prediction)
         }
     }
 

Sources/VimAssistant/Views/VimAssistantView.swift

@@ -71,7 +71,7 @@ struct VimPredictionView: View {
         for entity in prediction.entities {
             let entityText = text[entity.range]
             var attributedEntityString = AttributedString(entityText)
-            attributedEntityString.foregroundColor = .cyan
+            attributedEntityString.foregroundColor = entity.label.color
             attributedEntityString.underlineStyle = .single
             attributedEntityString.link = URL(string: "/\(entity.label)/\(entity.value)")!
             result.replaceSubrange(bounds: entity.range, with: attributedEntityString)
@@ -95,15 +95,17 @@ struct VimPredictionView: View {
             Divider()
                 .fixedSize()
 
-            Text("Recognized entities:")
-                .font(.subheadline).bold()
+            if prediction.entities.isNotEmpty {
+                Text("Recognized entities:")
+                    .font(.subheadline).bold()
+            }
             ForEach(prediction.entities) { entity in
                 HStack {
                     Text(text[entity.range])
                         .bold()
                     Text(entity.label.rawValue)
                         .padding(1)
-                        .background(Color.cyan)
+                        .background(entity.label.color)
                         .foregroundStyle(Color.black)
                         .cornerRadius(2)
                 }
@@ -133,3 +135,51 @@ struct VimPredictionView: View {
     let prediction = try! JSONDecoder().decode(VimPrediction.self, from: json.data(using: .utf8)!)
     VimPredictionView(prediction: prediction, explain: .constant(true))
 }
+
+extension VimPrediction.NerLabel {
+
+    var color: Color {
+        switch self {
+        case .person:
+            .purple
+        case .organization:
+            .blue
+        case .location:
+            .black
+        case .date:
+            .mint
+        case .time:
+            .cyan
+        case .event:
+            .cyan
+        case .workOfArt:
+            .cyan
+        case .fac:
+            .cyan
+        case .gpe:
+            .cyan
+        case .language:
+            .cyan
+        case .law:
+            .cyan
+        case .norp:
+            .cyan
+        case .product:
+            .cyan
+        case .cardinal:
+            .blue
+        case .bimCategory:
+            .cyan
+        case .bimFamily:
+            .orange
+        case .bimType:
+            .cyan
+        case .bimInstance:
+            .cyan
+        case .bimLevel:
+            .yellow
+        case .bimView:
+            .cyan
+        }
+    }
+}