fix make test

Author: yuan.wang

src/memos/llms/ollama.py

@@ -77,8 +77,9 @@ def generate(self, messages: MessageList, **kwargs) -> Any:
             tools=kwargs.get("tools"),
         )
         logger.info(f"Raw response from Ollama: {response.model_dump_json()}")
-        if response.message.tool_calls:
-            return self.tool_call_parser(response.message.tool_calls)
+        tool_calls = getattr(response.message, "tool_calls", None)
+        if isinstance(tool_calls, list) and len(tool_calls) > 0:
+            return self.tool_call_parser(tool_calls)
 
         str_thinking = (
             f"<think>{response.message.thinking}</think>"

src/memos/llms/openai.py

@@ -41,18 +41,18 @@ def generate(self, messages: MessageList, **kwargs) -> str:
             tools=kwargs.get("tools", NOT_GIVEN),
         )
         logger.info(f"Response from OpenAI: {response.model_dump_json()}")
-        if response.choices[0].message.tool_calls:
-            return self.tool_call_parser(response.choices[0].message.tool_calls)
-        reasoning_content = (
-            f"<think>{response.choices[0].message.reasoning_content}</think>"
-            if hasattr(response.choices[0].message, "reasoning_content")
-            else ""
-        )
+        tool_calls = getattr(response.choices[0].message, "tool_calls", None)
+        if isinstance(tool_calls, list) and len(tool_calls) > 0:
+            return self.tool_call_parser(tool_calls)
         response_content = response.choices[0].message.content
+        reasoning_content = getattr(response.choices[0].message, "reasoning_content", None)
+        if isinstance(reasoning_content, str) and reasoning_content:
+            reasoning_content = f"<think>{reasoning_content}</think>"
         if self.config.remove_think_prefix:
             return remove_thinking_tags(response_content)
-        else:
+        if reasoning_content:
             return reasoning_content + response_content
+        return response_content
 
     @timed(log=True, log_prefix="OpenAI LLM")
     def generate_stream(self, messages: MessageList, **kwargs) -> Generator[str, None, None]:

tests/configs/test_llm.py

@@ -19,7 +19,14 @@ def test_base_llm_config():
         required_fields=[
             "model_name_or_path",
         ],
-        optional_fields=["temperature", "max_tokens", "top_p", "top_k", "remove_think_prefix"],
+        optional_fields=[
+            "temperature",
+            "max_tokens",
+            "top_p",
+            "top_k",
+            "remove_think_prefix",
+            "default_headers",
+        ],
     )
 
     check_config_instantiation_valid(
@@ -48,6 +55,7 @@ def test_openai_llm_config():
             "api_base",
             "remove_think_prefix",
             "extra_body",
+            "default_headers",
         ],
     )
 
@@ -79,6 +87,8 @@ def test_ollama_llm_config():
             "top_k",
             "remove_think_prefix",
             "api_base",
+            "default_headers",
+            "enable_thinking",
         ],
     )
 
@@ -111,6 +121,7 @@ def test_hf_llm_config():
             "do_sample",
             "remove_think_prefix",
             "add_generation_prompt",
+            "default_headers",
         ],
     )
 

tests/llms/test_deepseek.py

@@ -12,12 +12,14 @@ def test_deepseek_llm_generate_with_and_without_think_prefix(self):
         """Test DeepSeekLLM generate method with and without <think> tag removal."""
 
         # Simulated full content including <think> tag
-        full_content = "<think>Thinking in progress...</think>Hello from DeepSeek!"
+        full_content = "Hello from DeepSeek!"
+        reasoning_content = "Thinking in progress..."
 
         # Mock response object
         mock_response = MagicMock()
         mock_response.model_dump_json.return_value = '{"mock": "true"}'
         mock_response.choices[0].message.content = full_content
+        mock_response.choices[0].message.reasoning_content = reasoning_content
 
         # Config with think prefix preserved
         config_with_think = DeepSeekLLMConfig.model_validate(
@@ -35,15 +37,15 @@ def test_deepseek_llm_generate_with_and_without_think_prefix(self):
         llm_with_think.client.chat.completions.create = MagicMock(return_value=mock_response)
 
         output_with_think = llm_with_think.generate([{"role": "user", "content": "Hello"}])
-        self.assertEqual(output_with_think, full_content)
+        self.assertEqual(output_with_think, f"<think>{reasoning_content}</think>{full_content}")
 
         # Config with think tag removed
         config_without_think = config_with_think.model_copy(update={"remove_think_prefix": True})
         llm_without_think = DeepSeekLLM(config_without_think)
         llm_without_think.client.chat.completions.create = MagicMock(return_value=mock_response)
 
         output_without_think = llm_without_think.generate([{"role": "user", "content": "Hello"}])
-        self.assertEqual(output_without_think, "Hello from DeepSeek!")
+        self.assertEqual(output_without_think, full_content)
 
     def test_deepseek_llm_generate_stream(self):
         """Test DeepSeekLLM generate_stream with reasoning_content and content chunks."""
@@ -84,5 +86,5 @@ def make_chunk(delta_dict):
 
         self.assertIn("Analyzing...", full_output)
         self.assertIn("Hello, DeepSeek!", full_output)
-        self.assertTrue(full_output.startswith("Analyzing..."))
+        self.assertTrue(full_output.startswith("<think>"))
         self.assertTrue(full_output.endswith("DeepSeek!"))

tests/llms/test_ollama.py

@@ -1,5 +1,6 @@
 import unittest
 
+from types import SimpleNamespace
 from unittest.mock import MagicMock
 
 from memos.configs.llm import LLMConfigFactory, OllamaLLMConfig
@@ -12,15 +13,15 @@ def test_llm_factory_with_mocked_ollama_backend(self):
         """Test LLMFactory with mocked Ollama backend."""
         mock_chat = MagicMock()
         mock_response = MagicMock()
-        mock_response.model_dump_json.return_value = '{"model":"qwen3:0.6b","created_at":"2025-05-13T18:07:04.508998134Z","done":true,"done_reason":"stop","total_duration":348924420,"load_duration":14321072,"prompt_eval_count":16,"prompt_eval_duration":16770943,"eval_count":21,"eval_duration":317395459,"message":{"role":"assistant","content":"Hello! How are you? I\'m here to help and smile!","images":null,"tool_calls":null}}'
-        mock_response.__getitem__.side_effect = lambda key: {
-            "message": {
-                "role": "assistant",
-                "content": "Hello! How are you? I'm here to help and smile!",
-                "images": None,
-                "tool_calls": None,
-            }
-        }[key]
+        mock_response.model_dump_json.return_value = '{"model":"qwen3:0.6b","created_at":"2025-05-13T18:07:04.508998134Z","done":true,"done_reason":"stop","total_duration":348924420,"load_duration":14321072,"prompt_eval_count":16,"prompt_eval_duration":16770943,"eval_count":21,"eval_duration":317395459,"message":{"role":"assistant","content":"Hello! How are you? I\'m here to help and smile!", "thinking":"Analyzing your request...","images":null,"tool_calls":null}}'
+
+        mock_response.message = SimpleNamespace(
+            role="assistant",
+            content="Hello! How are you? I'm here to help and smile!",
+            thinking="Analyzing your request...",
+            images=None,
+            tool_calls=None,
+        )
         mock_chat.return_value = mock_response
 
         config = LLMConfigFactory.model_validate(
@@ -32,6 +33,7 @@ def test_llm_factory_with_mocked_ollama_backend(self):
                     "max_tokens": 1024,
                     "top_p": 0.9,
                     "top_k": 50,
+                    "enable_thinking": True,
                 },
             }
         )
@@ -42,21 +44,23 @@ def test_llm_factory_with_mocked_ollama_backend(self):
         ]
         response = llm.generate(messages)
 
-        self.assertEqual(response, "Hello! How are you? I'm here to help and smile!")
+        self.assertEqual(
+            response,
+            "<think>Analyzing your request...</think>Hello! How are you? I'm here to help and smile!",
+        )
 
     def test_ollama_llm_with_mocked_backend(self):
         """Test OllamaLLM with mocked backend."""
         mock_chat = MagicMock()
         mock_response = MagicMock()
-        mock_response.model_dump_json.return_value = '{"model":"qwen3:0.6b","created_at":"2025-05-13T18:07:04.508998134Z","done":true,"done_reason":"stop","total_duration":348924420,"load_duration":14321072,"prompt_eval_count":16,"prompt_eval_duration":16770943,"eval_count":21,"eval_duration":317395459,"message":{"role":"assistant","content":"Hello! How are you? I\'m here to help and smile!","images":null,"tool_calls":null}}'
-        mock_response.__getitem__.side_effect = lambda key: {
-            "message": {
-                "role": "assistant",
-                "content": "Hello! How are you? I'm here to help and smile!",
-                "images": None,
-                "tool_calls": None,
-            }
-        }[key]
+        mock_response.model_dump_json.return_value = '{"model":"qwen3:0.6b","created_at":"2025-05-13T18:07:04.508998134Z","done":true,"done_reason":"stop","total_duration":348924420,"load_duration":14321072,"prompt_eval_count":16,"prompt_eval_duration":16770943,"eval_count":21,"eval_duration":317395459,"message":{"role":"assistant","content":"Hello! How are you? I\'m here to help and smile!","thinking":"Analyzing your request...","images":null,"tool_calls":null}}'
+        mock_response.message = SimpleNamespace(
+            role="assistant",
+            content="Hello! How are you? I'm here to help and smile!",
+            thinking="Analyzing your request...",
+            images=None,
+            tool_calls=None,
+        )
         mock_chat.return_value = mock_response
 
         config = OllamaLLMConfig(
@@ -73,4 +77,7 @@ def test_ollama_llm_with_mocked_backend(self):
         ]
         response = ollama.generate(messages)
 
-        self.assertEqual(response, "Hello! How are you? I'm here to help and smile!")
+        self.assertEqual(
+            response,
+            "<think>Analyzing your request...</think>Hello! How are you? I'm here to help and smile!",
+        )

tests/llms/test_openai.py

@@ -14,6 +14,7 @@ def test_llm_factory_with_mocked_openai_backend(self):
         mock_response = MagicMock()
         mock_response.model_dump_json.return_value = '{"id":"chatcmpl-BWoqIrvOeWdnFVZQUFzCcdVEpJ166","choices":[{"finish_reason":"stop","index":0,"message":{"content":"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?","role":"assistant"}}],"created":1747161634,"model":"gpt-4o-2024-08-06","object":"chat.completion"}'
         mock_response.choices[0].message.content = "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?"  # fmt: skip
+        mock_response.choices[0].message.reasoning_content = None
         mock_chat_completions_create.return_value = mock_response
 
         config = LLMConfigFactory.model_validate(

tests/llms/test_qwen.py

@@ -12,12 +12,14 @@ def test_qwen_llm_generate_with_and_without_think_prefix(self):
         """Test QwenLLM non-streaming response generation with and without <think> prefix removal."""
 
         # Simulated full response content with <think> tag
-        full_content = "<think>Analyzing your request...</think>Hello, world!"
+        full_content = "Hello from DeepSeek!"
+        reasoning_content = "Thinking in progress..."
 
         # Prepare the mock response object with expected structure
         mock_response = MagicMock()
         mock_response.model_dump_json.return_value = '{"mocked": "true"}'
         mock_response.choices[0].message.content = full_content
+        mock_response.choices[0].message.reasoning_content = reasoning_content
 
         # Create config with remove_think_prefix = False
         config_with_think = QwenLLMConfig.model_validate(
@@ -37,7 +39,7 @@ def test_qwen_llm_generate_with_and_without_think_prefix(self):
         llm_with_think.client.chat.completions.create = MagicMock(return_value=mock_response)
 
         response_with_think = llm_with_think.generate([{"role": "user", "content": "Hi"}])
-        self.assertEqual(response_with_think, full_content)
+        self.assertEqual(response_with_think, f"<think>{reasoning_content}</think>{full_content}")
 
         # Create config with remove_think_prefix = True
         config_without_think = config_with_think.model_copy(update={"remove_think_prefix": True})
@@ -47,7 +49,7 @@ def test_qwen_llm_generate_with_and_without_think_prefix(self):
         llm_without_think.client.chat.completions.create = MagicMock(return_value=mock_response)
 
         response_without_think = llm_without_think.generate([{"role": "user", "content": "Hi"}])
-        self.assertEqual(response_without_think, "Hello, world!")
+        self.assertEqual(response_without_think, full_content)
         self.assertNotIn("<think>", response_without_think)
 
     def test_qwen_llm_generate_stream(self):