OpenJarvis/rust/crates/openjarvis-python/src/lib.rs

//! PyO3 bridge — exposes ~50 Rust classes to Python via `openjarvis_rust`.

use once_cell::sync::Lazy;
use pyo3::prelude::*;

// Shared tokio runtime for async-to-sync bridge (agents, future async APIs).
pub(crate) static RUNTIME: Lazy<tokio::runtime::Runtime> = Lazy::new(|| {
    tokio::runtime::Runtime::new().expect("Failed to create tokio runtime")
});

pub mod agents;
pub mod core;
pub mod engine;
pub mod learning;
pub mod mcp;
pub mod security;
pub mod storage;
pub mod telemetry;
pub mod tools;
pub mod traces;

// Module-level functions

#[pyfunction]
#[pyo3(signature = (path=None))]
fn load_config(path: Option<&str>) -> PyResult<core::PyConfig> {
    let p = path.map(std::path::Path::new);
    let config = openjarvis_core::load_config(p)
        .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
    Ok(core::PyConfig { inner: config })
}

#[pyfunction]
fn detect_hardware() -> PyResult<String> {
    let hw = openjarvis_core::hardware::detect_hardware();
    Ok(serde_json::to_string(&hw).unwrap_or_default())
}

#[pyfunction]
fn check_ssrf(url: &str) -> Option<String> {
    openjarvis_security::check_ssrf(url)
}

#[pyfunction]
fn is_sensitive_file(path: &str) -> bool {
    openjarvis_security::is_sensitive_file(std::path::Path::new(path))
}

#[pymodule]
fn openjarvis_rust(m: &Bound<'_, PyModule>) -> PyResult<()> {
    // --- Core types ---
    m.add_class::<core::PyMessage>()?;
    m.add_class::<core::PyToolResult>()?;
    m.add_class::<core::PyToolCall>()?;
    m.add_class::<core::PyConfig>()?;
    m.add_class::<core::PyEventBus>()?;
    m.add_class::<core::PyModelSpec>()?;
    m.add_class::<core::PyRoutingContext>()?;
    m.add_class::<core::PyAgentContext>()?;
    m.add_class::<core::PyAgentResult>()?;

    // --- Engines ---
    m.add_class::<engine::PyEngine>()?;
    m.add_class::<engine::PyOllamaEngine>()?;

    // --- Agents ---
    m.add_class::<agents::PySimpleAgent>()?;
    m.add_class::<agents::PyOrchestratorAgent>()?;
    m.add_class::<agents::PyNativeReActAgent>()?;
    m.add_class::<agents::PyLoopGuard>()?;

    // --- Tools ---
    m.add_class::<tools::PyToolExecutor>()?;
    m.add_class::<tools::PyCalculatorTool>()?;
    m.add_class::<tools::PyThinkTool>()?;
    m.add_class::<tools::PyFileReadTool>()?;
    m.add_class::<tools::PyFileWriteTool>()?;
    m.add_class::<tools::PyShellExecTool>()?;
    m.add_class::<tools::PyHttpRequestTool>()?;
    m.add_class::<tools::PyGitStatusTool>()?;
    m.add_class::<tools::PyGitDiffTool>()?;
    m.add_class::<tools::PyGitLogTool>()?;

    // --- Storage / Memory ---
    m.add_class::<storage::PySQLiteMemory>()?;
    m.add_class::<storage::PyBM25Memory>()?;
    m.add_class::<storage::PyKnowledgeGraphMemory>()?;

    // --- Security ---
    m.add_class::<security::PySecretScanner>()?;
    m.add_class::<security::PyPIIScanner>()?;
    m.add_class::<security::PyGuardrailsEngine>()?;
    m.add_class::<security::PyAuditLogger>()?;
    m.add_class::<security::PyCapabilityPolicy>()?;
    m.add_class::<security::PyInjectionScanner>()?;
    m.add_class::<security::PyRateLimiter>()?;
    m.add_class::<security::PyTaintSet>()?;

    // --- Telemetry ---
    m.add_class::<telemetry::PyTelemetryStore>()?;
    m.add_class::<telemetry::PyTelemetryAggregator>()?;
    m.add_class::<telemetry::PyInstrumentedEngine>()?;
    // --- Telemetry (new session/phase/ITL/FLOPs classes) ---
    m.add_class::<telemetry::PyTelemetrySample>()?;
    m.add_class::<telemetry::PyTelemetrySessionCore>()?;
    m.add_class::<telemetry::PyItlStats>()?;
    m.add_class::<telemetry::PyFlopsEstimator>()?;
    m.add_class::<telemetry::PyPhaseMetrics>()?;

    // --- Traces ---
    m.add_class::<traces::PyTraceStore>()?;
    m.add_class::<traces::PyTraceCollector>()?;
    m.add_class::<traces::PyTraceAnalyzer>()?;

    // --- Learning ---
    m.add_class::<learning::PyHeuristicRouter>()?;
    m.add_class::<learning::PyBanditRouterPolicy>()?;
    m.add_class::<learning::PyGRPORouterPolicy>()?;

    // --- MCP ---
    m.add_class::<mcp::PyMcpServer>()?;

    // --- Module-level functions ---
    m.add_function(wrap_pyfunction!(load_config, m)?)?;
    m.add_function(wrap_pyfunction!(detect_hardware, m)?)?;
    m.add_function(wrap_pyfunction!(check_ssrf, m)?)?;
    m.add_function(wrap_pyfunction!(is_sensitive_file, m)?)?;

    Ok(())
}