import unittest
from deepsearcher.llm.base import BaseLLM, ChatResponse
from unittest.mock import patch
class TestBaseLLM(unittest.TestCase):
"""Tests for the BaseLLM abstract base class."""
def setUp(self):
"""Set up test fixtures."""
# Clear environment variables temporarily
self.env_patcher = patch.dict('os.environ', {}, clear=True)
self.env_patcher.start()
def tearDown(self):
"""Clean up test fixtures."""
self.env_patcher.stop()
def test_chat_response_init(self):
"""Test ChatResponse initialization and representation."""
content = "Test content"
total_tokens = 100
response = ChatResponse(content=content, total_tokens=total_tokens)
self.assertEqual(response.content, content)
self.assertEqual(response.total_tokens, total_tokens)
self.assertEqual(
repr(response),
f"ChatResponse(content={content}, total_tokens={total_tokens})"
)
def test_literal_eval_python_code_block(self):
"""Test literal_eval with Python code block."""
content = '''```python
{"key": "value", "number": 42}
```'''
result = BaseLLM.literal_eval(content)
self.assertEqual(result, {"key": "value", "number": 42})
def test_literal_eval_json_code_block(self):
"""Test literal_eval with JSON code block."""
content = '''```json
{"key": "value", "number": 42}
```'''
result = BaseLLM.literal_eval(content)
self.assertEqual(result, {"key": "value", "number": 42})
def test_literal_eval_str_code_block(self):
"""Test literal_eval with str code block."""
content = '''```str
{"key": "value", "number": 42}
```'''
result = BaseLLM.literal_eval(content)
self.assertEqual(result, {"key": "value", "number": 42})
def test_literal_eval_plain_code_block(self):
"""Test literal_eval with plain code block."""
content = '''```
{"key": "value", "number": 42}
```'''
result = BaseLLM.literal_eval(content)
self.assertEqual(result, {"key": "value", "number": 42})
def test_literal_eval_raw_dict(self):
"""Test literal_eval with raw dictionary string."""
content = '{"key": "value", "number": 42}'
result = BaseLLM.literal_eval(content)
self.assertEqual(result, {"key": "value", "number": 42})
def test_literal_eval_raw_list(self):
"""Test literal_eval with raw list string."""
content = '[1, 2, "three", {"four": 4}]'
result = BaseLLM.literal_eval(content)
self.assertEqual(result, [1, 2, "three", {"four": 4}])
def test_literal_eval_with_whitespace(self):
"""Test literal_eval with extra whitespace."""
content = '''
{"key": "value"}
'''
result = BaseLLM.literal_eval(content)
self.assertEqual(result, {"key": "value"})
def test_literal_eval_nested_structures(self):
"""Test literal_eval with nested data structures."""
content = '''
{
"string": "value",
"number": 42,
"list": [1, 2, 3],
"dict": {"nested": "value"},
"mixed": [1, {"key": "value"}, [2, 3]]
}
'''
result = BaseLLM.literal_eval(content)
expected = {
"string": "value",
"number": 42,
"list": [1, 2, 3],
"dict": {"nested": "value"},
"mixed": [1, {"key": "value"}, [2, 3]]
}
self.assertEqual(result, expected)
def test_literal_eval_invalid_format(self):
"""Test literal_eval with invalid format."""
invalid_contents = [
"Not a valid Python literal",
"{invalid: json}",
"[1, 2, 3", # Unclosed bracket
'{"key": undefined}', # undefined is not a valid Python literal
]
for content in invalid_contents:
with self.assertRaises(ValueError):
BaseLLM.literal_eval(content)
def test_remove_think_with_tags(self):
"""Test remove_think with think tags."""
content = '''
This is the reasoning process.
Multiple lines of thought.
This is the actual response.'''
result = BaseLLM.remove_think(content)
self.assertEqual(result.strip(), "This is the actual response.")
def test_remove_think_without_tags(self):
"""Test remove_think without think tags."""
content = "This is a response without think tags."
result = BaseLLM.remove_think(content)
self.assertEqual(result.strip(), content.strip())
def test_remove_think_multiple_tags(self):
"""Test remove_think with multiple think tags - should only remove first block."""
content = '''First think block
Actual response
Second think block'''
result = BaseLLM.remove_think(content)
self.assertEqual(
result.strip(),
"Actual response\n Second think block"
)
def test_remove_think_empty_tags(self):
"""Test remove_think with empty think tags."""
content = "Response"
result = BaseLLM.remove_think(content)
self.assertEqual(result.strip(), "Response")
if __name__ == "__main__":
unittest.main()