agent-learning-kit 0.1.0__py3-none-any.whl

fi/opt/optimizers/gepa.py

@@ -0,0 +1,322 @@
+import logging
+import time
+from typing import Any, Dict, List, Optional
+
+# Import GEPA's core components. A try/except block makes this a soft dependency.
+try:
+    import gepa
+    from gepa.core.adapter import GEPAAdapter, EvaluationBatch, DataInst
+except ImportError:
+    raise ImportError(
+        "To use GEPAOptimizer, please install the 'gepa' library with: pip install gepa"
+    )
+
+from ..base.base_optimizer import BaseOptimizer
+from ..datamappers.basic_mapper import BasicDataMapper
+from ..base.evaluator import Evaluator
+from ..generators.litellm import LiteLLMGenerator
+from ..types import OptimizationResult, IterationHistory
+from ..utils.early_stopping import (
+    EarlyStoppingConfig,
+    EarlyStoppingChecker,
+    EarlyStoppingException,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class _InternalGEPAAdapter(GEPAAdapter[DataInst, Dict[str, Any], Dict[str, Any]]):
+    """
+    An internal adapter that translates our framework's components (Evaluator,
+    DataMapper) into the interface GEPA's optimization engine expects.
+    """
+
+    def __init__(
+        self,
+        generator_model: str,
+        evaluator: Evaluator,
+        data_mapper: BasicDataMapper,
+        history_list: List[IterationHistory],
+        early_stopping_checker: Optional[EarlyStoppingChecker] = None,
+    ):
+        self.generator_model = generator_model
+        self.evaluator = evaluator
+        self.data_mapper = data_mapper
+        self.history_list = history_list
+        self.early_stopping_checker = early_stopping_checker
+        logger.info(f"Initialized with generator_model: {generator_model}")
+
+    def evaluate(
+        self,
+        batch: List[Dict[str, Any]],
+        candidate: Dict[str, str],
+        capture_traces: bool = False,
+    ) -> EvaluationBatch[Dict[str, Any], Dict[str, Any]]:
+        """
+        This method is called by GEPA during its optimization loop. It uses
+        our framework's components to perform the evaluation.
+        """
+        eval_start_time = time.time()
+        logger.info("Starting evaluation for a candidate prompt.")
+
+        # GEPA provides the prompt as the first (and only) value in the candidate dict.
+        prompt_text = next(iter(candidate.values()))
+        logger.info(f"Evaluating prompt: '{prompt_text[:100]}...'")
+        logger.info(f"Batch size: {len(batch)}")
+
+        temp_generator = LiteLLMGenerator(
+            model=self.generator_model, prompt_template=prompt_text
+        )
+
+        logger.info("Generating outputs...")
+        gen_start_time = time.time()
+        generated_outputs = [temp_generator.generate(example) for example in batch]
+        gen_end_time = time.time()
+        logger.info(
+            f"Output generation finished in {gen_end_time - gen_start_time:.2f}s."
+        )
+
+        logger.info("Mapping evaluation inputs...")
+        eval_inputs = [
+            self.data_mapper.map(gen_out, ex)
+            for gen_out, ex in zip(generated_outputs, batch)
+        ]
+
+        logger.info("Evaluating generated outputs...")
+        evaluator_start_time = time.time()
+        results = self.evaluator.evaluate(eval_inputs)
+        evaluator_end_time = time.time()
+        logger.info(
+            f"Evaluation with framework evaluator finished in {evaluator_end_time - evaluator_start_time:.2f}s."
+        )
+
+        scores = [res.score for res in results]
+        logger.info(f"Scores: {scores}")
+        outputs = [
+            {"generated_text": out, "full_result": res.model_dump()}
+            for out, res in zip(generated_outputs, results)
+        ]
+
+        # capture iteration history
+        avg_score = sum(scores) / len(scores) if scores else 0.0
+        self.history_list.append(
+            IterationHistory(
+                prompt=prompt_text,
+                average_score=avg_score,
+                individual_results=results,
+            )
+        )
+
+        # Check early stopping
+        if self.early_stopping_checker:
+            if self.early_stopping_checker.should_stop(avg_score, len(batch)):
+                reason = self.early_stopping_checker.get_state()["stop_reason"]
+                logger.info(f"Early stopping triggered: {reason}")
+                raise EarlyStoppingException(reason)
+
+        trajectories = []
+        if capture_traces:
+            logger.info("Capturing traces.")
+            for i in range(len(batch)):
+                trajectories.append(
+                    {
+                        "inputs": batch[i],
+                        "generated_output": generated_outputs[i],
+                        "evaluation_result": results[i].model_dump(),
+                    }
+                )
+
+        eval_end_time = time.time()
+        logger.info(f"Evaluation finished in {eval_end_time - eval_start_time:.2f}s.")
+        return EvaluationBatch(
+            outputs=outputs, scores=scores, trajectories=trajectories
+        )
+
+    def make_reflective_dataset(
+        self,
+        candidate: Dict[str, str],
+        eval_batch: EvaluationBatch,
+        components_to_update: List[str],
+    ) -> Dict[str, List[Dict[str, Any]]]:
+        """
+        Creates the dataset for GEPA's reflective LLM to analyze.
+        """
+        logger.info("Creating reflective dataset.")
+        reflective_data = {comp: [] for comp in components_to_update}
+
+        if not eval_batch.trajectories:
+            logger.warning("No trajectories found to create reflective dataset.")
+            return reflective_data
+
+        logger.info(f"Processing {len(eval_batch.trajectories)} trajectories.")
+        for trajectory in eval_batch.trajectories:
+            result = trajectory.get("evaluation_result", {})
+            score = result.get("score", 0.0)
+            reason = result.get("reason", "No reason provided.")
+
+            if score >= 0.8:
+                feedback = f"This output was successful (score={score:.2f}). The reasoning for this score was: {reason}"
+            else:
+                feedback = f"This output performed poorly (score={score:.2f}). The key reason for the failure was: {reason}. The prompt needs to be improved to avoid this specific failure mode."
+
+            example = {
+                "Inputs": trajectory.get("inputs", {}),
+                "Generated Outputs": trajectory.get("generated_output", ""),
+                "Feedback": feedback,
+            }
+
+            for comp in components_to_update:
+                reflective_data[comp].append(example)
+
+        logger.info(
+            f"Reflective dataset created for components: {components_to_update}"
+        )
+        return reflective_data
+
+
+class GEPAOptimizer(BaseOptimizer):
+    """
+    An adapter that integrates the powerful GEPA evolutionary optimization
+    algorithm into the prompt-optimizer framework.
+    """
+
+    def __init__(self, reflection_model: str, generator_model: str = "gpt-4o-mini"):
+        """
+        Initializes the GEPA Optimizer wrapper.
+
+        Args:
+            reflection_model (str): The name of a powerful LLM (e.g., "gpt-4-turbo")
+                that GEPA will use for its reflection and mutation steps.
+            generator_model (str): The name of the model that will be used by the
+                prompts being optimized (the "task language model").
+        """
+        self.reflection_model = reflection_model
+        self.generator_model = generator_model
+        logger.info(
+            f"Initialized with reflection_model: {reflection_model}, generator_model: {generator_model}"
+        )
+
+    def optimize(
+        self,
+        evaluator: Evaluator,
+        data_mapper: BasicDataMapper,
+        dataset: List[Dict[str, Any]],
+        initial_prompts: List[str],
+        max_metric_calls: Optional[int] = 150,
+        early_stopping: Optional[EarlyStoppingConfig] = None,
+    ) -> OptimizationResult:
+        opt_start_time = time.time()
+        logger.info("--- Starting GEPA Prompt Optimization ---")
+        logger.info(f"Dataset size: {len(dataset)}")
+        logger.info(f"Initial prompts: {initial_prompts}")
+        logger.info(f"Max metric calls: {max_metric_calls}")
+
+        # Initialize early stopping checker
+        checker = None
+        if early_stopping and early_stopping.is_enabled():
+            checker = EarlyStoppingChecker(early_stopping)
+            logger.info(f"Early stopping enabled: {early_stopping}")
+
+        if not initial_prompts:
+            raise ValueError("Initial prompts list cannot be empty for GEPAOptimizer.")
+        history: List[IterationHistory] = []
+        # 1. Create the internal adapter that bridges our framework to GEPA
+        logger.info("Creating internal GEPA adapter...")
+        adapter = _InternalGEPAAdapter(
+            generator_model=self.generator_model,
+            evaluator=evaluator,
+            data_mapper=data_mapper,
+            history_list=history,
+            early_stopping_checker=checker,
+        )
+
+        # 2. Prepare the inputs for gepa.optimize
+        seed_candidate = {"prompt": initial_prompts[0]}
+        logger.info(f"Seed candidate for GEPA: {seed_candidate}")
+
+        # 3. Call the external GEPA library's optimize function
+        logger.info("Calling gepa.optimize...")
+        gepa_start_time = time.time()
+
+        try:
+            gepa_result = gepa.optimize(
+                seed_candidate=seed_candidate,
+                trainset=dataset,
+                valset=dataset,
+                adapter=adapter,
+                reflection_lm=self.reflection_model,
+                max_metric_calls=max_metric_calls,
+                display_progress_bar=True,
+            )
+            gepa_end_time = time.time()
+            logger.info(
+                f"gepa.optimize finished in {gepa_end_time - gepa_start_time:.2f}s."
+            )
+            logger.info(
+                f"GEPA result best score: {gepa_result.val_aggregate_scores[gepa_result.best_idx]}"
+            )
+            logger.info(f"GEPA best candidate: {gepa_result.best_candidate}")
+
+            logger.info(f"Captured {len(history)} iterations in history.")
+            # 4. Translate GEPA's result back into our framework's standard format
+            logger.info("Translating GEPA result to OptimizationResult...")
+
+            final_best_generator = LiteLLMGenerator(
+                model=self.generator_model,
+                prompt_template=gepa_result.best_candidate.get("prompt", ""),
+            )
+
+            # Build result with early stopping metadata
+            result = OptimizationResult(
+                best_generator=final_best_generator,
+                history=history,
+                final_score=gepa_result.val_aggregate_scores[gepa_result.best_idx],
+                early_stopped=False,
+                stop_reason=None,
+                total_iterations=len(history),
+                total_evaluations=(
+                    checker.get_state()["total_evaluations"]
+                    if checker
+                    else sum(len(h.individual_results) for h in history)
+                ),
+            )
+
+        except EarlyStoppingException as e:
+            gepa_end_time = time.time()
+            logger.info(
+                f"GEPA stopped early after {gepa_end_time - gepa_start_time:.2f}s: {e.reason}"
+            )
+
+            # Use best from history
+            if not history:
+                raise RuntimeError(
+                    "Early stopping triggered before any evaluations completed"
+                )
+
+            best_history = max(history, key=lambda h: h.average_score)
+            final_best_generator = LiteLLMGenerator(
+                model=self.generator_model,
+                prompt_template=best_history.prompt,
+            )
+
+            result = OptimizationResult(
+                best_generator=final_best_generator,
+                history=history,
+                final_score=best_history.average_score,
+                early_stopped=True,
+                stop_reason=e.reason,
+                total_iterations=len(history),
+                total_evaluations=(
+                    checker.get_state()["total_evaluations"]
+                    if checker
+                    else sum(len(h.individual_results) for h in history)
+                ),
+            )
+
+        opt_end_time = time.time()
+        logger.info(
+            f"--- GEPA Prompt Optimization finished in {opt_end_time - opt_start_time:.2f}s ---"
+        )
+        logger.info(f"Final best score: {result.final_score}")
+
+        return result

fi/opt/optimizers/metaprompt.py

@@ -0,0 +1,243 @@
+import json
+import random
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel, Field, ValidationError
+
+from ..base.base_optimizer import BaseOptimizer
+from ..datamappers.basic_mapper import BasicDataMapper
+from ..base.evaluator import Evaluator
+from ..generators.litellm import LiteLLMGenerator
+from ..types import IterationHistory, OptimizationResult
+from ..utils.early_stopping import EarlyStoppingConfig, EarlyStoppingChecker
+import logging
+
+logger = logging.getLogger(__name__)
+# ==============================================================================
+# Prompts and Pydantic Models for the Teacher LLM (Meta-Model)
+# ==============================================================================
+
+
+META_PROMPT_TEMPLATE = """
+You are a world-class expert in prompt engineering. Your task is to diagnose and optimize a given prompt based on its performance on a set of test cases.
+
+### Current Prompt
+The following is the current prompt being evaluated:
+---
+{current_prompt}
+---
+
+### Previous Failed Attempts
+You have already tried the following prompts, but they performed worse than the current one. Analyze why they failed to avoid repeating mistakes.
+---
+{other_attempts}
+---
+
+### Performance Data
+The current prompt was run on a set of examples, and here are the results. Pay close attention to the examples with low scores.
+---
+{annotated_results}
+---
+
+### Task Description
+{task_description}
+
+### Your Task
+Think step-by-step to generate an improved prompt:
+1.  **Analyze Failures:** Deeply analyze the failing examples. What patterns do you see? Is the prompt too vague, too restrictive, or missing key instructions?
+2.  **Formulate a Hypothesis:** Based on your analysis, state a clear hypothesis for how to improve the prompt. For example, "My hypothesis is that adding a chain-of-thought instruction will improve reasoning on multi-step problems."
+3.  **Generate Improved Prompt:** Rewrite the *entire* prompt, implementing your hypothesis. The new prompt should be a complete replacement for the current one.
+
+Return ONLY a valid JSON object with two keys: "hypothesis" (your string hypothesis) and "improved_prompt" (the complete new prompt string).
+"""
+
+
+class MetaPromptOutput(BaseModel):
+    hypothesis: str = Field(
+        description="The hypothesis for why the new prompt will be better."
+    )
+    improved_prompt: str = Field(description="The complete, new, improved prompt.")
+
+
+# ==============================================================================
+# The MetaPrompt Optimizer Class
+# ==============================================================================
+
+
+class MetaPromptOptimizer(BaseOptimizer):
+    """
+    Optimizes a prompt by using a powerful "teacher" LLM to analyze its
+    performance and rewrite it. This is inspired by the `promptim` library.
+    """
+
+    def __init__(
+        self,
+        teacher_generator: LiteLLMGenerator,
+        task_model: Optional[str] = None,
+    ):
+        """
+        Initializes the MetaPrompt Optimizer.
+
+        Args:
+            teacher_generator: A powerful generator (e.g., GPT-4o, Claude 3 Opus)
+                used to analyze performance and generate new prompts.
+            task_model: Model used to run candidate prompts while scoring them.
+                Defaults to the teacher generator's model.
+        """
+        self.teacher = teacher_generator
+        self.task_model = task_model
+
+    def optimize(
+        self,
+        evaluator: Evaluator,
+        data_mapper: BasicDataMapper,
+        dataset: List[Dict[str, Any]],
+        initial_prompts: List[str],
+        task_description: str = "I want to improve my prompt.",
+        num_rounds: Optional[int] = 5,
+        eval_subset_size: Optional[int] = 40,
+        early_stopping: Optional[EarlyStoppingConfig] = None,
+    ) -> OptimizationResult:
+        logger.info("--- Starting Meta-Prompt Optimization ---")
+
+        # Initialize early stopping checker
+        checker = None
+        if early_stopping and early_stopping.is_enabled():
+            checker = EarlyStoppingChecker(early_stopping)
+            logger.info(f"Early stopping enabled: {early_stopping}")
+
+        if not initial_prompts:
+            raise ValueError("Initial prompts list cannot be empty.")
+
+        current_prompt = initial_prompts[0]
+        best_prompt = current_prompt
+        best_score = -1.0
+        history: List[IterationHistory] = []
+        previous_attempts = set()
+
+        for round_num in range(num_rounds):
+            logger.info(
+                f"\n--- Starting Optimization Round {round_num + 1}/{num_rounds} ---"
+            )
+            logger.info(f"Current best prompt:\n{current_prompt}")
+
+            # 1. Evaluate the current prompt on a subset of data
+            eval_subset = random.sample(dataset, min(len(dataset), eval_subset_size))
+            iteration_history = self._score_prompt(
+                current_prompt, evaluator, data_mapper, eval_subset
+            )
+
+            if not iteration_history:
+                logger.warning("Evaluation of current prompt failed. Skipping round.")
+                continue
+
+            history.append(iteration_history)
+            current_score = iteration_history.average_score
+
+            if current_score > best_score:
+                best_score = current_score
+                best_prompt = current_prompt
+                logger.info(f"New best score found: {best_score:.4f}")
+
+            # Check early stopping
+            if checker:
+                num_evals = len(eval_subset)
+                if checker.should_stop(current_score, num_evals):
+                    logger.info(
+                        f"Early stopping triggered: {checker.get_state()['stop_reason']}"
+                    )
+                    break
+
+            # 2. Use the teacher model to generate a new, improved prompt
+            annotated_results_str = self._format_results(iteration_history, eval_subset)
+
+            # Format previous attempts for the meta-prompt
+            other_attempts_str = (
+                "\n---\n".join(list(previous_attempts)) if previous_attempts else "N/A"
+            )
+
+            meta_prompt = META_PROMPT_TEMPLATE.format(
+                current_prompt=current_prompt,
+                other_attempts=other_attempts_str,
+                annotated_results=annotated_results_str,
+                task_description=task_description,
+            )
+
+            logger.debug("Generating new prompt with meta-prompt...")
+            new_prompt_json = self.teacher.generate(
+                prompt_vars={"prompt": meta_prompt},
+                response_format={"type": "json_object"},
+            )
+
+            try:
+                parsed_output = MetaPromptOutput.model_validate_json(new_prompt_json)
+                logger.info(f"Teacher's Hypothesis: {parsed_output.hypothesis}")
+                previous_attempts.add(current_prompt)
+                current_prompt = parsed_output.improved_prompt
+            except (ValidationError, json.JSONDecodeError) as e:
+                logger.error(
+                    f"Failed to parse new prompt from teacher model, keeping current prompt. Error: {e}"
+                )
+
+        final_best_generator = LiteLLMGenerator(self.teacher.model_name, best_prompt)
+
+        # Build result with early stopping metadata
+        return OptimizationResult(
+            best_generator=final_best_generator,
+            history=history,
+            final_score=best_score,
+            early_stopped=checker.get_state()["stopped"] if checker else False,
+            stop_reason=checker.get_state()["stop_reason"] if checker else None,
+            total_iterations=len(history),
+            total_evaluations=(
+                checker.get_state()["total_evaluations"]
+                if checker
+                else sum(len(h.individual_results) for h in history)
+            ),
+        )
+
+    def _score_prompt(
+        self,
+        prompt: str,
+        evaluator: Evaluator,
+        data_mapper: BasicDataMapper,
+        dataset: List[Dict[str, Any]],
+    ) -> IterationHistory | None:
+        """Scores a single prompt and returns its history."""
+        try:
+            temp_generator = LiteLLMGenerator(
+                self.task_model or self.teacher.model_name, prompt
+            )
+            generated_outputs = [
+                temp_generator.generate(example) for example in dataset
+            ]
+            eval_inputs = [
+                data_mapper.map(gen_out, ex)
+                for gen_out, ex in zip(generated_outputs, dataset)
+            ]
+            results = evaluator.evaluate(eval_inputs)
+            avg_score = (
+                sum(res.score for res in results) / len(results) if results else 0.0
+            )
+            return IterationHistory(
+                prompt=prompt, average_score=avg_score, individual_results=results
+            )
+        except Exception as e:
+            logger.error(f"Failed to score prompt: {e}")
+            return None
+
+    def _format_results(
+        self, iteration_history: IterationHistory, dataset: List[Dict[str, Any]]
+    ) -> str:
+        """Formats the evaluation results into a string for the meta-prompt."""
+        formatted_lines = []
+        for i, result in enumerate(iteration_history.individual_results):
+            example_input = dataset[i]
+            formatted_lines.append(f"Example {i + 1}:")
+            formatted_lines.append(
+                f"  Input: {json.dumps(example_input, ensure_ascii=False)}"
+            )
+            formatted_lines.append(f"  Score: {result.score:.2f}")
+            formatted_lines.append(f"  Reason: {result.reason}")
+            formatted_lines.append("---")
+        return "\n".join(formatted_lines)