Spaces:
Sleeping
Sleeping
File size: 4,441 Bytes
058ae1e | 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 | 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()
|