1use pyo3::prelude::*;
40use std::ffi::CString;
41
42mod cost;
44mod drift;
45mod models;
46mod replay;
47mod runtime;
48mod sanitization;
49mod storage;
50
51use cost::*;
52use drift::*;
53use models::*;
54use replay::*;
55use runtime::{init_runtime, RuntimeConfig};
56use sanitization::*;
57use storage::*;
58
59#[pyfunction]
61fn init(py: Python) -> PyResult<()> {
62 let config = RuntimeConfig::default();
63 init_runtime(config)?;
64 let import_code = CString::new("import atexit")?;
65 let register_code = CString::new("atexit.register(lambda: briefcase_ai._shutdown_runtime())")?;
66 py.run(&import_code, None, None)?;
67 py.run(®ister_code, None, None)?;
68 Ok(())
69}
70
71#[pyfunction]
73#[pyo3(signature = (worker_threads = 2))]
74fn init_with_config(worker_threads: usize) -> PyResult<()> {
75 let config = RuntimeConfig { worker_threads };
76 init_runtime(config)
77}
78
79#[pyfunction]
81fn is_initialized() -> bool {
82 runtime::is_initialized()
83}
84
85#[pyfunction]
87fn _shutdown_runtime() -> PyResult<()> {
88 runtime::shutdown_runtime()
89}
90
91#[pymodule]
93fn briefcase_ai(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
94 m.add_class::<PyInput>()?;
96 m.add_class::<PyOutput>()?;
97 m.add_class::<PyModelParameters>()?;
98 m.add_class::<PyExecutionContext>()?;
99 m.add_class::<PyDecisionSnapshot>()?;
100 m.add_class::<PySnapshot>()?;
101
102 m.add_class::<PySqliteBackend>()?;
104 m.add_class::<PyLakeFSBackend>()?;
105
106 m.add_class::<PyDriftCalculator>()?;
108 m.add_class::<PyCostCalculator>()?;
109 m.add_class::<PySanitizer>()?;
110 m.add_class::<PyReplayEngine>()?;
111
112 m.add_class::<PySnapshotQuery>()?;
114
115 m.add_function(wrap_pyfunction!(init, m)?)?;
117 m.add_function(wrap_pyfunction!(init_with_config, m)?)?;
118 m.add_function(wrap_pyfunction!(is_initialized, m)?)?;
119 m.add_function(wrap_pyfunction!(_shutdown_runtime, m)?)?;
120
121 m.add("__version__", "2.0.8")?;
123 m.add("__author__", "Briefcase AI Team")?;
124
125 Ok(())
126}