Add ModelNames class for model name constants

config/chain_configs.py

@@ -26,13 +26,22 @@ from config.llm_config import llms
 
 # Note: The default LLM here is GPT-4o (low temp); the UI can override this choice.
 
+# Model name constants - these match the keys in llm_config.llms
+class ModelNames:
+    GPT4_ZERO = "GPT-4o (zero temp)"
+    GPT4_LOW = "GPT-4o (low temp)"
+    GPT4_MID = "GPT-4o (mid temp)"
+    GPT4_MINI_ZERO = "GPT-4o-mini (zero temp)"
+    O1_HIGH_REASONING = "o1 (high reasoning_effort)"
+    O3_MINI_HIGH_REASONING = "o3-mini (high reasoning_effort)"
+
 chain_configs = {
     "diagnoser": {
         "class": DiagnoserChain,
         "template_standardize": template_standardize_exercise,
-        "llm_standardize": llms["GPT-4o-mini (zero temp)"],     # Always fixed
-        "llm_4o_mini": llms["GPT-4o-mini (low temp)"],
-        "llm_4o": llms["GPT-4o (low temp)"],
+        "llm_standardize": llms[ModelNames.GPT4_MINI_ZERO],     # Always fixed
+        "llm_4o_mini": llms[ModelNames.GPT4_LOW],
+        "llm_4o": llms[ModelNames.GPT4_LOW],
         # 4 different diagnosis templates (to run in parallel:
         "templates_diagnose": [
             template_diagnose_double_negation,
@@ -41,44 +50,44 @@ chain_configs = {
             template_diagnose_distractor_partially_correct,
         ],
         "template_diagnose_scorecard": diagnose_scorecard_template,
-        "llm_diagnose": llms["GPT-4o (low temp)"],             # Default; can be replaced in UI
+        "llm_diagnose": llms[ModelNames.GPT4_LOW],             # Default; can be replaced in UI
     },
     "distractors": {
         "class": DistractorsChain,
         "template_standardize": template_standardize_exercise,
-        "llm_standardize": llms["GPT-4o-mini (zero temp)"],     # Always fixed
+        "llm_standardize": llms[ModelNames.GPT4_MINI_ZERO],     # Always fixed
         "template_distractors_brainstorm_1": template_distractors_brainstorm_1,
         "template_distractors_brainstorm_2": template_distractors_brainstorm_2,
-        "llm_brainstorm_1": llms["GPT-4o (low temp)"],
-        "llm_brainstorm_2": llms["GPT-4o (mid temp)"],
+        "llm_brainstorm_1": llms[ModelNames.GPT4_LOW],
+        "llm_brainstorm_2": llms[ModelNames.GPT4_MID],
         "template_consolidate": template_consolidate_distractors,
-        "llm_consolidate": llms["GPT-4o (low temp)"],
+        "llm_consolidate": llms[ModelNames.GPT4_LOW],
     },
     "learning_objectives": {
         "class": LearningObjectivesChain,
         "template_standardize": template_standardize_studytext,
-        "llm_standardize": llms["GPT-4o-mini (zero temp)"],     # Always fixed
+        "llm_standardize": llms[ModelNames.GPT4_MINI_ZERO],     # Always fixed
         "template_gen_prompt_a": template_gen_prompt_a,
         "template_gen_prompt_b": template_gen_prompt_b,
-        "default_llm_a": llms["o1 (high reasoning_effort)"],
-        "default_llm_b": llms["o3-mini (high reasoning_effort)"],
+        "default_llm_a": llms[ModelNames.O1_HIGH_REASONING],
+        "default_llm_b": llms[ModelNames.O3_MINI_HIGH_REASONING],
         "template_sanitize": template_sanitize_learning_objectives,
-        "llm_sanitize": "GPT-4o-mini (zero temp)",
+        "llm_sanitize": ModelNames.GPT4_MINI_ZERO,
     },
     "fluster": {
         "class": FlusterWritingChain,
         "template_write_fluster_a": template_write_fluster_a,
         "template_write_fluster_b": template_write_fluster_b,
-        "default_llm_a": llms["o1 (high reasoning_effort)"],
-        "default_llm_b": llms["o3-mini (high reasoning_effort)"],
+        "default_llm_a": llms[ModelNames.O1_HIGH_REASONING],
+        "default_llm_b": llms[ModelNames.O3_MINI_HIGH_REASONING],
         # Prompt & LLM for the refine-distractors step
         "template_refine_fluster": template_refine_fluster,
-        "llm_refine": llms["GPT-4o (zero temp)"],
+        "llm_refine": llms[ModelNames.GPT4_ZERO],
         "template_sanitize": template_sanitize_fluster,
-        "llm_sanitize": llms["GPT-4o-mini (zero temp)"],
+        "llm_sanitize": ModelNames.GPT4_MINI_ZERO,
         "template_structurize": template_isolate_exercises,
-        "llm_structurize": llms["GPT-4o (zero temp)"],
+        "llm_structurize": llms[ModelNames.GPT4_ZERO],
         "template_fix_exercise": template_fix_exercise,
-        "llm_fix_exercise": llms["GPT-4o (low temp)"],
+        "llm_fix_exercise": llms[ModelNames.GPT4_LOW],
     },
 }