diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index baf742b5..1aae7005 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -59,6 +59,7 @@ jobs: python -m pytest tests/test_bon_candidate_filtering.py -v --tb=short python -m pytest tests/test_compact_plugin.py -v --tb=short python -m pytest tests/test_completions_endpoint.py -v --tb=short + python -m pytest tests/test_developer_role.py -v --tb=short python -m pytest tests/test_entropy_decoding_sampling.py -v --tb=short python -m pytest tests/test_mars_agent_routing.py -v --tb=short python -m pytest tests/test_mars_aggregation_sampling.py -v --tb=short diff --git a/optillm/server.py b/optillm/server.py index 6bc7b4e1..d25d65bc 100644 --- a/optillm/server.py +++ b/optillm/server.py @@ -677,7 +677,8 @@ def parse_conversation(messages): if role == 'assistant' and not text_content: continue - if role == 'system': + # 'developer' is the OpenAI replacement for 'system' on o-series and newer models + if role in ('system', 'developer'): system_prompt, optillm_approach = extract_optillm_approach(text_content) elif role == 'user': if not optillm_approach: diff --git a/tests/test_developer_role.py b/tests/test_developer_role.py new file mode 100644 index 00000000..2703a267 --- /dev/null +++ b/tests/test_developer_role.py @@ -0,0 +1,106 @@ +""" +Offline unit tests for `developer` role messages. + +OpenAI clients send `developer` messages in place of `system` messages for +o-series and newer models. optillm used to drop them when it split a request +into system prompt and conversation, so the instructions never reached the +approach. Uses the Flask test client with a mocked approach (no network, no model). +""" + +import sys +import os +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import unittest + +from optillm import server + + +CAPTURED = {} + + +def fake_chat_with_mcts(system_prompt, initial_query, client, model, + num_simulations, exploration_weight, depth, + request_config=None, request_id=None): + CAPTURED['system_prompt'] = system_prompt + CAPTURED['initial_query'] = initial_query + return "ok", 1 + + +class TestParseConversationDeveloperRole(unittest.TestCase): + + def test_developer_message_becomes_system_prompt(self): + messages = [ + {"role": "developer", "content": "Answer in French."}, + {"role": "user", "content": "Hello"}, + ] + system_prompt, initial_query, approach = server.parse_conversation(messages) + self.assertEqual(system_prompt, "Answer in French.") + self.assertEqual(initial_query, "User: Hello") + self.assertIsNone(approach) + + def test_developer_message_with_content_parts(self): + messages = [ + {"role": "developer", "content": [{"type": "text", "text": "Answer in French."}]}, + {"role": "user", "content": "Hello"}, + ] + system_prompt, _, _ = server.parse_conversation(messages) + self.assertEqual(system_prompt, "Answer in French.") + + def test_developer_message_can_select_approach(self): + messages = [ + {"role": "developer", "content": "re2 Be brief."}, + {"role": "user", "content": "Hello"}, + ] + system_prompt, _, approach = server.parse_conversation(messages) + self.assertEqual(system_prompt, "Be brief.") + self.assertEqual(approach, "re2") + + def test_system_message_unchanged(self): + messages = [ + {"role": "system", "content": "Answer in French."}, + {"role": "user", "content": "Hello"}, + ] + system_prompt, initial_query, _ = server.parse_conversation(messages) + self.assertEqual(system_prompt, "Answer in French.") + self.assertEqual(initial_query, "User: Hello") + + +class TestDeveloperRoleReachesApproach(unittest.TestCase): + + def setUp(self): + CAPTURED.clear() + self._originals = { + 'get_config': server.get_config, + 'chat_with_mcts': server.chat_with_mcts, + 'config': dict(server.server_config), + } + server.get_config = lambda: (object(), 'optillm') + server.chat_with_mcts = fake_chat_with_mcts + server.server_config['optillm_api_key'] = '' + server.app.config['TESTING'] = True + self.client = server.app.test_client() + + def tearDown(self): + server.get_config = self._originals['get_config'] + server.chat_with_mcts = self._originals['chat_with_mcts'] + server.server_config.clear() + server.server_config.update(self._originals['config']) + + def test_developer_message_is_passed_to_approach(self): + resp = self.client.post( + '/v1/chat/completions', + json={"model": "gpt-4o-mini", + "messages": [ + {"role": "developer", "content": "Answer in French."}, + {"role": "user", "content": "Hello"}, + ], + "optillm_approach": "mcts"}) + + self.assertEqual(resp.status_code, 200, resp.get_data(as_text=True)) + self.assertEqual(CAPTURED['system_prompt'], "Answer in French.") + self.assertEqual(CAPTURED['initial_query'], "User: Hello") + + +if __name__ == '__main__': + unittest.main()