config.sampler to turn on/off sampler (#496)

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": (

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

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(

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