Spaces:
Sleeping
Sleeping
| import unittest | |
| from unittest.mock import MagicMock, patch | |
| import sys | |
| import os | |
| import json | |
| # Add project root to path | |
| sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) | |
| # Mock Transformers | |
| with patch('transformers.AutoProcessor.from_pretrained'), \ | |
| patch('transformers.AutoModelForCausalLM.from_pretrained'), \ | |
| patch('transformers.AutoTokenizer.from_pretrained'): | |
| import agent_module | |
| import oracle_module | |
| class TestAgentScenarios(unittest.TestCase): | |
| def setUp(self): | |
| self.mock_llm_patcher = patch('agent_module.get_llm') | |
| self.mock_llm = self.mock_llm_patcher.start() | |
| self.mock_model = MagicMock() | |
| self.mock_model.device = "cpu" | |
| self.mock_processor = MagicMock() | |
| self.mock_llm.return_value = (self.mock_model, self.mock_processor) | |
| self.mock_oracle_patcher = patch('agent_module.get_oracle_data') | |
| self.mock_oracle = self.mock_oracle_patcher.start() | |
| self.mock_oracle.return_value = {"wisdom": "Mock Wisdom"} | |
| self.mock_streamer_patcher = patch('agent_module.TextIteratorStreamer') | |
| self.mock_streamer = self.mock_streamer_patcher.start() | |
| def tearDown(self): | |
| self.mock_llm_patcher.stop() | |
| self.mock_oracle_patcher.stop() | |
| self.mock_streamer_patcher.stop() | |
| def run_chat(self, user_input, llm_output_stream): | |
| """Helper to run chat with mocked LLM stream""" | |
| self.mock_streamer.return_value = iter(llm_output_stream) | |
| gen = agent_module.chat_agent_stream(user_input, []) | |
| return list(gen) | |
| def test_name_only(self): | |
| """Scenario: User says 'I am Julian'""" | |
| tool_call = { | |
| "name": "oracle_consultation", | |
| "arguments": {"topic": "General", "name": "Julian", "date_str": "today"} | |
| } | |
| # LLM output simulation | |
| llm_output = [ | |
| "Greetings Julian. ", | |
| f"<tool_call>{json.dumps(tool_call)}</tool_call>" | |
| ] | |
| results = self.run_chat("I am Julian", llm_output) | |
| # Verify Oracle called correctly | |
| self.mock_oracle.assert_called_with(name="Julian", topic="General", date_str="today") | |
| self.assertTrue("*(Consulting the Oracle...)*" in results) | |
| def test_topic_only(self): | |
| """Scenario: User says 'Thema: Liebe'""" | |
| tool_call = { | |
| "name": "oracle_consultation", | |
| "arguments": {"topic": "Liebe", "name": "Seeker", "date_str": "today"} | |
| } | |
| llm_output = [ | |
| "Let us check the Oracle for Love. ", | |
| f"<tool_call>{json.dumps(tool_call)}</tool_call>" | |
| ] | |
| results = self.run_chat("Thema: Liebe", llm_output) | |
| self.mock_oracle.assert_called_with(name="Seeker", topic="Liebe", date_str="today") | |
| self.assertTrue("*(Consulting the Oracle...)*" in results) | |
| def test_date_only(self): | |
| """Scenario: User says 'Reading for 2025-01-01'""" | |
| tool_call = { | |
| "name": "oracle_consultation", | |
| "arguments": {"topic": "General", "name": "Seeker", "date_str": "2025-01-01"} | |
| } | |
| llm_output = [ | |
| "Consulting for that date. ", | |
| f"<tool_call>{json.dumps(tool_call)}</tool_call>" | |
| ] | |
| results = self.run_chat("Reading for 2025-01-01", llm_output) | |
| self.mock_oracle.assert_called_with(name="Seeker", topic="General", date_str="2025-01-01") | |
| def test_name_and_date(self): | |
| """Scenario: User says 'I am Julian, date 2025-01-01'""" | |
| tool_call = { | |
| "name": "oracle_consultation", | |
| "arguments": {"topic": "General", "name": "Julian", "date_str": "2025-01-01"} | |
| } | |
| llm_output = [ | |
| f"<tool_call>{json.dumps(tool_call)}</tool_call>" | |
| ] | |
| results = self.run_chat("I am Julian, date 2025-01-01", llm_output) | |
| self.mock_oracle.assert_called_with(name="Julian", topic="General", date_str="2025-01-01") | |
| def test_general_conversation(self): | |
| """Scenario: User says 'Hi' (No tool)""" | |
| llm_output = ["Hello there! How can I help you?"] | |
| results = self.run_chat("Hi", llm_output) | |
| self.mock_oracle.assert_not_called() | |
| self.assertIn("Hello there! How can I help you?", results) | |
| if __name__ == '__main__': | |
| unittest.main() | |