Update agent.py

agent.py

@@ -380,19 +380,21 @@ class MagAgent:
                 user_prompts = yaml.safe_load(f)
                 
                 # Validate loaded prompts structure
+                # In MagAgent's __init__ method:
                 if isinstance(user_prompts, dict):
                     for key in self.prompt_templates:
                         if key in user_prompts:
-                            # Merge main template
-                            self.prompt_templates[key].update(
-                                {k:v for k,v in user_prompts[key].items() 
-                                 if k in ['template', 'variables']}
-                            )
-                            # Merge sub-templates
+                            # Merge main template properties
+                            self.prompt_templates[key].update({
+                                k: v for k, v in user_prompts[key].items()
+                                if k in ['template', 'variables']
+                            })
+                            
+                            # Merge sub-templates if present
                             if 'sub_templates' in user_prompts[key]:
                                 self.prompt_templates[key].setdefault('sub_templates', {}).update(
                                     user_prompts[key]['sub_templates']
-                                )
+                                )                
                 else:
                     print(f"Invalid prompts.yaml structure. Using defaults. Got type: {type(user_prompts)}")
                     
@@ -404,6 +406,11 @@ class MagAgent:
         if missing := required_keys - self.prompt_templates.keys():
             raise ValueError(f"Missing required prompt templates: {missing}")
 
+        # After loading templates, add:
+        planning_subtemplates = self.prompt_templates['planning'].get('sub_templates', {})
+        required_subtemplates = {'initial_plan', 'initial_facts'}
+        if missing := required_subtemplates - planning_subtemplates.keys():
+            raise ValueError(f"Missing required planning sub-templates: {missing}")
         
         self.agent = CodeAgent(
             model=self.model,