Restructure agent2 modules

Nathan Sobo created 5 months ago

Change summary

crates/agent2/Cargo.toml                     |  2 
crates/agent2/src/agent2.rs                  |  6 +++
crates/agent2/src/prompts.rs                 |  8 +++
crates/agent2/src/thread.rs                  | 19 +++------
crates/agent2/src/thread/tests.rs            | 44 +++++++++++----------
crates/agent2/src/thread/tests/test_tools.rs |  0 
crates/agent2/src/tools/glob.rs              |  2 
7 files changed, 44 insertions(+), 37 deletions(-)

Detailed changes

crates/agent2/Cargo.toml

@@ -6,7 +6,7 @@ license = "GPL-3.0-or-later"
 publish = false
 
 [lib]
-path = "src/agent.rs"
+path = "src/agent2.rs"
 
 [lints]
 workspace = true

crates/agent2/src/agent2.rs

@@ -0,0 +1,6 @@
+mod prompts;
+mod templates;
+mod thread;
+mod tools;
+
+pub use thread::*;

crates/agent2/src/prompts.rs

@@ -1,4 +1,10 @@
-use super::*;
+use crate::{
+    templates::{BaseTemplate, Template, Templates, WorktreeData},
+    thread::Prompt,
+};
+use anyhow::Result;
+use gpui::{App, Entity};
+use project::Project;
 
 struct BasePrompt {
     project: Entity<Project>,

crates/agent2/src/agent.rs → crates/agent2/src/thread.rs

@@ -1,24 +1,17 @@
-mod prompts;
-mod templates;
-#[cfg(test)]
-mod tests;
-mod tools;
-
-use anyhow::{Result, anyhow};
+use crate::templates::Templates;
+use anyhow::{anyhow, Result};
 use futures::{channel::mpsc, future};
-use gpui::{App, Context, Entity, SharedString, Task};
+use gpui::{App, Context, SharedString, Task};
 use language_model::{
     CompletionIntent, CompletionMode, LanguageModel, LanguageModelCompletionError,
     LanguageModelCompletionEvent, LanguageModelRequest, LanguageModelRequestMessage,
     LanguageModelRequestTool, LanguageModelToolResult, LanguageModelToolResultContent,
     LanguageModelToolSchemaFormat, LanguageModelToolUse, MessageContent, Role, StopReason,
 };
-use project::Project;
 use schemars::{JsonSchema, Schema};
 use serde::Deserialize;
 use smol::stream::StreamExt;
 use std::{collections::BTreeMap, sync::Arc};
-use templates::{BaseTemplate, Template, Templates, WorktreeData};
 use util::ResultExt;
 
 #[derive(Debug)]
@@ -29,11 +22,11 @@ pub struct AgentMessage {
 
 pub type AgentResponseEvent = LanguageModelCompletionEvent;
 
-trait Prompt {
+pub trait Prompt {
     fn render(&self, prompts: &Templates, cx: &App) -> Result<String>;
 }
 
-pub struct Agent {
+pub struct Thread {
     messages: Vec<AgentMessage>,
     completion_mode: CompletionMode,
     /// Holds the task that handles agent interaction until the end of the turn.
@@ -47,7 +40,7 @@ pub struct Agent {
     // action_log: Entity<ActionLog>,
 }
 
-impl Agent {
+impl Thread {
     pub fn new(templates: Arc<Templates>) -> Self {
         Self {
             messages: Vec::new(),

crates/agent2/src/tests.rs → crates/agent2/src/thread/tests.rs

@@ -1,12 +1,16 @@
 use super::*;
-use client::{Client, UserStore, proto::language_server_prompt_request};
+use client::{proto::language_server_prompt_request, Client, UserStore};
 use fs::FakeFs;
-use gpui::{AppContext, TestAppContext};
-use language_model::LanguageModelRegistry;
+use gpui::{AppContext, Entity, TestAppContext};
+use language_model::{
+    LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent,
+    LanguageModelRegistry, MessageContent, StopReason,
+};
 use reqwest_client::ReqwestClient;
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
-use std::time::Duration;
+use smol::stream::StreamExt;
+use std::{sync::Arc, time::Duration};
 
 mod test_tools;
 use test_tools::*;
@@ -69,21 +73,19 @@ async fn test_basic_tool_calls(cx: &mut TestAppContext) {
         vec![StopReason::ToolUse, StopReason::EndTurn]
     );
     agent.update(cx, |agent, _cx| {
-        assert!(
-            agent
-                .messages
-                .last()
-                .unwrap()
-                .content
-                .iter()
-                .any(|content| {
-                    if let MessageContent::Text(text) = content {
-                        text.contains("Ding")
-                    } else {
-                        false
-                    }
-                })
-        );
+        assert!(agent
+            .messages
+            .last()
+            .unwrap()
+            .content
+            .iter()
+            .any(|content| {
+                if let MessageContent::Text(text) = content {
+                    text.contains("Ding")
+                } else {
+                    false
+                }
+            }));
     });
 }
 
@@ -200,7 +202,7 @@ fn stop_events(
 
 struct AgentTest {
     model: Arc<dyn LanguageModel>,
-    agent: Entity<Agent>,
+    agent: Entity<Thread>,
 }
 
 async fn setup(cx: &mut TestAppContext) -> AgentTest {
@@ -210,7 +212,7 @@ async fn setup(cx: &mut TestAppContext) -> AgentTest {
     // let project = Project::test(fs.clone(), [], cx).await;
     // let action_log = cx.new(|_| ActionLog::new(project.clone()));
     let templates = Templates::new();
-    let agent = cx.new(|_| Agent::new(templates));
+    let agent = cx.new(|_| Thread::new(templates));
 
     let model = cx
         .update(|cx| {

crates/agent2/src/tests/test_tools.rs → crates/agent2/src/thread/tests/test_tools.rs

crates/agent2/src/tools/glob.rs

@@ -9,7 +9,7 @@ use worktree::Snapshot as WorktreeSnapshot;
 
 use crate::{
     templates::{GlobTemplate, Template, Templates},
-    AgentTool,
+    thread::AgentTool,
 };
 
 // Description is dynamic, see `fn description` below