orizu 0.5.19 → 0.5.21

package/vendor/gepa-python/src/gepa/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 gepa import optimize
+from gepa.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,
+        )

package/vendor/gepa-python/src/gepa/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

package/vendor/gepa-python/src/gepa/gskill/__init__.py

@@ -0,0 +1 @@
+# gskill: Learning Repository-Specific Skills for Coding Agents

package/vendor/gepa-python/src/gepa/gskill/gskill/cost_tracker.py

@@ -0,0 +1,194 @@
+"""
+Unified cost tracker for GEPA optimization.
+Uses LiteLLM's built-in cost tracking for accuracy.
+Thread-safe for parallel execution.
+"""
+
+import json
+import threading
+from datetime import datetime
+from pathlib import Path
+
+import litellm
+
+
+class UnifiedCostTracker:
+    """
+    Thread-safe cost tracker that uses LiteLLM's built-in cost calculation.
+    Tracks agent and reflection costs separately.
+    """
+
+    def __init__(self, log_dir: str = "gepa_results/logs"):
+        self.log_dir = Path(log_dir)
+        self.log_dir.mkdir(parents=True, exist_ok=True)
+
+        self.session_start = datetime.now()
+        self._lock = threading.Lock()
+
+        # Counters
+        self.total_cost = 0.0
+        self.agent_cost = 0.0
+        self.reflection_cost = 0.0
+        self.call_count = 0
+        self.agent_calls = 0
+        self.reflection_calls = 0
+        self.total_tokens = 0
+
+        # Log files
+        timestamp = self.session_start.strftime("%Y%m%d_%H%M%S")
+        self.log_file = self.log_dir / f"cost_log_{timestamp}.jsonl"
+        self.summary_file = self.log_dir / "cost_summary.txt"
+
+        # Register LiteLLM callback
+        self._register_callback()
+
+    def _register_callback(self):
+        """Register ourselves as a LiteLLM success callback."""
+        # Remove any existing callback to avoid duplicates
+        litellm.success_callback = [self._on_completion]
+
+    def _on_completion(self, kwargs, completion_response, start_time, end_time):
+        """Called by LiteLLM after each successful completion."""
+        with self._lock:
+            self.call_count += 1
+            call_num = self.call_count
+
+            # Get model name
+            model = kwargs.get("model", "unknown")
+
+            # Calculate cost using LiteLLM's built-in function
+            try:
+                cost = litellm.completion_cost(completion_response=completion_response)
+            except Exception:
+                cost = 0.0
+
+            # Get token usage
+            tokens = 0
+            if hasattr(completion_response, "usage") and completion_response.usage:
+                tokens = getattr(completion_response.usage, "total_tokens", 0)
+            self.total_tokens += tokens
+
+            # Categorize: reflection models typically have "pro" in name
+            is_reflection = "pro" in model.lower() or "reflection" in kwargs.get("metadata", {}).get("type", "")
+
+            if is_reflection:
+                self.reflection_cost += cost
+                self.reflection_calls += 1
+                call_type = "REFLECTION"
+            else:
+                self.agent_cost += cost
+                self.agent_calls += 1
+                call_type = "AGENT"
+
+            self.total_cost += cost
+
+            # Log entry
+            entry = {
+                "timestamp": datetime.now().isoformat(),
+                "call_number": call_num,
+                "type": call_type,
+                "model": model,
+                "tokens": tokens,
+                "cost": cost,
+                "cumulative_cost": self.total_cost,
+                "agent_cost": self.agent_cost,
+                "reflection_cost": self.reflection_cost,
+            }
+
+            # Write to log file
+            try:
+                with open(self.log_file, "a") as f:
+                    f.write(json.dumps(entry) + "\n")
+            except Exception:
+                pass
+
+            # Cost tracking is silent - check log file or call print_summary() for details
+
+    def write_summary(self) -> str:
+        """Write human-readable summary to file and return it."""
+        with self._lock:
+            elapsed = (datetime.now() - self.session_start).total_seconds()
+            elapsed_min = elapsed / 60
+
+            summary = f"""
+================================================================================
+GEPA COST TRACKER - Experiment Summary
+================================================================================
+
+Session Start: {self.session_start.strftime("%Y-%m-%d %H:%M:%S")}
+Elapsed Time:  {elapsed_min:.1f} minutes
+
+COST BREAKDOWN
+--------------
+Agent Calls:      {self.agent_calls:,} calls  |  ${self.agent_cost:.4f}
+Reflection Calls: {self.reflection_calls:,} calls  |  ${self.reflection_cost:.4f}
+Total Calls:      {self.call_count:,} calls  |  ${self.total_cost:.4f}
+
+Total Tokens: {self.total_tokens:,}
+
+RATES
+-----
+Cost Rate:  ${self.total_cost / elapsed_min:.4f}/min
+Call Rate:  {self.call_count / elapsed_min:.1f} calls/min
+
+PER-UNIT ESTIMATES
+------------------
+Per Agent Call:     ${self.agent_cost / max(self.agent_calls, 1):.4f}
+Per Reflection Call: ${self.reflection_cost / max(self.reflection_calls, 1):.4f}
+
+Log file: {self.log_file}
+================================================================================
+"""
+            with open(self.summary_file, "w") as f:
+                f.write(summary)
+
+            return summary
+
+    def get_stats(self) -> dict:
+        """Get current stats as a dictionary."""
+        with self._lock:
+            elapsed = (datetime.now() - self.session_start).total_seconds()
+            return {
+                "total_cost": self.total_cost,
+                "agent_cost": self.agent_cost,
+                "reflection_cost": self.reflection_cost,
+                "total_calls": self.call_count,
+                "agent_calls": self.agent_calls,
+                "reflection_calls": self.reflection_calls,
+                "total_tokens": self.total_tokens,
+                "elapsed_seconds": elapsed,
+                "cost_per_minute": self.total_cost / (elapsed / 60) if elapsed > 0 else 0,
+            }
+
+    def print_summary(self):
+        """Print and save final summary."""
+        summary = self.write_summary()
+        print(summary)
+
+
+# Global tracker instance
+_tracker: UnifiedCostTracker | None = None
+_tracker_lock = threading.Lock()
+
+
+def get_tracker(log_dir: str = "gepa_results/logs") -> UnifiedCostTracker:
+    """Get or create global tracker instance."""
+    global _tracker
+    with _tracker_lock:
+        if _tracker is None:
+            _tracker = UnifiedCostTracker(log_dir=log_dir)
+        return _tracker
+
+
+def reset_tracker(log_dir: str = "gepa_results/logs") -> UnifiedCostTracker:
+    """Reset and create a new tracker instance."""
+    global _tracker
+    with _tracker_lock:
+        _tracker = UnifiedCostTracker(log_dir=log_dir)
+        return _tracker
+
+
+# For backwards compatibility
+def log_call(model_name, input_tokens, output_tokens, operation="unknown"):
+    """Legacy function - now a no-op since LiteLLM callbacks handle tracking."""
+    pass  # Tracking is now automatic via LiteLLM callbacks