mantisdk 0.1.0__py3-none-any.whl

mantisdk/algorithm/gepa/lib/adapters/default_adapter/default_adapter.py

@@ -0,0 +1,209 @@
+# Copyright (c) 2025 Lakshya A Agrawal and the GEPA contributors
+# https://github.com/gepa-ai/gepa
+
+from collections.abc import Mapping, Sequence
+from typing import Any, NamedTuple, Protocol, TypedDict, cast
+
+from mantisdk.algorithm.gepa.lib.core.adapter import EvaluationBatch, GEPAAdapter
+
+
+# DataInst, Trajectory, RolloutOutput
+class DefaultDataInst(TypedDict):
+    input: str
+    additional_context: dict[str, str]
+    answer: str
+
+
+class EvaluationResult(NamedTuple):
+    score: float
+    feedback: str
+    objective_scores: dict[str, float] | None = None
+
+
+class DefaultTrajectory(TypedDict):
+    data: DefaultDataInst
+    full_assistant_response: str
+    feedback: str
+
+
+class DefaultRolloutOutput(TypedDict):
+    full_assistant_response: str
+
+
+DefaultReflectiveRecord = TypedDict(
+    "DefaultReflectiveRecord",
+    {
+        "Inputs": str,
+        "Generated Outputs": str,
+        "Feedback": str,
+    },
+)
+
+
+class ChatMessage(TypedDict):
+    role: str
+    content: str
+
+
+class ChatCompletionCallable(Protocol):
+    """Protocol for chat completion callables (duck typing for custom model wrappers)."""
+
+    def __call__(self, messages: Sequence[ChatMessage]) -> str: ...
+
+
+# Callable that evaluates a response and returns (score, feedback, optional objective_scores)
+class Evaluator(Protocol):
+    def __call__(self, data: DefaultDataInst, response: str) -> EvaluationResult:
+        """
+        Evaluates a response and returns a score, feedback, and optional objective scores.
+        """
+        ...
+
+
+class ContainsAnswerEvaluator:
+    """Default evaluator that checks if the expected answer is contained in the response."""
+
+    def __init__(self, failure_score: float = 0.0):
+        self.failure_score = failure_score
+
+    def __call__(self, data: DefaultDataInst, response: str) -> EvaluationResult:
+        is_correct = data["answer"] in response
+        score = 1.0 if is_correct else self.failure_score
+
+        if is_correct:
+            feedback = f"The generated response is correct. The response include the correct answer '{data['answer']}'"
+        else:
+            additional_context_str = "\n".join(f"{k}: {v}" for k, v in data["additional_context"].items())
+            feedback = (
+                f"The generated response is incorrect. The correct answer is '{data['answer']}'. "
+                "Ensure that the correct answer is included in the response exactly as it is."
+            )
+            if additional_context_str:
+                feedback += f" Here is some additional context that might be helpful:\n{additional_context_str}"
+
+        return EvaluationResult(score=score, feedback=feedback, objective_scores=None)
+
+
+class DefaultAdapter(GEPAAdapter[DefaultDataInst, DefaultTrajectory, DefaultRolloutOutput]):
+    def __init__(
+        self,
+        model: str | ChatCompletionCallable,
+        evaluator: Evaluator | None = None,
+        max_litellm_workers: int = 10,
+        litellm_batch_completion_kwargs: dict[str, Any] | None = None,
+    ):
+        if isinstance(model, str):
+            import litellm
+
+            self.litellm = litellm
+        self.model = model
+        self.evaluator = evaluator or ContainsAnswerEvaluator()
+        self.max_litellm_workers = max_litellm_workers
+        self.litellm_batch_completion_kwargs = litellm_batch_completion_kwargs or {}
+
+    def evaluate(
+        self,
+        batch: list[DefaultDataInst],
+        candidate: dict[str, str],
+        capture_traces: bool = False,
+    ) -> EvaluationBatch[DefaultTrajectory, DefaultRolloutOutput]:
+        outputs: list[DefaultRolloutOutput] = []
+        scores: list[float] = []
+        objective_scores: list[dict[str, float] | None] = []
+        trajectories: list[DefaultTrajectory] | None = [] if capture_traces else None
+
+        system_content = next(iter(candidate.values()))
+
+        litellm_requests = []
+
+        for data in batch:
+            user_content = f"{data['input']}"
+
+            messages: list[ChatMessage] = [
+                {"role": "system", "content": system_content},
+                {"role": "user", "content": user_content},
+            ]
+
+            litellm_requests.append(messages)
+
+        if isinstance(self.model, str):
+            responses = [
+                resp.choices[0].message.content.strip()
+                for resp in self.litellm.batch_completion(
+                    model=self.model,
+                    messages=litellm_requests,
+                    max_workers=self.max_litellm_workers,
+                    **self.litellm_batch_completion_kwargs,
+                )
+            ]
+        else:
+            responses = [self.model(messages) for messages in litellm_requests]
+
+        for data, assistant_response in zip(batch, responses, strict=True):
+            eval_result = self.evaluator(data, assistant_response)
+            score = eval_result.score
+            feedback = eval_result.feedback
+            obj_scores = eval_result.objective_scores
+
+            output: DefaultRolloutOutput = {"full_assistant_response": assistant_response}
+
+            outputs.append(output)
+            scores.append(score)
+            objective_scores.append(obj_scores)
+
+            if trajectories is not None:
+                trajectories.append(
+                    {
+                        "data": data,
+                        "full_assistant_response": assistant_response,
+                        "feedback": feedback,
+                    }
+                )
+
+        objective_scores_arg: list[dict[str, float]] | None = None
+        if objective_scores:
+            all_none = all(x is None for x in objective_scores)
+            all_not_none = all(x is not None for x in objective_scores)
+            if not (all_none or all_not_none):
+                raise ValueError("Objective scores must either be all None or all not None.")
+            if all_not_none:
+                objective_scores_arg = cast(list[dict[str, float]], objective_scores)
+
+        return EvaluationBatch(
+            outputs=outputs,
+            scores=scores,
+            trajectories=trajectories,
+            objective_scores=objective_scores_arg,
+        )
+
+    def make_reflective_dataset(
+        self,
+        candidate: dict[str, str],
+        eval_batch: EvaluationBatch[DefaultTrajectory, DefaultRolloutOutput],
+        components_to_update: list[str],
+    ) -> Mapping[str, Sequence[Mapping[str, Any]]]:
+        ret_d: dict[str, list[DefaultReflectiveRecord]] = {}
+
+        assert len(components_to_update) == 1
+        comp = components_to_update[0]
+
+        trajectories = eval_batch.trajectories
+        assert trajectories is not None, "Trajectories are required to build a reflective dataset."
+
+        items: list[DefaultReflectiveRecord] = []
+
+        for traj in trajectories:
+            d: DefaultReflectiveRecord = {
+                "Inputs": traj["data"]["input"],
+                "Generated Outputs": traj["full_assistant_response"],
+                "Feedback": traj["feedback"],
+            }
+
+            items.append(d)
+
+        ret_d[comp] = items
+
+        if len(items) == 0:
+            raise Exception("No valid predictions found for any module.")
+
+        return ret_d

mantisdk/algorithm/gepa/lib/adapters/dspy_adapter/README.md

@@ -0,0 +1,7 @@
+# DSPy <> GEPA
+
+[dspy_adapter.py](dspy_adapter.py) provides an example adapter to allow GEPA to optimize the signature instructions of any DSPy program. The most up-to-date version of this adapter is live in the DSPy repository at [gepa_utils.py](https://github.com/stanfordnlp/dspy/blob/main/dspy/teleprompt/gepa/gepa_utils.py).
+
+> The best way to use this adapter is from within DSPy itself.
+
+Extensive tutorials on using GEPA with DSPy are available at [dspy.GEPA tutorials](https://dspy.ai/tutorials/gepa_ai_program/).

mantisdk/algorithm/gepa/lib/adapters/dspy_adapter/dspy_adapter.py

@@ -0,0 +1,307 @@
+"""
+This file provides an example adapter allowing GEPA to optimize text components of DSPy programs (instructions and prompts).
+The most up-to-date version of this file is in the DSPy repository: https://github.com/stanfordnlp/dspy/blob/main/dspy/teleprompt/gepa/gepa_utils.py
+"""
+
+import logging
+import random
+from typing import Any, Callable, Protocol
+
+from dspy.adapters.chat_adapter import ChatAdapter
+from dspy.adapters.types import History
+from dspy.evaluate import Evaluate
+from dspy.primitives import Example, Prediction
+from dspy.teleprompt.bootstrap_trace import TraceData
+
+from mantisdk.algorithm.gepa.lib.core.adapter import EvaluationBatch, GEPAAdapter
+
+
+class LoggerAdapter:
+    def __init__(self, logger: logging.Logger):
+        self.logger = logger
+
+    def log(self, x: str):
+        self.logger.info(x)
+
+
+DSPyTrace = list[tuple[Any, dict[str, Any], Prediction]]
+
+
+class ScoreWithFeedback(Prediction):
+    score: float
+    feedback: str | None = None
+    subscores: dict[str, float] | None = None
+
+
+class PredictorFeedbackFn(Protocol):
+    def __call__(
+        self,
+        predictor_output: dict[str, Any],
+        predictor_inputs: dict[str, Any],
+        module_inputs: Example,
+        module_outputs: Prediction,
+        captured_trace: DSPyTrace,
+    ) -> ScoreWithFeedback:
+        """
+        This function is used to provide feedback to a specific predictor.
+        The function is called with the following arguments:
+        - predictor_output: The output of the predictor.
+        - predictor_inputs: The inputs to the predictor.
+        - module_inputs: The inputs to the whole program --- `Example`.
+        - module_outputs: The outputs of the whole program --- `Prediction`.
+        - captured_trace: The trace of the module's execution.
+        # Shape of trace is: [predictor_invocation_idx -> Tuple[Predictor, PredictorInputs, Prediction]]
+        # Each trace is a tuple of (Predictor, PredictorInputs, Prediction)
+
+        The function should return a `ScoreWithFeedback` object.
+        The feedback is a string that is used to guide the evolution of the predictor.
+        """
+        ...
+
+
+class DspyAdapter(GEPAAdapter[Example, TraceData, Prediction]):
+    def __init__(
+        self,
+        student_module,
+        metric_fn: Callable,
+        feedback_map: dict[str, Callable],
+        failure_score=0.0,
+        num_threads: int | None = None,
+        add_format_failure_as_feedback: bool = False,
+        rng: random.Random | None = None,
+    ):
+        self.student = student_module
+        self.metric_fn = metric_fn
+        self.feedback_map = feedback_map
+        self.failure_score = failure_score
+        self.num_threads = num_threads
+        self.add_format_failure_as_feedback = add_format_failure_as_feedback
+        self.rng = rng or random.Random(0)
+
+        # Cache predictor names/signatures
+        self.named_predictors = list(self.student.named_predictors())
+
+    def build_program(self, candidate: dict[str, str]):
+        new_prog = self.student.deepcopy()
+        for name, pred in new_prog.named_predictors():
+            if name in candidate:
+                pred.signature = pred.signature.with_instructions(candidate[name])
+        return new_prog
+
+    def evaluate(self, batch, candidate, capture_traces=False):
+        program = self.build_program(candidate)
+
+        outputs: list[Prediction] = []
+        scores: list[float] = []
+        subscores: list[dict[str, float]] = []
+        trajs: list[TraceData] | None = None
+
+        if capture_traces:
+            # bootstrap_trace_data-like flow with trace capture
+            from dspy.teleprompt.bootstrap_trace import bootstrap_trace_data
+
+            trajs = bootstrap_trace_data(
+                program=program,
+                dataset=batch,
+                metric=self.metric_fn,
+                num_threads=self.num_threads,
+                raise_on_error=False,
+                capture_failed_parses=True,
+                failure_score=self.failure_score,
+                format_failure_score=self.failure_score,
+            )
+            for t in trajs:
+                outputs.append(t["prediction"])
+                score_val, subscore_dict = self._extract_score_and_subscores(t.get("score"))
+                if score_val is None:
+                    score_val = self.failure_score
+                scores.append(score_val)
+                subscores.append(subscore_dict)
+        else:
+            evaluator = Evaluate(
+                devset=batch,
+                metric=self.metric_fn,
+                num_threads=self.num_threads,
+                return_all_scores=True,
+                return_outputs=True,
+                failure_score=self.failure_score,
+                provide_traceback=True,
+                max_errors=len(batch) * 100,
+            )
+            res = evaluator(program)
+            outputs = [r[1] for r in res.results]
+            raw_scores = [r[2] for r in res.results]
+            for raw_score in raw_scores:
+                score_val, subscore_dict = self._extract_score_and_subscores(raw_score)
+                if score_val is None:
+                    score_val = self.failure_score
+                scores.append(score_val)
+                subscores.append(subscore_dict)
+
+        has_subscores = any(subscores)
+        # Map DSPy "subscores" into GEPA objective score payloads.
+        return EvaluationBatch(
+            outputs=outputs,
+            scores=scores,
+            trajectories=trajs,
+            objective_scores=subscores if has_subscores else None,
+        )
+
+    def make_reflective_dataset(self, candidate, eval_batch, components_to_update):
+        from dspy.teleprompt.bootstrap_trace import FailedPrediction
+
+        program = self.build_program(candidate)
+
+        ret_d: dict[str, list[dict[str, Any]]] = {}
+        for pred_name in components_to_update:
+            module = None
+            for name, m in program.named_predictors():
+                if name == pred_name:
+                    module = m
+                    break
+            assert module is not None
+
+            items: list[dict[str, Any]] = []
+            for data in eval_batch.trajectories or []:
+                trace = data["trace"]
+                example = data["example"]
+                prediction = data["prediction"]
+                module_score_obj = data.get("score")
+                module_score, _ = self._extract_score_and_subscores(module_score_obj)
+
+                trace_instances = [t for t in trace if t[0].signature.equals(module.signature)]
+                if not self.add_format_failure_as_feedback:
+                    trace_instances = [t for t in trace_instances if not isinstance(t[2], FailedPrediction)]
+                if len(trace_instances) == 0:
+                    continue
+
+                selected = None
+                for t in trace_instances:
+                    if isinstance(t[2], FailedPrediction):
+                        selected = t
+                        break
+
+                if selected is None:
+                    if isinstance(prediction, FailedPrediction):
+                        continue
+                    selected = self.rng.choice(trace_instances)
+
+                inputs = selected[1]
+                outputs = selected[2]
+
+                new_inputs = {}
+                new_outputs = {}
+
+                contains_history = False
+                history_key_name = None
+                for input_key, input_val in inputs.items():
+                    if isinstance(input_val, History):
+                        contains_history = True
+                        assert history_key_name is None
+                        history_key_name = input_key
+
+                if contains_history:
+                    s = "```json\n"
+                    for i, message in enumerate(inputs[history_key_name].messages):
+                        s += f"  {i}: {message}\n"
+                    s += "```"
+                    new_inputs["Context"] = s
+
+                for input_key, input_val in inputs.items():
+                    if contains_history and input_key == history_key_name:
+                        continue
+                    new_inputs[input_key] = str(input_val)
+
+                if isinstance(outputs, FailedPrediction):
+                    s = "Couldn't parse the output as per the expected output format. The model's raw response was:\n"
+                    s += "```\n"
+                    s += outputs.completion_text + "\n"
+                    s += "```\n\n"
+                    new_outputs = s
+                else:
+                    for output_key, output_val in outputs.items():
+                        new_outputs[output_key] = str(output_val)
+
+                d = {"Inputs": new_inputs, "Generated Outputs": new_outputs}
+                if isinstance(outputs, FailedPrediction):
+                    adapter = ChatAdapter()
+                    structure_instruction = ""
+                    for dd in adapter.format(module.signature, [], {}):
+                        structure_instruction += dd["role"] + ": " + dd["content"] + "\n"
+                    d["Feedback"] = "Your output failed to parse. Follow this structure:\n" + structure_instruction
+                    # d['score'] = self.failure_score
+                else:
+                    feedback_fn = self.feedback_map[pred_name]
+                    fb = feedback_fn(
+                        predictor_output=outputs,
+                        predictor_inputs=inputs,
+                        module_inputs=example,
+                        module_outputs=prediction,
+                        captured_trace=trace,
+                    )
+                    if isinstance(fb, dict):
+                        feedback_score = fb.get("score")
+                        feedback_text = fb.get("feedback", "")
+                    else:
+                        feedback_score = getattr(fb, "score", None)
+                        feedback_text = getattr(fb, "feedback", "")
+                    d["Feedback"] = feedback_text
+                    if module_score is not None and feedback_score is not None:
+                        assert abs(feedback_score - module_score) < 1e-8, (
+                            "Currently, GEPA only supports feedback functions that return the same score as the module's score. "
+                            f"However, the module-level score is {module_score} and the feedback score is {feedback_score}."
+                        )
+                items.append(d)
+
+            if len(items) == 0:
+                # raise Exception(f"No valid predictions found for module {module.signature}.")
+                continue
+            ret_d[pred_name] = items
+
+        if len(ret_d) == 0:
+            raise Exception("No valid predictions found for any module.")
+
+        return ret_d
+
+    @staticmethod
+    def _extract_score_and_subscores(score_obj: Any) -> tuple[float | None, dict[str, float]]:
+        if score_obj is None:
+            return None, {}
+        if isinstance(score_obj, dict):
+            score_val = score_obj.get("score")
+            subscores = score_obj.get("subscores") or {}
+            return score_val, dict(subscores)
+        if hasattr(score_obj, "score"):
+            score_val = getattr(score_obj, "score", None)
+            subscores = getattr(score_obj, "subscores", None) or {}
+            return score_val, dict(subscores)
+        try:
+            return float(score_obj), {}
+        except (TypeError, ValueError):
+            return None, {}
+
+    # TODO: The current DSPyAdapter implementation uses the GEPA default propose_new_texts.
+    # We can potentially override this, to use the instruction proposal similar to MIPROv2.
+
+    # def propose_new_texts(
+    #     self,
+    #     candidate: Dict[str, str],
+    #     reflective_dataset: Dict[str, List[Dict[str, Any]]],
+    #     components_to_update: List[str]
+    # ) -> Dict[str, str]:
+    #     if self.adapter.propose_new_texts is not None:
+    #         return self.adapter.propose_new_texts(candidate, reflective_dataset, components_to_update)
+
+    #     from .instruction_proposal import InstructionProposalSignature
+    #     new_texts: Dict[str, str] = {}
+    #     for name in components_to_update:
+    #         base_instruction = candidate[name]
+    #         dataset_with_feedback = reflective_dataset[name]
+    #         new_texts[name] = InstructionProposalSignature.run(
+    #             lm=self.reflection_lm,
+    #             input_dict={
+    #                 "current_instruction_doc": base_instruction,
+    #                 "dataset_with_feedback": dataset_with_feedback
+    #             }
+    #         )['new_instruction']
+    #     return new_texts

mantisdk/algorithm/gepa/lib/adapters/dspy_full_program_adapter/README.md

@@ -0,0 +1,99 @@
+# DSPy Full Program Adapter
+
+This adapter lets GEPA evolve entire DSPy programs—including signatures, modules, and control flow—not just prompts or instructions.
+
+## Usage
+
+First, install DSPy (version 3.0.2 or higher) with: `pip install 'dspy>=3.0.2'`.
+
+Now, let's use GEPA to generate a DSPy program to solve MATH (benchmark). We start with a very very simple DSPy program `dspy.ChainOfThought("question -> answer")`:
+```python
+from gepa import optimize
+from gepa.adapters.dspy_full_program_adapter.full_program_adapter import DspyAdapter
+import dspy
+
+# Standard DSPy metric function
+def metric_fn(example, pred, trace=None):
+    ...
+
+# Start with a basic program. This code block must export a `program` that shows how the task should be performed
+seed_program = """import dspy
+program = dspy.ChainOfThought("question -> answer")"""
+
+# Run optimization
+reflection_lm = dspy.LM(model="openai/gpt-4.1", max_tokens=32000) # <-- This LM will only be used to propose new DSPy programs
+result = optimize(
+    seed_candidate={"program": seed_program},
+    trainset=train_data,
+    valset=val_data,
+    adapter=DspyAdapter(
+        task_lm=dspy.LM(model="openai/gpt-4.1-nano", max_tokens=32000), # <-- This LM will be used for the downstream task
+        metric_fn=metric_fn,
+        reflection_lm=lambda x: reflection_lm(x)[0],
+    ),
+    max_metric_calls=2000,
+)
+
+# Get the evolved program
+optimized_program_code = result.best_candidate["program"]
+print(optimized_program_code)
+```
+
+Using dspy.ChainOfThought with GPT-4.1 Nano achieves a score of **67%**, while the following GEPA-optimized program boosts performance to **93%**!
+```
+import dspy
+from typing import Optional
+
+class MathQAReasoningSignature(dspy.Signature):
+    """
+    Solve the given math word problem step by step, showing all necessary reasoning and calculations.
+    - First, provide a clear, detailed, and logically ordered reasoning chain, using equations and algebraic steps as needed.
+    - Then, extract the final answer in the required format, strictly following these rules:
+        * If the answer should be a number, output only the number (no units, unless explicitly requested).
+        * If the answer should be an algebraic expression, output it in LaTeX math mode (e.g., \frac{h^2}{m}).
+        * Do not include explanatory text, units, or extra formatting in the answer field unless the question explicitly requests it.
+    Common pitfalls:
+        - Including units when not required.
+        - Restating the answer with extra words or formatting.
+        - Failing to simplify expressions or extract the final answer.
+    Edge cases:
+        - If the answer is a sum or list, output only the final value(s) as required.
+        - If the answer is an expression, ensure it is fully simplified.
+    Successful strategies:
+        - Use step-by-step algebraic manipulation.
+        - Double-check the final answer for correct format and content.
+    """
+    question: str = dspy.InputField(desc="A math word problem to solve.")
+    reasoning: str = dspy.OutputField(desc="Step-by-step solution, with equations and logic.")
+    answer: str = dspy.OutputField(desc="Final answer, strictly in the required format (see instructions).")
+
+class MathQAExtractSignature(dspy.Signature):
+    """
+    Given a math word problem and a detailed step-by-step solution, extract ONLY the final answer in the required format.
+    - If the answer should be a number, output only the number (no units, unless explicitly requested).
+    - If the answer should be an algebraic expression, output it in LaTeX math mode (e.g., \frac{h^2}{m}).
+    - Do not include explanatory text, units, or extra formatting in the answer field unless the question explicitly requests it.
+    - If the answer is a sum or list, output only the final value(s) as required.
+    """
+    question: str = dspy.InputField(desc="The original math word problem.")
+    reasoning: str = dspy.InputField(desc="A detailed, step-by-step solution to the problem.")
+    answer: str = dspy.OutputField(desc="Final answer, strictly in the required format.")
+
+class MathQAModule(dspy.Module):
+    def __init__(self):
+        super().__init__()
+        self.reasoner = dspy.ChainOfThought(MathQAReasoningSignature)
+        self.extractor = dspy.Predict(MathQAExtractSignature)
+
+    def forward(self, question: str):
+        reasoning_pred = self.reasoner(question=question)
+        extract_pred = self.extractor(question=question, reasoning=reasoning_pred.reasoning)
+        return dspy.Prediction(
+            reasoning=reasoning_pred.reasoning,
+            answer=extract_pred.answer
+        )
+
+program = MathQAModule()
+```
+
+A fully executable notebook to run this example is in [src/gepa/examples/dspy_full_program_evolution/example.ipynb](../../examples/dspy_full_program_evolution/example.ipynb)