update py03 for 3.13 support

postgresml · montanalow · Jan 16, 2025 · Jan 11, 2025 · Jan 11, 2025 · Jan 11, 2025
commit fb91af4ea29dcaeb647a1388fc06927e8bdd83f0
diff --git a/pgml-extension/Cargo.lock b/pgml-extension/Cargo.lock
diff --git a/pgml-extension/Cargo.toml b/pgml-extension/Cargo.toml
@@ -45,7 +45,7 @@ ndarray-stats = "0.5"
 parking_lot = "0.12"
 pgrx = "=0.12.9"
 pgrx-pg-sys = "=0.12.9"
-pyo3 = { version = "0.20", features = ["anyhow", "auto-initialize"], optional = true }
+pyo3 = { version = "0.23", features = ["anyhow", "auto-initialize"], optional = true }
 rand = "0.8"
 rmp-serde = { version = "1.3" }
 signal-hook = "0.3"

diff --git a/pgml-extension/src/api.rs b/pgml-extension/src/api.rs
@@ -759,10 +759,11 @@ pub fn transform_stream_json(
    input: default!(&str, "''"),
    cache: default!(bool, false),
 ) -> SetOfIterator<'static, JsonB> {
-    // We can unwrap this becuase if there is an error the current transaction is aborted in the map_err call
+    // We can unwrap this because if there is an error the current transaction is aborted in the map_err call
    let python_iter = crate::bindings::transformers::transform_stream_iterator(&task.0, &args.0, input)
        .map_err(|e| error!("{e}"))
        .unwrap();
+
    SetOfIterator::new(python_iter)
 }

@@ -776,7 +777,7 @@ pub fn transform_stream_string(
    cache: default!(bool, false),
 ) -> SetOfIterator<'static, JsonB> {
    let task_json = json!({ "task": task });
-    // We can unwrap this becuase if there is an error the current transaction is aborted in the map_err call
+    // We can unwrap this because if there is an error the current transaction is aborted in the map_err call
    let python_iter = crate::bindings::transformers::transform_stream_iterator(&task_json, &args.0, input)
        .map_err(|e| error!("{e}"))
        .unwrap();
@@ -795,7 +796,7 @@ pub fn transform_stream_conversational_json(
    if !task.0["task"].as_str().is_some_and(|v| v == "conversational") {
        error!("ARRAY[]::JSONB inputs for transform_stream should only be used with a conversational task");
    }
-    // We can unwrap this becuase if there is an error the current transaction is aborted in the map_err call
+    // We can unwrap this because if there is an error the current transaction is aborted in the map_err call
    let python_iter = crate::bindings::transformers::transform_stream_iterator(&task.0, &args.0, inputs)
        .map_err(|e| error!("{e}"))
        .unwrap();
@@ -815,7 +816,7 @@ pub fn transform_stream_conversational_string(
        error!("ARRAY::JSONB inputs for transform_stream should only be used with a conversational task");
    }
    let task_json = json!({ "task": task });
-    // We can unwrap this becuase if there is an error the current transaction is aborted in the map_err call
+    // We can unwrap this because if there is an error the current transaction is aborted in the map_err call
    let python_iter = crate::bindings::transformers::transform_stream_iterator(&task_json, &args.0, inputs)
        .map_err(|e| error!("{e}"))
        .unwrap();

diff --git a/pgml-extension/src/bindings/langchain/mod.rs b/pgml-extension/src/bindings/langchain/mod.rs
@@ -1,22 +1,22 @@
 use anyhow::Result;
 use pyo3::prelude::*;
-use pyo3::types::PyTuple;
+use pyo3::ffi::c_str;
+use pyo3::types::PyString;

 use crate::create_pymodule;

 create_pymodule!("/src/bindings/langchain/langchain.py");

 pub fn chunk(splitter: &str, text: &str, kwargs: &serde_json::Value) -> Result<Vec<String>> {
-    let kwargs = serde_json::to_string(kwargs).unwrap();

    Python::with_gil(|py| -> Result<Vec<String>> {
        let chunk: Py<PyAny> = get_module!(PY_MODULE).getattr(py, "chunk")?;
+        let splitter = PyString::new(py, splitter);
+        let text = PyString::new(py, text);
+        let kwargs = PyString::new(py, serde_json::to_string(kwargs)?.as_str());

        Ok(chunk
-            .call1(
-                py,
-                PyTuple::new(py, &[splitter.into_py(py), text.into_py(py), kwargs.into_py(py)]),
-            )?
+            .call1(py,(splitter, text, kwargs))?
            .extract(py)?)
    })
 }
diff --git a/pgml-extension/src/bindings/mod.rs b/pgml-extension/src/bindings/mod.rs
@@ -3,7 +3,7 @@ use std::fmt::Debug;
 use anyhow::{anyhow, Result};
 #[allow(unused_imports)] // used for test macros
 use pgrx::*;
-use pyo3::{pyfunction, PyResult, Python};
+use pyo3::{pyfunction, PyResult, Python, prelude::PyTracebackMethods};

 use crate::orm::*;

@@ -42,11 +42,11 @@ macro_rules! create_pymodule {
            once_cell::sync::Lazy::new(|| {
                pyo3::Python::with_gil(|py| -> anyhow::Result<pyo3::Py<pyo3::types::PyModule>> {
                    use $crate::bindings::TracebackError;
-                    let src = include_str!(concat!(env!("CARGO_MANIFEST_DIR"), $pyfile));
-                    let module = pyo3::types::PyModule::from_code(py, src, "transformers.py", "__main__")
+                    let src = c_str!(include_str!(concat!(env!("CARGO_MANIFEST_DIR"), $pyfile)));
+                    let module = pyo3::types::PyModule::from_code(py, src, c_str!("transformers.py"), c_str!("__main__"))
                        .format_traceback(py)?;
-                    module.add_function(wrap_pyfunction!($crate::bindings::r_insert_logs, module)?)?;
-                    module.add_function(wrap_pyfunction!($crate::bindings::r_log, module)?)?;
+                    module.add_function(wrap_pyfunction!($crate::bindings::r_insert_logs, &module)?)?;
+                    module.add_function(wrap_pyfunction!($crate::bindings::r_log, &module)?)?;
                    Ok(module.into())
                })
            });

diff --git a/pgml-extension/src/bindings/python/mod.rs b/pgml-extension/src/bindings/python/mod.rs
@@ -4,7 +4,8 @@ use anyhow::Result;
 use pgrx::iter::TableIterator;
 use pgrx::*;
 use pyo3::prelude::*;
-use pyo3::types::PyTuple;
+use pyo3::types::PyString;
+use pyo3::ffi::c_str;

 use crate::config::PGML_VENV;
 use crate::create_pymodule;
@@ -13,8 +14,8 @@ create_pymodule!("/src/bindings/python/python.py");

 pub fn activate_venv(venv: &str) -> Result<bool> {
    Python::with_gil(|py| {
-        let activate_venv: Py<PyAny> = get_module!(PY_MODULE).getattr(py, "activate_venv")?;
-        let result: Py<PyAny> = activate_venv.call1(py, PyTuple::new(py, &[venv.to_string().into_py(py)]))?;
+        let activate_venv = get_module!(PY_MODULE).getattr(py, "activate_venv")?;
+        let result = activate_venv.call1(py, (PyString::new(py, venv),))?;

        Ok(result.extract(py)?)
    })