Skip to main content

max / makenotwork

4.6 KB · 163 lines History Blame Raw
1 //! End-to-end example:
2 //!
3 //! - Register two tiny tools (`add` = Read, `remember_project` = Write) into a
4 //! `ToolRegistry`.
5 //! - Spin up an MCP server on a random localhost port so external clients
6 //! (Claude Desktop, etc.) could connect.
7 //! - Drive the same registry in-process via an `Agent<OllamaProvider>`, with
8 //! the `remember_project` capability granted to the session.
9 //!
10 //! Run with a tool-calling model installed:
11 //!
12 //! ```
13 //! ollama pull qwen2.5
14 //! cargo run --example toy
15 //! ```
16 //!
17 //! Override the model with `KBERG_MODEL=llama3.2 cargo run --example toy`.
18
19 use std::collections::HashSet;
20
21 use async_trait::async_trait;
22 use kberg::{
23 Agent, Error, Tool, ToolCallResult, ToolKind, ToolRegistry, WriteCapability,
24 provider::ollama::OllamaProvider, server,
25 };
26 use serde_json::{Value, json};
27
28 struct Add;
29
30 #[async_trait]
31 impl Tool for Add {
32 fn name(&self) -> &'static str {
33 "add"
34 }
35 fn description(&self) -> &'static str {
36 "Add two integers and return the sum."
37 }
38 fn kind(&self) -> ToolKind {
39 ToolKind::Read
40 }
41 fn small_model_safe(&self) -> bool {
42 true
43 }
44 fn input_schema(&self) -> Value {
45 json!({
46 "type": "object",
47 "properties": {
48 "a": { "type": "integer" },
49 "b": { "type": "integer" }
50 },
51 "required": ["a", "b"]
52 })
53 }
54 async fn call(&self, args: Value) -> Result<ToolCallResult, Error> {
55 let a = args
56 .get("a")
57 .and_then(Value::as_i64)
58 .ok_or_else(|| Error::InvalidArgs {
59 tool: "add".into(),
60 message: "missing integer `a`".into(),
61 })?;
62 let b = args
63 .get("b")
64 .and_then(Value::as_i64)
65 .ok_or_else(|| Error::InvalidArgs {
66 tool: "add".into(),
67 message: "missing integer `b`".into(),
68 })?;
69 Ok(ToolCallResult::text((a + b).to_string()))
70 }
71 }
72
73 struct RememberProject;
74
75 #[async_trait]
76 impl Tool for RememberProject {
77 fn name(&self) -> &'static str {
78 "remember_project"
79 }
80 fn description(&self) -> &'static str {
81 "Record a new project the user is working on."
82 }
83 fn kind(&self) -> ToolKind {
84 ToolKind::Write(WriteCapability::new(
85 "remember_project",
86 "record new projects the user says they're working on",
87 ))
88 }
89 fn input_schema(&self) -> Value {
90 json!({
91 "type": "object",
92 "properties": { "name": { "type": "string" } },
93 "required": ["name"]
94 })
95 }
96 async fn call(&self, args: Value) -> Result<ToolCallResult, Error> {
97 let name = args
98 .get("name")
99 .and_then(Value::as_str)
100 .unwrap_or("<unnamed>");
101 Ok(ToolCallResult::text(format!("recorded project `{name}`")))
102 }
103 }
104
105 #[tokio::main]
106 async fn main() -> Result<(), Box<dyn std::error::Error>> {
107 tracing_subscriber::fmt()
108 .with_env_filter(
109 tracing_subscriber::EnvFilter::try_from_default_env()
110 .unwrap_or_else(|_| "info,kberg=debug".into()),
111 )
112 .init();
113
114 let mut registry = ToolRegistry::new();
115 registry.register(Add);
116 registry.register(RememberProject);
117
118 // Bring up an MCP server so external clients could reach the same tools.
119 let bound = server::serve(
120 registry.clone(),
121 "127.0.0.1:0",
122 server::ServeConfig {
123 projection: kberg::SurfaceProjection::Full,
124 grants: Some(HashSet::from(["remember_project".to_string()])),
125 resources: None,
126 },
127 )
128 .await?;
129 println!("MCP server: {}", bound.url());
130 println!(
131 "registered write capabilities: {:?}",
132 registry.write_capabilities()
133 );
134
135 let model = std::env::var("KBERG_MODEL").unwrap_or_else(|_| "qwen2.5".to_string());
136 let provider = OllamaProvider::new(&model);
137 let agent = Agent::new(provider).with_grants(HashSet::from(["remember_project".to_string()]));
138
139 let prompt = "What is 17 + 25? Also, remember that I'm working on a project called Kberg.";
140
141 println!("\n> {prompt}\n");
142 let outcome = agent
143 .run(
144 &registry,
145 Some(
146 "You are a terse assistant. Prefer tools over guessing. \
147 When you have the answer, reply directly and briefly.",
148 ),
149 prompt,
150 )
151 .await?;
152
153 println!(
154 "--- final ({} step{}) ---",
155 outcome.steps,
156 if outcome.steps == 1 { "" } else { "s" }
157 );
158 println!("{}", outcome.final_message);
159
160 bound.shutdown();
161 Ok(())
162 }
163