Make APO templates configurable via constructor arguments (#443)

Author: li2zhi

agentlightning/algorithm/apo/apo.py

@@ -112,6 +112,8 @@ def __init__(
         beam_rounds: int = 3,
         rollout_batch_timeout: float = 3600.0,
         run_initial_validation: bool = True,
+        gradient_prompt_files: Optional[List[Path]] = None,
+        apply_edit_prompt_files: Optional[List[Path]] = None,
         # Internal flags for debugging
         _poml_trace: bool = False,
     ):
@@ -132,6 +134,8 @@ def __init__(
             rollout_batch_timeout: Maximum time in seconds to wait for rollout batch completion.
             run_initial_validation: If True, runs validation on the seed prompt before starting
                 optimization to establish a baseline score. Defaults to True.
+            gradient_prompt_files: Prompt templates used to compute textual gradients (critiques).
+            apply_edit_prompt_files: Prompt templates used to apply edits based on critiques.
         """
         self.async_openai_client = async_openai_client
         self.gradient_model = gradient_model
@@ -144,6 +148,8 @@ def __init__(
         self.beam_rounds = beam_rounds
         self.rollout_batch_timeout = rollout_batch_timeout
         self.run_initial_validation = run_initial_validation
+        self.gradient_prompt_files = gradient_prompt_files or GRADIENT_PROMPT_FILES
+        self.apply_edit_prompt_files = apply_edit_prompt_files or APPLY_EDIT_PROMPT_FILES
 
         self._history_best_prompt: Optional[PromptTemplate] = None
         self._history_best_score: float = float("-inf")
@@ -270,7 +276,7 @@ async def compute_textual_gradient(
         Returns:
             A textual critique generated by the LLM, or None if generation fails.
         """
-        tg_template = random.choice(GRADIENT_PROMPT_FILES)
+        tg_template = random.choice(self.gradient_prompt_files)
 
         if len(rollout_results) < self.gradient_batch_size:
             self._log(
@@ -352,7 +358,7 @@ async def textual_gradient_and_apply_edit(
             return current_prompt.prompt_template.template
 
         # 2) Apply edit
-        ae_template = random.choice(APPLY_EDIT_PROMPT_FILES)
+        ae_template = random.choice(self.apply_edit_prompt_files)
         self._log(
             logging.INFO,
             f"Edit will be generated by {self.apply_edit_model} with template: {ae_template.name}",