mantisdk 0.1.0__py3-none-any.whl

mantisdk/algorithm/gepa/lib/examples/rag_adapter/requirements-rag.txt

@@ -0,0 +1,29 @@
+# RAG Adapter Dependencies
+# Install these dependencies based on which vector store you want to use
+
+# Common dependency for all vector stores
+litellm>=1.64.0
+
+# ChromaDB vector store
+chromadb>=0.4.0
+
+# Weaviate vector store  
+weaviate-client>=4.0.0
+
+# Qdrant vector store
+qdrant-client>=1.15.0
+
+# Milvus vector store
+pymilvus>=2.6.0
+
+# LanceDB vector store
+lancedb>=0.22.0
+pyarrow>=10.0.0
+
+# Installation examples:
+# For ChromaDB: pip install litellm>=1.64.0 chromadb>=0.4.0
+# For Weaviate: pip install litellm>=1.64.0 weaviate-client>=4.0.0  
+# For Qdrant: pip install litellm>=1.64.0 qdrant-client>=1.15.0
+# For Milvus: pip install litellm>=1.64.0 pymilvus>=2.6.0
+# For LanceDB: pip install litellm>=1.64.0 lancedb>=0.22.0 pyarrow>=10.0.0
+# For all: pip install -r requirements-rag.txt

mantisdk/algorithm/gepa/lib/examples/terminal-bench/prompt-templates/instruction_prompt.txt

@@ -0,0 +1,16 @@
+
+You are an AI assistant tasked with solving command-line tasks in a Linux environment. You will be given a task instruction and the output from previously executed commands. Your goal is to solve the task by providing batches of shell commands.
+
+For each response:
+1. Analyze the current state based on any terminal output provided
+2. Determine the next set of commands needed to make progress
+3. Decide if you need to see the output of these commands before proceeding
+
+Don't include markdown formatting.
+
+Note that you operate directly on the terminal from inside a tmux session. Use tmux keystrokes like `C-x` or `Escape` to interactively navigate the terminal. If you would like to execute a command that you have written you will need to append a newline character to the end of your command.
+
+For example, if you write "ls -la" you will need to append a newline character to the end of your command like this: `ls -la
+`.
+
+One thing to be very careful about is handling interactive sessions like less, vim, or git diff. In these cases, you should not wait for the output of the command. Instead, you should send the keystrokes to the terminal as if you were typing them.

mantisdk/algorithm/gepa/lib/examples/terminal-bench/prompt-templates/terminus.txt

@@ -0,0 +1,9 @@
+Instruction:
+{instruction}
+
+Your response must be a JSON object that matches this schema:
+
+{response_schema}
+
+The current terminal state is:
+{terminal_state}

mantisdk/algorithm/gepa/lib/examples/terminal-bench/train_terminus.py

@@ -0,0 +1,161 @@
+import argparse
+import json
+from pathlib import Path
+
+import litellm
+from terminal_bench.agents.terminus_1 import AgentResult, Chat, FailureMode, Terminus
+from terminal_bench.dataset.dataset import Dataset
+from terminal_bench.terminal.tmux_session import TmuxSession
+
+from mantisdk.algorithm.gepa.lib import optimize
+from mantisdk.algorithm.gepa.lib.adapters.terminal_bench_adapter.terminal_bench_adapter import (
+    TerminalBenchTask,
+    TerminusAdapter,
+)
+
+INSTRUCTION_PROMPT_PATH = Path(__file__).parent / "prompt-templates/instruction_prompt.txt"
+
+
+class TerminusWrapper(Terminus):
+    def __init__(
+        self,
+        model_name: str,
+        max_episodes: int = 50,
+        api_base: str | None = None,
+        **kwargs,
+    ):
+        self.PROMPT_TEMPLATE_PATH = Path(__file__).parent / "prompt-templates/terminus.txt"
+        self.instruction_prompt = INSTRUCTION_PROMPT_PATH.read_text()
+        super().__init__(model_name, max_episodes, api_base, **kwargs)
+
+    def perform_task(
+        self,
+        instruction: str,
+        session: TmuxSession,
+        logging_dir: Path | None = None,
+    ):
+        chat = Chat(self._llm)
+
+        initial_prompt = self.instruction_prompt + self._prompt_template.format(
+            response_schema=self._response_schema,
+            instruction=instruction,
+            history="",
+            terminal_state=session.capture_pane(),
+        )
+
+        self._run_agent_loop(initial_prompt, session, chat, logging_dir)
+
+        return AgentResult(
+            total_input_tokens=chat.total_input_tokens,
+            total_output_tokens=chat.total_output_tokens,
+            failure_mode=FailureMode.NONE,
+            timestamped_markers=self._timestamped_markers,
+        )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_name", type=str, default="gpt-4o-mini")
+    parser.add_argument("--n_concurrent", type=int, default=6)
+    args = parser.parse_args()
+
+    initial_prompt_from_terminus = """
+You are an AI assistant tasked with solving command-line tasks in a Linux environment. You will be given a task instruction and the output from previously executed commands. Your goal is to solve the task by providing batches of shell commands.
+
+For each response:
+1. Analyze the current state based on any terminal output provided
+2. Determine the next set of commands needed to make progress
+3. Decide if you need to see the output of these commands before proceeding
+
+Don't include markdown formatting.
+
+Note that you operate directly on the terminal from inside a tmux session. Use tmux keystrokes like `C-x` or `Escape` to interactively navigate the terminal. If you would like to execute a command that you have written you will need to append a newline character to the end of your command.
+
+For example, if you write "ls -la" you will need to append a newline character to the end of your command like this: `ls -la\n`.
+
+One thing to be very careful about is handling interactive sessions like less, vim, or git diff. In these cases, you should not wait for the output of the command. Instead, you should send the keystrokes to the terminal as if you were typing them.
+"""
+
+    terminal_bench_dataset = Dataset(name="terminal-bench-core", version="head")
+    terminal_bench_dataset.sort_by_duration()
+
+    terminal_bench_tasks = terminal_bench_dataset._tasks[::-1]
+
+    trainset = [
+        TerminalBenchTask(task_id=task.name, model_name=args.model_name) for task in terminal_bench_tasks[30:50]
+    ]
+    valset = [TerminalBenchTask(task_id=task.name, model_name=args.model_name) for task in terminal_bench_tasks[:30]]
+
+    testset = [
+        TerminalBenchTask(task_id=task.name, model_name=args.model_name)
+        for task in terminal_bench_tasks[50:]
+        if task.name != "chem-rf"
+    ]
+
+    reflection_lm_name = "openai/gpt-5"
+    reflection_lm = (
+        lambda prompt: litellm.completion(
+            model=reflection_lm_name,
+            messages=[{"role": "user", "content": prompt}],
+            reasoning_effort="high",
+        )
+        .choices[0]
+        .message.content
+    )
+
+    adapter = TerminusAdapter(n_concurrent=args.n_concurrent, instruction_prompt_path=INSTRUCTION_PROMPT_PATH)
+    testset_results_no_prompt = adapter.evaluate(testset, {"instruction_prompt": ""}, capture_traces=True)
+    testset_results_before_opt = adapter.evaluate(
+        testset,
+        {"instruction_prompt": initial_prompt_from_terminus},
+        capture_traces=True,
+    )
+
+    with open("gepa_terminus/testset_results_no_prompt.json", "w") as f:
+        json.dump(
+            {
+                "score": sum(trajectory["success"] for trajectory in testset_results_no_prompt.trajectories),
+                "trajectories": testset_results_no_prompt.trajectories,
+            },
+            f,
+            indent=4,
+        )
+    with open("gepa_terminus/testset_results_before_opt.json", "w") as f:
+        json.dump(
+            {
+                "score": sum(trajectory["success"] for trajectory in testset_results_before_opt.trajectories),
+                "trajectories": testset_results_before_opt.trajectories,
+            },
+            f,
+            indent=4,
+        )
+
+    optimized_results = optimize(
+        seed_candidate={"instruction_prompt": initial_prompt_from_terminus},
+        trainset=trainset,
+        valset=valset,
+        adapter=adapter,
+        reflection_lm=reflection_lm,
+        use_wandb=True,
+        max_metric_calls=400,
+        reflection_minibatch_size=3,
+        perfect_score=1,
+        skip_perfect_score=False,
+        run_dir="gepa_terminus",
+    )
+
+    testset_results_after_opt = adapter.evaluate(
+        testset,
+        {"instruction_prompt": optimized_results.best_candidate["instruction_prompt"]},
+        capture_traces=True,
+    )
+
+    with open("gepa_terminus/optimized_results.json", "w") as f:
+        json.dump(
+            {
+                "score": sum(trajectory["success"] for trajectory in testset_results_after_opt.trajectories),
+                "trajectories": testset_results_after_opt.trajectories,
+            },
+            f,
+            indent=4,
+        )

mantisdk/algorithm/gepa/lib/gepa_utils.py

@@ -0,0 +1,117 @@
+# Copyright (c) 2025 Lakshya A Agrawal and the GEPA contributors
+# https://github.com/gepa-ai/gepa
+
+
+import random
+from typing import Any, Mapping
+
+
+def json_default(x):
+    """Default JSON encoder for objects that are not serializable by default."""
+    try:
+        return {**x}
+    except Exception:
+        return repr(x)
+
+
+def idxmax(lst: list[float]) -> int:
+    """Return the index of the maximum value in a list."""
+    max_val = max(lst)
+    return lst.index(max_val)
+
+
+def is_dominated(y, programs, program_at_pareto_front_valset):
+    y_fronts = [front for front in program_at_pareto_front_valset.values() if y in front]
+    for front in y_fronts:
+        found_dominator_in_front = False
+        for other_prog in front:
+            if other_prog in programs:
+                found_dominator_in_front = True
+                break
+        if not found_dominator_in_front:
+            return False
+
+    return True
+
+
+def remove_dominated_programs(program_at_pareto_front_valset, scores=None):
+    freq = {}
+    for front in program_at_pareto_front_valset.values():
+        for p in front:
+            freq[p] = freq.get(p, 0) + 1
+
+    dominated = set()
+    programs = list(freq.keys())
+
+    if scores is None:
+        scores = dict.fromkeys(programs, 1)
+
+    programs = sorted(programs, key=lambda x: scores[x], reverse=False)
+
+    found_to_remove = True
+    while found_to_remove:
+        found_to_remove = False
+        for y in programs:
+            if y in dominated:
+                continue
+            if is_dominated(y, set(programs).difference({y}).difference(dominated), program_at_pareto_front_valset):
+                dominated.add(y)
+                found_to_remove = True
+                break
+
+    dominators = [p for p in programs if p not in dominated]
+    for front in program_at_pareto_front_valset.values():
+        if not front:
+            continue
+        assert any(p in front for p in dominators)
+
+    new_program_at_pareto_front_valset = {
+        val_id: {prog_idx for prog_idx in front if prog_idx in dominators}
+        for val_id, front in program_at_pareto_front_valset.items()
+    }
+    for val_id, front_new in new_program_at_pareto_front_valset.items():
+        assert front_new.issubset(program_at_pareto_front_valset[val_id])
+
+    return new_program_at_pareto_front_valset
+
+
+def find_dominator_programs(pareto_front_programs, train_val_weighted_agg_scores_for_all_programs):
+    train_val_pareto_front_programs = pareto_front_programs
+    new_program_at_pareto_front_valset = remove_dominated_programs(
+        train_val_pareto_front_programs, scores=train_val_weighted_agg_scores_for_all_programs
+    )
+    uniq_progs = []
+    for front in new_program_at_pareto_front_valset.values():
+        uniq_progs.extend(front)
+    uniq_progs = set(uniq_progs)
+    return list(uniq_progs)
+
+
+def select_program_candidate_from_pareto_front(
+    pareto_front_programs: Mapping[Any, set[int]],
+    train_val_weighted_agg_scores_for_all_programs: list[float],
+    rng: random.Random,
+) -> int:
+    train_val_pareto_front_programs = pareto_front_programs
+    new_program_at_pareto_front_valset = remove_dominated_programs(
+        train_val_pareto_front_programs, scores=train_val_weighted_agg_scores_for_all_programs
+    )
+    program_frequency_in_validation_pareto_front = {}
+    for testcase_pareto_front in new_program_at_pareto_front_valset.values():
+        for prog_idx in testcase_pareto_front:
+            if prog_idx not in program_frequency_in_validation_pareto_front:
+                program_frequency_in_validation_pareto_front[prog_idx] = 0
+            program_frequency_in_validation_pareto_front[prog_idx] += 1
+
+    sampling_list = [
+        prog_idx for prog_idx, freq in program_frequency_in_validation_pareto_front.items() for _ in range(freq)
+    ]
+
+    # TODO: Determine if we need this fallback
+    # if not sampling_list:
+    #     # No Pareto programs survived; fall back to the globally highest-scoring program.
+    #     return idxmax(train_val_weighted_agg_scores_for_all_programs)
+    assert len(sampling_list) > 0
+
+    curr_prog_id = rng.choice(sampling_list)
+    return curr_prog_id

mantisdk/algorithm/gepa/lib/logging/experiment_tracker.py

@@ -0,0 +1,187 @@
+# Copyright (c) 2025 Lakshya A Agrawal and the GEPA contributors
+# https://github.com/gepa-ai/gepa
+
+from typing import Any
+
+
+class ExperimentTracker:
+    """
+    Unified experiment tracking that supports both wandb and mlflow.
+    """
+
+    def __enter__(self):
+        """Context manager entry."""
+        self.initialize()
+        self.start_run()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit - always end the run."""
+        self.end_run()
+        return False  # Don't suppress exceptions
+
+    def __init__(
+        self,
+        use_wandb: bool = False,
+        wandb_api_key: str | None = None,
+        wandb_init_kwargs: dict[str, Any] | None = None,
+        use_mlflow: bool = False,
+        mlflow_tracking_uri: str | None = None,
+        mlflow_experiment_name: str | None = None,
+    ):
+        self.use_wandb = use_wandb
+        self.use_mlflow = use_mlflow
+
+        self.wandb_api_key = wandb_api_key
+        self.wandb_init_kwargs = wandb_init_kwargs or {}
+        self.mlflow_tracking_uri = mlflow_tracking_uri
+        self.mlflow_experiment_name = mlflow_experiment_name
+
+        self._created_mlflow_run = False
+
+    def initialize(self):
+        """Initialize the logging backends."""
+        if self.use_wandb:
+            self._initialize_wandb()
+        if self.use_mlflow:
+            self._initialize_mlflow()
+
+    def _initialize_wandb(self):
+        """Initialize wandb."""
+        try:
+            import wandb  # type: ignore
+
+            if self.wandb_api_key:
+                wandb.login(key=self.wandb_api_key, verify=True)
+            else:
+                wandb.login()
+        except ImportError:
+            raise ImportError("wandb is not installed. Please install it or set backend='mlflow' or 'none'.")
+        except Exception as e:
+            raise RuntimeError(f"Error logging into wandb: {e}")
+
+    def _initialize_mlflow(self):
+        """Initialize mlflow."""
+        try:
+            import mlflow  # type: ignore
+
+            if self.mlflow_tracking_uri:
+                mlflow.set_tracking_uri(self.mlflow_tracking_uri)
+            if self.mlflow_experiment_name:
+                mlflow.set_experiment(self.mlflow_experiment_name)
+        except ImportError:
+            raise ImportError("mlflow is not installed. Please install it or set backend='wandb' or 'none'.")
+        except Exception as e:
+            raise RuntimeError(f"Error setting up mlflow: {e}")
+
+    def start_run(self):
+        """Start a new run."""
+        if self.use_wandb:
+            import wandb  # type: ignore
+
+            wandb.init(**self.wandb_init_kwargs)
+        if self.use_mlflow:
+            import mlflow  # type: ignore
+
+            # Only start a new run if there's no active run
+            if mlflow.active_run() is None:
+                mlflow.start_run()
+                self._created_mlflow_run = True
+            else:
+                self._created_mlflow_run = False
+
+    def log_metrics(self, metrics: dict[str, Any], step: int | None = None):
+        """Log metrics to the active backends."""
+        if self.use_wandb:
+            try:
+                import wandb  # type: ignore
+
+                wandb.log(metrics, step=step)
+            except Exception as e:
+                print(f"Warning: Failed to log to wandb: {e}")
+
+        if self.use_mlflow:
+            try:
+                import mlflow  # type: ignore
+
+                mlflow.log_metrics(metrics, step=step)
+            except Exception as e:
+                print(f"Warning: Failed to log to mlflow: {e}")
+
+    def end_run(self):
+        """End the current run."""
+        if self.use_wandb:
+            try:
+                import wandb  # type: ignore
+
+                if wandb.run is not None:
+                    wandb.finish()
+            except Exception as e:
+                print(f"Warning: Failed to end wandb run: {e}")
+
+        if self.use_mlflow:
+            try:
+                import mlflow  # type: ignore
+
+                if self._created_mlflow_run and mlflow.active_run() is not None:
+                    mlflow.end_run()
+                    self._created_mlflow_run = False
+            except Exception as e:
+                print(f"Warning: Failed to end mlflow run: {e}")
+
+    def is_active(self) -> bool:
+        """Check if any backend has an active run."""
+        if self.use_wandb:
+            try:
+                import wandb  # type: ignore
+
+                if wandb.run is not None:
+                    return True
+            except Exception:
+                pass
+
+        if self.use_mlflow:
+            try:
+                import mlflow  # type: ignore
+
+                if mlflow.active_run() is not None:
+                    return True
+            except Exception:
+                pass
+
+        return False
+
+
+def create_experiment_tracker(
+    use_wandb: bool = False,
+    wandb_api_key: str | None = None,
+    wandb_init_kwargs: dict[str, Any] | None = None,
+    use_mlflow: bool = False,
+    mlflow_tracking_uri: str | None = None,
+    mlflow_experiment_name: str | None = None,
+) -> ExperimentTracker:
+    """
+    Create an experiment tracker based on the specified backends.
+
+    Args:
+        use_wandb: Whether to use wandb
+        use_mlflow: Whether to use mlflow
+        wandb_api_key: API key for wandb
+        wandb_init_kwargs: Additional kwargs for wandb.init()
+        mlflow_tracking_uri: Tracking URI for mlflow
+        mlflow_experiment_name: Experiment name for mlflow
+
+    Returns:
+        ExperimentTracker instance
+
+    Note:
+        Both wandb and mlflow can be used simultaneously if desired.
+    """
+    return ExperimentTracker(
+        use_wandb=use_wandb,
+        wandb_api_key=wandb_api_key,
+        wandb_init_kwargs=wandb_init_kwargs,
+        use_mlflow=use_mlflow,
+        mlflow_tracking_uri=mlflow_tracking_uri,
+        mlflow_experiment_name=mlflow_experiment_name,
+    )

mantisdk/algorithm/gepa/lib/logging/logger.py

@@ -0,0 +1,75 @@
+# Copyright (c) 2025 Lakshya A Agrawal and the GEPA contributors
+# https://github.com/gepa-ai/gepa
+
+import sys
+from typing import Protocol
+
+
+class LoggerProtocol(Protocol):
+    def log(self, message: str): ...
+
+
+class StdOutLogger(LoggerProtocol):
+    def log(self, message: str):
+        print(message)
+
+
+class Tee:
+    def __init__(self, *files):
+        self.files = files
+
+    def write(self, obj):
+        for f in self.files:
+            f.write(obj)
+
+    def flush(self):
+        for f in self.files:
+            if hasattr(f, "flush"):
+                f.flush()
+
+    def isatty(self):
+        # True if any of the files is a terminal
+        return any(hasattr(f, "isatty") and f.isatty() for f in self.files)
+
+    def close(self):
+        for f in self.files:
+            if hasattr(f, "close"):
+                f.close()
+
+    def fileno(self):
+        for f in self.files:
+            if hasattr(f, "fileno"):
+                return f.fileno()
+        raise OSError("No underlying file object with fileno")
+
+
+class Logger(LoggerProtocol):
+    def __init__(self, filename, mode="a"):
+        self.file_handle = open(filename, mode)
+        self.file_handle_stderr = open(filename.replace("run_log.", "run_log_stderr."), mode)
+        self.modified_sys = False
+
+    def __enter__(self):
+        self.original_stdout = sys.stdout
+        self.original_stderr = sys.stderr
+        sys.stdout = Tee(sys.stdout, self.file_handle)
+        sys.stderr = Tee(sys.stderr, self.file_handle_stderr)
+        self.modified_sys = True
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        sys.stdout = self.original_stdout
+        sys.stderr = self.original_stderr
+        self.file_handle.close()
+        self.file_handle_stderr.close()
+        self.modified_sys = False
+
+    def log(self, *args, **kwargs):
+        if self.modified_sys:
+            print(*args, **kwargs)
+        else:
+            # Emulate print(*args, **kwargs) behavior but write to the file
+            print(*args, **kwargs)
+            print(*args, file=self.file_handle_stderr, **kwargs)
+        self.file_handle.flush()
+        self.file_handle_stderr.flush()