Make runPython sync

Author: intitni

Core/Sources/ChatPlugins/MathChatPlugin/SolveMathProblem.swift

@@ -1,7 +1,7 @@
 import Foundation
 import LangChain
-import PythonKit
 import PythonHelper
+import PythonKit
 
 func solveMathProblem(_ problem: String) async throws -> String {
     #if DEBUG
@@ -16,16 +16,20 @@ func solveMathProblem(_ problem: String) async throws -> String {
         }
     }
 
-    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 }
+    let task = Task {
+        try 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,7 +1,7 @@
 import Foundation
 import Python
-import PythonKit
 import PythonHelper
+import PythonKit
 
 func initializePython() {
     guard let sitePackagePath = Bundle.main.path(forResource: "site-packages", ofType: nil),
@@ -14,24 +14,21 @@ func initializePython() {
 
     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()
-    
+    _ = PyEval_SaveThread()
+
     // Setup GIL state guard.
-    PythonHelper.PyGILState_Guard = { closure in
-        let gilState = PyGILState_Ensure()
-        try closure()
-        PyGILState_Release(gilState)
-    }
+    PythonHelper.gilStateEnsure = { PyGILState_Ensure() }
+    PythonHelper.gilStateRelease = { gilState in PyGILState_Release(gilState as! PyGILState_STATE) }
 
     Task {
         // All future task should run inside runPython.
-        try await runPython {
+        try runPython {
             let sys = Python.import("sys")
             print("Python Version: \(sys.version_info.major).\(sys.version_info.minor)")
         }

Tool/Sources/PythonHelper/RunPython.swift

@@ -1,40 +1,36 @@
 import Foundation
 import PythonKit
 
-public var PyGILState_Guard: ((() throws -> Void) throws -> Void)! = nil
+public var gilStateEnsure: (() -> Any)!
+public var gilStateRelease: ((Any) -> Void)!
+func gilStateGuard<T>(_ closure: @escaping () throws -> T) throws -> T {
+    let state = gilStateEnsure()
+    do {
+        let result = try closure()
+        gilStateRelease(state)
+        return result
+    } catch {
+        gilStateRelease(state)
+        throw error
+    }
+}
 
 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)
-                }
+) throws -> T {
+    if usePythonThread {
+        return try PythonThread.shared.runPythonAndWait {
+            return try gilStateGuard {
+                try closure()
             }
         }
+    } else {
+        return try gilStateGuard {
+            try closure()
+        }
     }
 }
 
@@ -66,5 +62,3 @@ public struct ReadablePythonError: Error, LocalizedError {
     }
 }
 
-
-