Skip to content

Commit e268b82

Browse files
CaralHsiCaralHsi
committed
test: add unit test for llms
1 parent 18f3bfc commit e268b82

File tree

3 files changed

+249
-0
lines changed

3 files changed

+249
-0
lines changed

‎tests/llms/test_deepseek.py‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
import unittest
2+
3+
from types import SimpleNamespace
4+
from unittest.mock import MagicMock
5+
6+
from memos.configs.llm import DeepSeekLLMConfig
7+
from memos.llms.deepseek import DeepSeekLLM
8+
9+
10+
class TestDeepSeekLLM(unittest.TestCase):
11+
def test_deepseek_llm_generate_with_and_without_think_prefix(self):
12+
"""Test DeepSeekLLM generate method with and without <think> tag removal."""
13+
14+
# Simulated full content including <think> tag
15+
full_content = "<think>Thinking in progress...</think>Hello from DeepSeek!"
16+
17+
# Mock response object
18+
mock_response = MagicMock()
19+
mock_response.model_dump_json.return_value = '{"mock": "true"}'
20+
mock_response.choices[0].message.content = full_content
21+
22+
# Config with think prefix preserved
23+
config_with_think = DeepSeekLLMConfig.model_validate(
24+
{
25+
"model_name_or_path": "deepseek-chat",
26+
"temperature": 0.7,
27+
"max_tokens": 512,
28+
"top_p": 0.9,
29+
"api_key": "sk-test",
30+
"api_base": "https://api.deepseek.com/v1",
31+
"remove_think_prefix": False,
32+
}
33+
)
34+
llm_with_think = DeepSeekLLM(config_with_think)
35+
llm_with_think.client.chat.completions.create = MagicMock(return_value=mock_response)
36+
37+
output_with_think = llm_with_think.generate([{"role": "user", "content": "Hello"}])
38+
self.assertEqual(output_with_think, full_content)
39+
40+
# Config with think tag removed
41+
config_without_think = config_with_think.copy(update={"remove_think_prefix": True})
42+
llm_without_think = DeepSeekLLM(config_without_think)
43+
llm_without_think.client.chat.completions.create = MagicMock(return_value=mock_response)
44+
45+
output_without_think = llm_without_think.generate([{"role": "user", "content": "Hello"}])
46+
self.assertEqual(output_without_think, "Hello from DeepSeek!")
47+
48+
def test_deepseek_llm_generate_stream(self):
49+
"""Test DeepSeekLLM generate_stream with reasoning_content and content chunks."""
50+
51+
def make_chunk(delta_dict):
52+
# Create a simulated stream chunk with delta fields
53+
delta = SimpleNamespace(**delta_dict)
54+
choice = SimpleNamespace(delta=delta)
55+
return SimpleNamespace(choices=[choice])
56+
57+
# Simulate chunks: reasoning + answer
58+
mock_stream_chunks = [
59+
make_chunk({"reasoning_content": "Analyzing..."}),
60+
make_chunk({"content": "Hello"}),
61+
make_chunk({"content": ", "}),
62+
make_chunk({"content": "DeepSeek!"}),
63+
]
64+
65+
mock_chat_completions_create = MagicMock(return_value=iter(mock_stream_chunks))
66+
67+
config = DeepSeekLLMConfig.model_validate(
68+
{
69+
"model_name_or_path": "deepseek-chat",
70+
"temperature": 0.7,
71+
"max_tokens": 512,
72+
"top_p": 0.9,
73+
"api_key": "sk-test",
74+
"api_base": "https://api.deepseek.com/v1",
75+
"remove_think_prefix": False,
76+
}
77+
)
78+
llm = DeepSeekLLM(config)
79+
llm.client.chat.completions.create = mock_chat_completions_create
80+
81+
messages = [{"role": "user", "content": "Say hello"}]
82+
streamed = list(llm.generate_stream(messages))
83+
full_output = "".join(streamed)
84+
85+
self.assertIn("Analyzing...", full_output)
86+
self.assertIn("Hello, DeepSeek!", full_output)
87+
self.assertTrue(full_output.startswith("Analyzing..."))
88+
self.assertTrue(full_output.endswith("DeepSeek!"))

‎tests/llms/test_openai.py‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import unittest
22

3+
from types import SimpleNamespace
34
from unittest.mock import MagicMock
45

56
from memos.configs.llm import LLMConfigFactory
@@ -40,3 +41,62 @@ def test_llm_factory_with_mocked_openai_backend(self):
4041
response,
4142
"Hello! I'm an AI language model created by OpenAI. I'm here to help answer questions, provide information, and assist with a wide range of topics. How can I assist you today?",
4243
)
44+
45+
def test_llm_factory_with_stream_openai_backend(self):
46+
"""Test LLMFactory stream generation with mocked OpenAI backend."""
47+
48+
def make_chunk(delta_dict):
49+
# Create a mock response chunk with a simulated delta dictionary
50+
delta = SimpleNamespace(**delta_dict)
51+
choice = SimpleNamespace(delta=delta, finish_reason="stop", index=0)
52+
return SimpleNamespace(choices=[choice])
53+
54+
# Simulate a stream response with both reasoning_content and content
55+
mock_stream_chunks = [
56+
make_chunk({"reasoning_content": "I am thinking"}),
57+
make_chunk({"content": "Hello"}),
58+
make_chunk({"content": ", "}),
59+
make_chunk({"content": "world!"}),
60+
]
61+
62+
# Mock the streaming chat completion call
63+
mock_chat_completions_create = MagicMock(return_value=iter(mock_stream_chunks))
64+
65+
# Create the LLM config with think prefix enabled
66+
config = LLMConfigFactory.model_validate(
67+
{
68+
"backend": "openai",
69+
"config": {
70+
"model_name_or_path": "gpt-4.1-nano",
71+
"temperature": 0.8,
72+
"max_tokens": 1024,
73+
"top_p": 0.9,
74+
"top_k": 50,
75+
"api_key": "sk-xxxx",
76+
"api_base": "https://api.openai.com/v1",
77+
"remove_think_prefix": False,
78+
# Ensure <think> tag is emitted
79+
},
80+
}
81+
)
82+
83+
# Instantiate the LLM and inject the mocked stream method
84+
llm = LLMFactory.from_config(config)
85+
llm.client.chat.completions.create = mock_chat_completions_create
86+
87+
# Input message to the model
88+
messages = [{"role": "user", "content": "Think and say hello"}]
89+
90+
# Collect streamed output as a list of chunks
91+
response_parts = list(llm.generate_stream(messages))
92+
response = "".join(response_parts)
93+
94+
# Assert the presence of the <think> tag and expected content
95+
self.assertIn("<think>", response)
96+
self.assertIn("I am thinking", response)
97+
self.assertIn("Hello, world!", response)
98+
99+
# Optional: check structure of stream response
100+
self.assertEqual(response_parts[0], "<think>")
101+
self.assertTrue(response.startswith("<think>I am thinking"))
102+
self.assertTrue(response.endswith("Hello, world!"))

‎tests/llms/test_qwen.py‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
import unittest
2+
3+
from types import SimpleNamespace
4+
from unittest.mock import MagicMock
5+
6+
from memos.configs.llm import QwenLLMConfig
7+
from memos.llms.qwen import QwenLLM
8+
9+
10+
class TestQwenLLM(unittest.TestCase):
11+
def test_qwen_llm_generate_with_and_without_think_prefix(self):
12+
"""Test QwenLLM non-streaming response generation with and without <think> prefix removal."""
13+
14+
# Simulated full response content with <think> tag
15+
full_content = "<think>Analyzing your request...</think>Hello, world!"
16+
17+
# Prepare the mock response object with expected structure
18+
mock_response = MagicMock()
19+
mock_response.model_dump_json.return_value = '{"mocked": "true"}'
20+
mock_response.choices[0].message.content = full_content
21+
22+
# Create config with remove_think_prefix = False
23+
config_with_think = QwenLLMConfig.model_validate(
24+
{
25+
"model_name_or_path": "qwen-test",
26+
"temperature": 0.7,
27+
"max_tokens": 100,
28+
"top_p": 0.9,
29+
"api_key": "sk-test",
30+
"api_base": "https://dashscope.aliyuncs.com/api/v1",
31+
"remove_think_prefix": False,
32+
}
33+
)
34+
35+
# Instance with think tag enabled
36+
llm_with_think = QwenLLM(config_with_think)
37+
llm_with_think.client.chat.completions.create = MagicMock(return_value=mock_response)
38+
39+
response_with_think = llm_with_think.generate([{"role": "user", "content": "Hi"}])
40+
self.assertEqual(response_with_think, full_content)
41+
42+
# Create config with remove_think_prefix = True
43+
config_without_think = config_with_think.copy(update={"remove_think_prefix": True})
44+
45+
# Instance with think tag removed
46+
llm_without_think = QwenLLM(config_without_think)
47+
llm_without_think.client.chat.completions.create = MagicMock(return_value=mock_response)
48+
49+
response_without_think = llm_without_think.generate([{"role": "user", "content": "Hi"}])
50+
self.assertEqual(response_without_think, "Hello, world!")
51+
self.assertNotIn("<think>", response_without_think)
52+
53+
def test_qwen_llm_generate_stream(self):
54+
"""Test QwenLLM stream generation with both reasoning_content and content."""
55+
56+
def make_chunk(delta_dict):
57+
# Construct a mock chunk with delta fields
58+
delta = SimpleNamespace(**delta_dict)
59+
choice = SimpleNamespace(delta=delta)
60+
return SimpleNamespace(choices=[choice])
61+
62+
# Simulate a sequence of streamed chunks
63+
mock_stream_chunks = [
64+
make_chunk({"reasoning_content": "Analyzing input..."}),
65+
make_chunk({"content": "Hello"}),
66+
make_chunk({"content": ", "}),
67+
make_chunk({"content": "world!"}),
68+
]
69+
70+
# Mock the client's streaming response
71+
mock_chat_completions_create = MagicMock(return_value=iter(mock_stream_chunks))
72+
73+
# Build QwenLLM config with think prefix enabled
74+
config = QwenLLMConfig.model_validate(
75+
{
76+
"model_name_or_path": "qwen-test",
77+
"temperature": 0.7,
78+
"max_tokens": 100,
79+
"top_p": 0.9,
80+
"api_key": "sk-test",
81+
"api_base": "https://dashscope.aliyuncs.com/api/v1",
82+
"remove_think_prefix": False,
83+
}
84+
)
85+
86+
# Create QwenLLM instance and inject mock client
87+
llm = QwenLLM(config)
88+
llm.client.chat.completions.create = mock_chat_completions_create
89+
90+
messages = [{"role": "user", "content": "Say hello"}]
91+
92+
# Collect the streamed output
93+
response_parts = list(llm.generate_stream(messages))
94+
response = "".join(response_parts)
95+
96+
# Assertions for structure and content
97+
self.assertIn("<think>", response)
98+
self.assertIn("Analyzing input...", response)
99+
self.assertIn("Hello, world!", response)
100+
self.assertTrue(response.startswith("<think>Analyzing input..."))
101+
self.assertTrue(response.endswith("Hello, world!"))

0 commit comments

Comments
 (0)