harn_vm/composition/
hosts.rs1use std::collections::BTreeMap;
2
3use serde_json::Value;
4
5use crate::agent_events::ToolCallErrorCategory;
6use crate::value::VmValue;
7use crate::vm::AsyncBuiltinCtx;
8use crate::VmClosure;
9
10use super::manifest::BindingManifestEntry;
11use super::types::{CompositionToolHost, CompositionToolOutput};
12
13pub struct StaticCompositionToolHost {
14 outputs: BTreeMap<String, Value>,
15}
16
17impl StaticCompositionToolHost {
18 pub fn new(outputs: BTreeMap<String, Value>) -> Self {
19 Self { outputs }
20 }
21}
22
23#[async_trait::async_trait(?Send)]
24impl CompositionToolHost for StaticCompositionToolHost {
25 async fn call(&self, binding: &BindingManifestEntry, input: Value) -> CompositionToolOutput {
26 if let Some(value) = self.outputs.get(&binding.name) {
27 return CompositionToolOutput::ok(value.clone());
28 }
29 if let Some(value) = binding.metadata.get("mock_output") {
30 return CompositionToolOutput::ok(value.clone());
31 }
32 CompositionToolOutput::ok(serde_json::json!({
33 "tool": binding.name,
34 "input": input,
35 }))
36 }
37}
38
39pub struct ClosureCompositionToolHost {
48 closure: VmClosure,
49 ctx: AsyncBuiltinCtx,
50}
51
52impl ClosureCompositionToolHost {
53 pub fn new(closure: VmClosure, ctx: AsyncBuiltinCtx) -> Self {
54 Self { closure, ctx }
55 }
56}
57
58#[async_trait::async_trait(?Send)]
59impl CompositionToolHost for ClosureCompositionToolHost {
60 async fn call(&self, binding: &BindingManifestEntry, input: Value) -> CompositionToolOutput {
61 let mut vm = self.ctx.child_vm();
62 let args = vec![
63 VmValue::String(std::sync::Arc::from(binding.name.as_str())),
64 crate::json_to_vm_value(&input),
65 ];
66 let result = vm.call_closure_pub(&self.closure, &args).await;
67 self.ctx.forward_output(&vm.take_output());
68 match result {
69 Ok(value) => {
70 let json = crate::llm::vm_value_to_json(&value);
71 CompositionToolOutput::ok(json)
72 }
73 Err(error) => {
74 CompositionToolOutput::error(error.to_string(), ToolCallErrorCategory::ToolError)
75 }
76 }
77 }
78}