diff --git a/mentat/config.py b/mentat/config.py index 574611d6e..b3666ed98 100644 --- a/mentat/config.py +++ b/mentat/config.py @@ -109,7 +109,7 @@ class Config: }, converter=converters.optional(converters.to_bool), ) - auto_save_snapshot: bool = attr.field( + sampler: bool = attr.field( default=False, metadata={ "description": ( diff --git a/mentat/sampler/sampler.py b/mentat/sampler/sampler.py index c4216c35a..e5c7cd5e8 100644 --- a/mentat/sampler/sampler.py +++ b/mentat/sampler/sampler.py @@ -17,7 +17,6 @@ class Sampler: - active: bool = True diff_active: str | None = None commit_active: str | None = None last_sample_id: str | None = None @@ -43,7 +42,7 @@ def set_active_diff(self): f"Sampler error setting active diff: {e}. Disabling sampler.", style="error", ) - self.active = False + ctx.config.sampler = False async def create_sample(self) -> Sample: # Check for repo and merge_base in config diff --git a/mentat/session.py b/mentat/session.py index 5e7f66749..d9a444da8 100644 --- a/mentat/session.py +++ b/mentat/session.py @@ -117,7 +117,7 @@ def __init__( ): for file in code_context.diff_context.diff_files(): code_context.include(file) - if config.auto_save_snapshot: + if config.sampler: sampler.set_active_diff() def _create_task(self, coro: Coroutine[None, None, Any]): @@ -185,7 +185,7 @@ async def _main(self): await get_user_feedback_on_edits(file_edits) ) - if session_context.sampler and session_context.sampler.active: + if session_context.config.sampler: session_context.sampler.set_active_diff() applied_edits = await code_file_manager.write_changes_to_files( diff --git a/tests/sampler_test.py b/tests/sampler_test.py index 72f512e59..60af89d60 100644 --- a/tests/sampler_test.py +++ b/tests/sampler_test.py @@ -37,6 +37,7 @@ async def test_sample_from_context( mock_collect_user_input, ): mock_session_context.config.sample_repo = "test_sample_repo" + mock_session_context.config.sampler = True mocker.patch( "mentat.conversation.Conversation.get_messages", @@ -382,6 +383,7 @@ async def test_sampler_integration( # Generate a sample using Mentat python_client = PythonClient(cwd=temp_testbed, paths=["."]) await python_client.startup() + python_client.session.ctx.config.sampler = True await python_client.call_mentat_auto_accept(dedent("""\ Make the following changes to "multifile_calculator/operations.py": 1. Add "# Inserted line 2" as the first line