Switch to RLM instead of ReAct

nanocode.py

@@ -234,16 +234,6 @@ class CodingAssistant(dspy.Signature):
     )
 
 
-tools = {
-    "readfile": read_file,
-    "writefile": write_file,
-    "editfile": edit_file,
-    "globfiles": glob_files,
-    "grepfiles": grep_files,
-    "runbash": run_bash,
-}
-
-
 class ToolLoggingCallback(BaseCallback):
     """Callback that logs tool calls as they happen."""
 
@@ -278,9 +268,10 @@ class RLMCodingConfig(PrecompiledConfig):
     lm: str = "openrouter/anthropic/claude-3.5-sonnet"  # Default fallback
     sub_lm: str = "openrouter/openai/gpt-4.1"  # Default fallback
     api_base: str = "https://openrouter.ai/api/v1"
-    max_tokens: int = 16000
+    max_tokens: int = 32000
     max_output_chars: int = 100000
     verbose: bool = False
+    track_usage: bool = True
 
 
 class RLMCodingProgram(PrecompiledProgram):
@@ -289,6 +280,14 @@ class RLMCodingProgram(PrecompiledProgram):
     def __init__(self, config: RLMCodingConfig, **kwargs):
         self.config = config
         super().__init__(config, **kwargs)
+        self.tools = {
+            "read_file": read_file,
+            "write_file": write_file,
+            "edit_file": edit_file,
+            "glob_files": glob_files,
+            "grep_files": grep_files,
+            "run_bash": run_bash,
+        }
 
         # tool logging for introspections on multi-turn conversations
         dspy.settings.configure(callbacks=[ToolLoggingCallback()])
@@ -296,16 +295,18 @@ class RLMCodingProgram(PrecompiledProgram):
             self.config.lm,
             api_base=self.config.api_base,
             max_tokens=self.config.max_tokens,
+            track_usage=self.config.track_usage,
         )
         sub_lm = dspy.LM(
             self.config.sub_lm,
             api_base=self.config.api_base,
             max_tokens=self.config.max_tokens,
+            track_usage=self.config.track_usage,
         )
         agent = dspy.RLM(
             CodingAssistant,
             sub_lm=sub_lm,
-            tools=tools,
+            tools=self.tools,
             max_output_chars=self.config.max_output_chars,
             max_iterations=self.config.max_iters,
             verbose=self.config.verbose,
@@ -318,6 +319,15 @@ class RLMCodingProgram(PrecompiledProgram):
         assert task, "Task cannot be empty"
         return self.agent(task=task)
 
+    def get_tools(self):
+        return self.tools
+    
+    def set_tool(self, name: str, tool: callable):
+        self.tools[name] = tool
+
+    def remove_tool(self, name: str):
+        del self.tools[name]
+
 def main():
     model = os.getenv("MODEL")
     if model is None:
@@ -332,7 +342,7 @@ def main():
 
     agent = RLMCodingProgram(config)
     print(
-        f"{BOLD}nanocode-dspy{RESET} | {DIM}{agent.config.lm} | {os.getcwd()}{RESET}\n"
+        f"{BOLD}NANOCODE DSPY{RESET} | {DIM}{agent.config.lm} | {os.getcwd()}{RESET}\n"
     )
 
     # Conversation history for context
@@ -352,6 +362,37 @@ def main():
                 history = []
                 print(f"{GREEN}⏺ Cleared conversation{RESET}")
                 continue
+            if user_input == "/model":
+                print(f"\n{BOLD}Current model: {agent.config.lm}{RESET}")
+                print(f"\n{BOLD}Select a new model:{RESET}")
+                for key, (name, model_id) in AVAILABLE_MODELS.items():
+                    print(f"  {BLUE}{key}{RESET}. {name} ({DIM}{model_id}{RESET})")
+                print(f"  {BLUE}c{RESET}. Custom model (enter manually)")
+                print(f"  {BLUE}k{RESET}. Keep current model")
+
+                choice = input(f"\n{BOLD}{BLUE}❯{RESET} Enter choice: ").strip().lower()
+
+                if choice == "k":
+                    print(f"{GREEN}⏺ Keeping current model: {agent.config.lm}{RESET}")
+                    continue
+                elif choice in AVAILABLE_MODELS:
+                    name, model_id = AVAILABLE_MODELS[choice]
+                    new_model = model_id if model_id.startswith("openrouter/") else f"openrouter/{model_id}"
+                    config.lm = new_model
+                    agent = RLMCodingProgram(config)
+                    print(f"{GREEN}⏺ Switched to: {name} ({new_model}){RESET}")
+                elif choice == "c":
+                    custom_model = input(f"{BOLD}{BLUE}❯{RESET} Enter model ID: ").strip()
+                    if custom_model:
+                        new_model = custom_model if custom_model.startswith("openrouter/") else f"openrouter/{custom_model}"
+                        config.lm = new_model
+                        agent = RLMCodingProgram(config)
+                        print(f"{GREEN}⏺ Switched to custom model: {new_model}{RESET}")
+                    else:
+                        print(f"{RED}⏺ Invalid model ID, keeping current model{RESET}")
+                else:
+                    print(f"{RED}⏺ Invalid choice, keeping current model{RESET}")
+                continue
 
             # Build context from history
             context = f"Working directory: {os.getcwd()}\n"
@@ -364,11 +405,14 @@ def main():
 
             print(f"\n{CYAN}⏺{RESET} Thinking...", flush=True)
 
-            # Run the ReAct agent
+            # Run the RLM agent
             result = agent(task=task)
 
             # Display the answer
             print(f"\n{CYAN}⏺{RESET} {render_markdown(result.answer)}")
+            
+            # Display usage
+            print(f"\n{MAGENTA}⏺ Debug Prediction: {result}{RESET}")
 
             # Save to history
             history.append({"user": user_input, "assistant": result.answer})
@@ -386,5 +430,5 @@ def main():
 
 if __name__ == "__main__":
     agent = RLMCodingProgram(RLMCodingConfig())
-    agent.push_to_hub(MODAIC_REPO_PATH, commit_message="Switch to RLM instead of ReAct", tag="v0.0.2")
+    agent.push_to_hub(MODAIC_REPO_PATH, commit_message="Switch to RLM instead of ReAct", tag="v0.0.3")
     #main()