Skip to main content

briefcase_ai/
lib.rs

1//! # Briefcase AI Python Bindings
2//!
3//! Python FFI bindings for the Briefcase AI core library, providing high-performance
4//! AI observability, replay, and decision tracking for Python applications.
5//!
6//! ## Features
7//!
8//! - **AI Decision Tracking**: Capture inputs, outputs, and context for every AI decision
9//! - **Deterministic Replay**: Reproduce AI decisions exactly with full context preservation
10//! - **Cost Management**: Track and optimize AI model usage costs
11//! - **Drift Detection**: Monitor model performance and behavior changes
12//! - **Data Sanitization**: Built-in privacy controls for sensitive data
13//! - **Multiple Storage Backends**: SQLite, cloud storage, and custom backends
14//!
15//! ## Installation
16//!
17//! ```bash
18//! pip install briefcase-ai
19//! ```
20//!
21//! ## Usage
22//!
23//! ```python
24//! import briefcase_ai
25//!
26//! # Initialize the library
27//! briefcase_ai.init()
28//!
29//! # Create a decision snapshot
30//! decision = briefcase_ai.DecisionSnapshot("ai_function")
31//! decision.add_input(briefcase_ai.Input("query", "Hello world", "string"))
32//! decision.add_output(briefcase_ai.Output("response", "Hello back!", "string"))
33//!
34//! # Save to storage
35//! storage = briefcase_ai.SqliteBackend.in_memory()
36//! decision_id = storage.save_decision(decision)
37//! ```
38
39use pyo3::prelude::*;
40use std::ffi::CString;
41
42// Re-export core types with Python wrappers
43mod 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/// Initialize the Briefcase AI library with default configuration
60#[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(&register_code, None, None)?;
68    Ok(())
69}
70
71/// Initialize with custom configuration
72#[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/// Check if the runtime is initialized
80#[pyfunction]
81fn is_initialized() -> bool {
82    runtime::is_initialized()
83}
84
85/// Internal function to shutdown the runtime (called by atexit)
86#[pyfunction]
87fn _shutdown_runtime() -> PyResult<()> {
88    runtime::shutdown_runtime()
89}
90
91/// Briefcase AI Python module
92#[pymodule]
93fn briefcase_ai(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
94    // Core data models
95    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    // Storage backends
103    m.add_class::<PySqliteBackend>()?;
104    m.add_class::<PyLakeFSBackend>()?;
105
106    // Core functionality
107    m.add_class::<PyDriftCalculator>()?;
108    m.add_class::<PyCostCalculator>()?;
109    m.add_class::<PySanitizer>()?;
110    m.add_class::<PyReplayEngine>()?;
111
112    // Query types
113    m.add_class::<PySnapshotQuery>()?;
114
115    // Runtime initialization functions
116    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    // Add module metadata
122    m.add("__version__", "2.0.8")?;
123    m.add("__author__", "Briefcase AI Team")?;
124
125    Ok(())
126}