Feature: Retrieval functionality to attach document to user prompt

Enchanted/Application/EnchantedApp.swift

@@ -18,6 +18,7 @@ extension KeyboardShortcuts.Name {
 @main
 struct EnchantedApp: App {
     @State private var appStore = AppStore.shared
+    @State private var retrievalStore = RetrievalStore.shared
 #if os(macOS)
     @NSApplicationDelegateAdaptor(PanelManager.self) var panelManager
 #endif
@@ -30,10 +31,20 @@ struct EnchantedApp: App {
                     print("heya")
                     panelManager.togglePanel()
                 }
+#endif
                 .onAppear {
+#if os(macOS)
                     NSWindow.allowsAutomaticWindowTabbing = false
-                }
 #endif
+                    Task{
+                        do {
+                            try await retrievalStore.fetchDatabases()
+                            retrievalStore.selectDatabase(database: nil)
+                        } catch {
+                            print("Error fetching databases: \(error)")
+                        }
+                    }
+                }
         }
 #if os(macOS)
         .commands {

Enchanted/Extensions/URL+Extension.swift

@@ -0,0 +1,21 @@
+//
+//  URL+Extension.swift
+//  Enchanted
+//
+//  Created by Daniel on 3/19/25.
+//
+
+import Foundation
+import Collections
+import UniformTypeIdentifiers
+
+extension URL {
+    public func mimeType() -> String {
+        if let mimeType = UTType(filenameExtension: self.pathExtension)?.preferredMIMEType {
+            return mimeType
+        }
+        else {
+            return "application/octet-stream"
+        }
+    }
+}

Enchanted/Models/DocumentStatus.swift

@@ -0,0 +1,37 @@
+//
+//  DocumentStatus.swift
+//  Enchanted
+//
+//  Created by Daniel on 3/19/25.
+//
+
+import Foundation
+import SwiftUI
+
+enum DocumentIndexStatus: String, Codable {
+    case indexing
+    case completed
+    case notStarted
+    case failed
+
+    var humanReadable: String {
+        switch self {
+        case .indexing: "Indexing"
+        case .completed: "Completed"
+        case .notStarted: "Not Started"
+        case .failed: "Failed"
+        }
+    }
+}
+
+// MARK: - Icon extension
+extension DocumentIndexStatus {
+    var icon: some View {
+        switch self {
+        case .indexing: return Image(systemName: "hourglass.circle").foregroundColor(.indigo)
+        case .completed: return Image(systemName: "checkmark.circle").foregroundColor(.green)
+        case .notStarted: return Image(systemName: "hourglass.circle").foregroundColor(.yellow)
+        case .failed: return Image(systemName: "circle.slash").foregroundColor(.red)
+        }
+    }
+}

Enchanted/Retrieval/DataLoader.swift

@@ -0,0 +1,63 @@
+//
+//  DataLoader.swift
+//  Enchanted
+//
+//  Created by Daniel on 3/19/25.
+//
+
+import Foundation
+import PDFKit
+
+struct DataLoader {
+    func fromPDF(_ url: URL) -> String? {
+        guard url.startAccessingSecurityScopedResource() else {
+            print("could not get access to file \(url.absoluteString)")
+            return nil
+        }
+
+        let pdfDocument = PDFDocument(url: url)
+
+        guard let document = pdfDocument else {
+            print("Could not load \(url.absoluteString)")
+            return nil
+        }
+
+        var contents = ""
+        for i in 0..<document.pageCount {
+            if let page = document.page(at: i) {
+                if let pageContent = page.string {
+                    contents += pageContent
+                }
+            }
+        }
+        return contents
+    }
+
+    func fromTextFile(_ url: URL) -> String? {
+        do {
+            let contents = try String(contentsOf: url, encoding: .utf8)
+            return contents
+        } catch {
+            print("Could not parse \(url.absoluteString). Error: \(error.localizedDescription)")
+        }
+
+        return nil
+    }
+
+    func from(_ url: URL) -> String? {
+        let mimeType = url.mimeType()
+
+        if mimeType.contains("text/") {
+            return fromTextFile(url)
+        } else if mimeType == "application/pdf" {
+            return fromPDF(url)
+        } else if mimeType == "application/octet-stream" {
+            return fromTextFile(url)
+        } else if mimeType == "application/x-yaml" {
+            return fromTextFile(url)
+        }
+
+        print("unhandled mime type \(mimeType) for \(url.absoluteString)")
+        return nil
+    }
+}

Enchanted/Retrieval/MessageEmbedding.swift

@@ -0,0 +1,48 @@
+//
+//  MessageEmbedding.swift
+//  Enchanted
+//
+//  Created by Now or Never on 3/19/25.
+//
+
+import Foundation
+
+protocol MessageEmbeddingProtocol {
+    func generateContext(_ prompt: String, _ model: LanguageModelSD) async -> String
+}
+
+struct MessageEmbedding: MessageEmbeddingProtocol {
+    private let languageModelStore: LanguageModelStore
+    private let retrievalStore: RetrievalStore
+    private let svdbService: SVDBService
+
+    init(
+        languageModelStore: LanguageModelStore = .shared,
+        retrievalStore: RetrievalStore = .shared,
+        svdbService: SVDBService = .shared
+    ) {
+        self.languageModelStore = languageModelStore
+        self.retrievalStore = retrievalStore
+        self.svdbService = svdbService
+    }
+
+    func generateContext(_ prompt: String, _ model: LanguageModelSD) async -> String {
+        // Obtain the embedding vector for the given prompt
+        guard let promptEmbedding = await languageModelStore.getEmbedding(model: model, prompt: prompt) else {
+            print("[MessageEmbedding] Error: Unable to generate embedding for prompt: \(prompt)")
+            return ""
+        }
+
+        print("[MessageEmbedding] Embedding vector generated: \(promptEmbedding)")
+
+        guard let databaseId = retrievalStore.selectedDatabase?.id else {
+            print("Database ID not found.")
+            return ""
+        }
+
+        // Perform similarity search to retrieve relevant context
+        let retrievedContext = await svdbService.search(query: promptEmbedding, databaseId: databaseId)
+
+        return retrievedContext
+    }
+}

Enchanted/Retrieval/TextSplitter.swift

@@ -0,0 +1,68 @@
+//
+//  TextSplitter.swift
+//  Enchanted
+//
+//  Created by Daniel on 3/19/25.
+//
+
+import Foundation
+
+protocol TextBatchSplitter {
+    func split(text: String, maxChunkSize: Int) -> [String]
+}
+
+struct SimpleTextSplitter: TextBatchSplitter {
+    func split(text: String, maxChunkSize: Int) -> [String] {
+        guard !text.isEmpty, maxChunkSize > 0 else { return [] }
+
+        let separators = ["\n\n", "\n", ".", " ", ""]
+
+        // Attempt to split using separators
+        for separator in separators {
+            let components = separator.isEmpty ? text.map { String($0) } : text.components(separatedBy: separator)
+
+            // If any single component is larger than maxChunkSize, skip this separator
+            if components.contains(where: { $0.count > maxChunkSize }) {
+                continue
+            }
+
+            var chunks: [String] = []
+            var currentChunk = ""
+
+            for component in components {
+                let separatorLength = currentChunk.isEmpty ? 0 : separator.count
+                if currentChunk.count + separatorLength + component.count <= maxChunkSize {
+                    currentChunk += (currentChunk.isEmpty ? "" : separator) + component
+                } else {
+                    if !currentChunk.isEmpty {
+                        chunks.append(currentChunk)
+                    }
+                    currentChunk = component
+                }
+            }
+
+            if !currentChunk.isEmpty {
+                chunks.append(currentChunk)
+            }
+
+            return chunks
+        }
+
+        // Fallback: split text into fixed-size chunks (by characters)
+        return fixedSizeCharacterChunks(text: text, maxChunkSize: maxChunkSize)
+    }
+
+    private func fixedSizeCharacterChunks(text: String, maxChunkSize: Int) -> [String] {
+        var chunks: [String] = []
+        var currentIndex = text.startIndex
+
+        while currentIndex < text.endIndex {
+            let endIndex = text.index(currentIndex, offsetBy: maxChunkSize, limitedBy: text.endIndex) ?? text.endIndex
+            let chunk = String(text[currentIndex..<endIndex])
+            chunks.append(chunk)
+            currentIndex = endIndex
+        }
+
+        return chunks
+    }
+}

Enchanted/Services/OllamaService.swift

@@ -50,3 +50,11 @@ class OllamaService: @unchecked Sendable {
         return await ollamaKit.reachable()
     }
 }
+
+// MARK: - Vector embeddings
+extension OllamaService {
+    func getEmbedding(prompt: String, model: LanguageModelSD) async -> [Double]? {
+        let res = try? await ollamaKit.generateEmbeddings(data: .init(model: model.name, prompt: prompt))
+        return res?.embedding
+    }
+}

Enchanted/Services/SVDBService.swift

@@ -0,0 +1,81 @@
+//
+//  SVDBService.swift
+//  Enchanted
+//
+//  Created by Now or Never on 3/24/25.
+//
+
+import SVDB
+import Foundation
+
+final actor SVDBService {
+    var svdb: SVDB
+    var collection: Collection?
+
+    static let shared = SVDBService()
+
+    private init() {
+        self.svdb = SVDB.shared
+    }
+
+    func createCollection(databaseId: UUID) async throws -> Collection {
+        do {
+            let newCollection = try svdb.collection("\(databaseId)")
+            self.collection = newCollection
+            return newCollection
+        } catch let error as SVDBError {
+            if(error == .collectionAlreadyExists) {
+                if let collection = self.collection {
+                    return collection
+                }
+            }
+            throw error
+        }  catch {
+            throw error
+        }
+    }
+
+    func releaseCollection(_ collectionName: String) async throws -> Void {
+        svdb.releaseCollection(collectionName)
+    }
+
+    func addDocument(text: String, embedding: [Double]) throws -> Void {
+        guard let collection = collection else {
+            throw NSError(
+                domain: "SVDBService",
+                code: 404,
+                userInfo: [NSLocalizedDescriptionKey: "Collection not initialized"]
+            )
+        }
+
+        collection.addDocument(text: text, embedding: embedding)
+    }
+
+    func search(query: [Double], databaseId: UUID) -> String {
+        do {
+            let collection = try svdb.collection(databaseId.uuidString)
+            return searchInCollection(collection, with: query)
+        } catch let error as SVDBError {
+            if(error == .collectionAlreadyExists) {
+                guard let collection = collection else {
+                    print("Collection not initialized")
+                    return ""
+                }
+                return searchInCollection(collection, with: query)
+            }
+        } catch {
+            print("Failed to load or search collection:", error)
+        }
+        return ""
+    }
+
+    private func searchInCollection(_ collection: Collection, with query: [Double]) -> String {
+        if let firstResult = collection.search(query: query, num_results: 1).first {
+            print("Result found: \(firstResult.text)")
+            return firstResult.text
+        } else {
+            print("No result found.")
+            return ""
+        }
+    }
+}