Add Generics, reduce dynamic dispatch, python integration

This commit is contained in:
krypticmouse
2026-03-04 18:13:45 -08:00
parent 896404b4e0
commit 77026f6803
40 changed files with 3563 additions and 537 deletions
+198
View File
@@ -0,0 +1,198 @@
//! PyO3 bindings for agent types.
//!
//! At the Python boundary, agents use `Box<dyn OjAgent>` for type erasure
//! since Python can't handle Rust generics. The shared tokio Runtime
//! bridges async→sync.
use crate::core::PyAgentResult;
use crate::RUNTIME;
use openjarvis_agents::OjAgent;
use pyo3::prelude::*;
use std::sync::Arc;
/// Python wrapper for SimpleAgent (type-erased via Box<dyn OjAgent>).
#[pyclass(name = "SimpleAgent")]
pub struct PySimpleAgent {
inner: Box<dyn OjAgent>,
}
#[pymethods]
impl PySimpleAgent {
/// Create a SimpleAgent backed by an Engine enum.
#[new]
#[pyo3(signature = (engine_key="ollama", host="http://localhost:11434", model="qwen3:8b", system_prompt="You are a helpful assistant.", temperature=0.7))]
fn new(
engine_key: &str,
host: &str,
model: &str,
system_prompt: &str,
temperature: f64,
) -> PyResult<Self> {
let config = openjarvis_core::JarvisConfig::default();
let engine = openjarvis_engine::get_engine_static(&config, Some(engine_key))
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
let adapter = openjarvis_engine::rig_adapter::RigModelAdapter::new(
Arc::new(engine),
model.to_string(),
);
let agent = openjarvis_agents::SimpleAgent::new(adapter, system_prompt, temperature);
Ok(Self {
inner: Box::new(agent),
})
}
fn agent_id(&self) -> &str {
self.inner.agent_id()
}
fn accepts_tools(&self) -> bool {
self.inner.accepts_tools()
}
fn run(&self, input: &str) -> PyResult<PyAgentResult> {
let result = RUNTIME
.block_on(self.inner.run(input, None))
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(PyAgentResult {
content: result.content,
turns: result.turns,
})
}
}
/// Python wrapper for OrchestratorAgent.
#[pyclass(name = "OrchestratorAgent")]
pub struct PyOrchestratorAgent {
inner: Box<dyn OjAgent>,
}
#[pymethods]
impl PyOrchestratorAgent {
#[new]
#[pyo3(signature = (engine_key="ollama", host="http://localhost:11434", model="qwen3:8b", system_prompt="You are a helpful orchestrator agent.", max_turns=10, temperature=0.7))]
fn new(
engine_key: &str,
host: &str,
model: &str,
system_prompt: &str,
max_turns: usize,
temperature: f64,
) -> PyResult<Self> {
let config = openjarvis_core::JarvisConfig::default();
let engine = openjarvis_engine::get_engine_static(&config, Some(engine_key))
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
let adapter = openjarvis_engine::rig_adapter::RigModelAdapter::new(
Arc::new(engine),
model.to_string(),
);
let executor = Arc::new(openjarvis_tools::ToolExecutor::new(None, None));
let agent = openjarvis_agents::OrchestratorAgent::new(
adapter,
system_prompt,
executor,
max_turns,
temperature,
);
Ok(Self {
inner: Box::new(agent),
})
}
fn agent_id(&self) -> &str {
self.inner.agent_id()
}
fn accepts_tools(&self) -> bool {
self.inner.accepts_tools()
}
fn run(&self, input: &str) -> PyResult<PyAgentResult> {
let result = RUNTIME
.block_on(self.inner.run(input, None))
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(PyAgentResult {
content: result.content,
turns: result.turns,
})
}
}
/// Python wrapper for NativeReActAgent.
#[pyclass(name = "NativeReActAgent")]
pub struct PyNativeReActAgent {
inner: Box<dyn OjAgent>,
}
#[pymethods]
impl PyNativeReActAgent {
#[new]
#[pyo3(signature = (engine_key="ollama", host="http://localhost:11434", model="qwen3:8b", max_turns=10, temperature=0.7))]
fn new(
engine_key: &str,
host: &str,
model: &str,
max_turns: usize,
temperature: f64,
) -> PyResult<Self> {
let config = openjarvis_core::JarvisConfig::default();
let engine = openjarvis_engine::get_engine_static(&config, Some(engine_key))
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
let adapter = openjarvis_engine::rig_adapter::RigModelAdapter::new(
Arc::new(engine),
model.to_string(),
);
let executor = Arc::new(openjarvis_tools::ToolExecutor::new(None, None));
let agent = openjarvis_agents::NativeReActAgent::new(
adapter,
executor,
max_turns,
temperature,
);
Ok(Self {
inner: Box::new(agent),
})
}
fn agent_id(&self) -> &str {
self.inner.agent_id()
}
fn accepts_tools(&self) -> bool {
self.inner.accepts_tools()
}
fn run(&self, input: &str) -> PyResult<PyAgentResult> {
let result = RUNTIME
.block_on(self.inner.run(input, None))
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(PyAgentResult {
content: result.content,
turns: result.turns,
})
}
}
/// Python wrapper for LoopGuard.
#[pyclass(name = "LoopGuard")]
pub struct PyLoopGuard {
inner: openjarvis_agents::LoopGuard,
}
#[pymethods]
impl PyLoopGuard {
#[new]
#[pyo3(signature = (max_identical=50, max_ping_pong=4, poll_budget=100))]
fn new(max_identical: usize, max_ping_pong: usize, poll_budget: usize) -> Self {
Self {
inner: openjarvis_agents::LoopGuard::new(max_identical, max_ping_pong, poll_budget),
}
}
fn check(&mut self, tool_name: &str, arguments: &str) -> Option<String> {
self.inner.check(tool_name, arguments)
}
fn reset(&mut self) {
self.inner.reset()
}
}
+212
View File
@@ -0,0 +1,212 @@
//! PyO3 bindings for core types.
use pyo3::prelude::*;
use std::collections::HashMap;
#[pyclass(name = "Message")]
#[derive(Clone)]
pub struct PyMessage {
#[pyo3(get, set)]
pub role: String,
#[pyo3(get, set)]
pub content: String,
#[pyo3(get, set)]
pub name: Option<String>,
#[pyo3(get, set)]
pub tool_call_id: Option<String>,
}
#[pymethods]
impl PyMessage {
#[new]
fn new(role: String, content: String) -> Self {
Self {
role,
content,
name: None,
tool_call_id: None,
}
}
fn __repr__(&self) -> String {
format!("Message(role='{}', content='{}')", self.role, &self.content[..self.content.len().min(50)])
}
}
impl PyMessage {
pub fn to_core(&self) -> openjarvis_core::Message {
let role = match self.role.as_str() {
"system" => openjarvis_core::Role::System,
"assistant" => openjarvis_core::Role::Assistant,
"tool" => openjarvis_core::Role::Tool,
_ => openjarvis_core::Role::User,
};
openjarvis_core::Message {
role,
content: self.content.clone(),
name: self.name.clone(),
tool_calls: None,
tool_call_id: self.tool_call_id.clone(),
metadata: HashMap::new(),
}
}
}
#[pyclass(name = "ToolResult")]
#[derive(Clone)]
pub struct PyToolResult {
#[pyo3(get)]
pub tool_name: String,
#[pyo3(get)]
pub content: String,
#[pyo3(get)]
pub success: bool,
}
#[pymethods]
impl PyToolResult {
#[new]
fn new(tool_name: String, content: String, success: bool) -> Self {
Self { tool_name, content, success }
}
fn __repr__(&self) -> String {
format!("ToolResult(tool='{}', success={})", self.tool_name, self.success)
}
}
#[pyclass(name = "ToolCall")]
#[derive(Clone)]
pub struct PyToolCall {
#[pyo3(get, set)]
pub id: String,
#[pyo3(get, set)]
pub name: String,
#[pyo3(get, set)]
pub arguments: String,
}
#[pymethods]
impl PyToolCall {
#[new]
fn new(id: String, name: String, arguments: String) -> Self {
Self { id, name, arguments }
}
}
#[pyclass(name = "Config")]
pub struct PyConfig {
pub inner: openjarvis_core::JarvisConfig,
}
#[pymethods]
impl PyConfig {
#[new]
fn new() -> Self {
Self {
inner: openjarvis_core::JarvisConfig::default(),
}
}
fn __repr__(&self) -> String {
format!(
"Config(engine={}, model={})",
self.inner.engine.default, self.inner.intelligence.default_model
)
}
#[getter]
fn engine_default(&self) -> String {
self.inner.engine.default.clone()
}
#[getter]
fn model_default(&self) -> String {
self.inner.intelligence.default_model.clone()
}
}
#[pyclass(name = "EventBus")]
pub struct PyEventBus {
pub inner: std::sync::Arc<openjarvis_core::EventBus>,
}
#[pymethods]
impl PyEventBus {
#[new]
fn new() -> Self {
Self {
inner: std::sync::Arc::new(openjarvis_core::EventBus::new(true)),
}
}
fn history_len(&self) -> usize {
self.inner.history().len()
}
}
#[pyclass(name = "ModelSpec")]
#[derive(Clone)]
pub struct PyModelSpec {
#[pyo3(get, set)]
pub name: String,
#[pyo3(get, set)]
pub params_b: f64,
#[pyo3(get, set)]
pub context_length: usize,
}
#[pymethods]
impl PyModelSpec {
#[new]
fn new(name: String, params_b: f64, context_length: usize) -> Self {
Self { name, params_b, context_length }
}
}
#[pyclass(name = "RoutingContext")]
#[derive(Clone)]
pub struct PyRoutingContext {
#[pyo3(get, set)]
pub query: String,
#[pyo3(get, set)]
pub query_class: String,
}
#[pymethods]
impl PyRoutingContext {
#[new]
fn new(query: String) -> Self {
Self { query, query_class: "general".into() }
}
}
#[pyclass(name = "AgentContext")]
pub struct PyAgentContext {
#[pyo3(get, set)]
pub session_id: String,
}
#[pymethods]
impl PyAgentContext {
#[new]
fn new(session_id: String) -> Self {
Self { session_id }
}
}
#[pyclass(name = "AgentResult")]
#[derive(Clone)]
pub struct PyAgentResult {
#[pyo3(get)]
pub content: String,
#[pyo3(get)]
pub turns: usize,
}
#[pymethods]
impl PyAgentResult {
fn __repr__(&self) -> String {
format!("AgentResult(turns={}, content='{}')", self.turns, &self.content[..self.content.len().min(50)])
}
}
+146
View File
@@ -0,0 +1,146 @@
//! PyO3 bindings for engine types.
use crate::core::PyMessage;
use openjarvis_engine::InferenceEngine;
use pyo3::prelude::*;
/// Wraps the Engine enum (static dispatch internally, opaque to Python).
#[pyclass(name = "Engine")]
pub struct PyEngine {
pub inner: openjarvis_engine::Engine,
}
#[pymethods]
impl PyEngine {
/// Create an engine by key (e.g. "ollama", "vllm", "sglang", "llamacpp", "mlx", "lmstudio").
#[new]
#[pyo3(signature = (engine_key="ollama", host=None))]
fn new(engine_key: &str, host: Option<&str>) -> PyResult<Self> {
let engine = match engine_key {
"ollama" => openjarvis_engine::Engine::Ollama(
openjarvis_engine::OllamaEngine::new(
host.unwrap_or("http://localhost:11434"),
120.0,
),
),
"vllm" => openjarvis_engine::Engine::Vllm(
openjarvis_engine::OpenAICompatEngine::vllm(
host.unwrap_or("http://localhost:8000"),
),
),
"sglang" => openjarvis_engine::Engine::Sglang(
openjarvis_engine::OpenAICompatEngine::sglang(
host.unwrap_or("http://localhost:30000"),
),
),
"llamacpp" => openjarvis_engine::Engine::LlamaCpp(
openjarvis_engine::OpenAICompatEngine::llamacpp(
host.unwrap_or("http://localhost:8080"),
),
),
"mlx" => openjarvis_engine::Engine::Mlx(
openjarvis_engine::OpenAICompatEngine::mlx(
host.unwrap_or("http://localhost:8080"),
),
),
"lmstudio" => openjarvis_engine::Engine::LmStudio(
openjarvis_engine::OpenAICompatEngine::lmstudio(
host.unwrap_or("http://localhost:1234"),
),
),
other => {
return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
format!("Unknown engine: {}", other),
));
}
};
Ok(Self { inner: engine })
}
fn engine_id(&self) -> &str {
self.inner.engine_id()
}
fn variant_key(&self) -> &str {
self.inner.variant_key()
}
fn health(&self) -> bool {
self.inner.health()
}
fn list_models(&self) -> PyResult<Vec<String>> {
self.inner
.list_models()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
}
#[pyo3(signature = (messages, model, temperature=0.7, max_tokens=1024))]
fn generate(
&self,
messages: Vec<PyMessage>,
model: &str,
temperature: f64,
max_tokens: i64,
) -> PyResult<String> {
let core_msgs: Vec<openjarvis_core::Message> =
messages.iter().map(|m| m.to_core()).collect();
let result = self
.inner
.generate(&core_msgs, model, temperature, max_tokens, None)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(serde_json::to_string(&result).unwrap_or_default())
}
fn __repr__(&self) -> String {
format!("Engine({})", self.inner.variant_key())
}
}
/// Convenience alias for backward compatibility.
#[pyclass(name = "OllamaEngine")]
pub struct PyOllamaEngine {
inner: openjarvis_engine::OllamaEngine,
}
#[pymethods]
impl PyOllamaEngine {
#[new]
#[pyo3(signature = (host="http://localhost:11434", timeout=120.0))]
fn new(host: &str, timeout: f64) -> Self {
Self {
inner: openjarvis_engine::OllamaEngine::new(host, timeout),
}
}
fn engine_id(&self) -> &str {
self.inner.engine_id()
}
fn health(&self) -> bool {
self.inner.health()
}
fn list_models(&self) -> PyResult<Vec<String>> {
self.inner
.list_models()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
}
#[pyo3(signature = (messages, model, temperature=0.7, max_tokens=1024))]
fn generate(
&self,
messages: Vec<PyMessage>,
model: &str,
temperature: f64,
max_tokens: i64,
) -> PyResult<String> {
let core_msgs: Vec<openjarvis_core::Message> =
messages.iter().map(|m| m.to_core()).collect();
let result = self
.inner
.generate(&core_msgs, model, temperature, max_tokens, None)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(serde_json::to_string(&result).unwrap_or_default())
}
}
@@ -0,0 +1,98 @@
//! PyO3 bindings for learning/router policy types.
use openjarvis_learning::RouterPolicy;
use pyo3::prelude::*;
#[pyclass(name = "HeuristicRouter")]
pub struct PyHeuristicRouter {
inner: openjarvis_learning::HeuristicRouter,
}
#[pymethods]
impl PyHeuristicRouter {
#[new]
#[pyo3(signature = (default_model="qwen3:8b", code_model=None, math_model=None, fast_model=None))]
fn new(
default_model: &str,
code_model: Option<String>,
math_model: Option<String>,
fast_model: Option<String>,
) -> Self {
Self {
inner: openjarvis_learning::HeuristicRouter::new(
default_model.to_string(),
code_model,
math_model,
fast_model,
),
}
}
fn select_model(&self, query: &str, has_code: bool, has_math: bool) -> String {
let ctx = openjarvis_core::RoutingContext {
query: query.to_string(),
query_length: query.len(),
has_code,
has_math,
..Default::default()
};
self.inner.select_model(&ctx)
}
}
#[pyclass(name = "BanditRouterPolicy")]
pub struct PyBanditRouterPolicy {
inner: openjarvis_learning::BanditRouterPolicy,
}
#[pymethods]
impl PyBanditRouterPolicy {
#[new]
#[pyo3(signature = (models, strategy="thompson"))]
fn new(models: Vec<String>, strategy: &str) -> Self {
let strat = match strategy {
"ucb1" | "UCB1" => openjarvis_learning::bandit::BanditStrategy::UCB1,
_ => openjarvis_learning::bandit::BanditStrategy::ThompsonSampling,
};
Self {
inner: openjarvis_learning::BanditRouterPolicy::new(models, strat),
}
}
fn select_model(&self) -> String {
let ctx = openjarvis_core::RoutingContext::default();
self.inner.select_model(&ctx)
}
fn update(&self, model: &str, reward: f64) {
self.inner.update(model, reward);
}
}
#[pyclass(name = "GRPORouterPolicy")]
pub struct PyGRPORouterPolicy {
inner: openjarvis_learning::GRPORouterPolicy,
}
#[pymethods]
impl PyGRPORouterPolicy {
#[new]
#[pyo3(signature = (models, temperature=1.0))]
fn new(models: Vec<String>, temperature: f64) -> Self {
Self {
inner: openjarvis_learning::GRPORouterPolicy::new(models, temperature),
}
}
fn select_model(&self) -> String {
let ctx = openjarvis_core::RoutingContext::default();
self.inner.select_model(&ctx)
}
fn update_weights(&self, rewards_json: &str) -> PyResult<()> {
let rewards: Vec<(String, f64)> = serde_json::from_str(rewards_json)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
self.inner.update_weights(&rewards);
Ok(())
}
}
+86 -210
View File
@@ -1,218 +1,33 @@
//! PyO3 bridge — exposes Rust backend to Python.
//! PyO3 bridge — exposes ~50 Rust classes to Python via `openjarvis_rust`.
use once_cell::sync::Lazy;
use pyo3::prelude::*;
use pyo3::types::PyDict;
use std::collections::HashMap;
// Re-export core types as Python classes
// 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")
});
#[pyclass(name = "Message")]
#[derive(Clone)]
struct PyMessage {
#[pyo3(get, set)]
role: String,
#[pyo3(get, set)]
content: String,
#[pyo3(get, set)]
name: Option<String>,
#[pyo3(get, set)]
tool_call_id: Option<String>,
}
#[pymethods]
impl PyMessage {
#[new]
fn new(role: String, content: String) -> Self {
Self {
role,
content,
name: None,
tool_call_id: None,
}
}
}
impl PyMessage {
fn to_core(&self) -> openjarvis_core::Message {
let role = match self.role.as_str() {
"system" => openjarvis_core::Role::System,
"assistant" => openjarvis_core::Role::Assistant,
"tool" => openjarvis_core::Role::Tool,
_ => openjarvis_core::Role::User,
};
openjarvis_core::Message {
role,
content: self.content.clone(),
name: self.name.clone(),
tool_calls: None,
tool_call_id: self.tool_call_id.clone(),
metadata: HashMap::new(),
}
}
}
#[pyclass(name = "ToolResult")]
#[derive(Clone)]
struct PyToolResult {
#[pyo3(get)]
tool_name: String,
#[pyo3(get)]
content: String,
#[pyo3(get)]
success: bool,
}
#[pyclass(name = "Config")]
struct PyConfig {
inner: openjarvis_core::JarvisConfig,
}
#[pymethods]
impl PyConfig {
#[new]
fn new() -> Self {
Self {
inner: openjarvis_core::JarvisConfig::default(),
}
}
fn __repr__(&self) -> String {
format!(
"Config(engine={}, model={})",
self.inner.engine.default, self.inner.intelligence.default_model
)
}
}
#[pyclass(name = "OllamaEngine")]
struct PyOllamaEngine {
inner: openjarvis_engine::OllamaEngine,
}
#[pymethods]
impl PyOllamaEngine {
#[new]
#[pyo3(signature = (host="http://localhost:11434", timeout=120.0))]
fn new(host: &str, timeout: f64) -> Self {
Self {
inner: openjarvis_engine::OllamaEngine::new(host, timeout),
}
}
fn engine_id(&self) -> &str {
use openjarvis_engine::InferenceEngine;
self.inner.engine_id()
}
fn health(&self) -> bool {
use openjarvis_engine::InferenceEngine;
self.inner.health()
}
fn list_models(&self) -> PyResult<Vec<String>> {
use openjarvis_engine::InferenceEngine;
self.inner
.list_models()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
}
#[pyo3(signature = (messages, model, temperature=0.7, max_tokens=1024))]
fn generate(
&self,
messages: Vec<PyMessage>,
model: &str,
temperature: f64,
max_tokens: i64,
) -> PyResult<String> {
use openjarvis_engine::InferenceEngine;
let core_msgs: Vec<openjarvis_core::Message> =
messages.iter().map(|m| m.to_core()).collect();
let result = self
.inner
.generate(&core_msgs, model, temperature, max_tokens, None)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(serde_json::to_string(&result).unwrap_or_default())
}
}
#[pyclass(name = "SecretScanner")]
struct PySecretScanner {
inner: openjarvis_security::SecretScanner,
}
#[pymethods]
impl PySecretScanner {
#[new]
fn new() -> Self {
Self {
inner: openjarvis_security::SecretScanner::new(),
}
}
fn scan(&self, text: &str) -> PyResult<String> {
let result = self.inner.scan(text);
Ok(serde_json::to_string(&result).unwrap_or_default())
}
fn redact(&self, text: &str) -> String {
self.inner.redact(text)
}
}
#[pyclass(name = "PIIScanner")]
struct PyPIIScanner {
inner: openjarvis_security::PIIScanner,
}
#[pymethods]
impl PyPIIScanner {
#[new]
fn new() -> Self {
Self {
inner: openjarvis_security::PIIScanner::new(),
}
}
fn scan(&self, text: &str) -> PyResult<String> {
let result = self.inner.scan(text);
Ok(serde_json::to_string(&result).unwrap_or_default())
}
fn redact(&self, text: &str) -> String {
self.inner.redact(text)
}
}
#[pyclass(name = "CalculatorTool")]
struct PyCalculatorTool;
#[pymethods]
impl PyCalculatorTool {
#[new]
fn new() -> Self {
Self
}
fn execute(&self, expression: &str) -> PyResult<String> {
use openjarvis_tools::traits::BaseTool;
let tool = openjarvis_tools::builtin::calculator::CalculatorTool;
let params = serde_json::json!({"expression": expression});
let result = tool
.execute(&params)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(result.content)
}
}
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<PyConfig> {
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(PyConfig { inner: config })
Ok(core::PyConfig { inner: config })
}
#[pyfunction]
@@ -233,16 +48,77 @@ fn is_sensitive_file(path: &str) -> bool {
#[pymodule]
fn openjarvis_rust(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<PyMessage>()?;
m.add_class::<PyToolResult>()?;
m.add_class::<PyConfig>()?;
m.add_class::<PyOllamaEngine>()?;
m.add_class::<PySecretScanner>()?;
m.add_class::<PyPIIScanner>()?;
m.add_class::<PyCalculatorTool>()?;
// --- 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>()?;
// --- 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(())
}
+25
View File
@@ -0,0 +1,25 @@
//! PyO3 bindings for the MCP server.
use crate::tools::PyToolExecutor;
use pyo3::prelude::*;
use std::sync::Arc;
#[pyclass(name = "McpServer")]
pub struct PyMcpServer {
inner: openjarvis_mcp::McpServer,
}
#[pymethods]
impl PyMcpServer {
#[new]
fn new(executor: &PyToolExecutor) -> Self {
Self {
inner: openjarvis_mcp::McpServer::new(Arc::clone(&executor.inner)),
}
}
/// Process a JSON-RPC request string and return a JSON-RPC response string.
fn handle_json(&self, json_str: &str) -> String {
self.inner.handle_json(json_str)
}
}
@@ -0,0 +1,263 @@
//! PyO3 bindings for security types.
use pyo3::prelude::*;
use std::sync::Arc;
#[pyclass(name = "SecretScanner")]
pub struct PySecretScanner {
inner: openjarvis_security::SecretScanner,
}
#[pymethods]
impl PySecretScanner {
#[new]
fn new() -> Self {
Self {
inner: openjarvis_security::SecretScanner::new(),
}
}
fn scan(&self, text: &str) -> PyResult<String> {
let result = self.inner.scan(text);
Ok(serde_json::to_string(&result).unwrap_or_default())
}
fn redact(&self, text: &str) -> String {
self.inner.redact(text)
}
}
#[pyclass(name = "PIIScanner")]
pub struct PyPIIScanner {
inner: openjarvis_security::PIIScanner,
}
#[pymethods]
impl PyPIIScanner {
#[new]
fn new() -> Self {
Self {
inner: openjarvis_security::PIIScanner::new(),
}
}
fn scan(&self, text: &str) -> PyResult<String> {
let result = self.inner.scan(text);
Ok(serde_json::to_string(&result).unwrap_or_default())
}
fn redact(&self, text: &str) -> String {
self.inner.redact(text)
}
}
#[pyclass(name = "GuardrailsEngine")]
pub struct PyGuardrailsEngine {
inner: openjarvis_security::GuardrailsEngine<openjarvis_engine::Engine>,
}
#[pymethods]
impl PyGuardrailsEngine {
#[new]
#[pyo3(signature = (engine_key="ollama", host="http://localhost:11434", mode="warn", scan_input=true, scan_output=true))]
fn new(
engine_key: &str,
host: &str,
mode: &str,
scan_input: bool,
scan_output: bool,
) -> PyResult<Self> {
let config = openjarvis_core::JarvisConfig::default();
let engine = openjarvis_engine::get_engine_static(&config, Some(engine_key))
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
let redaction_mode = match mode {
"redact" => openjarvis_security::RedactionMode::Redact,
"block" => openjarvis_security::RedactionMode::Block,
_ => openjarvis_security::RedactionMode::Warn,
};
Ok(Self {
inner: openjarvis_security::GuardrailsEngine::new(
engine,
redaction_mode,
scan_input,
scan_output,
None,
),
})
}
fn engine_id(&self) -> &str {
use openjarvis_engine::InferenceEngine;
self.inner.engine_id()
}
}
#[pyclass(name = "AuditLogger")]
pub struct PyAuditLogger {
inner: parking_lot::Mutex<openjarvis_security::AuditLogger>,
}
#[pymethods]
impl PyAuditLogger {
#[new]
#[pyo3(signature = (path=None))]
fn new(path: Option<&str>) -> PyResult<Self> {
let db_path = match path {
Some(p) => std::path::PathBuf::from(p),
None => std::path::PathBuf::from(":memory:"),
};
let inner = openjarvis_security::AuditLogger::new(&db_path)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(Self {
inner: parking_lot::Mutex::new(inner),
})
}
fn count(&self) -> i64 {
self.inner.lock().count()
}
fn verify_chain(&self) -> PyResult<(bool, Option<i64>)> {
self.inner
.lock()
.verify_chain()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
}
fn tail_hash(&self) -> String {
self.inner.lock().tail_hash()
}
}
#[pyclass(name = "CapabilityPolicy")]
pub struct PyCapabilityPolicy {
inner: openjarvis_security::CapabilityPolicy,
}
#[pymethods]
impl PyCapabilityPolicy {
#[new]
#[pyo3(signature = (default_deny=true))]
fn new(default_deny: bool) -> Self {
Self {
inner: openjarvis_security::CapabilityPolicy::new(default_deny),
}
}
fn check(&self, agent_id: &str, capability: &str, resource: &str) -> bool {
self.inner.check(agent_id, capability, resource)
}
fn grant(&mut self, agent_id: &str, capability: &str, pattern: &str) {
self.inner.grant(agent_id, capability, pattern);
}
fn deny(&mut self, agent_id: &str, capability: &str) {
self.inner.deny(agent_id, capability);
}
fn list_agents(&self) -> Vec<String> {
self.inner.list_agents()
}
}
#[pyclass(name = "InjectionScanner")]
pub struct PyInjectionScanner {
inner: openjarvis_security::InjectionScanner,
}
#[pymethods]
impl PyInjectionScanner {
#[new]
fn new() -> Self {
Self {
inner: openjarvis_security::InjectionScanner::new(),
}
}
fn scan(&self, text: &str) -> PyResult<String> {
let result = self.inner.scan(text);
// InjectionScanResult doesn't derive Serialize, so format manually.
Ok(serde_json::json!({
"is_clean": result.is_clean,
"threat_level": format!("{:?}", result.threat_level),
"findings_count": result.findings.len(),
})
.to_string())
}
}
#[pyclass(name = "RateLimiter")]
pub struct PyRateLimiter {
inner: openjarvis_security::RateLimiter,
}
#[pymethods]
impl PyRateLimiter {
#[new]
#[pyo3(signature = (requests_per_minute=60, burst_size=10))]
fn new(requests_per_minute: u32, burst_size: u32) -> Self {
Self {
inner: openjarvis_security::RateLimiter::new(
openjarvis_security::RateLimitConfig {
requests_per_minute,
burst_size,
enabled: true,
},
),
}
}
/// Returns (allowed, wait_seconds).
fn check(&self, key: &str) -> (bool, f64) {
self.inner.check(key)
}
fn reset(&self, key: Option<&str>) {
self.inner.reset(key);
}
}
#[pyclass(name = "TaintSet")]
pub struct PyTaintSet {
inner: openjarvis_security::TaintSet,
}
#[pymethods]
impl PyTaintSet {
#[new]
fn new() -> Self {
Self {
inner: openjarvis_security::TaintSet::new(),
}
}
fn add(&mut self, label: &str) {
let taint_label = match label {
"pii" => openjarvis_security::TaintLabel::Pii,
"secret" => openjarvis_security::TaintLabel::Secret,
"user_private" => openjarvis_security::TaintLabel::UserPrivate,
"external" => openjarvis_security::TaintLabel::External,
_ => openjarvis_security::TaintLabel::External,
};
// TaintSet is immutable-style; union with a single-label set.
self.inner = self.inner.union(
&openjarvis_security::TaintSet::from_labels(&[taint_label]),
);
}
fn has(&self, label: &str) -> bool {
let taint_label = match label {
"pii" => openjarvis_security::TaintLabel::Pii,
"secret" => openjarvis_security::TaintLabel::Secret,
"user_private" => openjarvis_security::TaintLabel::UserPrivate,
"external" => openjarvis_security::TaintLabel::External,
_ => return false,
};
self.inner.has(taint_label)
}
fn is_empty(&self) -> bool {
self.inner.is_empty()
}
}
@@ -0,0 +1,148 @@
//! PyO3 bindings for storage/memory backends.
use openjarvis_tools::storage::MemoryBackend;
use pyo3::prelude::*;
#[pyclass(name = "SQLiteMemory")]
pub struct PySQLiteMemory {
inner: openjarvis_tools::storage::SQLiteMemory,
}
#[pymethods]
impl PySQLiteMemory {
#[new]
#[pyo3(signature = (path=":memory:"))]
fn new(path: &str) -> PyResult<Self> {
let inner = openjarvis_tools::storage::SQLiteMemory::new(std::path::Path::new(path))
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(Self { inner })
}
fn backend_id(&self) -> &str {
self.inner.backend_id()
}
#[pyo3(signature = (content, source, metadata=None))]
fn store(&self, content: &str, source: &str, metadata: Option<&str>) -> PyResult<String> {
let meta = metadata
.map(|m| serde_json::from_str(m))
.transpose()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
self.inner
.store(content, source, meta.as_ref())
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
}
#[pyo3(signature = (query, top_k=5))]
fn retrieve(&self, query: &str, top_k: usize) -> PyResult<String> {
let results = self
.inner
.retrieve(query, top_k)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(serde_json::to_string(&results).unwrap_or_default())
}
fn count(&self) -> PyResult<usize> {
self.inner
.count()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
}
fn clear(&self) -> PyResult<()> {
self.inner
.clear()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
}
}
#[pyclass(name = "BM25Memory")]
pub struct PyBM25Memory {
inner: openjarvis_tools::storage::BM25Memory,
}
#[pymethods]
impl PyBM25Memory {
#[new]
#[pyo3(signature = (k1=1.2, b=0.75))]
fn new(k1: f64, b: f64) -> Self {
Self {
inner: openjarvis_tools::storage::BM25Memory::new(k1, b),
}
}
fn backend_id(&self) -> &str {
self.inner.backend_id()
}
#[pyo3(signature = (content, source, metadata=None))]
fn store(&self, content: &str, source: &str, metadata: Option<&str>) -> PyResult<String> {
let meta = metadata
.map(|m| serde_json::from_str(m))
.transpose()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
self.inner
.store(content, source, meta.as_ref())
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
}
#[pyo3(signature = (query, top_k=5))]
fn retrieve(&self, query: &str, top_k: usize) -> PyResult<String> {
let results = self
.inner
.retrieve(query, top_k)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(serde_json::to_string(&results).unwrap_or_default())
}
fn count(&self) -> PyResult<usize> {
self.inner
.count()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
}
}
#[pyclass(name = "KnowledgeGraphMemory")]
pub struct PyKnowledgeGraphMemory {
inner: openjarvis_tools::storage::KnowledgeGraphMemory,
}
#[pymethods]
impl PyKnowledgeGraphMemory {
#[new]
#[pyo3(signature = (path=":memory:"))]
fn new(path: &str) -> PyResult<Self> {
let inner = openjarvis_tools::storage::KnowledgeGraphMemory::new(std::path::Path::new(path))
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(Self { inner })
}
fn backend_id(&self) -> &str {
self.inner.backend_id()
}
#[pyo3(signature = (content, source, metadata=None))]
fn store(&self, content: &str, source: &str, metadata: Option<&str>) -> PyResult<String> {
let meta = metadata
.map(|m| serde_json::from_str(m))
.transpose()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
self.inner
.store(content, source, meta.as_ref())
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
}
#[pyo3(signature = (query, top_k=5))]
fn retrieve(&self, query: &str, top_k: usize) -> PyResult<String> {
let results = self
.inner
.retrieve(query, top_k)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(serde_json::to_string(&results).unwrap_or_default())
}
fn count(&self) -> PyResult<usize> {
self.inner
.count()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
}
}
@@ -0,0 +1,93 @@
//! PyO3 bindings for telemetry types.
use pyo3::prelude::*;
use std::sync::Arc;
#[pyclass(name = "TelemetryStore")]
pub struct PyTelemetryStore {
pub inner: Arc<openjarvis_telemetry::TelemetryStore>,
}
#[pymethods]
impl PyTelemetryStore {
#[new]
#[pyo3(signature = (path=None))]
fn new(path: Option<&str>) -> PyResult<Self> {
let inner = match path {
Some(p) => openjarvis_telemetry::TelemetryStore::new(std::path::Path::new(p)),
None => openjarvis_telemetry::TelemetryStore::in_memory(),
}
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(Self {
inner: Arc::new(inner),
})
}
fn count(&self) -> PyResult<usize> {
self.inner
.count()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
}
fn clear(&self) -> PyResult<()> {
self.inner
.clear()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
}
}
/// TelemetryAggregator computes aggregate stats from a TelemetryStore.
/// The Rust type is a unit struct with a static method.
#[pyclass(name = "TelemetryAggregator")]
pub struct PyTelemetryAggregator {
store: Arc<openjarvis_telemetry::TelemetryStore>,
}
#[pymethods]
impl PyTelemetryAggregator {
#[new]
fn new(store: &PyTelemetryStore) -> Self {
Self {
store: Arc::clone(&store.inner),
}
}
fn stats(&self) -> PyResult<String> {
let stats = openjarvis_telemetry::TelemetryAggregator::stats(&self.store)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(serde_json::to_string(&stats).unwrap_or_default())
}
}
#[pyclass(name = "InstrumentedEngine")]
pub struct PyInstrumentedEngine {
inner: openjarvis_telemetry::InstrumentedEngine<openjarvis_engine::Engine>,
}
#[pymethods]
impl PyInstrumentedEngine {
#[new]
#[pyo3(signature = (engine_key="ollama", host="http://localhost:11434", store_path=None, agent_name="default"))]
fn new(engine_key: &str, host: &str, store_path: Option<&str>, agent_name: &str) -> PyResult<Self> {
let config = openjarvis_core::JarvisConfig::default();
let engine = openjarvis_engine::get_engine_static(&config, Some(engine_key))
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
let store = Arc::new(match store_path {
Some(p) => openjarvis_telemetry::TelemetryStore::new(std::path::Path::new(p)),
None => openjarvis_telemetry::TelemetryStore::in_memory(),
}
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?);
Ok(Self {
inner: openjarvis_telemetry::InstrumentedEngine::new(
engine,
store,
agent_name.to_string(),
),
})
}
fn engine_id(&self) -> &str {
use openjarvis_engine::InferenceEngine;
self.inner.engine_id()
}
}
+237
View File
@@ -0,0 +1,237 @@
//! PyO3 bindings for tool types.
use openjarvis_tools::traits::BaseTool;
use pyo3::prelude::*;
use std::sync::Arc;
#[pyclass(name = "ToolExecutor")]
pub struct PyToolExecutor {
pub inner: Arc<openjarvis_tools::ToolExecutor>,
}
#[pymethods]
impl PyToolExecutor {
#[new]
fn new() -> Self {
Self {
inner: Arc::new(openjarvis_tools::ToolExecutor::new(None, None)),
}
}
fn list_tools(&self) -> Vec<String> {
self.inner.list_tools()
}
fn execute(&self, tool_name: &str, params_json: &str) -> PyResult<String> {
let params: serde_json::Value = serde_json::from_str(params_json)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
let result = self
.inner
.execute(tool_name, &params, None, None)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(serde_json::to_string(&result).unwrap_or_default())
}
}
#[pyclass(name = "CalculatorTool")]
pub struct PyCalculatorTool;
#[pymethods]
impl PyCalculatorTool {
#[new]
fn new() -> Self {
Self
}
fn execute(&self, expression: &str) -> PyResult<String> {
let tool = openjarvis_tools::builtin::calculator::CalculatorTool;
let params = serde_json::json!({"expression": expression});
let result = tool
.execute(&params)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(result.content)
}
}
#[pyclass(name = "ThinkTool")]
pub struct PyThinkTool;
#[pymethods]
impl PyThinkTool {
#[new]
fn new() -> Self {
Self
}
fn execute(&self, thought: &str) -> PyResult<String> {
let tool = openjarvis_tools::builtin::think::ThinkTool;
let params = serde_json::json!({"thought": thought});
let result = tool
.execute(&params)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(result.content)
}
}
#[pyclass(name = "FileReadTool")]
pub struct PyFileReadTool;
#[pymethods]
impl PyFileReadTool {
#[new]
fn new() -> Self {
Self
}
fn execute(&self, path: &str) -> PyResult<String> {
let tool = openjarvis_tools::builtin::file_tools::FileReadTool;
let params = serde_json::json!({"path": path});
let result = tool
.execute(&params)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(result.content)
}
}
#[pyclass(name = "FileWriteTool")]
pub struct PyFileWriteTool;
#[pymethods]
impl PyFileWriteTool {
#[new]
fn new() -> Self {
Self
}
fn execute(&self, path: &str, content: &str) -> PyResult<String> {
let tool = openjarvis_tools::builtin::file_tools::FileWriteTool;
let params = serde_json::json!({"path": path, "content": content});
let result = tool
.execute(&params)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(result.content)
}
}
#[pyclass(name = "ShellExecTool")]
pub struct PyShellExecTool;
#[pymethods]
impl PyShellExecTool {
#[new]
fn new() -> Self {
Self
}
#[pyo3(signature = (command, cwd=None))]
fn execute(&self, command: &str, cwd: Option<&str>) -> PyResult<String> {
let tool = openjarvis_tools::builtin::shell::ShellExecTool;
let mut params = serde_json::json!({"command": command});
if let Some(cwd) = cwd {
params["cwd"] = serde_json::Value::String(cwd.to_string());
}
let result = tool
.execute(&params)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(result.content)
}
}
#[pyclass(name = "HttpRequestTool")]
pub struct PyHttpRequestTool;
#[pymethods]
impl PyHttpRequestTool {
#[new]
fn new() -> Self {
Self
}
#[pyo3(signature = (url, method="GET", body=None))]
fn execute(&self, url: &str, method: &str, body: Option<&str>) -> PyResult<String> {
let tool = openjarvis_tools::builtin::http_tools::HttpRequestTool;
let mut params = serde_json::json!({"url": url, "method": method});
if let Some(body) = body {
params["body"] = serde_json::Value::String(body.to_string());
}
let result = tool
.execute(&params)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(result.content)
}
}
#[pyclass(name = "GitStatusTool")]
pub struct PyGitStatusTool;
#[pymethods]
impl PyGitStatusTool {
#[new]
fn new() -> Self {
Self
}
#[pyo3(signature = (cwd=None))]
fn execute(&self, cwd: Option<&str>) -> PyResult<String> {
let tool = openjarvis_tools::builtin::git_tools::GitStatusTool;
let mut params = serde_json::json!({});
if let Some(cwd) = cwd {
params["cwd"] = serde_json::Value::String(cwd.to_string());
}
let result = tool
.execute(&params)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(result.content)
}
}
#[pyclass(name = "GitDiffTool")]
pub struct PyGitDiffTool;
#[pymethods]
impl PyGitDiffTool {
#[new]
fn new() -> Self {
Self
}
#[pyo3(signature = (cwd=None))]
fn execute(&self, cwd: Option<&str>) -> PyResult<String> {
let tool = openjarvis_tools::builtin::git_tools::GitDiffTool;
let mut params = serde_json::json!({});
if let Some(cwd) = cwd {
params["cwd"] = serde_json::Value::String(cwd.to_string());
}
let result = tool
.execute(&params)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(result.content)
}
}
#[pyclass(name = "GitLogTool")]
pub struct PyGitLogTool;
#[pymethods]
impl PyGitLogTool {
#[new]
fn new() -> Self {
Self
}
#[pyo3(signature = (cwd=None, count=None))]
fn execute(&self, cwd: Option<&str>, count: Option<u32>) -> PyResult<String> {
let tool = openjarvis_tools::builtin::git_tools::GitLogTool;
let mut params = serde_json::json!({});
if let Some(cwd) = cwd {
params["cwd"] = serde_json::Value::String(cwd.to_string());
}
if let Some(count) = count {
params["count"] = serde_json::Value::Number(count.into());
}
let result = tool
.execute(&params)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(result.content)
}
}
@@ -0,0 +1,92 @@
//! PyO3 bindings for trace types.
use pyo3::prelude::*;
use std::sync::Arc;
#[pyclass(name = "TraceStore")]
pub struct PyTraceStore {
pub inner: Arc<openjarvis_traces::TraceStore>,
}
#[pymethods]
impl PyTraceStore {
#[new]
#[pyo3(signature = (path=None))]
fn new(path: Option<&str>) -> PyResult<Self> {
let inner = match path {
Some(p) => openjarvis_traces::TraceStore::new(std::path::Path::new(p)),
None => openjarvis_traces::TraceStore::in_memory(),
}
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(Self {
inner: Arc::new(inner),
})
}
fn count(&self) -> PyResult<usize> {
self.inner
.count()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
}
}
#[pyclass(name = "TraceCollector")]
pub struct PyTraceCollector {
inner: openjarvis_traces::TraceCollector,
}
#[pymethods]
impl PyTraceCollector {
#[new]
fn new(store: &PyTraceStore) -> Self {
Self {
inner: openjarvis_traces::TraceCollector::new(Arc::clone(&store.inner)),
}
}
fn active_count(&self) -> usize {
self.inner.active_count()
}
}
/// TraceAnalyzer wraps stats computation over a TraceStore.
/// Since the Rust TraceAnalyzer has a lifetime parameter, we own the store
/// and create the analyzer on each call.
#[pyclass(name = "TraceAnalyzer")]
pub struct PyTraceAnalyzer {
store: Arc<openjarvis_traces::TraceStore>,
}
#[pymethods]
impl PyTraceAnalyzer {
#[new]
fn new(store: &PyTraceStore) -> Self {
Self {
store: Arc::clone(&store.inner),
}
}
fn stats(&self) -> PyResult<String> {
let analyzer = openjarvis_traces::TraceAnalyzer::new(&self.store);
let stats = analyzer
.overall_stats()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(serde_json::to_string(&stats).unwrap_or_default())
}
fn stats_by_agent(&self) -> PyResult<String> {
let analyzer = openjarvis_traces::TraceAnalyzer::new(&self.store);
let stats = analyzer
.stats_by_agent()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(serde_json::to_string(&stats).unwrap_or_default())
}
fn stats_by_model(&self) -> PyResult<String> {
let analyzer = openjarvis_traces::TraceAnalyzer::new(&self.store);
let stats = analyzer
.stats_by_model()
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
Ok(serde_json::to_string(&stats).unwrap_or_default())
}
}