openevolve/tests/test_llm.py at d492df20495a92347a66b6a8eaa903823df89676 · algorithmicsuperintelligence/openevolve · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
"""
Tests for LLM implementations
"""

import asyncio
import unittest
from unittest.mock import AsyncMock, MagicMock, patch

from openevolve.config import LLMConfig
from openevolve.llm.anthropic import AnthropicLLM
from openevolve.llm.openai import OpenAILLM


class TestLLMImplementations(unittest.TestCase):
    """Tests for LLM implementations"""

    def setUp(self):
        """Set up test configuration"""
        self.config = LLMConfig(
            primary_model="test-model",
            api_key="test-key",
            api_base="https://test.api",
        )

    @patch("anthropic.Anthropic")
    async def test_anthropic_llm_generate(self, mock_anthropic):
        """Test Anthropic LLM generate method"""
        # Set up mock response
        mock_response = MagicMock()
        mock_response.content = [MagicMock(text="Test response")]
        mock_anthropic.return_value.messages.create.return_value = mock_response

        # Create LLM instance
        llm = AnthropicLLM(self.config)

        # Test generate
        response = await llm.generate("Test prompt")
        self.assertEqual(response, "Test response")

        # Verify API call
        mock_anthropic.return_value.messages.create.assert_called_once()
        call_args = mock_anthropic.return_value.messages.create.call_args[1]
        self.assertEqual(call_args["model"], "test-model")
        self.assertEqual(call_args["messages"][0]["role"], "user")
        self.assertEqual(call_args["messages"][0]["content"], "Test prompt")

    @patch("anthropic.Anthropic")
    async def test_anthropic_llm_generate_with_context(self, mock_anthropic):
        """Test Anthropic LLM generate_with_context method"""
        # Set up mock response
        mock_response = MagicMock()
        mock_response.content = [MagicMock(text="Test response")]
        mock_anthropic.return_value.messages.create.return_value = mock_response

        # Create LLM instance
        llm = AnthropicLLM(self.config)

        # Test generate_with_context
        messages = [
            {"role": "user", "content": "Test message 1"},
            {"role": "assistant", "content": "Test response 1"},
            {"role": "user", "content": "Test message 2"},
        ]
        response = await llm.generate_with_context("Test system", messages)
        self.assertEqual(response, "Test response")

        # Verify API call
        mock_anthropic.return_value.messages.create.assert_called_once()
        call_args = mock_anthropic.return_value.messages.create.call_args[1]
        self.assertEqual(call_args["model"], "test-model")
        self.assertEqual(call_args["system"], "Test system")
        self.assertEqual(len(call_args["messages"]), 3)
        self.assertEqual(call_args["messages"][0]["role"], "user")
        self.assertEqual(call_args["messages"][0]["content"], "Test message 1")

    @patch("openai.OpenAI")
    async def test_openai_llm_generate(self, mock_openai):
        """Test OpenAI LLM generate method"""
        # Set up mock response
        mock_response = MagicMock()
        mock_response.choices = [MagicMock(message=MagicMock(content="Test response"))]
        mock_openai.return_value.chat.completions.create.return_value = mock_response

        # Create LLM instance
        llm = OpenAILLM(self.config)

        # Test generate
        response = await llm.generate("Test prompt")
        self.assertEqual(response, "Test response")

        # Verify API call
        mock_openai.return_value.chat.completions.create.assert_called_once()
        call_args = mock_openai.return_value.chat.completions.create.call_args[1]
        self.assertEqual(call_args["model"], "test-model")
        self.assertEqual(call_args["messages"][0]["role"], "user")
        self.assertEqual(call_args["messages"][0]["content"], "Test prompt")

    @patch("openai.OpenAI")
    async def test_openai_llm_generate_with_context(self, mock_openai):
        """Test OpenAI LLM generate_with_context method"""
        # Set up mock response
        mock_response = MagicMock()
        mock_response.choices = [MagicMock(message=MagicMock(content="Test response"))]
        mock_openai.return_value.chat.completions.create.return_value = mock_response

        # Create LLM instance
        llm = OpenAILLM(self.config)

        # Test generate_with_context
        messages = [
            {"role": "user", "content": "Test message 1"},
            {"role": "assistant", "content": "Test response 1"},
            {"role": "user", "content": "Test message 2"},
        ]
        response = await llm.generate_with_context("Test system", messages)
        self.assertEqual(response, "Test response")

        # Verify API call
        mock_openai.return_value.chat.completions.create.assert_called_once()
        call_args = mock_openai.return_value.chat.completions.create.call_args[1]
        self.assertEqual(call_args["model"], "test-model")
        self.assertEqual(call_args["messages"][0]["role"], "system")
        self.assertEqual(call_args["messages"][0]["content"], "Test system")
        self.assertEqual(len(call_args["messages"]), 4)  # system + 3 messages

    def test_llm_config_model_detection(self):
        """Test LLM configuration model type detection"""
        # Test OpenAI model
        config = LLMConfig(primary_model="gpt-4")
        self.assertEqual(config.api_base, "https://api.openai.com/v1")

        # Test Claude model
        config = LLMConfig(primary_model="claude-3-sonnet")
        self.assertEqual(config.api_base, "https://api.anthropic.com/v1")

        # Test custom API base
        config = LLMConfig(
            primary_model="claude-3-sonnet",
            api_base="https://custom.api",
        )
        self.assertEqual(config.api_base, "https://custom.api")


if __name__ == "__main__":
    unittest.main()