use axllm::{agent, AxCodeRuntime, AxCodeSession, AxResult, RuntimeEnvelope};
use serde_json::{json, Value};
struct DemoSession {
globals: Value,
closed: bool,
}
impl AxCodeSession for DemoSession {
fn execute(&mut self, code: &str, _options: Value) -> AxResult<RuntimeEnvelope> {
if code == "timeout()" {
return Ok(RuntimeEnvelope::timeout("demo timeout"));
}
self.globals["answer"] = json!("runtime final");
Ok(RuntimeEnvelope::final_payload(
json!({"answer": self.globals["answer"]}),
))
}
fn snapshot_globals(&mut self, _options: Value) -> AxResult<Value> {
Ok(json!({"version": 1, "bindings": self.globals, "closed": self.closed}))
}
fn patch_globals(&mut self, snapshot: Value, _options: Value) -> AxResult<Value> {
self.globals = snapshot
.get("bindings")
.cloned()
.unwrap_or_else(|| json!({}));
self.snapshot_globals(json!({}))
}
fn close(&mut self) -> AxResult<Value> {
self.closed = true;
Ok(json!({"closed": true}))
}
}
struct DemoRuntime;
impl AxCodeRuntime for DemoRuntime {
fn language(&self) -> &str {
"Rust"
}
fn create_session(
&mut self,
globals: Value,
_options: Value,
) -> AxResult<Box<dyn AxCodeSession>> {
Ok(Box::new(DemoSession {
globals,
closed: false,
}))
}
}
fn main() -> AxResult<()> {
let mut runtime = DemoRuntime;
let mut runner = agent("question:string -> answer:string")?;
let step = runner.execute_actor_step(
&mut runtime,
"final()",
json!({"question": "adapter"}),
json!({}),
)?;
let snapshot = runner.export_session_state()?;
let timeout = runner.execute_actor_step(
&mut runtime,
"timeout()",
json!({"question": "adapter"}),
json!({}),
)?;
let closed = runner.close_runtime_session()?;
println!(
"{}",
serde_json::to_string_pretty(&json!({
"stepKind": step.payload["kind"],
"snapshotAnswer": snapshot["bindings"]["answer"],
"timeoutCategory": timeout.payload["error_category"],
"closed": closed
}))?
);
Ok(())
}