diff --git a/llm_cooperation/experiments/dilemma.py b/llm_cooperation/experiments/dilemma.py index 58d996a..bfa60b8 100755 --- a/llm_cooperation/experiments/dilemma.py +++ b/llm_cooperation/experiments/dilemma.py @@ -140,9 +140,9 @@ def strategy_t4t( def move_as_str(move: DilemmaEnum) -> str: if move == DilemmaEnum.D: - return "Project {COLOR_DEFECT}" + return f"Project {COLOR_DEFECT}" elif move == DilemmaEnum.C: - return "Project {COLOR_COOPERATE}" + return f"Project {COLOR_COOPERATE}" raise ValueError(f"Invalid choice {move}") diff --git a/tests/test_dilemma.py b/tests/test_dilemma.py index 5e612a9..f296b28 100644 --- a/tests/test_dilemma.py +++ b/tests/test_dilemma.py @@ -2,7 +2,7 @@ from unittest.mock import Mock import pytest -from openai_pygenerator import Completion +from openai_pygenerator import Completion, logger from llm_cooperation import DEFAULT_MODEL_SETUP, Payoffs from llm_cooperation.experiments.dilemma import ( @@ -18,6 +18,7 @@ T, extract_choice_pd, get_prompt_pd, + move_as_str, payoffs_pd, strategy_cooperate, strategy_defect, @@ -33,8 +34,9 @@ def test_get_instruction_prompt(condition: bool): role_prompt = "You are a helpful assistant." prompt = get_prompt_pd(condition, role_prompt) - assert COLOR_COOPERATE in prompt - assert COLOR_DEFECT in prompt + logger.debug("prompt = %s", prompt) + assert "COLOR_COOPERATE" not in prompt + assert "COLOR_DEFECT" not in prompt assert ("Explanation:" in prompt) == condition assert "Choice:" in prompt for payoff in [R, S, T, P]: @@ -110,6 +112,11 @@ def test_dilemma_choice(): assert Cooperate != Defect +def test_move_as_str(): + assert COLOR_COOPERATE in move_as_str(DilemmaEnum.C) + assert COLOR_DEFECT in move_as_str(DilemmaEnum.D) + + def test_run_repeated_game(mocker): completions = [ {"role": "assistant", "content": "project green"},