Skip to content

Commit

Permalink
[sidecar] complete integration of the tool based agent
Browse files Browse the repository at this point in the history
  • Loading branch information
theskcd committed Dec 20, 2024
1 parent 7c05932 commit 110ecb7
Show file tree
Hide file tree
Showing 6 changed files with 16 additions and 161 deletions.
21 changes: 10 additions & 11 deletions llm_client/src/clients/codestory.rs
Original file line number Diff line number Diff line change
Expand Up @@ -276,17 +276,16 @@ impl CodeStoryClient {
let access_token = self.access_token(api_key)?;

let request = OpenRouterRequest::from_chat_request(request, model.to_owned());
let mut response_stream = dbg!(
self.client
.post(endpoint)
.header("X-Accel-Buffering", "no")
.header("Authorization", format!("Bearer {}", access_token))
.json(&request)
.send()
.await
)?
.bytes_stream()
.eventsource();
let mut response_stream = self
.client
.post(endpoint)
.header("X-Accel-Buffering", "no")
.header("Authorization", format!("Bearer {}", access_token))
.json(&request)
.send()
.await?
.bytes_stream()
.eventsource();
let mut buffered_stream = "".to_owned();
// controls which tool we will be using if any
let mut tool_use_indication: Vec<(String, (String, String))> = vec![];
Expand Down
2 changes: 1 addition & 1 deletion sidecar/src/agentic/symbol/identifier.rs
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,7 @@ impl LLMProperties {
/// Only allow tool use when we are using anthropic since open-router does not
/// support the str_replace_editor tool natively
pub fn supports_midwit_and_tool_use(&self) -> bool {
self.llm() == &LLMType::ClaudeSonnet && matches!(&self.provider, &LLMProvider::Anthropic)
self.llm() == &LLMType::ClaudeSonnet && matches!(&self.provider, &LLMProvider::CodeStory(_))
}
}

Expand Down
147 changes: 0 additions & 147 deletions sidecar/src/agentic/tool/session/service.rs
Original file line number Diff line number Diff line change
Expand Up @@ -702,153 +702,6 @@ impl SessionService {
Ok(())
}

/// Use the tool based json mode over here
async fn tool_use_midwit_json_mode(
&self,
mut session: Session,
session_id: String,
storage_path: String,
user_message: String,
exchange_id: String,
all_files: Vec<String>,
open_files: Vec<String>,
shell: String,
project_labels: Vec<String>,
repo_ref: RepoRef,
root_directory: String,
tool_box: Arc<ToolBox>,
llm_broker: Arc<LLMBroker>,
user_context: UserContext,
is_midwit_tool_agent: bool,
mut message_properties: SymbolEventMessageProperties,
) -> Result<(), SymbolError> {
// os can be passed over here safely since we can assume the sidecar is running
// close to the vscode server
// we should ideally get this information from the vscode-server side setting
let tool_agent = ToolUseAgent::new(
llm_broker.clone(),
root_directory.to_owned(),
std::env::consts::OS.to_owned(),
shell.to_owned(),
None,
);

session = session.human_message_tool_use(
exchange_id.to_owned(),
user_message.to_owned(),
all_files,
open_files,
shell,
user_context,
);
let _ = self.save_to_storage(&session).await;

session = session.accept_open_exchanges_if_any(message_properties.clone());
let mut human_message_ticker = 0;
// now that we have saved it we can start the loop over here and look out for the cancellation
// token which will imply that we should end the current loop
loop {
let _ = self.save_to_storage(&session).await;
let tool_exchange_id = self
.tool_box
.create_new_exchange(session_id.to_owned(), message_properties.clone())
.await?;

println!("tool_exchange_id::({:?})", &tool_exchange_id);

let cancellation_token = tokio_util::sync::CancellationToken::new();

message_properties = message_properties
.set_request_id(tool_exchange_id.to_owned())
.set_cancellation_token(cancellation_token.clone());

// track the new exchange over here
self.track_exchange(&session_id, &tool_exchange_id, cancellation_token.clone())
.await;

let tool_use_output = session
// the clone here is pretty bad but its the easiest and the sanest
// way to keep things on the happy path
.clone()
.get_tool_to_use(
tool_box.clone(),
tool_exchange_id.to_owned(),
exchange_id.to_owned(),
tool_agent.clone(),
message_properties.clone(),
)
.await;

println!("tool_use_output::{:?}", tool_use_output);

match tool_use_output {
Ok(AgentToolUseOutput::Success((tool_input_partial, new_session))) => {
// update our session
session = new_session;
// store to disk
let _ = self.save_to_storage(&session).await;
let tool_type = tool_input_partial.to_tool_type();
session = session
.invoke_tool(
tool_type.clone(),
tool_input_partial,
tool_box.clone(),
root_directory.to_owned(),
message_properties.clone(),
)
.await?;

let _ = self.save_to_storage(&session).await;
if matches!(tool_type, ToolType::AskFollowupQuestions)
|| matches!(tool_type, ToolType::AttemptCompletion)
{
// we break if it is any of these 2 events, since these
// require the user to intervene
println!("session_service::tool_use_agentic::reached_terminating_tool");
break;
}
}
Ok(AgentToolUseOutput::Cancelled) => {
// if it is cancelled then we should break
break;
}
Ok(AgentToolUseOutput::Failed(failed_to_parse_output)) => {
let human_message = format!(
r#"Your output was incorrect, please give me the output in the correct format:
{}"#,
failed_to_parse_output.to_owned()
);
human_message_ticker = human_message_ticker + 1;
session = session.human_message(
human_message_ticker.to_string(),
human_message,
UserContext::default(),
vec![],
repo_ref.clone(),
);
let _ = message_properties
.ui_sender()
.send(UIEventWithID::tool_not_found(
session_id.to_owned(),
tool_exchange_id.to_owned(),
"Failed to get tool output".to_owned(),
));
}
Err(e) => {
let _ = message_properties
.ui_sender()
.send(UIEventWithID::tool_not_found(
session_id.to_owned(),
tool_exchange_id.to_owned(),
e.to_string(),
));
Err(e)?
}
}
}
Ok(())
}

pub async fn code_edit_agentic(
&self,
session_id: String,
Expand Down
3 changes: 3 additions & 0 deletions sidecar/src/agentic/tool/session/session.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1174,6 +1174,9 @@ impl Session {
"{}",
format!("inference::engine::retrying_tool_call::erroredbefore")
);
if matches!(tool_use_output, Err(SymbolError::CancelledResponseStream)) {
return Ok(AgentToolUseOutput::Cancelled);
}
tokio::time::sleep(Duration::from_secs(1)).await;
// just give it a plain retry and call it a day
tool_retry_index = tool_retry_index + 1;
Expand Down
2 changes: 1 addition & 1 deletion sidecar/src/agentic/tool/session/tool_use_agent.rs
Original file line number Diff line number Diff line change
Expand Up @@ -799,7 +799,7 @@ You accomplish a given task iteratively, breaking it down into clear steps and w
thinking.trim().to_owned(),
)))
} else {
Ok(ToolUseAgentOutputWithTools::Failure(None))
Err(SymbolError::CancelledResponseStream)
}
}

Expand Down
2 changes: 1 addition & 1 deletion sidecar/src/webserver/agentic.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1442,7 +1442,7 @@ pub async fn agent_tool_use(
let (sender, receiver) = tokio::sync::mpsc::unbounded_channel();
// check if model and provider combo supports tool use and midwit agent use
// let is_midwit_tool_agent = llm_provider.supports_midwit_and_tool_use();
let is_midwit_tool_agent = true;
let is_midwit_tool_agent = llm_provider.supports_midwit_and_tool_use();
let message_properties = SymbolEventMessageProperties::new(
SymbolEventRequestId::new(exchange_id.to_owned(), session_id.to_string()),
sender.clone(),
Expand Down

0 comments on commit 110ecb7

Please sign in to comment.