1use std::sync::Arc;
2
3use anyhow::Result;
4use assistant_tool::{ToolId, ToolWorkingSet};
5use collections::HashMap;
6use context_server::manager::ContextServerManager;
7use context_server::{ContextServerFactoryRegistry, ContextServerTool};
8use gpui::{prelude::*, AppContext, Model, ModelContext, Task};
9use project::Project;
10use unindent::Unindent;
11use util::ResultExt as _;
12
13use crate::thread::{Thread, ThreadId};
14
15pub struct ThreadStore {
16 #[allow(unused)]
17 project: Model<Project>,
18 tools: Arc<ToolWorkingSet>,
19 context_server_manager: Model<ContextServerManager>,
20 context_server_tool_ids: HashMap<Arc<str>, Vec<ToolId>>,
21 threads: Vec<Model<Thread>>,
22}
23
24impl ThreadStore {
25 pub fn new(
26 project: Model<Project>,
27 tools: Arc<ToolWorkingSet>,
28 cx: &mut AppContext,
29 ) -> Task<Result<Model<Self>>> {
30 cx.spawn(|mut cx| async move {
31 let this = cx.new_model(|cx: &mut ModelContext<Self>| {
32 let context_server_factory_registry =
33 ContextServerFactoryRegistry::default_global(cx);
34 let context_server_manager = cx.new_model(|cx| {
35 ContextServerManager::new(context_server_factory_registry, project.clone(), cx)
36 });
37
38 let mut this = Self {
39 project,
40 tools,
41 context_server_manager,
42 context_server_tool_ids: HashMap::default(),
43 threads: Vec::new(),
44 };
45 this.mock_recent_threads(cx);
46 this.register_context_server_handlers(cx);
47
48 this
49 })?;
50
51 Ok(this)
52 })
53 }
54
55 pub fn recent_threads(&self, limit: usize, cx: &ModelContext<Self>) -> Vec<Model<Thread>> {
56 self.threads
57 .iter()
58 .filter(|thread| !thread.read(cx).is_empty())
59 .take(limit)
60 .cloned()
61 .collect()
62 }
63
64 pub fn create_thread(&mut self, cx: &mut ModelContext<Self>) -> Model<Thread> {
65 let thread = cx.new_model(|cx| Thread::new(self.tools.clone(), cx));
66 self.threads.push(thread.clone());
67 thread
68 }
69
70 pub fn open_thread(&self, id: &ThreadId, cx: &mut ModelContext<Self>) -> Option<Model<Thread>> {
71 self.threads
72 .iter()
73 .find(|thread| thread.read(cx).id() == id)
74 .cloned()
75 }
76
77 fn register_context_server_handlers(&self, cx: &mut ModelContext<Self>) {
78 cx.subscribe(
79 &self.context_server_manager.clone(),
80 Self::handle_context_server_event,
81 )
82 .detach();
83 }
84
85 fn handle_context_server_event(
86 &mut self,
87 context_server_manager: Model<ContextServerManager>,
88 event: &context_server::manager::Event,
89 cx: &mut ModelContext<Self>,
90 ) {
91 let tool_working_set = self.tools.clone();
92 match event {
93 context_server::manager::Event::ServerStarted { server_id } => {
94 if let Some(server) = context_server_manager.read(cx).get_server(server_id) {
95 let context_server_manager = context_server_manager.clone();
96 cx.spawn({
97 let server = server.clone();
98 let server_id = server_id.clone();
99 |this, mut cx| async move {
100 let Some(protocol) = server.client() else {
101 return;
102 };
103
104 if protocol.capable(context_server::protocol::ServerCapability::Tools) {
105 if let Some(tools) = protocol.list_tools().await.log_err() {
106 let tool_ids = tools
107 .tools
108 .into_iter()
109 .map(|tool| {
110 log::info!(
111 "registering context server tool: {:?}",
112 tool.name
113 );
114 tool_working_set.insert(Arc::new(
115 ContextServerTool::new(
116 context_server_manager.clone(),
117 server.id(),
118 tool,
119 ),
120 ))
121 })
122 .collect::<Vec<_>>();
123
124 this.update(&mut cx, |this, _cx| {
125 this.context_server_tool_ids.insert(server_id, tool_ids);
126 })
127 .log_err();
128 }
129 }
130 }
131 })
132 .detach();
133 }
134 }
135 context_server::manager::Event::ServerStopped { server_id } => {
136 if let Some(tool_ids) = self.context_server_tool_ids.remove(server_id) {
137 tool_working_set.remove(&tool_ids);
138 }
139 }
140 }
141 }
142}
143
144impl ThreadStore {
145 /// Creates some mocked recent threads for testing purposes.
146 fn mock_recent_threads(&mut self, cx: &mut ModelContext<Self>) {
147 use language_model::Role;
148
149 self.threads.push(cx.new_model(|cx| {
150 let mut thread = Thread::new(self.tools.clone(), cx);
151 thread.insert_user_message("Hello! Can you help me understand quantum computing?", cx);
152 thread.insert_message(Role::Assistant, "Of course! I'd be happy to help you understand quantum computing. Quantum computing is a fascinating field that uses the principles of quantum mechanics to process information. Unlike classical computers that use bits (0s and 1s), quantum computers use quantum bits or 'qubits'. These qubits can exist in multiple states simultaneously, a property called superposition. This allows quantum computers to perform certain calculations much faster than classical computers. What specific aspect of quantum computing would you like to know more about?", cx);
153 thread.insert_user_message("That's interesting! Can you explain how quantum entanglement is used in quantum computing?", cx);
154 thread.insert_message(Role::Assistant, "Certainly! Quantum entanglement is a key principle used in quantum computing. When two qubits become entangled, the state of one qubit is directly related to the state of the other, regardless of the distance between them. This property is used in quantum computing to create complex quantum states and to perform operations on multiple qubits simultaneously. Entanglement allows quantum computers to process information in ways that classical computers cannot, potentially solving certain problems much more efficiently. For example, it's crucial in quantum error correction and in algorithms like quantum teleportation, which is important for quantum communication.", cx);
155 thread
156 }));
157
158 self.threads.push(cx.new_model(|cx| {
159 let mut thread = Thread::new(self.tools.clone(), cx);
160 thread.insert_user_message("Can you show me an example of Rust code for a simple web server?", cx);
161 thread.insert_message(Role::Assistant, "Certainly! Here's an example of a simple web server in Rust using the `actix-web` framework:
162
163 ```rust
164 use actix_web::{web, App, HttpResponse, HttpServer, Responder};
165
166 async fn hello() -> impl Responder {
167 HttpResponse::Ok().body(\"Hello, World!\")
168 }
169
170 #[actix_web::main]
171 async fn main() -> std::io::Result<()> {
172 HttpServer::new(|| {
173 App::new()
174 .route(\"/\", web::get().to(hello))
175 })
176 .bind(\"127.0.0.1:8080\")?
177 .run()
178 .await
179 }
180 ```
181
182 This code creates a basic web server that responds with 'Hello, World!' when you access the root URL. Here's a breakdown of what's happening:
183
184 1. We import necessary items from the `actix-web` crate.
185 2. We define an async `hello` function that returns a simple HTTP response.
186 3. In the `main` function, we set up the server to listen on `127.0.0.1:8080`.
187 4. We configure the app to respond to GET requests on the root path with our `hello` function.
188
189 To run this, you'd need to add `actix-web` to your `Cargo.toml` dependencies:
190
191 ```toml
192 [dependencies]
193 actix-web = \"4.0\"
194 ```
195
196 Then you can run the server with `cargo run` and access it at `http://localhost:8080`.".unindent(), cx);
197 thread.insert_user_message("That's great! Can you explain more about async functions in Rust?", cx);
198 thread.insert_message(Role::Assistant, "Certainly! Async functions are a key feature in Rust for writing efficient, non-blocking code, especially for I/O-bound operations. Here's an overview:
199
200 1. **Syntax**: Async functions are declared using the `async` keyword:
201
202 ```rust
203 async fn my_async_function() -> Result<(), Error> {
204 // Asynchronous code here
205 }
206 ```
207
208 2. **Futures**: Async functions return a `Future`. A `Future` represents a value that may not be available yet but will be at some point.
209
210 3. **Await**: Inside an async function, you can use the `.await` syntax to wait for other async operations to complete:
211
212 ```rust
213 async fn fetch_data() -> Result<String, Error> {
214 let response = make_http_request().await?;
215 let data = process_response(response).await?;
216 Ok(data)
217 }
218 ```
219
220 4. **Non-blocking**: Async functions allow the runtime to work on other tasks while waiting for I/O or other operations to complete, making efficient use of system resources.
221
222 5. **Runtime**: To execute async code, you need a runtime like `tokio` or `async-std`. Actix-web, which we used in the previous example, includes its own runtime.
223
224 6. **Error Handling**: Async functions work well with Rust's `?` operator for error handling.
225
226 Async programming in Rust provides a powerful way to write concurrent code that's both safe and efficient. It's particularly useful for servers, network programming, and any application that deals with many concurrent operations.".unindent(), cx);
227 thread
228 }));
229 }
230}