Add full loop test for run_tournament

README.md

@@ -11,6 +11,9 @@ This project provides a small interface for running "tournaments" between langua
    - `NUM_GENERATIONS`
    - `OPENAI_API_BASE`
    - `OPENAI_API_KEY`
+   - `GENERATE_MODEL`
+   - `SCORE_MODEL`
+   - `PAIRWISE_MODEL`
    - `ENABLE_SCORE_FILTER`
    - `ENABLE_PAIRWISE_FILTER`
 2. Install dependencies (example with `pip`):

main.py

@@ -4,7 +4,7 @@ import os, json, re, ast, gradio as gr
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from tqdm import tqdm
 import matplotlib.pyplot as plt
-from tournament_utils import generate_players, prompt_score, prompt_play
+from tournament_utils import generate_players, prompt_score, prompt_pairwise
 
 NUM_TOP_PICKS_DEFAULT = int(os.getenv("NUM_TOP_PICKS", 3))
 POOL_SIZE_DEFAULT = int(os.getenv("POOL_SIZE", 5))
@@ -14,6 +14,9 @@ API_BASE_DEFAULT = os.getenv("OPENAI_API_BASE", "")
 API_TOKEN_DEFAULT = os.getenv("OPENAI_API_KEY", "")
 SCORE_FILTER_DEFAULT = os.getenv("ENABLE_SCORE_FILTER", "true").lower() == "true"
 PAIRWISE_FILTER_DEFAULT = os.getenv("ENABLE_PAIRWISE_FILTER", "true").lower() == "true"
+GENERATE_MODEL_DEFAULT = os.getenv("GENERATE_MODEL", "gpt-4o-mini")
+SCORE_MODEL_DEFAULT = os.getenv("SCORE_MODEL", "gpt-4o-mini")
+PAIRWISE_MODEL_DEFAULT = os.getenv("PAIRWISE_MODEL", "gpt-4o-mini")
 
 def _clean_json(txt):
     txt = re.sub(r"^```.*?\n|```$", "", txt, flags=re.DOTALL).strip()
@@ -25,6 +28,9 @@ def _clean_json(txt):
 def run_tournament(
     api_base,
     api_token,
+    generate_model,
+    score_model,
+    pairwise_model,
     instruction_input,
     criteria_input,
     n_gen,
@@ -40,10 +46,16 @@ def run_tournament(
     num_top_picks = int(num_top_picks)
     pool_size = int(pool_size)
     max_workers = int(max_workers)
-    if api_base:
-        os.environ["OPENAI_API_BASE"] = api_base
-    if api_token:
-        os.environ["OPENAI_API_KEY"] = api_token
+    if not api_base:
+        api_base = API_BASE_DEFAULT
+    if not api_token:
+        api_token = API_TOKEN_DEFAULT
+    if not generate_model:
+        generate_model = GENERATE_MODEL_DEFAULT
+    if not score_model:
+        score_model = SCORE_MODEL_DEFAULT
+    if not pairwise_model:
+        pairwise_model = PAIRWISE_MODEL_DEFAULT
     enable_score_filter = bool(enable_score_filter)
     enable_pairwise_filter = bool(enable_pairwise_filter)
     process_log = []
@@ -54,7 +66,13 @@ def run_tournament(
         tqdm.write(msg)
         yield "\n".join(process_log), hist_fig, top_picks_str
     yield from log("Generating players …")
-    all_players = generate_players(instruction, n_gen)
+    all_players = generate_players(
+        instruction,
+        n_gen,
+        model=generate_model,
+        api_base=api_base,
+        api_key=api_token,
+    )
     yield from log(f"{len(all_players)} players generated")
     def criteria_block():
         return "\n".join(f"{i + 1}) {c}" for i, c in enumerate(criteria_list))
@@ -62,7 +80,15 @@ def run_tournament(
     if enable_score_filter:
         def score(player):
             data = _clean_json(
-                prompt_score(instruction, criteria_list, criteria_block(), player)
+                prompt_score(
+                    instruction,
+                    criteria_list,
+                    criteria_block(),
+                    player,
+                    model=score_model,
+                    api_base=api_base,
+                    api_key=api_token,
+                )
             )
             if "scores" in data and isinstance(data["scores"], list):
                 vals = data["scores"]
@@ -88,7 +114,15 @@ def run_tournament(
     if enable_pairwise_filter:
         def play(a, b):
             winner_label = _clean_json(
-                prompt_play(instruction, criteria_block(), a, b)
+                prompt_pairwise(
+                    instruction,
+                    criteria_block(),
+                    a,
+                    b,
+                    model=pairwise_model,
+                    api_base=api_base,
+                    api_key=api_token,
+                )
             ).get("winner", "A")
             return a if winner_label == "A" else b
 
@@ -151,6 +185,9 @@ demo = gr.Interface(
     inputs=[
         gr.Textbox(value=API_BASE_DEFAULT, label="API Base Path"),
         gr.Textbox(value="", label="API Token", type="password"),
+        gr.Textbox(value=GENERATE_MODEL_DEFAULT, label="Generation Model"),
+        gr.Textbox(value=SCORE_MODEL_DEFAULT, label="Score Model"),
+        gr.Textbox(value=PAIRWISE_MODEL_DEFAULT, label="Pairwise Model"),
         gr.Textbox(lines=10, label="Instruction"),
         gr.Textbox(lines=5, label="Criteria (comma separated)"),
         gr.Number(value=NUM_GENERATIONS_DEFAULT, label="Number of Generations"),

tests/test_main.py

@@ -0,0 +1,112 @@
+import sys, os, types, json
+from unittest.mock import patch, MagicMock
+
+# Ensure project root in path
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+# Provide dummy litellm module so import succeeds
+fake_litellm = types.ModuleType('litellm')
+fake_litellm.completion = MagicMock()
+sys.modules.setdefault('litellm', fake_litellm)
+
+# Provide dummy dotenv module
+fake_dotenv = types.ModuleType('dotenv')
+fake_dotenv.load_dotenv = MagicMock()
+sys.modules.setdefault('dotenv', fake_dotenv)
+
+# Dummy gradio module so import succeeds
+fake_gradio = types.ModuleType('gradio')
+fake_gradio.Interface = MagicMock(return_value=MagicMock(launch=MagicMock()))
+fake_gradio.Textbox = MagicMock
+fake_gradio.Number = MagicMock
+fake_gradio.Checkbox = MagicMock
+fake_gradio.Plot = MagicMock
+sys.modules.setdefault('gradio', fake_gradio)
+
+# Dummy tqdm module for write method
+class FakeTqdmModule(types.ModuleType):
+    def __init__(self):
+        super().__init__('tqdm')
+        self.write = MagicMock()
+    def __call__(self, iterable=None, total=None):
+        return iterable
+
+fake_tqdm_mod = FakeTqdmModule()
+fake_tqdm_mod.tqdm = fake_tqdm_mod
+sys.modules.setdefault('tqdm', fake_tqdm_mod)
+
+# Dummy matplotlib module
+fake_plt = types.ModuleType('matplotlib.pyplot')
+fake_plt.figure = MagicMock(return_value='fig')
+fake_plt.hist = MagicMock()
+fake_matplotlib = types.ModuleType('matplotlib')
+fake_matplotlib.pyplot = fake_plt
+sys.modules.setdefault('matplotlib', fake_matplotlib)
+sys.modules.setdefault('matplotlib.pyplot', fake_plt)
+
+import main
+
+class DummyFuture:
+    def __init__(self, func, *args):
+        self._func = func
+        self._args = args
+    def result(self):
+        return self._func(*self._args)
+
+class DummyExecutor:
+    def __init__(self, *args, **kwargs):
+        pass
+    def __enter__(self):
+        return self
+    def __exit__(self, exc_type, exc, tb):
+        pass
+    def submit(self, func, *args):
+        return DummyFuture(func, *args)
+    def map(self, func, iterable):
+        for item in iterable:
+            yield func(item)
+
+class DummyTqdm:
+    def __call__(self, iterable=None, total=None):
+        return iterable
+    def write(self, msg):
+        pass
+
+def test_run_tournament_full_loop():
+    dummy_tqdm = DummyTqdm()
+    with patch('main.generate_players') as mock_gen, \
+         patch('main.prompt_score') as mock_score, \
+         patch('main.prompt_pairwise') as mock_pair, \
+         patch('main.ThreadPoolExecutor', return_value=DummyExecutor()) as MockExec, \
+         patch('main.as_completed', new=lambda futs: futs), \
+         patch('main.tqdm', new=dummy_tqdm), \
+         patch('main.plt.figure', return_value='fig'), \
+         patch('main.plt.hist'):
+        mock_gen.return_value = ['p1', 'p2', 'p3', 'p4']
+        scores = {'p1':3, 'p2':2, 'p3':1, 'p4':0}
+        mock_score.side_effect = lambda instr, cl, block, player, **kw: json.dumps({'score': scores[player]})
+        mock_pair.side_effect = lambda instr, block, a, b, **kw: json.dumps({'winner': 'A'})
+
+        results = list(main.run_tournament(
+            api_base='b',
+            api_token='k',
+            generate_model='gm',
+            score_model='sm',
+            pairwise_model='pm',
+            instruction_input='instr',
+            criteria_input='c1,c2',
+            n_gen=4,
+            pool_size=2,
+            num_top_picks=1,
+            max_workers=1,
+            enable_score_filter=True,
+            enable_pairwise_filter=True,
+        ))
+
+    process_log, hist_fig, top_picks = results[-1]
+    assert 'Done' in process_log
+    assert hist_fig == 'fig'
+    assert top_picks.strip() in {'p1', 'p2'}
+    mock_gen.assert_called_once_with('instr', 4, model='gm', api_base='b', api_key='k')
+    assert mock_score.call_count == 4
+    assert mock_pair.called

tests/test_tournament_utils.py

@@ -25,22 +25,26 @@ def make_response(contents):
 def test_generate_players():
     resp = make_response([" player1 ", "player2\n"])
     with patch('tournament_utils.completion', return_value=resp) as mock_comp:
-        players = tu.generate_players('instr', 2, model='m')
-        mock_comp.assert_called_once_with(model='m', messages=[{'role': 'user', 'content': 'instr'}], n=2)
+        players = tu.generate_players('instr', 2, model='m', api_base='b', api_key='k')
+        mock_comp.assert_called_once_with(model='m', messages=[{'role': 'user', 'content': 'instr'}], n=2, api_base='b', api_key='k')
         assert players == ['player1', 'player2']
 
 
 def test_prompt_score():
     resp = make_response([" {\"score\": [5]} "])
     with patch('tournament_utils.completion', return_value=resp) as mock_comp:
-        result = tu.prompt_score('instr', ['c1'], 'block', 'pl', model='m')
+        result = tu.prompt_score('instr', ['c1'], 'block', 'pl', model='m', api_base='b', api_key='k')
         mock_comp.assert_called_once()
+        assert mock_comp.call_args.kwargs['api_base'] == 'b'
+        assert mock_comp.call_args.kwargs['api_key'] == 'k'
         assert result == '{"score": [5]}'
 
 
-def test_prompt_play():
+def test_prompt_pairwise():
     resp = make_response([" {\"winner\": \"A\"} "])
     with patch('tournament_utils.completion', return_value=resp) as mock_comp:
-        result = tu.prompt_play('instr', 'block', 'A text', 'B text', model='m')
+        result = tu.prompt_pairwise('instr', 'block', 'A text', 'B text', model='m', api_base='b', api_key='k')
         mock_comp.assert_called_once()
+        assert mock_comp.call_args.kwargs['api_base'] == 'b'
+        assert mock_comp.call_args.kwargs['api_key'] == 'k'
         assert result == '{"winner": "A"}'

tournament_utils.py

@@ -1,12 +1,30 @@
 from litellm import completion
 
 
-def generate_players(instruction: str, n: int, model: str = "gpt-4o-mini"):
+def _completion_kwargs(api_base: str | None, api_key: str | None) -> dict:
+    """Build kwargs for litellm.completion from api settings."""
+    kwargs: dict = {}
+    if api_base:
+        kwargs["api_base"] = api_base
+    if api_key:
+        kwargs["api_key"] = api_key
+    return kwargs
+
+
+def generate_players(
+    instruction: str,
+    n: int,
+    model: str = "gpt-4o-mini",
+    *,
+    api_base: str | None = None,
+    api_key: str | None = None,
+):
     """Request `n` completions for the instruction using the given model."""
     response = completion(
         model=model,
         messages=[{"role": "user", "content": instruction}],
         n=n,
+        **_completion_kwargs(api_base, api_key),
     )
     return [c.message.content.strip() for c in response.choices]
 
@@ -17,6 +35,9 @@ def prompt_score(
     criteria_block: str,
     player: str,
     model: str = "gpt-4o-mini",
+    *,
+    api_base: str | None = None,
+    api_key: str | None = None,
 ) -> str:
     """Return a JSON score string evaluating `player` on the criteria."""
     example_scores = ", ".join(["1-10"] * len(criteria_list)) or "1-10"
@@ -30,11 +51,24 @@ Instruction:
 
 Output:
 {player}"""
-    response = completion(model=model, messages=[{"role": "system", "content": prompt}])
+    response = completion(
+        model=model,
+        messages=[{"role": "system", "content": prompt}],
+        **_completion_kwargs(api_base, api_key),
+    )
     return response.choices[0].message.content.strip()
 
 
-def prompt_play(instruction: str, criteria_block: str, a: str, b: str, model: str = "gpt-4o-mini") -> str:
+def prompt_pairwise(
+    instruction: str,
+    criteria_block: str,
+    a: str,
+    b: str,
+    model: str = "gpt-4o-mini",
+    *,
+    api_base: str | None = None,
+    api_key: str | None = None,
+) -> str:
     """Return which player wins in JSON using the given criteria."""
     prompt = f"""Compare the two players below using:
 {criteria_block}
@@ -47,5 +81,9 @@ Instruction:
 Players:
 <A>{a}</A>
 <B>{b}</B>"""
-    response = completion(model=model, messages=[{"role": "system", "content": prompt}])
+    response = completion(
+        model=model,
+        messages=[{"role": "system", "content": prompt}],
+        **_completion_kwargs(api_base, api_key),
+    )
     return response.choices[0].message.content.strip()