Some cleanup

Author: brandonskiser

crates/agent/src/agent/agent_config/parse.rs

@@ -151,8 +151,6 @@ pub enum ToolParseErrorKind {
     SchemaFailure(String),
     #[error("The tool arguments failed validation: {}", .0)]
     InvalidArgs(String),
-    #[error("The tool name could not be resolved: {}", .0)]
-    AmbiguousToolName(String),
     #[error("An unexpected error occurred parsing the tools: {}", .0)]
     Other(#[from] AgentError),
 }

crates/agent/src/agent/agent_loop/mod.rs

@@ -23,7 +23,6 @@ use protocol::{
     StreamMetadata,
     UserTurnMetadata,
 };
-use rand::seq::IndexedRandom;
 use serde::{
     Deserialize,
     Serialize,
@@ -77,10 +76,6 @@ impl AgentLoopId {
             rand: rand::random::<u32>(),
         }
     }
-
-    pub fn agent_id(&self) -> &AgentId {
-        &self.agent_id
-    }
 }
 
 impl std::fmt::Display for AgentLoopId {
@@ -89,23 +84,6 @@ impl std::fmt::Display for AgentLoopId {
     }
 }
 
-// impl FromStr for AgentLoopId {
-//     type Err = String;
-//
-//     fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
-//         match s.find("/") {
-//             Some(i) => Ok(Self {
-//                 agent_id: s[..i].to_string(),
-//                 rand: match s[i + 1..].to_string().parse() {
-//                     Ok(v) => v,
-//                     Err(_) => return Err(s.to_string()),
-//                 },
-//             }),
-//             None => Err(s.to_string()),
-//         }
-//     }
-// }
-
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize, strum::Display, strum::EnumString)]
 #[serde(rename_all = "camelCase")]
 #[strum(serialize_all = "camelCase")]
@@ -126,14 +104,6 @@ pub enum LoopState {
     Errored,
 }
 
-// #[derive(Debug)]
-// struct StreamRequest {
-//     model: Box<dyn AgentLoopModel>,
-//     messages: Vec<Message>,
-//     tool_specs: Option<Vec<ToolSpec>>,
-//     system_prompt: Option<String>,
-// }
-
 /// Tracks the execution of a user turn, ending when either the model returns a response with no
 /// tool uses, or a non-retryable error is encountered.
 pub struct AgentLoop {
@@ -147,6 +117,7 @@ pub struct AgentLoop {
     cancel_token: CancellationToken,
 
     /// The current response stream future being received along with it's associated parse state
+    #[allow(clippy::type_complexity)]
     curr_stream: Option<(
         StreamParseState,
         Pin<Box<dyn Stream<Item = Result<StreamEvent, StreamError>> + Send>>,
@@ -201,15 +172,14 @@ impl AgentLoop {
     /// the spawned task.
     pub fn spawn(mut self) -> AgentLoopHandle {
         let id_clone = self.id.clone();
-        let cancel_token_clone = self.cancel_token.clone();
         let loop_event_rx = self.loop_event_rx.take().expect("loop_event_rx should exist");
         let loop_req_tx = self.loop_req_tx.take().expect("loop_req_tx should exist");
         let handle = tokio::spawn(async move {
             info!("agent loop start");
             self.run().await;
             info!("agent loop end");
         });
-        AgentLoopHandle::new(id_clone, loop_req_tx, loop_event_rx, cancel_token_clone, handle)
+        AgentLoopHandle::new(id_clone, loop_req_tx, loop_event_rx, handle)
     }
 
     async fn run(mut self) {
@@ -349,15 +319,6 @@ impl AgentLoop {
 
                 Ok(AgentLoopResponse::Metadata(metadata))
             },
-
-            AgentLoopRequest::GetPendingToolUses => {
-                if self.execution_state != LoopState::PendingToolUseResults {
-                    return Ok(AgentLoopResponse::PendingToolUses(None));
-                }
-                let tool_uses = self.stream_states.last().map(|s| s.tool_uses.clone());
-                debug_assert!(tool_uses.as_ref().is_some_and(|v| !v.is_empty()));
-                Ok(AgentLoopResponse::PendingToolUses(tool_uses))
-            },
         }
     }
 
@@ -648,8 +609,6 @@ pub struct AgentLoopHandle {
     /// Sender for sending requests to the agent loop
     sender: RequestSender<AgentLoopRequest, AgentLoopResponse, AgentLoopResponseError>,
     loop_event_rx: mpsc::Receiver<AgentLoopEventKind>,
-    /// A [CancellationToken] used for gracefully closing the agent loop.
-    cancel_token: CancellationToken,
     /// The [JoinHandle] to the task executing the agent loop.
     handle: JoinHandle<()>,
 }
@@ -659,14 +618,12 @@ impl AgentLoopHandle {
         id: AgentLoopId,
         sender: RequestSender<AgentLoopRequest, AgentLoopResponse, AgentLoopResponseError>,
         loop_event_rx: mpsc::Receiver<AgentLoopEventKind>,
-        cancel_token: CancellationToken,
         handle: JoinHandle<()>,
     ) -> Self {
         Self {
             id,
             sender,
             loop_event_rx,
-            cancel_token,
             handle,
         }
     }
@@ -676,19 +633,6 @@ impl AgentLoopHandle {
         &self.id
     }
 
-    /// Id of the agent this loop was created for.
-    pub fn agent_id(&self) -> &AgentId {
-        self.id.agent_id()
-    }
-
-    pub fn clone_weak(&self) -> AgentLoopWeakHandle {
-        AgentLoopWeakHandle {
-            id: self.id.clone(),
-            sender: self.sender.clone(),
-            cancel_token: self.cancel_token.clone(),
-        }
-    }
-
     pub async fn recv(&mut self) -> Option<AgentLoopEventKind> {
         self.loop_event_rx.recv().await
     }
@@ -722,21 +666,6 @@ impl AgentLoopHandle {
         }
     }
 
-    pub async fn get_pending_tool_uses(&self) -> Result<Option<Vec<ToolUseBlock>>, AgentLoopResponseError> {
-        match self
-            .sender
-            .send_recv(AgentLoopRequest::GetPendingToolUses)
-            .await
-            .unwrap_or(Err(AgentLoopResponseError::AgentLoopExited))?
-        {
-            AgentLoopResponse::PendingToolUses(v) => Ok(v),
-            other => Err(AgentLoopResponseError::Custom(format!(
-                "unknown response getting stream metadata: {:?}",
-                other,
-            ))),
-        }
-    }
-
     /// Ends the agent loop
     pub async fn close(&self) -> Result<UserTurnMetadata, AgentLoopResponseError> {
         match self
@@ -760,105 +689,3 @@ impl Drop for AgentLoopHandle {
         self.handle.abort();
     }
 }
-
-/// A weak handle to an executing agent loop.
-///
-/// Where [AgentLoopHandle] can receive agent loop events and abort the task on drop,
-/// [AgentLoopWeakHandle] is only used for sending messages to the agent loop.
-#[derive(Debug, Clone)]
-pub struct AgentLoopWeakHandle {
-    id: AgentLoopId,
-    sender: RequestSender<AgentLoopRequest, AgentLoopResponse, AgentLoopResponseError>,
-    cancel_token: CancellationToken,
-}
-
-impl AgentLoopWeakHandle {
-    pub async fn send_request<M: AgentLoopModel>(
-        &self,
-        model: M,
-        args: SendRequestArgs,
-    ) -> Result<AgentLoopResponse, AgentLoopResponseError> {
-        self.sender
-            .send_recv(AgentLoopRequest::SendRequest {
-                model: Box::new(model),
-                args,
-            })
-            .await
-            .unwrap_or(Err(AgentLoopResponseError::AgentLoopExited))
-    }
-
-    pub async fn get_loop_state(&self) -> Result<LoopState, AgentLoopResponseError> {
-        match self
-            .sender
-            .send_recv(AgentLoopRequest::GetExecutionState)
-            .await
-            .unwrap_or(Err(AgentLoopResponseError::AgentLoopExited))?
-        {
-            AgentLoopResponse::ExecutionState(state) => Ok(state),
-            other => Err(AgentLoopResponseError::Custom(format!(
-                "unknown response getting execution state: {:?}",
-                other,
-            ))),
-        }
-    }
-
-    pub async fn get_pending_tool_uses(&self) -> Result<Option<Vec<ToolUseBlock>>, AgentLoopResponseError> {
-        match self
-            .sender
-            .send_recv(AgentLoopRequest::GetPendingToolUses)
-            .await
-            .unwrap_or(Err(AgentLoopResponseError::AgentLoopExited))?
-        {
-            AgentLoopResponse::PendingToolUses(v) => Ok(v),
-            other => Err(AgentLoopResponseError::Custom(format!(
-                "unknown response getting stream metadata: {:?}",
-                other,
-            ))),
-        }
-    }
-
-    /// Ends the agent loop
-    pub async fn close(&self) -> Result<UserTurnMetadata, AgentLoopResponseError> {
-        match self
-            .sender
-            .send_recv(AgentLoopRequest::Close)
-            .await
-            .unwrap_or(Err(AgentLoopResponseError::AgentLoopExited))?
-        {
-            AgentLoopResponse::Metadata(md) => Ok(md),
-            other => Err(AgentLoopResponseError::Custom(format!(
-                "unknown response getting execution state: {:?}",
-                other,
-            ))),
-        }
-    }
-
-    /// Cancel the executing loop for graceful shutdown.
-    fn cancel(&self) {
-        self.cancel_token.cancel();
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use std::sync::Arc;
-
-    use super::*;
-    use crate::api_client::error::{
-        ConverseStreamError,
-        ConverseStreamErrorKind,
-    };
-
-    #[test]
-    fn test_other_stream_err_downcasting() {
-        let err = StreamError::new(StreamErrorKind::Interrupted).with_source(Arc::new(ConverseStreamError::new(
-            ConverseStreamErrorKind::ModelOverloadedError,
-            None::<aws_smithy_types::error::operation::BuildError>, /* annoying type inference
-                                                                     * required */
-        )));
-        assert!(
-            err.as_rts_error()
-                .is_some_and(|r| matches!(r.kind, ConverseStreamErrorKind::ModelOverloadedError))
-        );
-    }
-}

crates/agent/src/agent/agent_loop/model.rs

@@ -42,13 +42,6 @@ pub enum Models {
 }
 
 impl Models {
-    pub fn supported_model(&self) -> SupportedModel {
-        match self {
-            Models::Rts(_) => SupportedModel::Rts,
-            Models::Test(_) => SupportedModel::Test,
-        }
-    }
-
     pub fn state(&self) -> ModelsState {
         match self {
             Models::Rts(v) => ModelsState::Rts {
@@ -79,17 +72,6 @@ impl Default for ModelsState {
     }
 }
 
-/// Identifier for the models we support.
-///
-/// TODO - probably not required, use [ModelsState] instead
-#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, strum::Display, strum::EnumString)]
-#[serde(rename_all = "camelCase")]
-#[strum(serialize_all = "camelCase")]
-pub enum SupportedModel {
-    Rts,
-    Test,
-}
-
 impl Model for Models {
     fn stream(
         &self,
@@ -100,7 +82,7 @@ impl Model for Models {
     ) -> Pin<Box<dyn Stream<Item = Result<StreamEvent, StreamError>> + Send + 'static>> {
         match self {
             Models::Rts(rts_model) => rts_model.stream(messages, tool_specs, system_prompt, cancel_token),
-            Models::Test(test_model) => todo!(),
+            Models::Test(test_model) => test_model.stream(messages, tool_specs, system_prompt, cancel_token),
         }
     }
 }
@@ -113,3 +95,15 @@ impl TestModel {
         Self {}
     }
 }
+
+impl Model for TestModel {
+    fn stream(
+        &self,
+        messages: Vec<Message>,
+        tool_specs: Option<Vec<ToolSpec>>,
+        system_prompt: Option<String>,
+        cancel_token: CancellationToken,
+    ) -> Pin<Box<dyn Stream<Item = Result<StreamEvent, StreamError>> + Send + 'static>> {
+        todo!()
+    }
+}

crates/agent/src/agent/agent_loop/protocol.rs

@@ -24,7 +24,6 @@ use super::{
     InvalidToolUse,
     LoopState,
 };
-use crate::agent::types::AgentId;
 
 #[derive(Debug)]
 pub enum AgentLoopRequest {
@@ -33,7 +32,6 @@ pub enum AgentLoopRequest {
         model: Box<dyn AgentLoopModel>,
         args: SendRequestArgs,
     },
-    GetPendingToolUses,
     /// Ends the agent loop
     Close,
 }
@@ -93,11 +91,6 @@ impl AgentLoopEvent {
     pub fn new(id: AgentLoopId, kind: AgentLoopEventKind) -> Self {
         Self { id, kind }
     }
-
-    /// Id of the agent this loop event is associated with
-    pub fn agent_id(&self) -> &AgentId {
-        self.id.agent_id()
-    }
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]