feat: Agent sessions with multi

Cargo.toml

@@ -93,6 +93,7 @@ copypasta = "0.10.1"
 strip-ansi-escapes = "0.2.1"
 inquire = "0.7.5"
 config = { version = "0.15.6", features = ["toml", "convert-case"] }
+dyn-clone = "1.0.17"
 
 # Something is still pulling in libssl, this is a quickfix and should be investigated
 [target.'cfg(linux)'.dependencies]

src/agent/agent_session.rs

@@ -0,0 +1,102 @@
+use std::sync::Arc;
+
+use anyhow::{Context as _, Result};
+use swiftide::{
+    agents::{
+        system_prompt::SystemPrompt, tools::local_executor::LocalExecutor, Agent, DefaultContext,
+    },
+    chat_completion::{self, errors::ToolError, ChatCompletion, Tool, ToolOutput},
+    prompt::Prompt,
+    traits::{AgentContext, Command, SimplePrompt, ToolExecutor},
+};
+use swiftide_macros::Tool;
+use uuid::Uuid;
+
+use crate::{commands::Responder, git::github::GithubSession, repository::Repository};
+
+use super::{env_setup::AgentEnvironment, RunningAgent};
+// TODO: No need for this to be an interface, just making it easier to test
+// pub trait AgentSession: Send + Sync {
+//     fn init(&self) -> Result<()>;
+//
+//     fn session_id(&self) -> Uuid;
+//
+//     fn repository(&self) -> &Repository;
+//     fn github_session(&self) -> Option<&GithubSession>;
+//     fn branch_name(&self) -> &str;
+//     fn executor(&self) -> Arc<dyn ToolExecutor>;
+//     fn agent_environment(&self) -> &AgentEnvironment;
+//     fn available_tools(&self) -> &[Box<dyn Tool>];
+//
+//     fn responder(&self) -> &dyn Responder;
+//     fn responder_clone(&self) -> Box<dyn Responder>;
+// }
+
+pub struct AgentSession {
+    session_id: Uuid,
+    repository: Repository,
+    agent_environment: Option<AgentEnvironment>,
+    responder: Arc<dyn Responder>,
+
+    // After calling init
+    github_session: Option<GithubSession>,
+    executor: Option<Arc<dyn ToolExecutor>>,
+    available_tools: Option<Vec<Box<dyn Tool>>>,
+    // The agent that is currently running
+    // active_agent: RunningAgent,
+}
+
+impl AgentSession {
+    pub fn new(session_id: Uuid, repository: Repository, responder: Arc<dyn Responder>) -> Self {
+        Self {
+            session_id,
+            repository,
+            responder,
+            agent_environment: None,
+            github_session: None,
+            executor: None,
+            available_tools: None,
+        }
+    }
+
+    pub async fn init(&self) -> Result<()> {
+        Ok(())
+    }
+
+    pub fn repository(&self) -> &Repository {
+        &self.repository
+    }
+
+    pub fn github_session(&self) -> Option<&GithubSession> {
+        self.github_session.as_ref()
+    }
+
+    pub fn executor(&self) -> Arc<dyn ToolExecutor> {
+        Arc::clone(
+            self.executor
+                .as_ref()
+                .expect("Agent session not initialized"),
+        )
+    }
+
+    pub fn agent_environment(&self) -> &AgentEnvironment {
+        self.agent_environment
+            .as_ref()
+            .expect("Agent session not initialized")
+    }
+
+    pub fn available_tools(&self) -> &[Box<dyn Tool>] {
+        self.available_tools
+            .as_ref()
+            .expect("Agent session not initialized")
+            .as_slice()
+    }
+
+    pub fn responder(&self) -> &dyn Responder {
+        self.responder.as_ref()
+    }
+
+    pub fn responder_clone(&self) -> Arc<dyn Responder> {
+        Arc::clone(&self.responder)
+    }
+}

src/agent/delegating_agent.rs

@@ -0,0 +1,186 @@
+use anyhow::{Context as _, Result};
+use std::sync::Arc;
+use swiftide_macros::{tool, Tool};
+use uuid::Uuid;
+
+use super::{
+    agent_session::AgentSession,
+    conversation_summarizer::ConversationSummarizer,
+    env_setup::{self, EnvSetup},
+    tool_summarizer::ToolSummarizer,
+    tools, RunningAgent,
+};
+use swiftide::{
+    agents::{
+        system_prompt::SystemPrompt, tools::local_executor::LocalExecutor, Agent, DefaultContext,
+    },
+    chat_completion::{self, errors::ToolError, ChatCompletion, Tool, ToolOutput},
+    prompt::Prompt,
+    traits::{AgentContext, Command, SimplePrompt, ToolExecutor},
+};
+
+use crate::{commands::Responder, git::github::GithubSession, indexing, repository::Repository};
+
+use super::env_setup::AgentEnvironment;
+
+#[derive(Clone, Tool)]
+#[tool(
+    description = "Delegate to the coding agent",
+    param(
+        name = "task",
+        description = "A thorough description of the task to be completed"
+    )
+)]
+pub struct RunCodingAgent {
+    session: Arc<AgentSession>,
+}
+
+impl RunCodingAgent {
+    pub fn new(session: Arc<AgentSession>) -> Self {
+        Self { session }
+    }
+
+    pub async fn run_coding_agent(
+        &self,
+        context: &dyn AgentContext,
+        task: &str,
+    ) -> Result<ToolOutput, ToolError> {
+        // Start the agent on this session
+        // TODO:
+        // - How do we deal with output
+        // - Any additional prompting etc
+        // - Track the agent; only allow one agent per session
+        // - Keep it simple -> Delegating to this agent will make this the active agent
+        //
+        Ok("".into())
+    }
+}
+
+pub async fn start(query: &str, session: AgentSession) -> Result<RunningAgent> {
+    // Ensure the session is set up
+    // tools, etc => Session should provide:
+    // - providers
+    // - github session
+    // - branch_name / executor / initial context
+    // - env_setup / agent env
+    // - all tools
+    // - Tracks running agent?
+    // - Interaction with command responder?
+    // - uuid = session id
+    //
+    let executor = Arc::clone(&session.executor());
+    let mut context = Arc::new(DefaultContext::from_executor(Arc::clone(&executor)));
+    let initial_context = generate_initial_context(&session.repository(), query).await?;
+
+    let query_provider: Box<dyn ChatCompletion> =
+        session.repository().config().query_provider().try_into()?;
+    let fast_query_provider: Box<dyn SimplePrompt> = session
+        .repository()
+        .config()
+        .indexing_provider()
+        .try_into()?;
+
+    let system_prompt = SystemPrompt::builder().build()?;
+    let tools = session
+        .available_tools()
+        .iter()
+        .filter_map(|tool| {
+            if [
+                "search_file",
+                "search_code",
+                "fetch_url",
+                "explain_code",
+                "read_file",
+                "github_search_code",
+                "search_web",
+                "run_tests",
+                "run_coverage",
+            ]
+            .contains(&tool.name())
+            {
+                Some(tool.clone())
+            } else {
+                None
+            }
+        })
+        .collect::<Vec<_>>();
+
+    let tool_summarizer = ToolSummarizer::new(
+        fast_query_provider,
+        &["run_tests", "run_coverage"],
+        &tools,
+        &session.agent_environment().start_ref,
+    );
+    let conversation_summarizer = ConversationSummarizer::new(
+        query_provider.clone(),
+        &tools,
+        &session.agent_environment().start_ref,
+    );
+    let responder = session.responder_clone();
+
+    // tmp
+    let tx1 = responder.clone();
+    let tx2 = responder.clone();
+    let tx3 = responder.clone();
+
+    let agent = Agent::builder()
+        .context(Arc::clone(&context) as Arc<dyn AgentContext>)
+        .system_prompt(system_prompt)
+        .tools(tools)
+        .before_all(move |context| {
+            let initial_context = initial_context.clone();
+
+            Box::pin(async move {
+                context
+                    .add_message(chat_completion::ChatMessage::new_user(initial_context))
+                    .await;
+
+                let top_level_project_overview = context.exec_cmd(&Command::shell("fd -iH -d2 -E '.git/*'")).await?.output;
+                context.add_message(chat_completion::ChatMessage::new_user(format!("The following is a max depth 2, high level overview of the directory structure of the project: \n ```{top_level_project_overview}```"))).await;
+
+                Ok(())
+            })
+        })
+        .on_new_message(move |_, message| {
+            let command_responder = tx1.clone();
+            let message = message.clone();
+
+            Box::pin(async move {
+                command_responder.agent_message(message);
+
+                Ok(())
+            })
+        })
+        .before_completion(move |_, _| {
+            let command_responder = tx2.clone();
+            Box::pin(async move {
+                command_responder.update("running completions");
+                Ok(())
+            })
+        })
+        .before_tool(move |_, tool| {
+            let command_responder = tx3.clone();
+            let tool = tool.clone();
+            Box::pin(async move {
+                command_responder.update(&format!("running tool {}", tool.name()));
+                Ok(())
+            })
+        })
+        .after_tool(tool_summarizer.summarize_hook())
+        .after_each(conversation_summarizer.summarize_hook())
+        .llm(&query_provider)
+        .build()?;
+
+    RunningAgent::builder()
+        .agent(agent)
+        .executor(executor)
+        .agent_environment(session.agent_environment().clone())
+        .agent_context(context as Arc<dyn AgentContext>)
+        .build()
+}
+
+async fn generate_initial_context(repository: &Repository, query: &str) -> Result<String> {
+    let retrieved_context = indexing::query(repository, &query).await?;
+    let formatted_context = format!("Additional information:\n\n{retrieved_context}");
+    Ok(formatted_context)
+}

src/agent/mod.rs

@@ -1,4 +1,6 @@
+mod agent_session;
 mod conversation_summarizer;
+mod delegating_agent;
 mod env_setup;
 mod running_agent;
 mod tool_summarizer;

src/commands/responder.rs

@@ -1,7 +1,8 @@
 use std::sync::Arc;
 
+use dyn_clone::DynClone;
 #[cfg(test)]
-use mockall::automock;
+use mockall::mock;
 use swiftide::chat_completion;
 use uuid::Uuid;
 
@@ -46,8 +47,7 @@ impl CommandResponse {
 ///
 /// TODO: Consider, perhaps with the new structure, less concrete methods are needed
 /// and the frontend just uses a oneoff handler for each command
-#[cfg_attr(test, automock)]
-pub trait Responder: std::fmt::Debug + Send + Sync {
+pub trait Responder: std::fmt::Debug + Send + Sync + DynClone {
     /// Generic handler for command responses
     fn send(&self, response: CommandResponse);
 
@@ -67,6 +67,28 @@ pub trait Responder: std::fmt::Debug + Send + Sync {
     fn rename_branch(&self, name: &str);
 }
 
+dyn_clone::clone_trait_object!(Responder);
+
+#[cfg(test)]
+mock! {
+    #[derive(Debug)]
+    pub Responder {}
+
+    impl Responder for Responder {
+        fn send(&self, response: CommandResponse);
+        fn agent_message(&self, message: chat_completion::ChatMessage);
+        fn system_message(&self, message: &str);
+        fn update(&self, state: &str);
+        fn rename_chat(&self, name: &str);
+        fn rename_branch(&self, name: &str);
+    }
+
+    impl Clone for Responder {
+        fn clone(&self) -> Self;
+
+    }
+}
+
 // TODO: Naming should be identical to command response
 impl Responder for tokio::sync::mpsc::UnboundedSender<CommandResponse> {
     fn send(&self, response: CommandResponse) {