mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-07-31 03:12:16 +00:00
init commit
This commit is contained in:
Jon Saad-Falcon
@@ -0,0 +1,111 @@
|
||||
//! PyO3 bindings for A2A (Agent-to-Agent) protocol types and task store.
|
||||
|
||||
use pyo3::prelude::*;
|
||||
|
||||
#[pyclass(name = "AgentCard")]
|
||||
pub struct PyAgentCard {
|
||||
inner: openjarvis_a2a::AgentCard,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyAgentCard {
|
||||
#[new]
|
||||
#[pyo3(signature = (name, description, version, url))]
|
||||
fn new(name: &str, description: &str, version: &str, url: &str) -> Self {
|
||||
Self {
|
||||
inner: openjarvis_a2a::AgentCard::new(name, description, version, url),
|
||||
}
|
||||
}
|
||||
|
||||
fn with_skills(&mut self, skills: Vec<String>) {
|
||||
let refs: Vec<&str> = skills.iter().map(|s| s.as_str()).collect();
|
||||
self.inner = self.inner.clone().with_skills(&refs);
|
||||
}
|
||||
|
||||
fn with_modes(&mut self, modes: Vec<String>) {
|
||||
let refs: Vec<&str> = modes.iter().map(|s| s.as_str()).collect();
|
||||
self.inner = self.inner.clone().with_modes(&refs);
|
||||
}
|
||||
|
||||
fn to_json(&self) -> String {
|
||||
serde_json::to_string(&self.inner).unwrap_or_default()
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn name(&self) -> &str {
|
||||
&self.inner.name
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn description(&self) -> &str {
|
||||
&self.inner.description
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn version(&self) -> &str {
|
||||
&self.inner.version
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn url(&self) -> &str {
|
||||
&self.inner.url
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn skills(&self) -> Vec<String> {
|
||||
self.inner.skills.clone()
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass(name = "A2ATaskStore")]
|
||||
pub struct PyA2ATaskStore {
|
||||
inner: openjarvis_a2a::A2ATaskStore,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyA2ATaskStore {
|
||||
#[new]
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
inner: openjarvis_a2a::A2ATaskStore::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn create_task(&mut self, input: &str) -> String {
|
||||
let task = self.inner.create_task(input);
|
||||
serde_json::to_string(&task).unwrap_or_default()
|
||||
}
|
||||
|
||||
fn get_task(&self, id: &str) -> Option<String> {
|
||||
self.inner
|
||||
.get_task(id)
|
||||
.map(|t| serde_json::to_string(t).unwrap_or_default())
|
||||
}
|
||||
|
||||
fn update_state(&mut self, id: &str, state: &str) -> bool {
|
||||
let s = match state {
|
||||
"pending" => openjarvis_a2a::TaskState::Pending,
|
||||
"active" => openjarvis_a2a::TaskState::Active,
|
||||
"completed" => openjarvis_a2a::TaskState::Completed,
|
||||
"cancelled" => openjarvis_a2a::TaskState::Cancelled,
|
||||
"failed" => openjarvis_a2a::TaskState::Failed,
|
||||
_ => return false,
|
||||
};
|
||||
self.inner.update_state(id, s)
|
||||
}
|
||||
|
||||
fn set_output(&mut self, id: &str, output: &str) -> bool {
|
||||
self.inner.set_output(id, output)
|
||||
}
|
||||
|
||||
fn list_tasks(&self) -> String {
|
||||
serde_json::to_string(self.inner.list_tasks()).unwrap_or_default()
|
||||
}
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
pub fn parse_a2a_request(json_str: &str) -> PyResult<String> {
|
||||
let req = openjarvis_a2a::parse_request(json_str)
|
||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e))?;
|
||||
Ok(serde_json::to_string(&req).unwrap_or_default())
|
||||
}
|
||||
@@ -0,0 +1,378 @@
|
||||
//! PyO3 bindings for agent types.
|
||||
//!
|
||||
//! Uses `AgentEnum` for static dispatch instead of `Box<dyn OjAgent>`.
|
||||
|
||||
use crate::core::PyAgentResult;
|
||||
use crate::RUNTIME;
|
||||
use openjarvis_agents::OjAgent;
|
||||
use openjarvis_engine::rig_adapter::RigModelAdapter;
|
||||
use openjarvis_engine::Engine;
|
||||
use pyo3::prelude::*;
|
||||
use std::sync::Arc;
|
||||
|
||||
type DefaultAdapter = RigModelAdapter<Engine>;
|
||||
|
||||
enum AgentEnum {
|
||||
Simple(openjarvis_agents::SimpleAgent<DefaultAdapter>),
|
||||
Orchestrator(openjarvis_agents::OrchestratorAgent<DefaultAdapter>),
|
||||
NativeReAct(openjarvis_agents::NativeReActAgent<DefaultAdapter>),
|
||||
}
|
||||
|
||||
impl AgentEnum {
|
||||
fn agent_id(&self) -> &str {
|
||||
match self {
|
||||
AgentEnum::Simple(a) => a.agent_id(),
|
||||
AgentEnum::Orchestrator(a) => a.agent_id(),
|
||||
AgentEnum::NativeReAct(a) => a.agent_id(),
|
||||
}
|
||||
}
|
||||
|
||||
fn accepts_tools(&self) -> bool {
|
||||
match self {
|
||||
AgentEnum::Simple(a) => a.accepts_tools(),
|
||||
AgentEnum::Orchestrator(a) => a.accepts_tools(),
|
||||
AgentEnum::NativeReAct(a) => a.accepts_tools(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn run(
|
||||
&self,
|
||||
input: &str,
|
||||
context: Option<&openjarvis_core::AgentContext>,
|
||||
) -> Result<openjarvis_core::AgentResult, openjarvis_core::OpenJarvisError> {
|
||||
match self {
|
||||
AgentEnum::Simple(a) => a.run(input, context).await,
|
||||
AgentEnum::Orchestrator(a) => a.run(input, context).await,
|
||||
AgentEnum::NativeReAct(a) => a.run(input, context).await,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn make_adapter(engine_key: &str, model: &str) -> PyResult<DefaultAdapter> {
|
||||
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()))?;
|
||||
Ok(RigModelAdapter::new(Arc::new(engine), model.to_string()))
|
||||
}
|
||||
|
||||
#[pyclass(name = "SimpleAgent")]
|
||||
pub struct PySimpleAgent {
|
||||
inner: AgentEnum,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PySimpleAgent {
|
||||
#[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 adapter = make_adapter(engine_key, model)?;
|
||||
let agent = openjarvis_agents::SimpleAgent::new(adapter, system_prompt, temperature);
|
||||
Ok(Self { inner: AgentEnum::Simple(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,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass(name = "OrchestratorAgent")]
|
||||
pub struct PyOrchestratorAgent {
|
||||
inner: AgentEnum,
|
||||
}
|
||||
|
||||
#[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 adapter = make_adapter(engine_key, model)?;
|
||||
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: AgentEnum::Orchestrator(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,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass(name = "NativeReActAgent")]
|
||||
pub struct PyNativeReActAgent {
|
||||
inner: AgentEnum,
|
||||
}
|
||||
|
||||
#[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 adapter = make_adapter(engine_key, model)?;
|
||||
let executor = Arc::new(openjarvis_tools::ToolExecutor::new(None, None));
|
||||
let agent = openjarvis_agents::NativeReActAgent::new(
|
||||
adapter, executor, max_turns, temperature,
|
||||
);
|
||||
Ok(Self { inner: AgentEnum::NativeReAct(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 NativeOpenHandsAgent.
|
||||
#[pyclass(name = "NativeOpenHandsAgent")]
|
||||
pub struct PyNativeOpenHandsAgent {
|
||||
inner: Box<dyn OjAgent>,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyNativeOpenHandsAgent {
|
||||
#[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::NativeOpenHandsAgent::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 MonitorOperativeAgent.
|
||||
#[pyclass(name = "MonitorOperativeAgent")]
|
||||
pub struct PyMonitorOperativeAgent {
|
||||
inner: Box<dyn OjAgent>,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyMonitorOperativeAgent {
|
||||
/// Create a MonitorOperativeAgent.
|
||||
///
|
||||
/// Strategy parameters are strings:
|
||||
/// - `memory_extraction`: "causality_graph" | "scratchpad" | "structured_json" | "none"
|
||||
/// - `observation_compression`: "summarize" | "truncate" | "none"
|
||||
/// - `retrieval_strategy`: "hybrid_with_self_eval" | "keyword" | "semantic" | "none"
|
||||
/// - `task_decomposition`: "phased" | "monolithic" | "hierarchical"
|
||||
#[new]
|
||||
#[pyo3(signature = (
|
||||
engine_key="ollama",
|
||||
host="http://localhost:11434",
|
||||
model="qwen3:8b",
|
||||
max_turns=10,
|
||||
temperature=0.7,
|
||||
memory_extraction="causality_graph",
|
||||
observation_compression="summarize",
|
||||
retrieval_strategy="hybrid_with_self_eval",
|
||||
task_decomposition="phased",
|
||||
compression_threshold=2000,
|
||||
truncation_limit=2000
|
||||
))]
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn new(
|
||||
engine_key: &str,
|
||||
host: &str,
|
||||
model: &str,
|
||||
max_turns: usize,
|
||||
temperature: f64,
|
||||
memory_extraction: &str,
|
||||
observation_compression: &str,
|
||||
retrieval_strategy: &str,
|
||||
task_decomposition: &str,
|
||||
compression_threshold: usize,
|
||||
truncation_limit: usize,
|
||||
) -> 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 mem_ext = match memory_extraction {
|
||||
"scratchpad" => openjarvis_agents::MemoryExtraction::Scratchpad,
|
||||
"structured_json" => openjarvis_agents::MemoryExtraction::StructuredJson,
|
||||
"none" => openjarvis_agents::MemoryExtraction::None,
|
||||
_ => openjarvis_agents::MemoryExtraction::CausalityGraph,
|
||||
};
|
||||
let obs_comp = match observation_compression {
|
||||
"truncate" => openjarvis_agents::ObservationCompression::Truncate,
|
||||
"none" => openjarvis_agents::ObservationCompression::None,
|
||||
_ => openjarvis_agents::ObservationCompression::Summarize,
|
||||
};
|
||||
let ret_strat = match retrieval_strategy {
|
||||
"keyword" => openjarvis_agents::RetrievalStrategy::Keyword,
|
||||
"semantic" => openjarvis_agents::RetrievalStrategy::Semantic,
|
||||
"none" => openjarvis_agents::RetrievalStrategy::None,
|
||||
_ => openjarvis_agents::RetrievalStrategy::HybridWithSelfEval,
|
||||
};
|
||||
let task_dec = match task_decomposition {
|
||||
"monolithic" => openjarvis_agents::TaskDecomposition::Monolithic,
|
||||
"hierarchical" => openjarvis_agents::TaskDecomposition::Hierarchical,
|
||||
_ => openjarvis_agents::TaskDecomposition::Phased,
|
||||
};
|
||||
|
||||
let monitor_config = openjarvis_agents::MonitorConfig {
|
||||
memory_extraction: mem_ext,
|
||||
observation_compression: obs_comp,
|
||||
retrieval_strategy: ret_strat,
|
||||
task_decomposition: task_dec,
|
||||
compression_threshold,
|
||||
truncation_limit,
|
||||
};
|
||||
|
||||
let agent = openjarvis_agents::MonitorOperativeAgent::new(
|
||||
adapter,
|
||||
executor,
|
||||
max_turns,
|
||||
temperature,
|
||||
monitor_config,
|
||||
);
|
||||
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()
|
||||
}
|
||||
}
|
||||
@@ -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)])
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,189 @@
|
||||
//! 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", "exo", "nexa", "uzu", "apple_fm").
|
||||
#[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"),
|
||||
),
|
||||
),
|
||||
"exo" => openjarvis_engine::Engine::Exo(
|
||||
openjarvis_engine::OpenAICompatEngine::exo(
|
||||
host.unwrap_or("http://localhost:52415"),
|
||||
),
|
||||
),
|
||||
"nexa" => openjarvis_engine::Engine::Nexa(
|
||||
openjarvis_engine::OpenAICompatEngine::nexa(
|
||||
host.unwrap_or("http://localhost:18181"),
|
||||
),
|
||||
),
|
||||
"uzu" => openjarvis_engine::Engine::Uzu(
|
||||
openjarvis_engine::OpenAICompatEngine::uzu(
|
||||
host.unwrap_or("http://localhost:8080"),
|
||||
),
|
||||
),
|
||||
"apple_fm" => openjarvis_engine::Engine::AppleFm(
|
||||
openjarvis_engine::OpenAICompatEngine::apple_fm(
|
||||
host.unwrap_or("http://localhost:8079"),
|
||||
),
|
||||
),
|
||||
"vllm_native" => openjarvis_engine::Engine::VLLM(
|
||||
openjarvis_engine::VLLMEngine::new(
|
||||
host.unwrap_or("http://localhost"),
|
||||
8000,
|
||||
None,
|
||||
120.0,
|
||||
),
|
||||
),
|
||||
"sglang_native" => openjarvis_engine::Engine::SGLang(
|
||||
openjarvis_engine::SGLangEngine::new(
|
||||
host.unwrap_or("http://localhost"),
|
||||
30000,
|
||||
120.0,
|
||||
),
|
||||
),
|
||||
"llamacpp_native" => openjarvis_engine::Engine::LlamaCppNative(
|
||||
openjarvis_engine::LlamaCppEngine::new(
|
||||
host.unwrap_or("http://localhost"),
|
||||
8080,
|
||||
120.0,
|
||||
),
|
||||
),
|
||||
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,573 @@
|
||||
//! PyO3 bindings for learning/router policy types and optimization.
|
||||
|
||||
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(())
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Optimization store
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[pyclass(name = "OptimizationStore")]
|
||||
pub struct PyOptimizationStore {
|
||||
inner: parking_lot::Mutex<openjarvis_learning::optimize::OptimizationStore>,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyOptimizationStore {
|
||||
/// Open or create a store at the given database path.
|
||||
/// Use `":memory:"` for an in-memory store.
|
||||
#[new]
|
||||
#[pyo3(signature = (path=":memory:"))]
|
||||
fn new(path: &str) -> PyResult<Self> {
|
||||
let store = if path == ":memory:" {
|
||||
openjarvis_learning::optimize::OptimizationStore::in_memory()
|
||||
} else {
|
||||
openjarvis_learning::optimize::OptimizationStore::open(path)
|
||||
};
|
||||
let store =
|
||||
store.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
|
||||
Ok(Self {
|
||||
inner: parking_lot::Mutex::new(store),
|
||||
})
|
||||
}
|
||||
|
||||
/// Save a trial result (JSON string) for a given run_id.
|
||||
fn save_trial(&self, run_id: &str, trial_json: &str) -> PyResult<()> {
|
||||
let trial: openjarvis_learning::optimize::TrialResult =
|
||||
serde_json::from_str(trial_json)
|
||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
|
||||
self.inner
|
||||
.lock()
|
||||
.save_trial(run_id, &trial)
|
||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
|
||||
}
|
||||
|
||||
/// Retrieve an optimization run by id. Returns JSON string or None.
|
||||
fn get_run(&self, run_id: &str) -> PyResult<Option<String>> {
|
||||
let run = self
|
||||
.inner
|
||||
.lock()
|
||||
.get_run(run_id)
|
||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
|
||||
match run {
|
||||
Some(r) => Ok(Some(
|
||||
serde_json::to_string(&r).unwrap_or_else(|_| "{}".into()),
|
||||
)),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
/// List recent optimization runs. Returns JSON string.
|
||||
#[pyo3(signature = (limit=50))]
|
||||
fn list_runs(&self, limit: usize) -> PyResult<String> {
|
||||
let summaries = self
|
||||
.inner
|
||||
.lock()
|
||||
.list_runs(limit)
|
||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
|
||||
// RunSummary doesn't derive Serialize, so manually build JSON.
|
||||
let items: Vec<serde_json::Value> = summaries
|
||||
.iter()
|
||||
.map(|s| {
|
||||
serde_json::json!({
|
||||
"run_id": s.run_id,
|
||||
"status": s.status,
|
||||
"optimizer_model": s.optimizer_model,
|
||||
"benchmark": s.benchmark,
|
||||
"best_trial_id": s.best_trial_id,
|
||||
"best_recipe_path": s.best_recipe_path,
|
||||
"created_at": s.created_at,
|
||||
"updated_at": s.updated_at,
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
Ok(serde_json::to_string(&items).unwrap_or_else(|_| "[]".into()))
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// LLM Optimizer
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[pyclass(name = "LLMOptimizer")]
|
||||
pub struct PyLLMOptimizer {
|
||||
inner: openjarvis_learning::optimize::LLMOptimizer,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyLLMOptimizer {
|
||||
/// Create a new LLM optimizer.
|
||||
///
|
||||
/// `search_space_json` is a JSON string representing the SearchSpace.
|
||||
/// `optimizer_model` is the model name to use for optimization proposals.
|
||||
#[new]
|
||||
#[pyo3(signature = (search_space_json, optimizer_model="gpt-4o"))]
|
||||
fn new(search_space_json: &str, optimizer_model: &str) -> PyResult<Self> {
|
||||
let search_space: openjarvis_learning::optimize::SearchSpace =
|
||||
serde_json::from_str(search_space_json)
|
||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
|
||||
let inner = openjarvis_learning::optimize::LLMOptimizer::new(
|
||||
search_space,
|
||||
optimizer_model.to_string(),
|
||||
);
|
||||
Ok(Self { inner })
|
||||
}
|
||||
|
||||
/// Propose an initial configuration. Returns JSON string.
|
||||
fn propose_initial(&self) -> String {
|
||||
let config = self.inner.propose_initial();
|
||||
serde_json::to_string(&config).unwrap_or_else(|_| "{}".into())
|
||||
}
|
||||
|
||||
/// Propose the next configuration based on trial history (JSON string).
|
||||
/// Returns JSON string.
|
||||
fn propose_next(&self, history_json: &str) -> PyResult<String> {
|
||||
let history: Vec<openjarvis_learning::optimize::TrialResult> =
|
||||
serde_json::from_str(history_json)
|
||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
|
||||
let config = self.inner.propose_next(&history);
|
||||
Ok(serde_json::to_string(&config).unwrap_or_else(|_| "{}".into()))
|
||||
}
|
||||
}
|
||||
|
||||
// --- SFT Router Policy ---
|
||||
|
||||
#[pyclass(name = "SFTRouterPolicy")]
|
||||
pub struct PySFTRouterPolicy {
|
||||
inner: openjarvis_learning::SFTRouterPolicy,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PySFTRouterPolicy {
|
||||
#[new]
|
||||
#[pyo3(signature = (min_samples=5))]
|
||||
fn new(min_samples: usize) -> Self {
|
||||
Self {
|
||||
inner: openjarvis_learning::SFTRouterPolicy::new(min_samples),
|
||||
}
|
||||
}
|
||||
|
||||
fn policy_map(&self) -> std::collections::HashMap<String, String> {
|
||||
self.inner.policy_map()
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn classify_query(query: &str) -> &'static str {
|
||||
openjarvis_learning::SFTRouterPolicy::classify_query(query)
|
||||
}
|
||||
|
||||
fn update_from_data(&self, traces_json: &str) -> PyResult<String> {
|
||||
let traces: Vec<(String, String, String, Option<f64>)> =
|
||||
serde_json::from_str(traces_json)
|
||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
|
||||
let result = self.inner.update_from_data(&traces);
|
||||
Ok(serde_json::to_string(&result).unwrap_or_default())
|
||||
}
|
||||
}
|
||||
|
||||
// --- HeuristicRewardFunction ---
|
||||
|
||||
#[pyclass(name = "HeuristicRewardFunction")]
|
||||
pub struct PyHeuristicRewardFunction {
|
||||
inner: openjarvis_learning::HeuristicRewardFunction,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyHeuristicRewardFunction {
|
||||
#[new]
|
||||
#[pyo3(signature = (weight_latency=0.4, weight_cost=0.3, weight_efficiency=0.3, max_latency=30.0, max_cost=0.01))]
|
||||
fn new(
|
||||
weight_latency: f64,
|
||||
weight_cost: f64,
|
||||
weight_efficiency: f64,
|
||||
max_latency: f64,
|
||||
max_cost: f64,
|
||||
) -> Self {
|
||||
Self {
|
||||
inner: openjarvis_learning::HeuristicRewardFunction::new(
|
||||
weight_latency,
|
||||
weight_cost,
|
||||
weight_efficiency,
|
||||
max_latency,
|
||||
max_cost,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
fn compute(
|
||||
&self,
|
||||
latency_seconds: f64,
|
||||
cost_usd: f64,
|
||||
prompt_tokens: u64,
|
||||
completion_tokens: u64,
|
||||
) -> f64 {
|
||||
self.inner
|
||||
.compute(latency_seconds, cost_usd, prompt_tokens, completion_tokens)
|
||||
}
|
||||
}
|
||||
|
||||
// --- SkillDiscovery ---
|
||||
|
||||
#[pyclass(name = "SkillDiscovery")]
|
||||
pub struct PySkillDiscovery {
|
||||
inner: openjarvis_learning::SkillDiscovery,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PySkillDiscovery {
|
||||
#[new]
|
||||
#[pyo3(signature = (min_frequency=3, min_sequence_length=2, max_sequence_length=4, min_outcome=0.5))]
|
||||
fn new(
|
||||
min_frequency: usize,
|
||||
min_sequence_length: usize,
|
||||
max_sequence_length: usize,
|
||||
min_outcome: f64,
|
||||
) -> Self {
|
||||
Self {
|
||||
inner: openjarvis_learning::SkillDiscovery::new(
|
||||
min_frequency,
|
||||
min_sequence_length,
|
||||
max_sequence_length,
|
||||
min_outcome,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
fn analyze(&mut self, traces_json: &str) -> PyResult<String> {
|
||||
let traces: Vec<(Vec<String>, f64, String)> = serde_json::from_str(traces_json)
|
||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
|
||||
self.inner.analyze(&traces);
|
||||
Ok(serde_json::to_string(&self.inner.to_manifests()).unwrap_or_default())
|
||||
}
|
||||
}
|
||||
|
||||
// --- TraceDrivenPolicy ---
|
||||
|
||||
#[pyclass(name = "TraceDrivenPolicy")]
|
||||
pub struct PyTraceDrivenPolicy {
|
||||
inner: openjarvis_learning::TraceDrivenPolicy,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyTraceDrivenPolicy {
|
||||
#[new]
|
||||
#[pyo3(signature = (available_models=vec![], default_model="", fallback_model=""))]
|
||||
fn new(available_models: Vec<String>, default_model: &str, fallback_model: &str) -> Self {
|
||||
Self {
|
||||
inner: openjarvis_learning::TraceDrivenPolicy::new(
|
||||
available_models,
|
||||
default_model.to_string(),
|
||||
fallback_model.to_string(),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
fn select_model(&self, query: &str) -> String {
|
||||
self.inner.select_model(query)
|
||||
}
|
||||
|
||||
fn policy_map(&self) -> std::collections::HashMap<String, String> {
|
||||
self.inner.policy_map()
|
||||
}
|
||||
|
||||
#[pyo3(signature = (query, model, outcome=None, feedback=None))]
|
||||
fn observe(&self, query: &str, model: &str, outcome: Option<String>, feedback: Option<f64>) {
|
||||
self.inner
|
||||
.observe(query, model, outcome.as_deref(), feedback);
|
||||
}
|
||||
}
|
||||
|
||||
// --- AgentAdvisorPolicy ---
|
||||
|
||||
#[pyclass(name = "AgentAdvisorPolicy")]
|
||||
pub struct PyAgentAdvisorPolicy {
|
||||
inner: openjarvis_learning::AgentAdvisorPolicy,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyAgentAdvisorPolicy {
|
||||
#[new]
|
||||
#[pyo3(signature = (max_traces=50))]
|
||||
fn new(max_traces: usize) -> Self {
|
||||
Self {
|
||||
inner: openjarvis_learning::AgentAdvisorPolicy::new(max_traces),
|
||||
}
|
||||
}
|
||||
|
||||
fn analyze_patterns(&self, traces_json: &str) -> PyResult<String> {
|
||||
let traces: Vec<openjarvis_learning::TraceInfo> = serde_json::from_str(traces_json)
|
||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
|
||||
let recs = self.inner.analyze_patterns(&traces);
|
||||
Ok(serde_json::to_string(&recs).unwrap_or_default())
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn classify(query: &str) -> &'static str {
|
||||
openjarvis_learning::AgentAdvisorPolicy::classify(query)
|
||||
}
|
||||
}
|
||||
|
||||
// --- ICLUpdaterPolicy ---
|
||||
|
||||
#[pyclass(name = "ICLUpdaterPolicy")]
|
||||
pub struct PyICLUpdaterPolicy {
|
||||
inner: openjarvis_learning::ICLUpdaterPolicy,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyICLUpdaterPolicy {
|
||||
#[new]
|
||||
#[pyo3(signature = (min_score=0.7, max_examples=20, min_skill_occurrences=3))]
|
||||
fn new(min_score: f64, max_examples: usize, min_skill_occurrences: usize) -> Self {
|
||||
Self {
|
||||
inner: openjarvis_learning::ICLUpdaterPolicy::new(
|
||||
min_score,
|
||||
max_examples,
|
||||
min_skill_occurrences,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
fn add_example(&mut self, query: &str, response: &str, outcome: f64) -> bool {
|
||||
self.inner.add_example(
|
||||
query.to_string(),
|
||||
response.to_string(),
|
||||
outcome,
|
||||
std::collections::HashMap::new(),
|
||||
)
|
||||
}
|
||||
|
||||
fn rollback(&mut self, version: u32) {
|
||||
self.inner.rollback(version);
|
||||
}
|
||||
|
||||
fn get_examples(&self, query_class: &str, top_k: usize) -> String {
|
||||
serde_json::to_string(&self.inner.get_examples(query_class, top_k)).unwrap_or_default()
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn version(&self) -> u32 {
|
||||
self.inner.version()
|
||||
}
|
||||
}
|
||||
|
||||
// --- TrainingDataMiner ---
|
||||
|
||||
#[pyclass(name = "TrainingDataMiner")]
|
||||
pub struct PyTrainingDataMiner {
|
||||
inner: openjarvis_learning::TrainingDataMiner,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyTrainingDataMiner {
|
||||
#[new]
|
||||
#[pyo3(signature = (min_quality=0.7, min_samples_per_class=1))]
|
||||
fn new(min_quality: f64, min_samples_per_class: usize) -> Self {
|
||||
Self {
|
||||
inner: openjarvis_learning::TrainingDataMiner::new(min_quality, min_samples_per_class),
|
||||
}
|
||||
}
|
||||
|
||||
fn extract_sft_pairs(&self, traces_json: &str) -> PyResult<String> {
|
||||
let traces: Vec<openjarvis_learning::MinerTraceData> = serde_json::from_str(traces_json)
|
||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
|
||||
let pairs = self.inner.extract_sft_pairs(&traces);
|
||||
Ok(serde_json::to_string(&pairs).unwrap_or_default())
|
||||
}
|
||||
|
||||
fn extract_routing_pairs(&self, traces_json: &str) -> PyResult<String> {
|
||||
let traces: Vec<openjarvis_learning::MinerTraceData> = serde_json::from_str(traces_json)
|
||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
|
||||
let pairs = self.inner.extract_routing_pairs(&traces);
|
||||
Ok(serde_json::to_string(&pairs).unwrap_or_default())
|
||||
}
|
||||
}
|
||||
|
||||
// --- AgentConfigEvolver ---
|
||||
|
||||
#[pyclass(name = "AgentConfigEvolver")]
|
||||
pub struct PyAgentConfigEvolver {
|
||||
inner: openjarvis_learning::AgentConfigEvolver,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyAgentConfigEvolver {
|
||||
#[new]
|
||||
#[pyo3(signature = (min_quality=0.5))]
|
||||
fn new(min_quality: f64) -> Self {
|
||||
Self {
|
||||
inner: openjarvis_learning::AgentConfigEvolver::new(min_quality),
|
||||
}
|
||||
}
|
||||
|
||||
fn analyze(&self, traces_json: &str) -> PyResult<String> {
|
||||
let traces: Vec<openjarvis_learning::EvolutionTraceData> =
|
||||
serde_json::from_str(traces_json)
|
||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
|
||||
let recs = self.inner.analyze(&traces);
|
||||
Ok(serde_json::to_string(&recs).unwrap_or_default())
|
||||
}
|
||||
}
|
||||
|
||||
// --- MultiObjectiveReward ---
|
||||
|
||||
#[pyclass(name = "MultiObjectiveReward")]
|
||||
pub struct PyMultiObjectiveReward {
|
||||
inner: openjarvis_learning::MultiObjectiveReward,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyMultiObjectiveReward {
|
||||
#[new]
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
inner: openjarvis_learning::MultiObjectiveReward::new(
|
||||
openjarvis_learning::RewardWeights::default(),
|
||||
openjarvis_learning::Normalizers::default(),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
fn compute(&self, episode_json: &str) -> PyResult<f64> {
|
||||
let ep: openjarvis_learning::Episode = serde_json::from_str(episode_json)
|
||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))?;
|
||||
Ok(self.inner.compute(&ep))
|
||||
}
|
||||
}
|
||||
|
||||
// --- LearningOrchestrator ---
|
||||
|
||||
#[pyclass(name = "LearningOrchestrator")]
|
||||
pub struct PyLearningOrchestrator {
|
||||
inner: openjarvis_learning::LearningOrchestrator,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyLearningOrchestrator {
|
||||
#[new]
|
||||
#[pyo3(signature = (min_improvement=0.02, min_sft_pairs=10, min_quality=0.7))]
|
||||
fn new(min_improvement: f64, min_sft_pairs: usize, min_quality: f64) -> Self {
|
||||
Self {
|
||||
inner: openjarvis_learning::LearningOrchestrator::new(
|
||||
min_improvement,
|
||||
min_sft_pairs,
|
||||
min_quality,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
#[pyo3(signature = (sft_pairs, routing_count, agent_count, recs_count, baseline=None, post=None))]
|
||||
fn evaluate_cycle(
|
||||
&self,
|
||||
sft_pairs: usize,
|
||||
routing_count: usize,
|
||||
agent_count: usize,
|
||||
recs_count: usize,
|
||||
baseline: Option<f64>,
|
||||
post: Option<f64>,
|
||||
) -> String {
|
||||
let result = self.inner.evaluate_cycle(
|
||||
sft_pairs,
|
||||
routing_count,
|
||||
agent_count,
|
||||
recs_count,
|
||||
baseline,
|
||||
post,
|
||||
);
|
||||
serde_json::to_string(&result).unwrap_or_default()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,195 @@
|
||||
//! PyO3 bridge — exposes ~50 Rust classes to Python via `openjarvis_rust`.
|
||||
#![allow(clippy::redundant_closure, unused_variables)]
|
||||
|
||||
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 a2a;
|
||||
pub mod agents;
|
||||
pub mod core;
|
||||
pub mod engine;
|
||||
pub mod learning;
|
||||
pub mod mcp;
|
||||
pub mod recipes;
|
||||
pub mod scheduler;
|
||||
pub mod security;
|
||||
pub mod sessions;
|
||||
pub mod skills;
|
||||
pub mod storage;
|
||||
pub mod telemetry;
|
||||
pub mod templates;
|
||||
pub mod tools;
|
||||
pub mod traces;
|
||||
pub mod workflow;
|
||||
|
||||
// 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))
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
fn register_builtin_models() {
|
||||
openjarvis_core::model_catalog::register_builtin_models();
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
fn classify_query(query: &str) -> &'static str {
|
||||
openjarvis_learning::classify_query(query)
|
||||
}
|
||||
|
||||
#[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::PyNativeOpenHandsAgent>()?;
|
||||
m.add_class::<agents::PyMonitorOperativeAgent>()?;
|
||||
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::PyFAISSMemory>()?;
|
||||
m.add_class::<storage::PyColBERTMemory>()?;
|
||||
m.add_class::<storage::PyHybridMemory>()?;
|
||||
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>()?;
|
||||
m.add_class::<learning::PyOptimizationStore>()?;
|
||||
m.add_class::<learning::PyLLMOptimizer>()?;
|
||||
m.add_class::<learning::PySFTRouterPolicy>()?;
|
||||
m.add_class::<learning::PyHeuristicRewardFunction>()?;
|
||||
m.add_class::<learning::PySkillDiscovery>()?;
|
||||
m.add_class::<learning::PyTraceDrivenPolicy>()?;
|
||||
m.add_class::<learning::PyAgentAdvisorPolicy>()?;
|
||||
m.add_class::<learning::PyICLUpdaterPolicy>()?;
|
||||
m.add_class::<learning::PyTrainingDataMiner>()?;
|
||||
m.add_class::<learning::PyAgentConfigEvolver>()?;
|
||||
m.add_class::<learning::PyMultiObjectiveReward>()?;
|
||||
m.add_class::<learning::PyLearningOrchestrator>()?;
|
||||
|
||||
// --- MCP ---
|
||||
m.add_class::<mcp::PyMcpServer>()?;
|
||||
|
||||
// --- Sessions ---
|
||||
m.add_class::<sessions::PySessionStore>()?;
|
||||
|
||||
// --- Workflow ---
|
||||
m.add_class::<workflow::PyWorkflowGraph>()?;
|
||||
m.add_class::<workflow::PyWorkflowBuilder>()?;
|
||||
|
||||
// --- Skills ---
|
||||
m.add_class::<skills::PySkillManifest>()?;
|
||||
|
||||
// --- Recipes ---
|
||||
m.add_class::<recipes::PyRecipe>()?;
|
||||
|
||||
// --- Templates ---
|
||||
m.add_class::<templates::PyAgentTemplate>()?;
|
||||
|
||||
// --- A2A ---
|
||||
m.add_class::<a2a::PyAgentCard>()?;
|
||||
m.add_class::<a2a::PyA2ATaskStore>()?;
|
||||
|
||||
// --- Scheduler ---
|
||||
m.add_class::<scheduler::PySchedulerStore>()?;
|
||||
|
||||
// --- 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)?)?;
|
||||
m.add_function(wrap_pyfunction!(register_builtin_models, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(classify_query, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(skills::load_skill, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(recipes::load_recipe, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(templates::load_template, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(a2a::parse_a2a_request, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(scheduler::parse_cron_next, m)?)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
@@ -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,67 @@
|
||||
//! PyO3 bindings for composable recipes.
|
||||
|
||||
use pyo3::prelude::*;
|
||||
|
||||
#[pyclass(name = "Recipe")]
|
||||
pub struct PyRecipe {
|
||||
inner: openjarvis_recipes::Recipe,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyRecipe {
|
||||
#[getter]
|
||||
fn name(&self) -> &str {
|
||||
&self.inner.name
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn description(&self) -> Option<&str> {
|
||||
self.inner.description.as_deref()
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn model(&self) -> Option<&str> {
|
||||
self.inner.model.as_deref()
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn engine_key(&self) -> Option<&str> {
|
||||
self.inner.engine_key.as_deref()
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn agent_type(&self) -> Option<&str> {
|
||||
self.inner.agent_type.as_deref()
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn max_turns(&self) -> Option<usize> {
|
||||
self.inner.max_turns
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn temperature(&self) -> Option<f64> {
|
||||
self.inner.temperature
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn tools(&self) -> Option<Vec<String>> {
|
||||
self.inner.tools.clone()
|
||||
}
|
||||
|
||||
fn to_builder_kwargs(&self) -> String {
|
||||
let kwargs = self.inner.to_builder_kwargs();
|
||||
serde_json::to_string(&kwargs).unwrap_or_default()
|
||||
}
|
||||
|
||||
fn to_json(&self) -> String {
|
||||
serde_json::to_string(&self.inner).unwrap_or_default()
|
||||
}
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
pub fn load_recipe(toml_str: &str) -> PyResult<PyRecipe> {
|
||||
let recipe = openjarvis_recipes::load_recipe(toml_str)
|
||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e))?;
|
||||
Ok(PyRecipe { inner: recipe })
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
//! PyO3 bindings for task scheduler.
|
||||
|
||||
use pyo3::prelude::*;
|
||||
|
||||
#[pyclass(name = "SchedulerStore", unsendable)]
|
||||
pub struct PySchedulerStore {
|
||||
inner: openjarvis_scheduler::SchedulerStore,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PySchedulerStore {
|
||||
#[new]
|
||||
#[pyo3(signature = (db_path=":memory:"))]
|
||||
fn new(db_path: &str) -> Self {
|
||||
Self {
|
||||
inner: openjarvis_scheduler::SchedulerStore::new(db_path),
|
||||
}
|
||||
}
|
||||
|
||||
fn create_task(&self, name: &str, schedule_type: &str, schedule_value: &str) -> PyResult<String> {
|
||||
let st = openjarvis_scheduler::ScheduleType::parse(schedule_type).ok_or_else(|| {
|
||||
PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
|
||||
"invalid schedule_type '{}', expected cron/interval/once",
|
||||
schedule_type
|
||||
))
|
||||
})?;
|
||||
let task = self.inner.create_task(name, st, schedule_value);
|
||||
Ok(serde_json::to_string(&task).unwrap_or_default())
|
||||
}
|
||||
|
||||
fn get_task(&self, id: &str) -> Option<String> {
|
||||
self.inner
|
||||
.get_task(id)
|
||||
.map(|t| serde_json::to_string(&t).unwrap_or_default())
|
||||
}
|
||||
|
||||
fn list_tasks(&self) -> String {
|
||||
serde_json::to_string(&self.inner.list_tasks()).unwrap_or_default()
|
||||
}
|
||||
|
||||
fn update_status(&self, id: &str, status: &str) -> PyResult<bool> {
|
||||
let s = openjarvis_scheduler::TaskStatus::parse(status).ok_or_else(|| {
|
||||
PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
|
||||
"invalid status '{}', expected active/paused/cancelled/completed",
|
||||
status
|
||||
))
|
||||
})?;
|
||||
Ok(self.inner.update_status(id, s))
|
||||
}
|
||||
|
||||
fn record_run(&self, id: &str, timestamp: f64) -> bool {
|
||||
self.inner.record_run(id, timestamp)
|
||||
}
|
||||
|
||||
fn delete_task(&self, id: &str) -> bool {
|
||||
self.inner.delete_task(id)
|
||||
}
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
pub fn parse_cron_next(expr: &str, after: f64) -> Option<f64> {
|
||||
openjarvis_scheduler::parse_cron_next(expr, after)
|
||||
}
|
||||
@@ -0,0 +1,258 @@
|
||||
//! PyO3 bindings for security types.
|
||||
|
||||
use pyo3::prelude::*;
|
||||
|
||||
#[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);
|
||||
serde_json::to_string(&result)
|
||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.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)
|
||||
}
|
||||
|
||||
#[pyo3(signature = (key=None))]
|
||||
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,66 @@
|
||||
//! PyO3 bindings for session management.
|
||||
|
||||
use pyo3::prelude::*;
|
||||
|
||||
#[pyclass(name = "SessionStore", unsendable)]
|
||||
pub struct PySessionStore {
|
||||
inner: openjarvis_sessions::SessionStore,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PySessionStore {
|
||||
#[new]
|
||||
#[pyo3(signature = (db_path=":memory:", max_age_hours=24.0, consolidation_threshold=100))]
|
||||
fn new(db_path: &str, max_age_hours: f64, consolidation_threshold: usize) -> Self {
|
||||
Self {
|
||||
inner: openjarvis_sessions::SessionStore::new(
|
||||
db_path,
|
||||
max_age_hours,
|
||||
consolidation_threshold,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
fn get_or_create(
|
||||
&self,
|
||||
user_id: &str,
|
||||
channel: &str,
|
||||
channel_user_id: &str,
|
||||
display_name: &str,
|
||||
) -> String {
|
||||
let session = self
|
||||
.inner
|
||||
.get_or_create(user_id, channel, channel_user_id, display_name);
|
||||
serde_json::to_string(&session).unwrap_or_default()
|
||||
}
|
||||
|
||||
fn save_message(
|
||||
&self,
|
||||
session_id: &str,
|
||||
role: &str,
|
||||
content: &str,
|
||||
channel: &str,
|
||||
) -> PyResult<()> {
|
||||
self.inner
|
||||
.save_message(session_id, role, content, channel)
|
||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
|
||||
}
|
||||
|
||||
fn consolidate(&self, session_id: &str) {
|
||||
self.inner.consolidate(session_id);
|
||||
}
|
||||
|
||||
fn link_channel(&self, session_id: &str, channel: &str, channel_user_id: &str) {
|
||||
self.inner.link_channel(session_id, channel, channel_user_id);
|
||||
}
|
||||
|
||||
fn list_sessions(&self, active_only: bool, limit: usize) -> String {
|
||||
let sessions = self.inner.list_sessions(active_only, limit);
|
||||
serde_json::to_string(&sessions).unwrap_or_default()
|
||||
}
|
||||
|
||||
#[pyo3(signature = (max_age_hours=None))]
|
||||
fn decay(&self, max_age_hours: Option<f64>) -> usize {
|
||||
self.inner.decay(max_age_hours)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
//! PyO3 bindings for skill manifests and verification.
|
||||
|
||||
use pyo3::prelude::*;
|
||||
|
||||
#[pyclass(name = "SkillManifest")]
|
||||
pub struct PySkillManifest {
|
||||
inner: openjarvis_skills::SkillManifest,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PySkillManifest {
|
||||
#[getter]
|
||||
fn name(&self) -> &str {
|
||||
&self.inner.name
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn version(&self) -> &str {
|
||||
&self.inner.version
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn description(&self) -> &str {
|
||||
&self.inner.description
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn author(&self) -> &str {
|
||||
&self.inner.author
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn steps_count(&self) -> usize {
|
||||
self.inner.steps.len()
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn required_capabilities(&self) -> Vec<String> {
|
||||
self.inner.required_capabilities.clone()
|
||||
}
|
||||
|
||||
fn to_json(&self) -> String {
|
||||
serde_json::to_string(&self.inner).unwrap_or_default()
|
||||
}
|
||||
|
||||
fn manifest_bytes(&self) -> Vec<u8> {
|
||||
self.inner.manifest_bytes()
|
||||
}
|
||||
|
||||
fn verify_signature(&self, public_key_hex: &str) -> bool {
|
||||
let key_bytes: Vec<u8> = (0..public_key_hex.len())
|
||||
.step_by(2)
|
||||
.filter_map(|i| u8::from_str_radix(&public_key_hex[i..i + 2], 16).ok())
|
||||
.collect();
|
||||
openjarvis_skills::verify_signature(&self.inner, &key_bytes)
|
||||
}
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
pub fn load_skill(toml_str: &str) -> PyResult<PySkillManifest> {
|
||||
let manifest = openjarvis_skills::load_skill(toml_str)
|
||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e))?;
|
||||
Ok(PySkillManifest { inner: manifest })
|
||||
}
|
||||
@@ -0,0 +1,356 @@
|
||||
//! 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 delete(&self, doc_id: &str) -> PyResult<bool> {
|
||||
self.inner
|
||||
.delete(doc_id)
|
||||
.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 = "FAISSMemory")]
|
||||
pub struct PyFAISSMemory {
|
||||
inner: openjarvis_tools::storage::FAISSMemory,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyFAISSMemory {
|
||||
#[new]
|
||||
#[pyo3(signature = (path=":memory:", dim=128))]
|
||||
fn new(path: &str, dim: usize) -> PyResult<Self> {
|
||||
let inner = openjarvis_tools::storage::FAISSMemory::new(std::path::Path::new(path), dim)
|
||||
.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 delete(&self, doc_id: &str) -> PyResult<bool> {
|
||||
self.inner
|
||||
.delete(doc_id)
|
||||
.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 = "ColBERTMemory")]
|
||||
pub struct PyColBERTMemory {
|
||||
inner: openjarvis_tools::storage::ColBERTMemory,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyColBERTMemory {
|
||||
#[new]
|
||||
#[pyo3(signature = (path=":memory:", token_dim=64))]
|
||||
fn new(path: &str, token_dim: usize) -> PyResult<Self> {
|
||||
let inner =
|
||||
openjarvis_tools::storage::ColBERTMemory::new(std::path::Path::new(path), token_dim)
|
||||
.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 delete(&self, doc_id: &str) -> PyResult<bool> {
|
||||
self.inner
|
||||
.delete(doc_id)
|
||||
.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 = "HybridMemory")]
|
||||
pub struct PyHybridMemory {
|
||||
inner: openjarvis_tools::storage::HybridMemory,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyHybridMemory {
|
||||
/// Create a HybridMemory combining named backends.
|
||||
///
|
||||
/// `backend_keys` is a list of strings like `["sqlite", "bm25"]`.
|
||||
/// Supported keys: "sqlite", "bm25", "faiss", "colbert".
|
||||
#[new]
|
||||
#[pyo3(signature = (backend_keys))]
|
||||
fn new(backend_keys: Vec<String>) -> PyResult<Self> {
|
||||
let mut backends: Vec<Box<dyn MemoryBackend>> = Vec::new();
|
||||
for key in &backend_keys {
|
||||
let backend: Box<dyn MemoryBackend> = match key.as_str() {
|
||||
"sqlite" => {
|
||||
let m = openjarvis_tools::storage::SQLiteMemory::new(
|
||||
std::path::Path::new(":memory:"),
|
||||
)
|
||||
.map_err(|e| {
|
||||
PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string())
|
||||
})?;
|
||||
Box::new(m)
|
||||
}
|
||||
"bm25" => Box::new(openjarvis_tools::storage::BM25Memory::default()),
|
||||
"faiss" => {
|
||||
let m = openjarvis_tools::storage::FAISSMemory::in_memory().map_err(|e| {
|
||||
PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string())
|
||||
})?;
|
||||
Box::new(m)
|
||||
}
|
||||
"colbert" => {
|
||||
let m =
|
||||
openjarvis_tools::storage::ColBERTMemory::in_memory().map_err(|e| {
|
||||
PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string())
|
||||
})?;
|
||||
Box::new(m)
|
||||
}
|
||||
other => {
|
||||
return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
|
||||
"Unknown backend key: {}. Supported: sqlite, bm25, faiss, colbert",
|
||||
other
|
||||
)));
|
||||
}
|
||||
};
|
||||
backends.push(backend);
|
||||
}
|
||||
Ok(Self {
|
||||
inner: openjarvis_tools::storage::HybridMemory::new(backends),
|
||||
})
|
||||
}
|
||||
|
||||
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,382 @@
|
||||
//! 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()
|
||||
}
|
||||
}
|
||||
|
||||
// --- New telemetry session classes ---
|
||||
|
||||
/// Python wrapper for TelemetrySample.
|
||||
#[pyclass(name = "TelemetrySample")]
|
||||
#[derive(Clone)]
|
||||
pub struct PyTelemetrySample {
|
||||
pub timestamp_ns: u64,
|
||||
pub gpu_power_w: f64,
|
||||
pub cpu_power_w: f64,
|
||||
pub gpu_energy_j: f64,
|
||||
pub cpu_energy_j: f64,
|
||||
pub gpu_util_pct: f64,
|
||||
pub gpu_temp_c: f64,
|
||||
pub gpu_mem_gb: f64,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyTelemetrySample {
|
||||
#[new]
|
||||
#[pyo3(signature = (timestamp_ns, gpu_power_w=0.0, cpu_power_w=0.0, gpu_energy_j=0.0, cpu_energy_j=0.0, gpu_util_pct=0.0, gpu_temp_c=0.0, gpu_mem_gb=0.0))]
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn new(
|
||||
timestamp_ns: u64,
|
||||
gpu_power_w: f64,
|
||||
cpu_power_w: f64,
|
||||
gpu_energy_j: f64,
|
||||
cpu_energy_j: f64,
|
||||
gpu_util_pct: f64,
|
||||
gpu_temp_c: f64,
|
||||
gpu_mem_gb: f64,
|
||||
) -> Self {
|
||||
Self {
|
||||
timestamp_ns,
|
||||
gpu_power_w,
|
||||
cpu_power_w,
|
||||
gpu_energy_j,
|
||||
cpu_energy_j,
|
||||
gpu_util_pct,
|
||||
gpu_temp_c,
|
||||
gpu_mem_gb,
|
||||
}
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn timestamp_ns(&self) -> u64 {
|
||||
self.timestamp_ns
|
||||
}
|
||||
#[getter]
|
||||
fn gpu_power_w(&self) -> f64 {
|
||||
self.gpu_power_w
|
||||
}
|
||||
#[getter]
|
||||
fn cpu_power_w(&self) -> f64 {
|
||||
self.cpu_power_w
|
||||
}
|
||||
#[getter]
|
||||
fn gpu_energy_j(&self) -> f64 {
|
||||
self.gpu_energy_j
|
||||
}
|
||||
#[getter]
|
||||
fn cpu_energy_j(&self) -> f64 {
|
||||
self.cpu_energy_j
|
||||
}
|
||||
#[getter]
|
||||
fn gpu_util_pct(&self) -> f64 {
|
||||
self.gpu_util_pct
|
||||
}
|
||||
#[getter]
|
||||
fn gpu_temp_c(&self) -> f64 {
|
||||
self.gpu_temp_c
|
||||
}
|
||||
#[getter]
|
||||
fn gpu_mem_gb(&self) -> f64 {
|
||||
self.gpu_mem_gb
|
||||
}
|
||||
}
|
||||
|
||||
/// Python wrapper for TelemetrySessionCore (ring buffer).
|
||||
#[pyclass(name = "TelemetrySessionCore")]
|
||||
pub struct PyTelemetrySessionCore {
|
||||
inner: openjarvis_telemetry::session::TelemetrySessionCore,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyTelemetrySessionCore {
|
||||
#[new]
|
||||
#[pyo3(signature = (capacity=100000, sampling_interval_ms=100))]
|
||||
fn new(capacity: usize, sampling_interval_ms: u64) -> Self {
|
||||
Self {
|
||||
inner: openjarvis_telemetry::session::TelemetrySessionCore::new(
|
||||
capacity,
|
||||
sampling_interval_ms,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
fn add_sample(&self, sample: &PyTelemetrySample) {
|
||||
let s = openjarvis_telemetry::session::TelemetrySample {
|
||||
timestamp_ns: sample.timestamp_ns,
|
||||
gpu_power_w: sample.gpu_power_w,
|
||||
cpu_power_w: sample.cpu_power_w,
|
||||
gpu_energy_j: sample.gpu_energy_j,
|
||||
cpu_energy_j: sample.cpu_energy_j,
|
||||
gpu_util_pct: sample.gpu_util_pct,
|
||||
gpu_temp_c: sample.gpu_temp_c,
|
||||
gpu_mem_gb: sample.gpu_mem_gb,
|
||||
};
|
||||
self.inner.add_sample(s);
|
||||
}
|
||||
|
||||
fn window(&self, start_ns: u64, end_ns: u64) -> Vec<PyTelemetrySample> {
|
||||
self.inner
|
||||
.window(start_ns, end_ns)
|
||||
.into_iter()
|
||||
.map(|s| PyTelemetrySample {
|
||||
timestamp_ns: s.timestamp_ns,
|
||||
gpu_power_w: s.gpu_power_w,
|
||||
cpu_power_w: s.cpu_power_w,
|
||||
gpu_energy_j: s.gpu_energy_j,
|
||||
cpu_energy_j: s.cpu_energy_j,
|
||||
gpu_util_pct: s.gpu_util_pct,
|
||||
gpu_temp_c: s.gpu_temp_c,
|
||||
gpu_mem_gb: s.gpu_mem_gb,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn compute_energy_delta(&self, start_ns: u64, end_ns: u64) -> (f64, f64) {
|
||||
self.inner.compute_energy_delta(start_ns, end_ns)
|
||||
}
|
||||
|
||||
fn compute_avg_power(&self, start_ns: u64, end_ns: u64) -> (f64, f64) {
|
||||
self.inner.compute_avg_power(start_ns, end_ns)
|
||||
}
|
||||
|
||||
fn len(&self) -> usize {
|
||||
self.inner.len()
|
||||
}
|
||||
|
||||
fn clear(&self) {
|
||||
self.inner.clear();
|
||||
}
|
||||
}
|
||||
|
||||
/// Python wrapper for ITL stats computation.
|
||||
#[pyclass(name = "ItlStats")]
|
||||
pub struct PyItlStats;
|
||||
|
||||
#[pymethods]
|
||||
impl PyItlStats {
|
||||
#[new]
|
||||
fn new() -> Self {
|
||||
Self
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn compute(token_timestamps_ms: Vec<f64>) -> PyResult<pyo3::Py<pyo3::types::PyDict>> {
|
||||
let stats = openjarvis_telemetry::itl::compute_itl_stats(&token_timestamps_ms);
|
||||
pyo3::Python::with_gil(|py| {
|
||||
let dict = pyo3::types::PyDict::new(py);
|
||||
dict.set_item("p50_ms", stats.p50_ms)?;
|
||||
dict.set_item("p90_ms", stats.p90_ms)?;
|
||||
dict.set_item("p95_ms", stats.p95_ms)?;
|
||||
dict.set_item("p99_ms", stats.p99_ms)?;
|
||||
dict.set_item("mean_ms", stats.mean_ms)?;
|
||||
dict.set_item("min_ms", stats.min_ms)?;
|
||||
dict.set_item("max_ms", stats.max_ms)?;
|
||||
Ok(dict.into())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Python wrapper for FLOPs estimation.
|
||||
#[pyclass(name = "FlopsEstimator")]
|
||||
pub struct PyFlopsEstimator;
|
||||
|
||||
#[pymethods]
|
||||
impl PyFlopsEstimator {
|
||||
#[new]
|
||||
fn new() -> Self {
|
||||
Self
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn estimate_flops(model: &str, input_tokens: u64, output_tokens: u64) -> (f64, f64) {
|
||||
openjarvis_telemetry::flops::estimate_flops(model, input_tokens, output_tokens)
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
#[pyo3(signature = (flops, duration_s, gpu_name, num_gpus=1))]
|
||||
fn compute_mfu(flops: f64, duration_s: f64, gpu_name: &str, num_gpus: u32) -> f64 {
|
||||
openjarvis_telemetry::flops::compute_mfu(flops, duration_s, gpu_name, num_gpus)
|
||||
}
|
||||
}
|
||||
|
||||
/// Python wrapper for phase metrics.
|
||||
#[pyclass(name = "PhaseMetrics")]
|
||||
pub struct PyPhaseMetrics;
|
||||
|
||||
#[pymethods]
|
||||
impl PyPhaseMetrics {
|
||||
#[new]
|
||||
fn new() -> Self {
|
||||
Self
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn compute(
|
||||
samples: Vec<PyTelemetrySample>,
|
||||
start_ns: u64,
|
||||
end_ns: u64,
|
||||
tokens: u64,
|
||||
) -> PyResult<pyo3::Py<pyo3::types::PyDict>> {
|
||||
let rust_samples: Vec<openjarvis_telemetry::session::TelemetrySample> = samples
|
||||
.iter()
|
||||
.map(|s| openjarvis_telemetry::session::TelemetrySample {
|
||||
timestamp_ns: s.timestamp_ns,
|
||||
gpu_power_w: s.gpu_power_w,
|
||||
cpu_power_w: s.cpu_power_w,
|
||||
gpu_energy_j: s.gpu_energy_j,
|
||||
cpu_energy_j: s.cpu_energy_j,
|
||||
gpu_util_pct: s.gpu_util_pct,
|
||||
gpu_temp_c: s.gpu_temp_c,
|
||||
gpu_mem_gb: s.gpu_mem_gb,
|
||||
})
|
||||
.collect();
|
||||
let metrics =
|
||||
openjarvis_telemetry::phase::compute_phase_metrics(&rust_samples, start_ns, end_ns, tokens);
|
||||
pyo3::Python::with_gil(|py| {
|
||||
let dict = pyo3::types::PyDict::new(py);
|
||||
dict.set_item("energy_j", metrics.energy_j)?;
|
||||
dict.set_item("mean_power_w", metrics.mean_power_w)?;
|
||||
dict.set_item("duration_s", metrics.duration_s)?;
|
||||
dict.set_item("energy_per_token_j", metrics.energy_per_token_j)?;
|
||||
dict.set_item("tokens", metrics.tokens)?;
|
||||
Ok(dict.into())
|
||||
})
|
||||
}
|
||||
|
||||
#[staticmethod]
|
||||
fn split_at_ttft(
|
||||
samples: Vec<PyTelemetrySample>,
|
||||
start_ns: u64,
|
||||
ttft_ns: u64,
|
||||
end_ns: u64,
|
||||
input_tokens: u64,
|
||||
output_tokens: u64,
|
||||
) -> PyResult<(pyo3::Py<pyo3::types::PyDict>, pyo3::Py<pyo3::types::PyDict>)> {
|
||||
let rust_samples: Vec<openjarvis_telemetry::session::TelemetrySample> = samples
|
||||
.iter()
|
||||
.map(|s| openjarvis_telemetry::session::TelemetrySample {
|
||||
timestamp_ns: s.timestamp_ns,
|
||||
gpu_power_w: s.gpu_power_w,
|
||||
cpu_power_w: s.cpu_power_w,
|
||||
gpu_energy_j: s.gpu_energy_j,
|
||||
cpu_energy_j: s.cpu_energy_j,
|
||||
gpu_util_pct: s.gpu_util_pct,
|
||||
gpu_temp_c: s.gpu_temp_c,
|
||||
gpu_mem_gb: s.gpu_mem_gb,
|
||||
})
|
||||
.collect();
|
||||
let (prefill, decode) = openjarvis_telemetry::phase::split_at_ttft(
|
||||
&rust_samples,
|
||||
start_ns,
|
||||
ttft_ns,
|
||||
end_ns,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
);
|
||||
pyo3::Python::with_gil(|py| {
|
||||
let prefill_dict = pyo3::types::PyDict::new(py);
|
||||
prefill_dict.set_item("energy_j", prefill.energy_j)?;
|
||||
prefill_dict.set_item("mean_power_w", prefill.mean_power_w)?;
|
||||
prefill_dict.set_item("duration_s", prefill.duration_s)?;
|
||||
prefill_dict.set_item("energy_per_token_j", prefill.energy_per_token_j)?;
|
||||
prefill_dict.set_item("tokens", prefill.tokens)?;
|
||||
|
||||
let decode_dict = pyo3::types::PyDict::new(py);
|
||||
decode_dict.set_item("energy_j", decode.energy_j)?;
|
||||
decode_dict.set_item("mean_power_w", decode.mean_power_w)?;
|
||||
decode_dict.set_item("duration_s", decode.duration_s)?;
|
||||
decode_dict.set_item("energy_per_token_j", decode.energy_per_token_j)?;
|
||||
decode_dict.set_item("tokens", decode.tokens)?;
|
||||
|
||||
Ok((prefill_dict.into(), decode_dict.into()))
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
//! PyO3 bindings for agent templates.
|
||||
|
||||
use pyo3::prelude::*;
|
||||
|
||||
#[pyclass(name = "AgentTemplate")]
|
||||
pub struct PyAgentTemplate {
|
||||
inner: openjarvis_templates::AgentTemplate,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyAgentTemplate {
|
||||
#[getter]
|
||||
fn name(&self) -> &str {
|
||||
&self.inner.name
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn description(&self) -> &str {
|
||||
&self.inner.description
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn system_prompt(&self) -> &str {
|
||||
&self.inner.system_prompt
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn agent_type(&self) -> &str {
|
||||
&self.inner.agent_type
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn tools(&self) -> Vec<String> {
|
||||
self.inner.tools.clone()
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn max_turns(&self) -> usize {
|
||||
self.inner.max_turns
|
||||
}
|
||||
|
||||
#[getter]
|
||||
fn temperature(&self) -> f64 {
|
||||
self.inner.temperature
|
||||
}
|
||||
|
||||
fn to_json(&self) -> String {
|
||||
serde_json::to_string(&self.inner).unwrap_or_default()
|
||||
}
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
pub fn load_template(toml_str: &str) -> PyResult<PyAgentTemplate> {
|
||||
let tpl = openjarvis_templates::load_template(toml_str)
|
||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e))?;
|
||||
Ok(PyAgentTemplate { inner: tpl })
|
||||
}
|
||||
@@ -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, ¶ms, 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(¶ms)
|
||||
.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(¶ms)
|
||||
.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(¶ms)
|
||||
.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(¶ms)
|
||||
.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(¶ms)
|
||||
.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(¶ms)
|
||||
.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(¶ms)
|
||||
.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(¶ms)
|
||||
.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(¶ms)
|
||||
.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())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,127 @@
|
||||
//! PyO3 bindings for workflow engine.
|
||||
|
||||
use pyo3::prelude::*;
|
||||
use std::collections::HashMap;
|
||||
|
||||
#[pyclass(name = "WorkflowGraph")]
|
||||
pub struct PyWorkflowGraph {
|
||||
inner: openjarvis_workflow::WorkflowGraph,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyWorkflowGraph {
|
||||
#[new]
|
||||
#[pyo3(signature = (name=""))]
|
||||
fn new(name: &str) -> Self {
|
||||
Self {
|
||||
inner: openjarvis_workflow::WorkflowGraph::new(name),
|
||||
}
|
||||
}
|
||||
|
||||
fn add_node(
|
||||
&mut self,
|
||||
id: &str,
|
||||
node_type: &str,
|
||||
agent: &str,
|
||||
tools: Vec<String>,
|
||||
) -> PyResult<()> {
|
||||
let nt = match node_type {
|
||||
"tool" => openjarvis_workflow::NodeType::Tool,
|
||||
"condition" => openjarvis_workflow::NodeType::Condition,
|
||||
"parallel" => openjarvis_workflow::NodeType::Parallel,
|
||||
"loop" => openjarvis_workflow::NodeType::Loop,
|
||||
"transform" => openjarvis_workflow::NodeType::Transform,
|
||||
_ => openjarvis_workflow::NodeType::Agent,
|
||||
};
|
||||
let node = openjarvis_workflow::WorkflowNode {
|
||||
id: id.to_string(),
|
||||
node_type: nt,
|
||||
agent: agent.to_string(),
|
||||
tools,
|
||||
config: HashMap::new(),
|
||||
condition_expr: String::new(),
|
||||
max_iterations: 10,
|
||||
transform_expr: String::new(),
|
||||
};
|
||||
self.inner
|
||||
.add_node(node)
|
||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e))
|
||||
}
|
||||
|
||||
fn add_edge(&mut self, source: &str, target: &str) -> PyResult<()> {
|
||||
let edge = openjarvis_workflow::WorkflowEdge {
|
||||
source: source.to_string(),
|
||||
target: target.to_string(),
|
||||
condition: String::new(),
|
||||
};
|
||||
self.inner
|
||||
.add_edge(edge)
|
||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e))
|
||||
}
|
||||
|
||||
fn validate(&self) -> (bool, String) {
|
||||
self.inner.validate()
|
||||
}
|
||||
|
||||
fn topological_sort(&self) -> PyResult<Vec<String>> {
|
||||
self.inner
|
||||
.topological_sort()
|
||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e))
|
||||
}
|
||||
|
||||
fn execution_stages(&self) -> Vec<Vec<String>> {
|
||||
self.inner.execution_stages()
|
||||
}
|
||||
|
||||
fn predecessors(&self, node_id: &str) -> Vec<String> {
|
||||
self.inner.predecessors(node_id)
|
||||
}
|
||||
|
||||
fn successors(&self, node_id: &str) -> Vec<String> {
|
||||
self.inner.successors(node_id)
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass(name = "WorkflowBuilder")]
|
||||
pub struct PyWorkflowBuilder {
|
||||
inner: Option<openjarvis_workflow::WorkflowBuilder>,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl PyWorkflowBuilder {
|
||||
#[new]
|
||||
#[pyo3(signature = (name=""))]
|
||||
fn new(name: &str) -> Self {
|
||||
Self {
|
||||
inner: Some(openjarvis_workflow::WorkflowBuilder::new(name)),
|
||||
}
|
||||
}
|
||||
|
||||
fn add_agent_node(&mut self, id: &str, agent: &str) {
|
||||
if let Some(ref mut b) = self.inner {
|
||||
b.add_agent_node(id, agent);
|
||||
}
|
||||
}
|
||||
|
||||
fn add_tool_node(&mut self, id: &str, tools: Vec<String>) {
|
||||
if let Some(ref mut b) = self.inner {
|
||||
b.add_tool_node(id, tools);
|
||||
}
|
||||
}
|
||||
|
||||
fn connect(&mut self, source: &str, target: &str) {
|
||||
if let Some(ref mut b) = self.inner {
|
||||
b.connect(source, target);
|
||||
}
|
||||
}
|
||||
|
||||
fn build(&mut self) -> PyResult<PyWorkflowGraph> {
|
||||
let builder = self.inner.take().ok_or_else(|| {
|
||||
PyErr::new::<pyo3::exceptions::PyRuntimeError, _>("builder already consumed")
|
||||
})?;
|
||||
let graph = builder
|
||||
.build()
|
||||
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e))?;
|
||||
Ok(PyWorkflowGraph { inner: graph })
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user