use async_trait::async_trait;
use oharness_core::{
CompletionRequest, CompletionResponse, Content, LlmCapabilities, Message, ModelId, StopReason,
Task, Usage,
};
use oharness_llm::{ChunkStream, Llm, LlmError};
use oharness_loop::{Agent, ReactLoop};
use oharness_tools::fs::FsToolSet;
use oharness_trace::InMemorySink;
use serde_json::json;
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc;
struct ScriptedLlm {
responses: Vec<CompletionResponse>,
cursor: AtomicU32,
}
#[async_trait]
impl Llm for ScriptedLlm {
fn name(&self) -> &str {
"scripted"
}
fn capabilities(&self) -> LlmCapabilities {
LlmCapabilities::default()
}
async fn complete(&self, _req: CompletionRequest) -> Result<CompletionResponse, LlmError> {
let idx = self.cursor.fetch_add(1, Ordering::SeqCst) as usize;
self.responses
.get(idx)
.cloned()
.ok_or(LlmError::Unsupported("script exhausted"))
}
async fn stream(&self, _req: CompletionRequest) -> Result<ChunkStream, LlmError> {
Err(LlmError::Unsupported("stream"))
}
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let script = vec![
CompletionResponse {
id: "msg_001".into(),
model: ModelId::new("scripted-tools-example"),
content: vec![
Content::text("Let me list the current directory to see what's here."),
Content::ToolUse {
id: "tu_1".into(),
name: "fs_list".into(),
input: json!({"path": "."}),
},
],
stop_reason: StopReason::ToolUse,
usage: Usage {
tokens_input: 12,
tokens_output: 40,
..Default::default()
},
},
CompletionResponse {
id: "msg_002".into(),
model: ModelId::new("scripted-tools-example"),
content: vec![Content::text(
"I can see the repository layout. There's a `crates/` \
directory, which is the Cargo workspace root.",
)],
stop_reason: StopReason::EndTurn,
usage: Usage {
tokens_input: 80,
tokens_output: 30,
..Default::default()
},
},
];
let llm = Arc::new(ScriptedLlm {
responses: script,
cursor: AtomicU32::new(0),
});
let sink = Arc::new(InMemorySink::new());
let agent = Agent::builder()
.with_llm(llm)
.with_tools(Arc::new(FsToolSet::new()))
.with_event_sink(sink.clone())
.with_loop(Box::new(ReactLoop::new()))
.with_max_turns(5)
.build()?;
let outcome = agent.run(Task::new("inspect the repo")).await?;
println!("Termination: {:?}", outcome.termination);
println!(
"Turns: {} | tool calls: {} | tokens in/out: {}/{}",
outcome.usage.turns,
outcome.usage.tool_calls,
outcome.usage.tokens_input,
outcome.usage.tokens_output,
);
if let Some(Message::Assistant { content, .. }) = outcome.final_messages.last() {
for c in content {
if let Content::Text { text } = c {
println!("Assistant: {text}");
}
}
}
println!("Trajectory events captured: {}", sink.events().len());
Ok(())
}