Save the last assistant message to a specified file and directory

This commit is contained in:
easyrider 2024-12-31 22:00:50 +01:00
parent f292e01980
commit b48539d98c
2 changed files with 75 additions and 0 deletions

View file

@ -1379,6 +1379,51 @@ class Commands:
except Exception as e:
self.io.tool_error(f"An unexpected error occurred while copying to clipboard: {str(e)}")
def cmd_md(self, args):
"Save the last assistant message to a specified file and directory"
if not args.strip():
self.io.tool_error("Please provide a filename to save the message to.")
return
try:
# Parse the filename from the arguments
filename = args.strip()
default_dir = Path(self.coder.root) / "notes"
# Ensure the default directory exists
default_dir.mkdir(parents=True, exist_ok=True)
# Set the filepath to the default directory if only a filename is provided
if not Path(filename).parent or Path(filename).parent == Path('.'):
filepath = default_dir / filename
else:
filepath = Path(expanduser(filename))
# Get the last assistant message
all_messages = self.coder.done_messages + self.coder.cur_messages
assistant_messages = [msg for msg in reversed(all_messages) if msg["role"] == "assistant"]
if not assistant_messages:
self.io.tool_error("No assistant messages found to save.")
return
last_assistant_message = assistant_messages[0]["content"]
# Save the message to the specified file
with open(filepath, "w", encoding=self.io.encoding) as f:
f.write(last_assistant_message)
self.io.tool_output(f"Saved last assistant message to {filepath}")
except PermissionError as e:
self.io.tool_error(f"Permission denied: {e}")
self.io.tool_output("Please ensure you have write permissions for the specified directory.")
except ValueError:
self.io.tool_error("Please provide a valid filename.")
except Exception as e:
self.io.tool_error(f"An unexpected error occurred while saving the message: {str(e)}")
def cmd_report(self, args):
"Report a problem by opening a GitHub Issue"
from aider.report import report_github_issue

View file

@ -145,6 +145,36 @@ class TestCommands(TestCase):
# Assert that tool_error was called with the clipboard error message
mock_tool_error.assert_called_once_with("Failed to copy to clipboard: Clipboard error")
def test_cmd_md(self):
# Initialize InputOutput and Coder instances
io = InputOutput(pretty=False, fancy_input=False, yes=True)
coder = Coder.create(self.GPT35, None, io)
commands = Commands(io, coder)
# Add some assistant messages to the chat history
coder.done_messages = [
{"role": "assistant", "content": "First assistant message"},
{"role": "user", "content": "User message"},
{"role": "assistant", "content": "Second assistant message"},
]
# Create a temporary directory for the notes
notes_dir = Path(self.tempdir) / "notes"
notes_dir.mkdir()
# Define the filename to save the message
filename = "test_note.md"
filepath = notes_dir / filename
# Invoke the /md command
commands.cmd_md(filename)
# Check if the file was created and contains the last assistant message
self.assertTrue(filepath.exists())
with open(filepath, "r", encoding=io.encoding) as f:
content = f.read()
self.assertEqual(content, "Second assistant message")
def test_cmd_add_bad_glob(self):
# https://github.com/Aider-AI/aider/issues/293