Update package

PreternaturalAI · Feb 28, 2024 · e0b1365 · e0b1365
1 parent 9ee2157
commit e0b1365
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 1 deletion.
diff --git a/Sources/Lite/Intramodular/Stores/LTDocumentStore+Embeddings.swift b/Sources/Lite/Intramodular/Stores/LTDocumentStore+Embeddings.swift
@@ -5,6 +5,12 @@
 import Cataphyl
 
 extension LTDocumentStore {
+    public func embedAllDocumentsIfNeeded() async throws {
+        if textEmbeddings.isEmpty {
+            try await embedAllDocuments()
+        }
+    }
+
     public func embedAllDocuments() async throws {
         try await internalTasks.perform {
             try await _embedAllDocuments()

diff --git a/Sources/Lite/Intramodular/Stores/LTDocumentStore+Search.swift b/Sources/Lite/Intramodular/Stores/LTDocumentStore+Search.swift
@@ -18,7 +18,7 @@ extension LTDocumentStore {
         public var id: ID {
             fragmentID
         }
-
+        
         @MainActor
         public var fragment: LTDocument.RetrievedFragment {
             get {
@@ -38,6 +38,36 @@ extension LTDocumentStore {
     }
 }
 
+extension LTDocumentStore {
+    public func search(
+        query: String,
+        maximumNumberOfResults: Int? = nil
+    ) async throws -> [SearchResult] {
+        try await embedAllDocumentsIfNeeded()
+
+        let query: _RawTextEmbedding = try await withTaskTimeout(.seconds(2)) {
+            return try await configuration.lite.textEmbedding(for: query)
+        }
+
+        let result: [SearchResult] = try textEmbeddings
+            .query(
+                .topMatches(
+                    for: query,
+                    maximumNumberOfResults: maximumNumberOfResults ?? self.documents.count
+                )
+            )
+            .map { (item: VectorIndexSearchResult<NaiveVectorIndex<LTDocumentFragmentIdentifier>>) in
+                SearchResult(
+                    store: self,
+                    fragmentID: item.item,
+                    score: item.score
+                )
+            }
+
+        return result
+    }
+}
+
 extension LTDocumentStore {
     public func relevantMatches(
         for query: String