Fix config override bug by recreating LMs after load_state

nanocode.py

@@ -23,7 +23,6 @@ MAGENTA = "\033[35m"
 
 # --- File operations ---
 
-
 def read_file(path: str, offset: int = 0, limit: int = None) -> str:
     """Read file contents with line numbers.
 
@@ -163,6 +162,9 @@ def run_bash(cmd: str) -> str:
         output_lines.append("\n(timed out after 30s)")
     return "".join(output_lines).strip() or "(empty output)"
 
+
+# -- Program ---
+
 class CodingAssistant(dspy.Signature):
     """You are a concise coding assistant with access to sub agents."""
 
@@ -171,6 +173,7 @@ class CodingAssistant(dspy.Signature):
         desc="Your response to the user after completing the task"
     )
 
+
 class RLMCodingConfig(PrecompiledConfig):
     max_iters: int = 50
     lm: str = "openrouter/openai/gpt-5.2-codex"
@@ -210,7 +213,7 @@ class RLMCodingProgram(PrecompiledProgram):
             max_tokens=self.config.max_tokens,
             track_usage=self.config.track_usage,
         )
-        agent = dspy.RLM(
+        self.agent = dspy.RLM(
             CodingAssistant,
             sub_lm=self.sub_lm,
             tools=self.tools,
@@ -218,12 +221,13 @@ class RLMCodingProgram(PrecompiledProgram):
             max_iterations=self.config.max_iters,
             verbose=self.config.verbose,
         )
-
-        agent.set_lm(self.lm)
-        self.agent = agent
+        self.agent.set_lm(self.lm)
+        self.set_lm(self.lm)
 
     def forward(self, task: str) -> str:
-        assert task, "Task cannot be empty"
+        if not task:
+            return dspy.Prediction(answer="No Task Given.")
+
         return self.agent(task=task)
 
     def get_tools(self):
@@ -231,14 +235,14 @@ class RLMCodingProgram(PrecompiledProgram):
 
     def set_tool(self, name: str, tool: callable):
         self.tools[name] = tool
-        self.reload_repl_tools()
+        self.reload_repl()
 
     def remove_tool(self, name: str):
         if name in self.tools:
             del self.tools[name]
-            self.reload_repl_tools()
+            self.reload_repl()
 
-    def reload_repl_tools(
+    def reload_repl(
         self,
     ):  # we need to create a new instance for tool mutations to be passed back into the REPL
         new_instance = dspy.RLM(
@@ -252,6 +256,36 @@ class RLMCodingProgram(PrecompiledProgram):
         new_instance.set_lm(self.lm)
         self.agent = new_instance
 
+    def reload_lms(self):
+        """Recreate LM objects from current config. Call this after changing config.lm or config.sub_lm."""
+        self.lm = dspy.LM(
+            model=self.config.lm,
+            api_base=self.config.api_base,
+            max_tokens=self.config.max_tokens,
+            track_usage=self.config.track_usage,
+        )
+        self.sub_lm = dspy.LM(
+            model=self.config.sub_lm,
+            api_base=self.config.api_base,
+            max_tokens=self.config.max_tokens,
+            track_usage=self.config.track_usage,
+        )
+        self.reload_repl()
+        if os.getenv("MODAIC_ENV") == "dev":
+            print(f"{BLUE}LMs RELOADED: {self.lm.model}, {self.sub_lm.model}{RESET}")
+
+    def load_state(self, state):
+        """Override to recreate LMs from config after loading state.
+
+        PrecompiledProgram.from_precompiled() calls load_state() AFTER __init__,
+        which overwrites our LMs with saved state. We fix this by recreating
+        the LMs from self.config after the parent load_state runs. Modaic will 
+        fix this in a later patch for future devs.
+        """
+        super().load_state(state) 
+        # recreate LMs from config (not from saved state)
+        self.reload_lms()
+
 if __name__ == "__main__":
     agent = RLMCodingProgram(RLMCodingConfig())
-    agent.push_to_hub(MODAIC_REPO_PATH, commit_message="change signature")
+    agent.push_to_hub(MODAIC_REPO_PATH, commit_message="Fix config override bug by recreating LMs after load_state")