Fix importing large packages like langchain

Author: intitni

Core/Sources/ChatPlugins/MathChatPlugin/SolveMathProblem.swift

@@ -1,6 +1,7 @@
 import Foundation
 import LangChain
 import PythonKit
+import PythonHelper
 
 func solveMathProblem(_ problem: String) async throws -> String {
     #if DEBUG
@@ -15,20 +16,16 @@ func solveMathProblem(_ problem: String) async throws -> String {
         }
     }
 
-    let task = Task {
-        try withReadableThrowingPython {
-            let langchain = try Python.attemptImport("langchain")
-            let LLMMathChain = langchain.LLMMathChain
-            let llm = try LangChainChatModel.DynamicChatOpenAI(temperature: 0)
-            let llmMath = LLMMathChain.from_llm(llm, verbose: verbose)
-            let result = try llmMath.run.throwing.dynamicallyCall(withArguments: problem)
-            let answer = String(result)
-            if let answer { return answer }
+    return try await runPython {
+        let langchain = try Python.attemptImportOnPythonThread("langchain")
+        let LLMMathChain = langchain.LLMMathChain
+        let llm = try LangChainChatModel.DynamicChatOpenAI(temperature: 0)
+        let llmMath = LLMMathChain.from_llm(llm, verbose: verbose)
+        let result = try llmMath.run.throwing.dynamicallyCall(withArguments: problem)
+        let answer = String(result)
+        if let answer { return answer }
 
-            throw E()
-        }
+        throw E()
     }
-
-    return try await task.value
 }
 

ExtensionService/InitializePython.swift

@@ -1,6 +1,7 @@
 import Foundation
 import Python
 import PythonKit
+import PythonHelper
 
 func initializePython() {
     guard let sitePackagePath = Bundle.main.path(forResource: "site-packages", ofType: nil),
@@ -10,8 +11,32 @@ func initializePython() {
               ofType: nil
           )
     else { return }
+
     setenv("PYTHONHOME", stdLibPath, 1)
     setenv("PYTHONPATH", "\(stdLibPath):\(libDynloadPath):\(sitePackagePath)", 1)
+    
+    // Initialize python
     Py_Initialize()
+    
+    // Immediately release the thread, so that we can ensure the GIL state later.
+    // We may not recover the thread because all future tasks will be done in the Python Thread.
+    let _ = PyEval_SaveThread()
+    
+    // Setup GIL state guard.
+    PythonHelper.PyGILState_Guard = { closure in
+        let gilState = PyGILState_Ensure()
+        try closure()
+        PyGILState_Release(gilState)
+    }
+
+    Task {
+        // All future task should run inside runPython.
+        try await runPython {
+            let sys = Python.import("sys")
+            print("Python Version: \(sys.version_info.major).\(sys.version_info.minor)")
+        }
+    }
 }
 
+let queue = DispatchQueue(label: "")
+

Tool/Package.swift

@@ -8,7 +8,7 @@ let package = Package(
     platforms: [.macOS(.v12)],
     products: [
         .library(name: "Terminal", targets: ["Terminal"]),
-        .library(name: "LangChain", targets: ["LangChain"]),
+        .library(name: "LangChain", targets: ["LangChain", "PythonHelper"]),
         .library(name: "Preferences", targets: ["Preferences", "Configs"]),
     ],
     dependencies: [
@@ -27,6 +27,14 @@ let package = Package(
 
         .target(
             name: "LangChain",
+            dependencies: [
+                "PythonHelper",
+                .product(name: "PythonKit", package: "PythonKit"),
+            ]
+        ),
+
+        .target(
+            name: "PythonHelper",
             dependencies: [
                 .product(name: "PythonKit", package: "PythonKit"),
             ]

Tool/Sources/LangChain/LangChainChatModel.swift

@@ -1,5 +1,6 @@
 import Foundation
 import Preferences
+import PythonHelper
 import PythonKit
 
 public enum LangChainChatModel {
@@ -11,32 +12,28 @@ public enum LangChainChatModel {
             let model = UserDefaults.shared.value(for: \.chatGPTModel)
             let apiBaseURL = UserDefaults.shared.value(for: \.openAIBaseURL)
             let apiKey = UserDefaults.shared.value(for: \.openAIAPIKey)
-            return try withReadableThrowingPython {
-                let chatModels = try Python.attemptImport("langchain.chat_models")
-                let ChatOpenAI = chatModels.ChatOpenAI
-                return ChatOpenAI(
-                    temperature: temperature,
-                    model: model,
-                    openai_api_base: "\(apiBaseURL)/v1",
-                    openai_api_key: apiKey
-                )
-            }
+            let chatModels = try Python.attemptImportOnPythonThread("langchain.chat_models")
+            let ChatOpenAI = chatModels.ChatOpenAI
+            return ChatOpenAI(
+                temperature: temperature,
+                model: model,
+                openai_api_base: "\(apiBaseURL)/v1",
+                openai_api_key: apiKey
+            )
         case .azureOpenAI:
             let apiBaseURL = UserDefaults.shared.value(for: \.azureOpenAIBaseURL)
             let apiKey = UserDefaults.shared.value(for: \.azureOpenAIAPIKey)
             let deployment = UserDefaults.shared.value(for: \.azureChatGPTDeployment)
-            return try withReadableThrowingPython {
-                let chatModels = try Python.attemptImport("langchain.chat_models")
-                let ChatOpenAI = chatModels.AzureChatOpenAI
-                return ChatOpenAI(
-                    temperature: temperature,
-                    openai_api_type: "azure",
-                    openai_api_version: "2023-03-15-preview",
-                    deployment_name: deployment,
-                    openai_api_base: apiBaseURL,
-                    openai_api_key: apiKey
-                )
-            }
+            let chatModels = try Python.attemptImportOnPythonThread("langchain.chat_models")
+            let ChatOpenAI = chatModels.AzureChatOpenAI
+            return ChatOpenAI(
+                temperature: temperature,
+                openai_api_type: "azure",
+                openai_api_version: "2023-03-15-preview",
+                deployment_name: deployment,
+                openai_api_base: apiBaseURL,
+                openai_api_key: apiKey
+            )
         }
     }
 }

Tool/Sources/LangChain/ReadablePythonError.swift

@@ -0,0 +1,90 @@
+import Foundation
+import PythonKit
+
+final class PythonThread: Thread {
+    static let shared = {
+        let thread = PythonThread(
+            target: PythonThread.self,
+            selector: #selector(PythonThread.setup),
+            object: nil
+        )
+        thread.name = "Python Thread"
+        thread.stackSize = 1_048_576 // so that langchain can be correctly imported.
+        return thread
+    }()
+
+    @objc static func setup() {
+        CFRunLoopRun()
+    }
+
+    @objc static func runPythonJob(_ job: PythonJob) {
+        job.run()
+    }
+
+    func runPython(_ closure: @escaping () -> Void) {
+        if !isExecuting {
+            start()
+        }
+
+        if Thread.current.isEqual(self) {
+            closure()
+        } else {
+            PythonThread.perform(
+                #selector(PythonThread.runPythonJob),
+                on: self,
+                with: PythonJob(closure: closure),
+                waitUntilDone: false
+            )
+        }
+    }
+
+    func runPythonAndWait<T>(_ closure: @escaping () throws -> T) throws -> T {
+        if !isExecuting {
+            start()
+        }
+
+        if Thread.current.isEqual(self) {
+            return try closure()
+        } else {
+            let job = PythonJob(closure: closure)
+            PythonThread.perform(
+                #selector(PythonThread.runPythonJob),
+                on: self,
+                with: job,
+                waitUntilDone: true
+            )
+            guard let result = job.result else {
+                throw FailedToGetPythonJobResultError()
+            }
+            switch result {
+            case let .success(value):
+                if let value = value as? T {
+                    return value
+                } else {
+                    throw FailedToGetPythonJobResultError()
+                }
+            case let .failure(error):
+                throw error
+            }
+        }
+    }
+}
+
+struct FailedToGetPythonJobResultError: Error, LocalizedError {
+    var errorDescription: String? {
+        "Failed to get PythonJob result."
+    }
+}
+
+final class PythonJob: NSObject {
+    let closure: () throws -> Any
+    var result: Result<Any, Error>?
+    init(closure: @escaping () throws -> Any) {
+        self.closure = closure
+    }
+
+    func run() {
+        result = Result(catching: closure)
+    }
+}
+

Tool/Sources/PythonHelper/PythonThread.swift

@@ -0,0 +1,70 @@
+import Foundation
+import PythonKit
+
+public var PyGILState_Guard: ((() throws -> Void) throws -> Void)! = nil
+
+let pythonQueue = DispatchQueue(label: "Python Queue")
+
+public func runPython<T>(
+    usePythonThread: Bool = false,
+    _ closure: @escaping () throws -> T
+) async throws -> T {
+    return try await withUnsafeThrowingContinuation { con in
+        if usePythonThread {
+            PythonThread.shared.runPython {
+                do {
+                    try PyGILState_Guard {
+                        con.resume(returning: try closure())
+                    }
+                } catch let error as PythonError {
+                    con.resume(throwing: ReadablePythonError(error))
+                } catch {
+                    con.resume(throwing: error)
+                }
+            }
+        } else {
+            pythonQueue.async {
+                do {
+                    try PyGILState_Guard {
+                        con.resume(returning: try closure())
+                    }
+                } catch let error as PythonError {
+                    con.resume(throwing: ReadablePythonError(error))
+                } catch {
+                    con.resume(throwing: error)
+                }
+            }
+        }
+    }
+}
+
+public extension PythonInterface {
+    func attemptImportOnPythonThread(_ name: String) throws -> PythonObject {
+        try PythonThread.shared.runPythonAndWait {
+            let module = try Python.attemptImport(name)
+            return module
+        }
+    }
+}
+
+public struct ReadablePythonError: Error, LocalizedError {
+    public var error: PythonError
+
+    public init(_ error: PythonError) {
+        self.error = error
+    }
+
+    public var errorDescription: String? {
+        switch error {
+        case let .exception(object, _):
+            return "\(object)"
+        case let .invalidCall(object):
+            return "Invalid call: \(object)"
+        case let .invalidModule(module):
+            return "Invalid module: \(module)"
+        }
+    }
+}
+
+
+