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::*;
40
41// Re-export core types with Python wrappers
42mod 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/// Initialize the Briefcase AI library with default configuration
61#[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/// Initialize with custom configuration
73#[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/// Check if the runtime is initialized
81#[pyfunction]
82fn is_initialized() -> bool {
83    runtime::is_initialized()
84}
85
86/// Internal function to shutdown the runtime (called by atexit)
87#[pyfunction]
88fn _shutdown_runtime() -> PyResult<()> {
89    runtime::shutdown_runtime()
90}
91
92/// Briefcase AI Python module
93#[pymodule]
94fn _core(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
95    // Core data models
96    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    // Storage backends
110    m.add_class::<PySqliteBackend>()?;
111    m.add_class::<PyLakeFSBackend>()?;
112    m.add_class::<PyBufferedBackend>()?;
113
114    // Core functionality
115    m.add_class::<PyDriftCalculator>()?;
116    m.add_class::<PyCostCalculator>()?;
117    m.add_class::<PySanitizer>()?;
118    m.add_class::<PyReplayEngine>()?;
119
120    // Query types
121    m.add_class::<PySnapshotQuery>()?;
122
123    // Authenticated client
124    m.add_class::<client::PyBriefcaseClient>()?;
125    m.add_class::<client::PyValidatedClient>()?;
126
127    // Runtime initialization functions
128    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    // Token estimation
134    m.add_function(wrap_pyfunction!(tokens::count_tokens, m)?)?;
135    m.add_function(wrap_pyfunction!(tokens::count_tokens_batch, m)?)?;
136
137    // Add module metadata
138    m.add("__version__", "2.5.1")?;
139    m.add("__author__", "Briefcase AI Team")?;
140
141    Ok(())
142}