1use pyo3::prelude::*;
40
41mod client;
43mod cost;
44mod drift;
45mod models;
46mod replay;
47mod runtime;
48mod sanitization;
49mod storage;
50mod tokens;
51
52use cost::*;
53use drift::*;
54use models::*;
55use replay::*;
56use runtime::{init_runtime, RuntimeConfig};
57use sanitization::*;
58use storage::*;
59
60#[pyfunction]
62fn init(py: Python) -> PyResult<()> {
63 let config = RuntimeConfig::default();
64 init_runtime(config)?;
65 let atexit = py.import("atexit")?;
66 let core_module = py.import("briefcase_ai._core")?;
67 let shutdown = core_module.getattr("_shutdown_runtime")?;
68 atexit.call_method1("register", (shutdown,))?;
69 Ok(())
70}
71
72#[pyfunction]
74#[pyo3(signature = (worker_threads = 2))]
75fn init_with_config(worker_threads: usize) -> PyResult<()> {
76 let config = RuntimeConfig { worker_threads };
77 init_runtime(config)
78}
79
80#[pyfunction]
82fn is_initialized() -> bool {
83 runtime::is_initialized()
84}
85
86#[pyfunction]
88fn _shutdown_runtime() -> PyResult<()> {
89 runtime::shutdown_runtime()
90}
91
92#[pymodule]
94fn _core(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
95 m.add_class::<PyScorecard>()?;
97 m.add_class::<PyExperimentMetadata>()?;
98 m.add_class::<PyAgentMetadata>()?;
99 m.add_class::<PyToolCallMetadata>()?;
100 m.add_class::<PyHardwareMetadata>()?;
101 m.add_class::<PyDataRef>()?;
102 m.add_class::<PyInput>()?;
103 m.add_class::<PyOutput>()?;
104 m.add_class::<PyModelParameters>()?;
105 m.add_class::<PyExecutionContext>()?;
106 m.add_class::<PyDecisionSnapshot>()?;
107 m.add_class::<PySnapshot>()?;
108
109 m.add_class::<PySqliteBackend>()?;
111 m.add_class::<PyLakeFSBackend>()?;
112 m.add_class::<PyBufferedBackend>()?;
113
114 m.add_class::<PyDriftCalculator>()?;
116 m.add_class::<PyCostCalculator>()?;
117 m.add_class::<PySanitizer>()?;
118 m.add_class::<PyReplayEngine>()?;
119
120 m.add_class::<PySnapshotQuery>()?;
122
123 m.add_class::<client::PyBriefcaseClient>()?;
125 m.add_class::<client::PyValidatedClient>()?;
126
127 m.add_function(wrap_pyfunction!(init, m)?)?;
129 m.add_function(wrap_pyfunction!(init_with_config, m)?)?;
130 m.add_function(wrap_pyfunction!(is_initialized, m)?)?;
131 m.add_function(wrap_pyfunction!(_shutdown_runtime, m)?)?;
132
133 m.add_function(wrap_pyfunction!(tokens::count_tokens, m)?)?;
135 m.add_function(wrap_pyfunction!(tokens::count_tokens_batch, m)?)?;
136
137 m.add("__version__", "2.5.1")?;
139 m.add("__author__", "Briefcase AI Team")?;
140
141 Ok(())
142}