Improve chatbot agent: fix config key overrides, update RAG docs, add…

server_api/chatbot/agent_eval.py

@@ -0,0 +1,101 @@
+"""
+End-to-end agent evaluation for training command generation.
+
+Sends 10 realistic user requests through the full agent pipeline and
+prints tool call logs + full responses for manual review.
+"""
+
+import sys
+import re
+from io import StringIO
+
+TRAINING_REQUESTS = [
+    # Basic config selection
+    "Give me the training command for mitochondria segmentation on the Lucchi dataset",
+    "Give me the command to train synapse detection on CREMI",
+    "Write the training command for neuron segmentation on SNEMI3D",
+    "Generate a command to train on the MitoEM-R dataset",
+    "Give me the command to train nucleus segmentation on zebrafish data",
+    "Write the command for vesicle segmentation training",
+    "Give me a quick demo training command to test if everything works",
+    # Single override
+    "Give me the command to train CREMI synapse detection with learning rate 0.0001",
+    "Write the training command for Lucchi mitochondria with 300 epochs",
+    "Generate the command to train SNEMI3D neuron segmentation with batch size 4",
+    "Give me the command to train MitoEM-H with AdamW optimizer",
+    "Write the command for Lucchi mito training with bf16 mixed precision",
+    # Multiple overrides
+    "Generate the training command for vesicle segmentation with learning rate 0.001 and gradient clipping of 1.0",
+    "Give me the command to train Lucchi mito with 200 epochs and learning rate 0.0005",
+    "Write the command to train CREMI synapse detection, saving checkpoints every 5 epochs with early stopping patience 20",
+    # Advanced / tricky
+    "Give me the command for neuron instance segmentation training at 40nm resolution with gradient clipping 1.0",
+    "Generate the training command for MitoEM-R with cosine learning rate schedule and 500 epochs",
+    "Write the command to train on the MitoEM human+rat combined dataset with 4 GPUs",
+    "Give me the training command for Lucchi mitochondria, but I want to save only the top 1 checkpoint",
+    "Generate a command to train SNEMI3D neurons with learning rate 0.0003, batch size 2, and gradient clipping 0.5",
+]
+
+INFERENCE_REQUESTS = [
+    # Basic inference
+    "Give me the inference command for my trained Lucchi mito model. The checkpoint is at outputs/mito_lucchi/checkpoints/epoch=100.ckpt",
+    "Write the command to run inference on CREMI with my checkpoint at outputs/syn_cremi/checkpoints/last.ckpt",
+    "Generate the inference command for SNEMI3D using checkpoint outputs/neuron_snemi/checkpoints/epoch=050.ckpt",
+    # With overrides
+    "Give me the inference command for Lucchi mito with checkpoint outputs/mito_lucchi/checkpoints/epoch=100.ckpt and Gaussian blending",
+    "Write the command to run inference on CREMI with checkpoint outputs/syn_cremi/checkpoints/last.ckpt and test-time augmentation enabled",
+    "Generate the inference command for MitoEM-R with checkpoint outputs/mitoem_r/checkpoints/epoch=200.ckpt and output saved to results/mitoem_predictions/",
+    # Advanced
+    "Give me the command to run inference on vesicle data with checkpoint outputs/vesicle_xm/checkpoints/epoch=150.ckpt, sliding window overlap 0.75, and Gaussian blending",
+    "Write the inference command for SNEMI3D neuron segmentation using outputs/neuron_snemi/checkpoints/epoch=100.ckpt with TTA enabled and ensemble mode set to mean",
+    "Generate the command to run inference on the Lucchi dataset with checkpoint outputs/mito_lucchi/checkpoints/best.ckpt and batch size 2",
+    "Give me the inference command for nucleus segmentation on zebrafish with checkpoint outputs/nuc_nucmm/checkpoints/epoch=080.ckpt and save output to results/nuc_output/",
+]
+
+
+def main():
+    from server_api.chatbot.chatbot import build_chain
+
+    print("Building agent chain...")
+    supervisor, reset_search = build_chain()
+    print("Agent chain ready.\n")
+
+    all_requests = [("TRAIN", r) for r in TRAINING_REQUESTS]
+
+    for i, (category, request) in enumerate(all_requests):
+        print(f"\n{'=' * 80}")
+        print(f"[{category}] TEST {i + 1}/{len(all_requests)}")
+        print(f"USER: {request}")
+        print("=" * 80)
+
+        # Capture stdout to get tool call logs
+        old_stdout = sys.stdout
+        captured = StringIO()
+        sys.stdout = captured
+
+        try:
+            reset_search()
+            result = supervisor.invoke(
+                {"messages": [{"role": "user", "content": request}]}
+            )
+            messages = result.get("messages", [])
+            response = messages[-1].content if messages else "(no response)"
+        except Exception as e:
+            response = f"ERROR: {e}"
+        finally:
+            sys.stdout = old_stdout
+            tool_log = captured.getvalue()
+
+        # Print tool calls
+        tool_calls = re.findall(r"\[TOOL\] .*", tool_log)
+        if tool_calls:
+            print("\nTOOL CALLS:")
+            for tc in tool_calls:
+                print(f"  {tc}")
+
+        print(f"\nRESPONSE:\n{response}")
+        print()
+
+
+if __name__ == "__main__":
+    main()

server_api/chatbot/chatbot.py

@@ -9,6 +9,15 @@
 import os
 from pathlib import Path
 
+
+def _format_admin_llm_error(error):
+    """Format an LLM connection error for display to the user."""
+    return (
+        "The AI assistant could not connect to its configured language model. "
+        "Please contact your system administrator with this error: "
+        f"{str(error).strip() or error.__class__.__name__}"
+    )
+
 from langchain_ollama import OllamaEmbeddings, ChatOllama
 from langchain_community.vectorstores import FAISS
 from langchain_core.tools import tool
@@ -30,24 +39,16 @@
 4. Be concise. State the facts, generate the command, stop.
 
 WORKFLOW: The available configs are provided in your task message. Pick the best match, then:
-1. Call read_config on the chosen config path to see its YAML overrides.
-2. For common parameters (learning rate, batch size, iterations, optimizer, checkpoint interval), ALWAYS use the keys listed below. DO NOT search for these.
-3. For specialized parameters (augmentation settings, loss functions, architecture details), call search_documentation.
-4. Build the command with overrides using the SECTION.KEY=value format.
+1. ALWAYS call read_config on the chosen config path BEFORE generating any command. Never skip this.
+2. Use ONLY the exact key paths you see in the returned YAML to build overrides (Hydra dot-path format: key=value).
+3. For parameters not visible in the config, call search_documentation.
+4. Build the command.
 
-IMPORTANT: YAML configs only show overrides — many valid keys exist in the defaults but are not shown in read_config output.
+CRITICAL: PyTC has its own config key names — they are NOT PyTorch Lightning CLI names.
+Do NOT guess key names. You MUST read the config first to see the actual keys.
 
-Common override keys (ALWAYS use these exact keys, never search for alternatives):
-- SOLVER.BASE_LR, SOLVER.SAMPLES_PER_BATCH, SOLVER.ITERATION_TOTAL
-- SOLVER.ITERATION_SAVE (checkpoint save interval), SOLVER.ITERATION_STEP (LR decay steps)
-- SOLVER.NAME (values: SGD, Adam, AdamW)
-- SOLVER.LR_SCHEDULER_NAME (values: WarmupMultiStepLR, WarmupCosineLR)
-- SOLVER.CLIP_GRADIENTS.ENABLED (True/False), SOLVER.CLIP_GRADIENTS.CLIP_VALUE
-- MODEL.ARCHITECTURE, MODEL.BLOCK_TYPE, MODEL.FILTERS
-
-NEVER invent keys like TRAIN.MAX_ITER, TRAINING.BATCH_SIZE, or CLI flags like --batch-size, --checkpoint-interval — these do not exist.
-
-Command format: `python scripts/main.py --config-file <path> [SECTION.KEY=value ...]`
+Command format: `python scripts/main.py --config <path> [key=value ...]`
+Example: `python scripts/main.py --config tutorials/syn_cremi.yaml optimization.optimizer.lr=0.0001 optimization.max_epochs=300`
 Always generate commands for the user to run — never execute directly."""
 
 
@@ -58,26 +59,13 @@
 2. Be concise. State the facts, generate the command, stop.
 
 WORKFLOW:
-1. If the user did NOT provide a checkpoint path, call list_checkpoints first to see available checkpoints.
-2. If the user DID provide a checkpoint path (e.g., outputs/model/checkpoint.pth.tar), skip list_checkpoints.
-3. Call read_config to see the INFERENCE section keys.
+1. If the user did NOT provide a checkpoint path, call list_checkpoints first.
+2. If the user DID provide one, skip list_checkpoints.
+3. Call read_config to see the config’s keys, then use those for overrides (Hydra dot-path format).
 4. For specialized inference parameters, call search_documentation if needed.
 
-Here is the correct override key mapping (use these exact keys):
-- Output path → INFERENCE.OUTPUT_PATH
-- TTA augmentation count → INFERENCE.AUG_NUM
-- TTA mode → INFERENCE.AUG_MODE (values: mean, max)
-- Blending → INFERENCE.BLENDING (values: gaussian, bump)
-- Stride → INFERENCE.STRIDE
-- Process volumes one at a time → INFERENCE.DO_SINGLY
-- Batch size → INFERENCE.SAMPLES_PER_BATCH
-
-Command format: `python scripts/main.py --config-file <path> --inference --checkpoint <ckpt> [SECTION.KEY=value ...]`
-
-IMPORTANT: Overrides use SECTION.KEY=value format (NO -- prefix). Example:
-  ✅ CORRECT: INFERENCE.AUG_NUM=8
-  ❌ WRONG: --inference.AUG_NUM=8
-
+Command format: `python scripts/main.py --config <path> --mode test --checkpoint <checkpoint.ckpt> [key=value ...]`
+Example: `python scripts/main.py --config tutorials/syn_cremi.yaml --mode test --checkpoint outputs/model/checkpoints/epoch=100.ckpt test.output_path=results/`
 Always generate commands for the user to run — never execute directly."""
 
 
@@ -224,12 +212,53 @@ def delegate_to_training_agent(task: str) -> str:
         # Auto-inject available configs so the agent doesn't need to call list_training_configs
         configs = list_training_configs.invoke({})
         config_summary = "\n".join(
-            f"- {c['name']} ({c['model']}) → {c['path']}" for c in configs if isinstance(c, dict) and 'name' in c
+            f"- {c['name']} ({c['model']}): {c.get('description', '')} → {c['path']}"
+            for c in configs if isinstance(c, dict) and 'name' in c
         )
+
+        # Auto-read a representative config and flatten grouping keys (default/train/test)
+        # so the agent sees the actual Hydra override paths, not YAML grouping structure
+        ref_config = next(
+            (c for c in configs if isinstance(c, dict) and 'lucchi' in c.get('name', '').lower()),
+            configs[0] if configs and isinstance(configs[0], dict) else None,
+        )
+        key_structure_section = ""
+        if ref_config:
+            ref_content = read_config.invoke({"config_path": ref_config['path']})
+            if isinstance(ref_content, dict) and 'error' not in ref_content:
+                import yaml as _yaml
+
+                def _deep_merge(base, override):
+                    """Merge override into base dict recursively."""
+                    for k, v in override.items():
+                        if k in base and isinstance(base[k], dict) and isinstance(v, dict):
+                            _deep_merge(base[k], v)
+                        else:
+                            base[k] = v
+
+                # Flatten: merge default, train, test into one level
+                flat = {}
+                for group_key in ("default", "train", "test"):
+                    group = ref_content.get(group_key, {})
+                    if isinstance(group, dict):
+                        _deep_merge(flat, group)
+
+                # Also include top-level non-grouping keys
+                for k, v in ref_content.items():
+                    if k not in ("default", "train", "test", "_base_", "description", "experiment_name"):
+                        flat[k] = v
+
+                key_structure_section = (
+                    f"\n\nOVERRIDE KEY PATHS (from {ref_config['name']} — all configs share this structure):\n"
+                    f"```yaml\n{_yaml.dump(flat, default_flow_style=False)[:3000]}```\n"
+                    f"These are the exact Hydra override paths. Use them as-is (e.g. optimization.max_epochs=200)."
+                )
+
         enriched_task = (
             f"{task}\n\n"
-            f"AVAILABLE CONFIGS (already fetched for you):\n{config_summary}\n\n"
-            f"Pick the best match and call read_config on its path to see the exact YAML keys before generating the command."
+            f"AVAILABLE CONFIGS (already fetched for you):\n{config_summary}"
+            f"{key_structure_section}\n\n"
+            f"Pick the best match. Use the key paths shown above for overrides — do NOT guess key names."
         )
         result = training_agent.invoke(
             {"messages": [{"role": "user", "content": enriched_task}]}