use std::{
io::{BufRead, IsTerminal},
path::{Path, PathBuf},
};
use clap::Args;
use llmy_agent::tool::ToolBox;
use llmy_agent_tools::bash::{BashTool, BashToolConfig};
use llmy_agent_tools::files::{
DeleteFileTool, EditFileTool, FindFileTool, ListDirectoryTool, ReadFileTool, WriteFileTool,
};
use llmy_agent_tools::memory::{
AgentMemory, AgentMemoryContext,
embed::{SimilarityModel, SimilarityModelConfig},
};
use llmy_clap::OpenAISetup;
use llmy_client::client::LLM;
use llmy_harness::{Agent, memory::AgentMemorySystemPromptCriteria};
use rustyline::{DefaultEditor, error::ReadlineError};
use super::chat_commands::{ChatInput, parse_chat_input, run_chat_command};
#[derive(Args)]
pub struct ChatArgs {
#[command(flatten)]
openai: OpenAISetup,
#[arg(long)]
system: Option<String>,
#[arg(long, value_name = "ROOT", num_args = 0..=1, default_missing_value = ".")]
agent_files: Option<PathBuf>,
#[arg(long, default_value_t = false)]
agent_bash: bool,
#[arg(long, default_value_t = false)]
memory: bool,
#[arg(long, default_value_t = default_memory_embed_model())]
memory_embed_model: String,
#[arg(long)]
memory_cache_dir: Option<PathBuf>,
}
pub async fn run_chat(args: ChatArgs) -> color_eyre::Result<()> {
let settings = args.openai.settings();
let llm: LLM = args.openai.clone().to_llm();
let system = args
.system
.as_deref()
.unwrap_or("You are a helpful assistant.");
let files_root = resolve_files_root(args.agent_files.clone())?;
let bash_root = resolve_bash_root(args.agent_bash)?;
let system_prompt = build_system_prompt(system, files_root.as_deref(), bash_root.as_deref());
let tools = build_toolbox(files_root.clone(), bash_root);
let mut agent = build_agent(&args, system_prompt, tools).await?;
let stdin = std::io::stdin();
let is_tty = stdin.is_terminal();
let mut reader = ChatReader::new(is_tty)?;
while let Some(input) = reader.read_next()? {
match parse_chat_input(&input) {
Ok(ChatInput::User(input)) => {
agent
.step_with_user(input, &llm, Some("chat"), Some(settings.clone()))
.await?;
while print_last_step(&agent, is_tty) {
agent
.step(&llm, Some("chat"), Some(settings.clone()))
.await?;
}
}
Ok(ChatInput::Command(command)) => {
run_chat_command(
command,
&mut agent,
&llm,
Some("chat"),
Some(settings.clone()),
is_tty,
)
.await?;
}
Err(error) => {
eprintln!("{error}");
}
}
}
Ok(())
}
async fn build_agent(
args: &ChatArgs,
system_prompt: String,
tools: ToolBox,
) -> color_eyre::Result<Agent> {
if !args.memory {
return Ok(Agent::new(
system_prompt,
tools,
"llmy-cli-chat".to_string(),
));
}
let memory = AgentMemoryContext::new(
AgentMemory::default(),
SimilarityModel::new(build_memory_config(args)?).await?,
);
let criteria = AgentMemorySystemPromptCriteria::default();
Ok(Agent::with_memory(
system_prompt,
tools,
"llmy-cli-chat".to_string(),
&memory,
&criteria,
)
.await)
}
fn build_memory_config(args: &ChatArgs) -> color_eyre::Result<SimilarityModelConfig> {
let mut config = SimilarityModelConfig::default();
config.model = args.memory_embed_model.parse().map_err(|error| {
color_eyre::eyre::eyre!(
"invalid memory embed model {:?}: {}",
args.memory_embed_model,
error
)
})?;
config.cache_dir = args.memory_cache_dir.clone();
Ok(config)
}
fn default_memory_embed_model() -> String {
SimilarityModelConfig::default().model.to_string()
}
fn resolve_files_root(root: Option<PathBuf>) -> color_eyre::Result<Option<PathBuf>> {
root.map(|path| {
if path == Path::new(".") {
std::env::current_dir().map_err(Into::into)
} else {
path.canonicalize().map_err(Into::into)
}
})
.transpose()
}
fn resolve_bash_root(enabled: bool) -> color_eyre::Result<Option<PathBuf>> {
enabled
.then(std::env::current_dir)
.transpose()
.map_err(Into::into)
}
fn build_toolbox(files_root: Option<PathBuf>, bash_root: Option<PathBuf>) -> ToolBox {
let mut toolbox = ToolBox::new();
if let Some(root) = bash_root {
toolbox.add_tool(BashTool::new(root, BashToolConfig::default()));
}
if let Some(root) = files_root {
toolbox.add_tool(ReadFileTool::new(root.clone()));
toolbox.add_tool(ListDirectoryTool::new_root(root.clone()));
toolbox.add_tool(FindFileTool::new(root.clone()));
toolbox.add_tool(WriteFileTool::new(root.clone()));
toolbox.add_tool(DeleteFileTool::new(root.clone()));
toolbox.add_tool(EditFileTool::new(root));
}
toolbox
}
fn build_system_prompt(base: &str, files_root: Option<&Path>, bash_root: Option<&Path>) -> String {
let mut sections = vec![base.to_string()];
if let Some(root) = files_root {
sections.push(format!(
"You can use sandboxed file tools rooted at {}. All tool paths must be relative to this root. Use the available file tools when you need to inspect or modify files.",
root.display()
));
}
if let Some(root) = bash_root {
sections.push(format!(
"You can use the dangerous bash tool to execute shell commands. Commands start in {} unless `working_directory` is provided. This tool can modify files, run programs, and access the network, so use it carefully.",
root.display()
));
}
sections.join("\n\n")
}
fn print_last_step(agent: &Agent, is_tty: bool) -> bool {
let last_step = agent
.last_step()
.as_ref()
.expect("agent step completed without recording last_step");
if let Some(msg) = last_step.assistant_message() {
if is_tty {
println!("\nAssistant: {}\n", msg);
} else {
println!("{}", msg);
}
}
last_step.did_tool_call()
}
enum ChatReader {
Interactive(DefaultEditor),
Plain(std::io::Stdin),
}
impl ChatReader {
fn new(is_tty: bool) -> color_eyre::Result<Self> {
if is_tty {
Ok(Self::Interactive(DefaultEditor::new()?))
} else {
Ok(Self::Plain(std::io::stdin()))
}
}
fn read_next(&mut self) -> color_eyre::Result<Option<String>> {
loop {
match self {
Self::Interactive(editor) => match editor.readline("You: ") {
Ok(line) => {
let input = line.trim();
if input.is_empty() {
continue;
}
let _ = editor.add_history_entry(input);
return Ok(Some(input.to_string()));
}
Err(ReadlineError::Interrupted) | Err(ReadlineError::Eof) => {
return Ok(None);
}
Err(error) => return Err(error.into()),
},
Self::Plain(stdin) => {
let mut input = String::new();
if stdin.lock().read_line(&mut input)? == 0 {
return Ok(None);
}
let input = input.trim();
if input.is_empty() {
continue;
}
return Ok(Some(input.to_string()));
}
}
}
}
}