Merge pull request #41981 from BerriAI/litellm_rust_typed_pyo3

refactor(rust): use typed pyo3 APIs instead of getattr/import strings

litellm-rust/crates/callbacks-legacy/src/adapter.rs

@@ -12,7 +12,7 @@ use pyo3::{
     exceptions::{PyBaseException, PyException},
     gc::{PyTraverseError, PyVisit},
     prelude::*,
-    types::{PyDict, PyList},
+    types::{PyDateTime, PyDict, PyList},
 };
 use serde_json::Value;
 
@@ -73,10 +73,7 @@ pub struct LegacyLogging {
 }
 
 fn datetime(py: Python<'_>, epoch_seconds: f64) -> PyResult<Py<PyAny>> {
-    py.import("datetime")?
-        .getattr("datetime")?
-        .call_method1("fromtimestamp", (epoch_seconds,))
-        .map(Bound::unbind)
+    PyDateTime::from_timestamp(py, epoch_seconds, None).map(|value| value.into_any().unbind())
 }
 
 fn is_cancellation(py: Python<'_>, error: &PyErr) -> bool {

litellm-rust/crates/host-python/src/callable.rs

@@ -73,13 +73,7 @@ abort = KeyboardInterrupt('cancelled')
             let wrapped = wrap_failure(py, TEMPLATE, failure(&original)).unwrap_err();
             assert!(wrapped.is_instance_of::<PyRuntimeError>(py));
             assert!(wrapped.cause(py).unwrap().value(py).is(&original));
-            assert!(
-                wrapped
-                    .value(py)
-                    .getattr("__context__")
-                    .unwrap()
-                    .is(&original)
-            );
+            assert!(wrapped.context(py).unwrap().value(py).is(&original));
             assert_eq!(
                 wrapped.value(py).str().unwrap().to_str().unwrap(),
                 "Failed to reach the caller: unavailable"
@@ -115,13 +109,7 @@ original = Unformattable('cannot render')
             let original = raised(&locals, "original");
             let error = wrap_failure(py, TEMPLATE, failure(&original)).unwrap_err();
             assert!(error.is_instance_of::<pyo3::exceptions::PyValueError>(py));
-            assert!(
-                error
-                    .value(py)
-                    .getattr("__context__")
-                    .unwrap()
-                    .is(&original)
-            );
+            assert!(error.context(py).unwrap().value(py).is(&original));
         });
     }
 

litellm-rust/crates/host-python/src/driver.rs

@@ -445,14 +445,8 @@ where
             Ok(failure) => return failure.into(),
             Err(classifier_error) => classifier_error,
         };
-        let attached = classifier_error.value(py).setattr(
-            "__context__",
-            PyRuntimeError::new_err(native).into_value(py),
-        );
-        match attached {
-            Ok(()) => classifier_error,
-            Err(error) => error,
-        }
+        classifier_error.set_context(py, Some(PyRuntimeError::new_err(native)));
+        classifier_error
     }
 
     fn succeeded(&mut self, py: Python<'_>, response: Py<PyAny>) -> PyResult<ExecutionStep> {
@@ -1071,9 +1065,9 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri
             let error = result.unwrap_err();
             assert!(error.is_instance_of::<pyo3::exceptions::PyTypeError>(py));
             assert_eq!(error.value(py).to_string(), "classifier failed");
-            let context = error.value(py).getattr("__context__").unwrap();
-            assert!(context.is_instance_of::<PyRuntimeError>());
-            assert_eq!(context.str().unwrap().to_string(), "provider exploded");
+            let context = error.context(py).unwrap();
+            assert!(context.is_instance_of::<PyRuntimeError>(py));
+            assert_eq!(context.value(py).to_string(), "provider exploded");
             assert_eq!(
                 log,
                 [
@@ -1186,20 +1180,12 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri
                 type Failure = Classified;
                 fn invoke(
                     &mut self,
-                    py: Python<'_>,
+                    _: Python<'_>,
                     _: &Bound<'_, PyDict>,
                     _: &'static str,
                 ) -> Result<String, InvokeError<Error>> {
                     self.0.push("route");
-                    Err(PyErr::from_value(
-                        py.import("asyncio")
-                            .unwrap()
-                            .getattr("CancelledError")
-                            .unwrap()
-                            .call0()
-                            .unwrap(),
-                    )
-                    .into())
+                    Err(pyo3::exceptions::asyncio::CancelledError::new_err(()).into())
                 }
                 fn chunk(
                     &mut self,

litellm-rust/crates/python-bridge/src/routes/ocr/document.rs

@@ -7,7 +7,8 @@ use pyo3::{
     gc::{PyTraverseError, PyVisit},
     prelude::*,
     pybacked::PyBackedBytes,
-    types::{PyBytes, PyString},
+    sync::PyOnceLock,
+    types::{PyBytes, PyString, PyType},
 };
 
 #[derive(Debug)]
@@ -84,7 +85,8 @@ impl FromPyObject<'_, '_> for FileDocumentInput {
                 "OCR file input does not accept bare str values. Pass bytes, a pathlib.Path, or a file-like object.",
             ));
         }
-        if file.is_instance(&py.import("os")?.getattr("PathLike")?)? {
+        static PATH_LIKE: PyOnceLock<Py<PyType>> = PyOnceLock::new();
+        if file.is_instance(PATH_LIKE.import(py, "os", "PathLike")?)? {
             return Ok(Self {
                 input: OcrDocumentInput::Path {
                     path: file.extract::<PathBuf>()?,