mantisdk 0.1.0__py3-none-any.whl

mantisdk/algorithm/gepa/lib/adapters/dspy_full_program_adapter/dspy_program_proposal_signature.py

@@ -0,0 +1,137 @@
+# Copyright (c) 2025 Lakshya A Agrawal and the GEPA contributors
+# https://github.com/gepa-ai/gepa
+
+from typing import Any, ClassVar
+
+import yaml
+
+from mantisdk.algorithm.gepa.lib.proposer.reflective_mutation.base import Signature
+
+
+class DSPyProgramProposalSignature(Signature):
+    prompt_template = """I am trying to solve a task using the DSPy framework. Here's a comprehensive overview of DSPy concepts to guide your improvements:
+
+Signatures:
+- Signatures define tasks declaratively through input/output fields and explicit instructions.
+- They serve as blueprints for what the LM needs to accomplish.
+
+Signature Types:
+- Simple signatures: Specified as strings like "input1, ..., inputN -> output1, ..., outputM" (e.g., "topic -> tweet").
+- Typed signatures: Create a subclass of dspy.Signature with a detailed docstring that includes task instructions, common pitfalls, edge cases, and successful strategies. Define fields using dspy.InputField(desc="...", type=...) and dspy.OutputField(desc="...", type=...) with pydantic types such as str, List[str], Literal["option1", "option2"], or custom classes.
+
+Modules:
+- Modules specify __how__ to solve the task defined by a signature.
+- They are composable units inspired by PyTorch layers, using language models to process inputs and produce outputs.
+- Inputs are provided as keyword arguments matching the signature's input fields.
+- Outputs are returned as dspy.Prediction objects containing the signature's output fields.
+- Key built-in modules:
+  - dspy.Predict(signature): Performs a single LM call to directly generate the outputs from the inputs.
+  - dspy.ChainOfThought(signature): Performs a single LM call that first generates a reasoning chain, then the outputs (adds a 'reasoning' field to the prediction).
+  - Other options: dspy.ReAct(signature) for reasoning and acting, or custom chains.
+- Custom modules: Subclass dspy.Module. In __init__, compose sub-modules (e.g., other Predict or ChainOfThought instances). In forward(self, **kwargs), define the data flow: call sub-modules, execute Python logic if needed, and return dspy.Prediction with the output fields.
+
+Example Usage:
+```
+# Simple signature
+simple_signature = "question -> answer"
+
+# Typed signature
+class ComplexSignature(dspy.Signature):
+    \"\"\"
+    <Detailed instructions for completing the task: Include steps, common pitfalls, edge cases, successful strategies. Include domain knowledge...>
+    \"\"\"
+    question: str = dspy.InputField(desc="The question to answer")
+    answer: str = dspy.OutputField(desc="Concise and accurate answer")
+
+# Built-in module
+simple_program = dspy.Predict(simple_signature)  # or dspy.ChainOfThought(ComplexSignature)
+
+# Custom module
+class ComplexModule(dspy.Module):
+    def __init__(self):
+        self.reasoner = dspy.ChainOfThought("question -> intermediate_answer")
+        self.finalizer = dspy.Predict("intermediate_answer -> answer")
+
+    def forward(self, question: str):
+        intermediate = self.reasoner(question=question)
+        final = self.finalizer(intermediate_answer=intermediate.intermediate_answer)
+        return dspy.Prediction(answer=final.answer, reasoning=intermediate.reasoning) # dspy.ChainOfThought returns 'reasoning' in addition to the signature outputs.
+
+complex_program = ComplexModule()
+```
+
+DSPy Improvement Strategies:
+1. Analyze traces for LM overload: If a single call struggles (e.g., skips steps or hallucinates), decompose into multi-step modules with ChainOfThought or custom logic for stepwise reasoning.
+2. Avoid over-decomposition: If the program is too fragmented, consolidate related steps into fewer modules for efficiency and coherence.
+3. Refine signatures: Enhance docstrings with actionable guidance from traces—address specific errors, incorporate domain knowledge, document edge cases, and suggest reasoning patterns. Ensure docstrings are self-contained, as the LM won't have access external traces during runtime.
+4. Balance LM and Python: Use Python for symbolic/logical operations (e.g., loops, conditionals); delegate complex reasoning or generation to LM calls.
+5. Incorporate control flow: Add loops, conditionals, sub-modules in custom modules if the task requires iteration (e.g., multi-turn reasoning, selection, voting, etc.).
+6. Leverage LM strengths: For code-heavy tasks, define signatures with 'code' outputs, extract and execute the generated code in the module's forward pass.
+
+Here's my current code:
+```
+<curr_program>
+```
+
+Here is the execution trace of the current code on example inputs, their outputs, and detailed feedback on improvements:
+```
+<dataset_with_feedback>
+```
+
+Assignment:
+- Think step-by-step: First, deeply analyze the current code, traces, and feedback to identify failure modes, strengths, and opportunities.
+- Create a concise checklist (3-7 bullets) outlining your high-level improvement plan, focusing on conceptual changes (e.g., "Decompose step X into a multi-stage module").
+- Then, propose a drop-in replacement code that instantiates an improved 'program' object.
+- Ensure the code is modular, efficient, and directly addresses feedback.
+- Output everything in a single code block using triple backticks—no additional explanations, comments, or language markers outside the block.
+- The code must be a valid, self-contained Python script with all necessary imports, definitions, and assignment to 'program'.
+
+Output Format:
+- Start with the checklist in plain text (3-7 short bullets).
+- Follow immediately with one code block in triple backticks containing the complete Python code, including assigning a `program` object."""
+    input_keys: ClassVar[list[str]] = ["curr_program", "dataset_with_feedback"]
+    output_keys: ClassVar[list[str]] = ["new_program"]
+
+    @classmethod
+    def prompt_renderer(cls, input_dict: dict[str, Any]) -> str:
+        curr_program = input_dict["curr_program"]
+        if not isinstance(curr_program, str):
+            raise TypeError("curr_program must be a string")
+
+        dataset = input_dict["dataset_with_feedback"]
+        if not isinstance(dataset, list):
+            raise TypeError("dataset_with_feedback must be a list")
+
+        def format_samples(samples):
+            # Serialize the samples list to YAML for concise, structured representation
+            yaml_str = yaml.dump(samples, sort_keys=False, default_flow_style=False, indent=2)
+            # Optionally, wrap or label it for clarity in the prompt
+            return yaml_str
+
+        prompt = cls.prompt_template
+        prompt = prompt.replace("<curr_program>", curr_program)
+        prompt = prompt.replace("<dataset_with_feedback>", format_samples(dataset))
+        return prompt
+
+    @staticmethod
+    def output_extractor(lm_out: str) -> dict[str, str]:
+        # Extract ``` blocks
+        new_instruction = None
+        if lm_out.count("```") >= 2:
+            start = lm_out.find("```")
+            end = lm_out.rfind("```")
+            if start >= end:
+                new_instruction = lm_out
+            if start == -1 or end == -1:
+                new_instruction = lm_out
+            else:
+                new_instruction = lm_out[start + 3 : end].strip()
+        else:
+            lm_out = lm_out.strip()
+            if lm_out.startswith("```"):
+                lm_out = lm_out[3:]
+            if lm_out.endswith("```"):
+                lm_out = lm_out[:-3]
+            new_instruction = lm_out
+
+        return {"new_program": new_instruction}

mantisdk/algorithm/gepa/lib/adapters/dspy_full_program_adapter/full_program_adapter.py

@@ -0,0 +1,266 @@
+import random
+from typing import Any, Callable
+
+import dspy
+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 import EvaluationBatch, GEPAAdapter
+
+
+class DspyAdapter(GEPAAdapter[Example, TraceData, Prediction]):
+    def __init__(
+        self,
+        task_lm: dspy.LM,
+        metric_fn: Callable,
+        reflection_lm: dspy.LM,
+        failure_score=0.0,
+        num_threads: int | None = None,
+        add_format_failure_as_feedback: bool = False,
+        rng: random.Random | None = None,
+    ):
+        self.task_lm = task_lm
+        self.metric_fn = metric_fn
+        assert reflection_lm is not None, (
+            "DspyAdapter for full-program evolution requires a reflection_lm to be provided"
+        )
+        self.reflection_lm = reflection_lm
+        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)
+
+    def build_program(self, candidate: dict[str, str]) -> tuple[dspy.Module, None] | tuple[None, str]:
+        candidate_src = candidate["program"]
+        context = {}
+        o = self.load_dspy_program_from_code(candidate_src, context)
+        return o
+
+    def load_dspy_program_from_code(
+        self,
+        candidate_src: str,
+        context: dict,
+    ):
+        try:
+            compile(candidate_src, "<string>", "exec")
+        except SyntaxError as e:
+            # print(f"Syntax Error in original code {e}")
+            # return None
+            import traceback
+
+            tb = traceback.format_exc()
+            return None, f"Syntax Error in code: {e}\n{tb}"
+
+        try:
+            exec(candidate_src, context)  # expose to current namespace
+        except Exception as e:
+            import traceback
+
+            tb = traceback.format_exc()
+            return None, f"Error in executing code: {e}\n{tb}"
+
+        dspy_program = context.get("program")
+
+        if dspy_program is None:
+            return (
+                None,
+                "Your code did not define a `program` object. Please define a `program` object which is an instance of `dspy.Module`, either directly by dspy.Predict or dspy.ChainOfThought, or by instantiating a class that inherits from `dspy.Module`.",
+            )
+        else:
+            if not isinstance(dspy_program, dspy.Module):
+                return (
+                    None,
+                    f"Your code defined a `program` object, but it is an instance of {type(dspy_program)}, not `dspy.Module`. Please define a `program` object which is an instance of `dspy.Module`, either directly by dspy.Predict or dspy.ChainOfThought, or by instantiating a class that inherits from `dspy.Module`.",
+                )
+
+        dspy_program.set_lm(self.task_lm)
+
+        return dspy_program, None
+
+    def evaluate(self, batch, candidate, capture_traces=False):
+        program, feedback = self.build_program(candidate)
+
+        if program is None:
+            return EvaluationBatch(outputs=None, scores=[self.failure_score for _ in batch], trajectories=feedback)
+
+        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,
+            )
+            scores = []
+            outputs = []
+            for t in trajs:
+                outputs.append(t["prediction"])
+                if hasattr(t["prediction"], "__class__") and t.get("score") is None:
+                    scores.append(self.failure_score)
+                else:
+                    score = t["score"]
+                    if hasattr(score, "score"):
+                        score = score["score"]
+                    scores.append(score)
+            return EvaluationBatch(outputs=outputs, scores=scores, trajectories=trajs)
+        else:
+            evaluator = Evaluate(
+                devset=batch,
+                metric=self.metric_fn,
+                num_threads=self.num_threads,
+                return_all_scores=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]
+            scores = [r[2] for r in res.results]
+            scores = [s["score"] if hasattr(s, "score") else s for s in scores]
+            return EvaluationBatch(outputs=outputs, scores=scores, trajectories=None)
+
+    def make_reflective_dataset(self, candidate, eval_batch, components_to_update):
+        proposed_program, _ = self.build_program(candidate)
+
+        assert set(components_to_update) == {"program"}, f"set(components_to_update) = {set(components_to_update)}"
+        from dspy.teleprompt.bootstrap_trace import FailedPrediction
+
+        ret_d: dict[str, list[dict[str, Any]]] = {}
+
+        if isinstance(eval_batch.trajectories, str):
+            feedback = eval_batch.trajectories
+            return {"program": {"Feedback": feedback}}
+
+        ########
+        items: list[dict[str, Any]] = []
+        for data in eval_batch.trajectories or []:
+            example_data = {}
+            trace = data["trace"]
+            example = data["example"]
+            example_data["Program Inputs"] = {**example.inputs()}
+            prediction = data["prediction"]
+            example_data["Program Outputs"] = {**prediction}
+            module_score = data["score"]
+
+            if hasattr(module_score, "feedback"):
+                feedback_text = module_score["feedback"]
+            else:
+                feedback_text = None
+
+            if hasattr(module_score, "score"):
+                module_score = module_score["score"]
+
+            trace_instances = trace
+
+            if len(trace_instances) == 0:
+                continue
+
+            selected = None
+            for t in trace_instances:
+                if isinstance(t[2], FailedPrediction):
+                    selected = t
+                    break
+
+            if selected is not None:
+                trace_instances = [selected]
+
+            trace_d = []
+            example_data["Program Trace"] = trace_d
+            for selected in trace_instances:
+                inputs = selected[1]
+                outputs = selected[2]
+
+                pred_name = None
+                for name, predictor in proposed_program.named_predictors():
+                    if predictor.signature.equals(selected[0].signature):
+                        pred_name = name
+                        break
+                assert pred_name is not None, f"Could not find predictor for {selected[0].signature}"
+
+                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 = {"Called Module": pred_name, "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:
+                # assert fb["score"] == module_score, f"Currently, GEPA only supports feedback functions that return the same score as the module's score. However, the module-level score is {module_score} and the feedback score is {fb.score}."
+                # d['score'] = fb.score
+                trace_d.append(d)
+
+            if feedback_text is not None:
+                example_data["Feedback"] = feedback_text
+
+            items.append(example_data)
+
+        if len(items) == 0:
+            raise Exception("No valid predictions found for program.")
+
+        ret_d["program"] = items
+
+        ########
+        if len(ret_d) == 0:
+            raise Exception("No valid predictions found for any module.")
+
+        return ret_d
+
+    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]:
+        from mantisdk.algorithm.gepa.lib.adapters.dspy_full_program_adapter.dspy_program_proposal_signature import DSPyProgramProposalSignature
+
+        new_texts: dict[str, str] = {}
+        for name in components_to_update:
+            base_instruction = candidate[name]
+            dataset_with_feedback = reflective_dataset[name]
+            new_texts[name] = DSPyProgramProposalSignature.run(
+                lm=self.reflection_lm,
+                input_dict={"curr_program": base_instruction, "dataset_with_feedback": dataset_with_feedback},
+            )["new_program"]
+        return new_texts