Fix PR 166 runtime cleanup and chatbot test collection

Author: akgohain

server_api/chatbot/test_agent.py server_api/chatbot/agent_cli.pyserver_api/chatbot/test_agent.py renamed to server_api/chatbot/agent_cli.py

@@ -4,9 +4,9 @@
 Shows all RAG retrievals, tool calls, and final responses.
 
 Usage:
-  python test_agent.py                     # interactive mode
-  python test_agent.py -b                  # batch: run 20-question PyTC eval
-  python test_agent.py "your question"     # single question
+  python agent_cli.py                     # interactive mode
+  python agent_cli.py -b                  # batch: run 20-question PyTC eval
+  python agent_cli.py "your question"     # single question
 """
 
 import os
@@ -69,7 +69,7 @@ def run_batch():
     print(f"{'#'*80}")
 
 
-def test_single(question: str):
+def run_single(question: str):
     """Test the agent with a single question."""
     print(f"\n{'='*80}")
     print(f"QUESTION: {question}")
@@ -127,6 +127,6 @@ def interactive_mode():
     elif args.interactive:
         interactive_mode()
     elif args.question:
-        test_single(" ".join(args.question))
+        run_single(" ".join(args.question))
     else:
         interactive_mode()

server_pytc/services/model.py

@@ -444,12 +444,18 @@ def _matches_pytc_mode_process(cmdline: list[str], mode: str) -> bool:
     if script_path not in normalized:
         return False
 
-    try:
+    # Support both the legacy "--mode train/test" CLI and the newer
+    # "--inference" flag that PyTC now uses to switch the entrypoint into test mode.
+    if "--mode" in normalized:
         mode_index = normalized.index("--mode")
-    except ValueError:
-        return False
+        return mode_index + 1 < len(normalized) and normalized[mode_index + 1] == mode
 
-    return mode_index + 1 < len(normalized) and normalized[mode_index + 1] == mode
+    is_inference = "--inference" in normalized
+    if mode == "test":
+        return is_inference
+    if mode == "train":
+        return not is_inference
+    return False
 
 
 def stop_pytc_processes(mode: str):

tests/test_pytc_runtime_routes.py

@@ -129,11 +129,11 @@ def test_write_temp_config_uses_origin_parent_for_relative_bases(self):
         config_path = model_service._write_temp_config(
             "foo: bar\n",
             "training",
-            config_origin_path="tutorials/neuron_snemi.yaml",
+            config_origin_path="configs/SNEMI/SNEMI-Base.yaml",
         )
         written_path = pathlib.Path(config_path)
         expected_parent = (
-            model_service._project_root() / "pytorch_connectomics" / "tutorials"
+            model_service._project_root() / "pytorch_connectomics" / "configs" / "SNEMI"
         )
 
         self.assertTrue(written_path.exists())

tests/test_worker_model_service.py

@@ -9,7 +9,7 @@ class WorkerModelServiceTests(unittest.TestCase):
     def tearDown(self):
         model_service.cleanup_temp_files()
 
-    def test_matches_pytc_mode_process_with_config_between_script_and_mode(self):
+    def test_matches_pytc_mode_process_supports_legacy_cli(self):
         script_path = str(model_service._pytc_script_path())
         cmdline = [
             "/usr/bin/python",
@@ -23,6 +23,31 @@ def test_matches_pytc_mode_process_with_config_between_script_and_mode(self):
         self.assertTrue(model_service._matches_pytc_mode_process(cmdline, "train"))
         self.assertFalse(model_service._matches_pytc_mode_process(cmdline, "test"))
 
+    def test_matches_pytc_mode_process_supports_current_cli(self):
+        script_path = str(model_service._pytc_script_path())
+        train_cmdline = [
+            "/usr/bin/python",
+            script_path,
+            "--config-file",
+            "/tmp/runtime.yaml",
+        ]
+        inference_cmdline = [
+            "/usr/bin/python",
+            script_path,
+            "--config-file",
+            "/tmp/runtime.yaml",
+            "--inference",
+        ]
+
+        self.assertTrue(model_service._matches_pytc_mode_process(train_cmdline, "train"))
+        self.assertFalse(model_service._matches_pytc_mode_process(train_cmdline, "test"))
+        self.assertTrue(
+            model_service._matches_pytc_mode_process(inference_cmdline, "test")
+        )
+        self.assertFalse(
+            model_service._matches_pytc_mode_process(inference_cmdline, "train")
+        )
+
     def test_cleanup_temp_files_is_scoped_by_kind(self):
         training_file = tempfile.NamedTemporaryFile(delete=False)
         inference_file = tempfile.NamedTemporaryFile(delete=False)