1//! This example creates a basic Chat UI with a function for rolling a die.
2
3use anyhow::{Context as _, Result};
4use assets::Assets;
5use assistant2::AssistantPanel;
6use assistant_tooling::{LanguageModelTool, ToolRegistry};
7use client::{Client, UserStore};
8use fs::Fs;
9use futures::StreamExt as _;
10use gpui::{actions, AnyElement, App, AppContext, KeyBinding, Model, Task, View, WindowOptions};
11use language::LanguageRegistry;
12use project::Project;
13use rand::Rng;
14use schemars::JsonSchema;
15use serde::{Deserialize, Serialize};
16use settings::{KeymapFile, DEFAULT_KEYMAP_PATH};
17use std::{path::PathBuf, sync::Arc};
18use theme::LoadThemes;
19use ui::{div, prelude::*, Render};
20use util::ResultExt as _;
21
22actions!(example, [Quit]);
23
24struct RollDiceTool {}
25
26impl RollDiceTool {
27 fn new() -> Self {
28 Self {}
29 }
30}
31
32#[derive(Serialize, Deserialize, JsonSchema, Clone)]
33#[serde(rename_all = "snake_case")]
34enum Die {
35 D6 = 6,
36 D20 = 20,
37}
38
39impl Die {
40 fn into_str(&self) -> &'static str {
41 match self {
42 Die::D6 => "d6",
43 Die::D20 => "d20",
44 }
45 }
46}
47
48#[derive(Serialize, Deserialize, JsonSchema, Clone)]
49struct DiceParams {
50 /// The number of dice to roll.
51 num_dice: u8,
52 /// Which die to roll. Defaults to a d6 if not provided.
53 die_type: Option<Die>,
54}
55
56#[derive(Serialize, Deserialize)]
57struct DieRoll {
58 die: Die,
59 roll: u8,
60}
61
62impl DieRoll {
63 fn render(&self) -> AnyElement {
64 match self.die {
65 Die::D6 => {
66 let face = match self.roll {
67 6 => div().child("⚅"),
68 5 => div().child("⚄"),
69 4 => div().child("⚃"),
70 3 => div().child("⚂"),
71 2 => div().child("⚁"),
72 1 => div().child("⚀"),
73 _ => div().child("😅"),
74 };
75 face.text_3xl().into_any_element()
76 }
77 _ => div()
78 .child(format!("{}", self.roll))
79 .text_3xl()
80 .into_any_element(),
81 }
82 }
83}
84
85#[derive(Serialize, Deserialize)]
86struct DiceRoll {
87 rolls: Vec<DieRoll>,
88}
89
90pub struct DiceView {
91 result: Result<DiceRoll>,
92}
93
94impl Render for DiceView {
95 fn render(&mut self, _cx: &mut ViewContext<Self>) -> impl IntoElement {
96 let output = match &self.result {
97 Ok(output) => output,
98 Err(_) => return "Somehow dice failed 🎲".into_any_element(),
99 };
100
101 h_flex()
102 .children(
103 output
104 .rolls
105 .iter()
106 .map(|roll| div().p_2().child(roll.render())),
107 )
108 .into_any_element()
109 }
110}
111
112impl LanguageModelTool for RollDiceTool {
113 type Input = DiceParams;
114 type Output = DiceRoll;
115 type View = DiceView;
116
117 fn name(&self) -> String {
118 "roll_dice".to_string()
119 }
120
121 fn description(&self) -> String {
122 "Rolls N many dice and returns the results.".to_string()
123 }
124
125 fn execute(&self, input: &Self::Input, _cx: &AppContext) -> Task<gpui::Result<Self::Output>> {
126 let rolls = (0..input.num_dice)
127 .map(|_| {
128 let die_type = input.die_type.as_ref().unwrap_or(&Die::D6).clone();
129
130 DieRoll {
131 die: die_type.clone(),
132 roll: rand::thread_rng().gen_range(1..=die_type as u8),
133 }
134 })
135 .collect();
136
137 return Task::ready(Ok(DiceRoll { rolls }));
138 }
139
140 fn output_view(
141 _tool_call_id: String,
142 _input: Self::Input,
143 result: Result<Self::Output>,
144 cx: &mut WindowContext,
145 ) -> gpui::View<Self::View> {
146 cx.new_view(|_cx| DiceView { result })
147 }
148
149 fn format(_: &Self::Input, output: &Result<Self::Output>) -> String {
150 let output = match output {
151 Ok(output) => output,
152 Err(_) => return "Somehow dice failed 🎲".to_string(),
153 };
154
155 let mut result = String::new();
156 for roll in &output.rolls {
157 let die = &roll.die;
158 result.push_str(&format!("{}: {}\n", die.into_str(), roll.roll));
159 }
160 result
161 }
162}
163
164struct FileBrowserTool {
165 fs: Arc<dyn Fs>,
166 root_dir: PathBuf,
167}
168
169impl FileBrowserTool {
170 fn new(fs: Arc<dyn Fs>, root_dir: PathBuf) -> Self {
171 Self { fs, root_dir }
172 }
173}
174
175#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
176struct FileBrowserParams {
177 command: FileBrowserCommand,
178}
179
180#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
181enum FileBrowserCommand {
182 Ls { path: PathBuf },
183 Cat { path: PathBuf },
184}
185
186#[derive(Serialize, Deserialize)]
187enum FileBrowserOutput {
188 Ls { entries: Vec<String> },
189 Cat { content: String },
190}
191
192pub struct FileBrowserView {
193 result: Result<FileBrowserOutput>,
194}
195
196impl Render for FileBrowserView {
197 fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
198 let Ok(output) = self.result.as_ref() else {
199 return h_flex().child("Failed to perform operation");
200 };
201
202 match output {
203 FileBrowserOutput::Ls { entries } => v_flex().children(
204 entries
205 .into_iter()
206 .map(|entry| h_flex().text_ui(cx).child(entry.clone())),
207 ),
208 FileBrowserOutput::Cat { content } => h_flex().child(content.clone()),
209 }
210 }
211}
212
213impl LanguageModelTool for FileBrowserTool {
214 type Input = FileBrowserParams;
215 type Output = FileBrowserOutput;
216 type View = FileBrowserView;
217
218 fn name(&self) -> String {
219 "file_browser".to_string()
220 }
221
222 fn description(&self) -> String {
223 "A tool for browsing the filesystem.".to_string()
224 }
225
226 fn execute(&self, input: &Self::Input, cx: &AppContext) -> Task<gpui::Result<Self::Output>> {
227 cx.spawn({
228 let fs = self.fs.clone();
229 let root_dir = self.root_dir.clone();
230 let input = input.clone();
231 |_cx| async move {
232 match input.command {
233 FileBrowserCommand::Ls { path } => {
234 let path = root_dir.join(path);
235
236 let mut output = fs.read_dir(&path).await?;
237
238 let mut entries = Vec::new();
239 while let Some(entry) = output.next().await {
240 let entry = entry?;
241 entries.push(entry.display().to_string());
242 }
243
244 Ok(FileBrowserOutput::Ls { entries })
245 }
246 FileBrowserCommand::Cat { path } => {
247 let path = root_dir.join(path);
248
249 let output = fs.load(&path).await?;
250
251 Ok(FileBrowserOutput::Cat { content: output })
252 }
253 }
254 }
255 })
256 }
257
258 fn output_view(
259 _tool_call_id: String,
260 _input: Self::Input,
261 result: Result<Self::Output>,
262 cx: &mut WindowContext,
263 ) -> gpui::View<Self::View> {
264 cx.new_view(|_cx| FileBrowserView { result })
265 }
266
267 fn format(_input: &Self::Input, output: &Result<Self::Output>) -> String {
268 let Ok(output) = output else {
269 return "Failed to perform command: {input:?}".to_string();
270 };
271
272 match output {
273 FileBrowserOutput::Ls { entries } => entries.join("\n"),
274 FileBrowserOutput::Cat { content } => content.to_owned(),
275 }
276 }
277}
278
279fn main() {
280 env_logger::init();
281 App::new().with_assets(Assets).run(|cx| {
282 cx.bind_keys(Some(KeyBinding::new("cmd-q", Quit, None)));
283 cx.on_action(|_: &Quit, cx: &mut AppContext| {
284 cx.quit();
285 });
286
287 settings::init(cx);
288 language::init(cx);
289 Project::init_settings(cx);
290 editor::init(cx);
291 theme::init(LoadThemes::JustBase, cx);
292 Assets.load_fonts(cx).unwrap();
293 KeymapFile::load_asset(DEFAULT_KEYMAP_PATH, cx).unwrap();
294 client::init_settings(cx);
295 release_channel::init("0.130.0", cx);
296
297 let client = Client::production(cx);
298 {
299 let client = client.clone();
300 cx.spawn(|cx| async move { client.authenticate_and_connect(false, &cx).await })
301 .detach_and_log_err(cx);
302 }
303 assistant2::init(client.clone(), cx);
304
305 let language_registry = Arc::new(LanguageRegistry::new(
306 Task::ready(()),
307 cx.background_executor().clone(),
308 ));
309
310 let user_store = cx.new_model(|cx| UserStore::new(client.clone(), cx));
311 let node_runtime = node_runtime::RealNodeRuntime::new(client.http_client());
312 languages::init(language_registry.clone(), node_runtime, cx);
313
314 cx.spawn(|cx| async move {
315 cx.update(|cx| {
316 let fs = Arc::new(fs::RealFs::new(None));
317 let cwd = std::env::current_dir().expect("Failed to get current working directory");
318
319 cx.open_window(WindowOptions::default(), |cx| {
320 let mut tool_registry = ToolRegistry::new();
321 tool_registry
322 .register(RollDiceTool::new(), cx)
323 .context("failed to register DummyTool")
324 .log_err();
325
326 tool_registry
327 .register(FileBrowserTool::new(fs, cwd), cx)
328 .context("failed to register FileBrowserTool")
329 .log_err();
330
331 let tool_registry = Arc::new(tool_registry);
332
333 println!("Tools registered");
334 for definition in tool_registry.definitions() {
335 println!("{}", definition);
336 }
337
338 cx.new_view(|cx| Example::new(language_registry, tool_registry, user_store, cx))
339 });
340 cx.activate(true);
341 })
342 })
343 .detach_and_log_err(cx);
344 })
345}
346
347struct Example {
348 assistant_panel: View<AssistantPanel>,
349}
350
351impl Example {
352 fn new(
353 language_registry: Arc<LanguageRegistry>,
354 tool_registry: Arc<ToolRegistry>,
355 user_store: Model<UserStore>,
356 cx: &mut ViewContext<Self>,
357 ) -> Self {
358 Self {
359 assistant_panel: cx.new_view(|cx| {
360 AssistantPanel::new(language_registry, tool_registry, user_store, cx)
361 }),
362 }
363 }
364}
365
366impl Render for Example {
367 fn render(&mut self, _cx: &mut ViewContext<Self>) -> impl ui::prelude::IntoElement {
368 div().size_full().child(self.assistant_panel.clone())
369 }
370}