Spaces:
Runtime error
Runtime error
File size: 5,241 Bytes
4962437 |
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 |
import pytest
from unittest.mock import Mock, patch
from swarms.agents.agents import AgentNodeInitializer, AgentNode, agent # replace with actual import
# For initializing AgentNodeInitializer in multiple tests
@pytest.fixture
def mock_agent_node_initializer():
with patch('swarms.agents.agents.ChatOpenAI') as mock_llm, \
patch('swarms.agents.agents.AutoGPT') as mock_agent:
initializer = AgentNodeInitializer(model_type='openai', model_id='test', openai_api_key='test_key', temperature=0.5)
initializer.llm = mock_llm
initializer.tools = [Mock(spec=BaseTool)]
initializer.vectorstore = Mock()
initializer.agent = mock_agent
return initializer
# Test initialize_llm method of AgentNodeInitializer class
@pytest.mark.parametrize("model_type", ['openai', 'huggingface', 'invalid'])
def test_agent_node_initializer_initialize_llm(model_type, mock_agent_node_initializer):
with patch('swarms.agents.agents.ChatOpenAI') as mock_openai, \
patch('swarms.agents.agents.HuggingFaceLLM') as mock_huggingface:
if model_type == 'invalid':
with pytest.raises(ValueError):
mock_agent_node_initializer.initialize_llm(model_type, 'model_id', 'openai_api_key', 0.5)
else:
mock_agent_node_initializer.initialize_llm(model_type, 'model_id', 'openai_api_key', 0.5)
if model_type == 'openai':
mock_openai.assert_called_once()
elif model_type == 'huggingface':
mock_huggingface.assert_called_once()
# Test add_tool method of AgentNodeInitializer class
def test_agent_node_initializer_add_tool(mock_agent_node_initializer):
with patch('swarms.agents.agents.BaseTool') as mock_base_tool:
mock_agent_node_initializer.add_tool(mock_base_tool)
assert mock_base_tool in mock_agent_node_initializer.tools
# Test run method of AgentNodeInitializer class
@pytest.mark.parametrize("prompt", ['valid prompt', ''])
def test_agent_node_initializer_run(prompt, mock_agent_node_initializer):
if prompt == '':
with pytest.raises(ValueError):
mock_agent_node_initializer.run(prompt)
else:
assert mock_agent_node_initializer.run(prompt) == "Task completed by AgentNode"
# For initializing AgentNode in multiple tests
@pytest.fixture
def mock_agent_node():
with patch('swarms.agents.agents.ChatOpenAI') as mock_llm, \
patch('swarms.agents.agents.AgentNodeInitializer') as mock_agent_node_initializer:
mock_agent_node = AgentNode('test_key')
mock_agent_node.llm_class = mock_llm
mock_agent_node.vectorstore = Mock()
mock_agent_node_initializer.llm = mock_llm
return mock_agent_node
# Test initialize_llm method of AgentNode class
@pytest.mark.parametrize("llm_class", ['openai', 'huggingface'])
def test_agent_node_initialize_llm(llm_class, mock_agent_node):
with patch('swarms.agents.agents.ChatOpenAI') as mock_openai, \
patch('swarms.agents.agents.HuggingFaceLLM') as mock_huggingface:
mock_agent_node.initialize_llm(llm_class)
if llm_class == 'openai':
mock_openai.assert_called_once()
elif llm_class == 'huggingface':
mock_huggingface.assert_called_once()
# Test initialize_tools method of AgentNode class
def test_agent_node_initialize_tools(mock_agent_node):
with patch('swarms.agents.agents.DuckDuckGoSearchRun') as mock_ddg, \
patch('swarms.agents.agents.WriteFileTool') as mock_write_file, \
patch('swarms.agents.agents.ReadFileTool') as mock_read_file, \
patch('swarms.agents.agents.process_csv') as mock_process_csv, \
patch('swarms.agents.agents.WebpageQATool') as mock_webpage_qa:
mock_agent_node.initialize_tools('openai')
assert mock_ddg.called
assert mock_write_file.called
assert mock_read_file.called
assert mock_process_csv.called
assert mock_webpage_qa.called
# Test create_agent method of AgentNode class
def test_agent_node_create_agent(mock_agent_node):
with patch.object(mock_agent_node, 'initialize_llm'), \
patch.object(mock_agent_node, 'initialize_tools'), \
patch.object(mock_agent_node, 'initialize_vectorstore'), \
patch('swarms.agents.agents.AgentNodeInitializer') as mock_agent_node_initializer:
mock_agent_node.create_agent()
mock_agent_node_initializer.assert_called_once()
mock_agent_node_initializer.return_value.create_agent.assert_called_once()
# Test agent function
@pytest.mark.parametrize("openai_api_key,objective", [('valid_key', 'valid_objective'), ('', 'valid_objective'), ('valid_key', '')])
def test_agent(openai_api_key, objective):
if openai_api_key == '' or objective == '':
with pytest.raises(ValueError):
agent(openai_api_key, objective)
else:
with patch('swarms.agents.agents.AgentNodeInitializer') as mock_agent_node_initializer:
mock_agent_node = mock_agent_node_initializer.return_value.create_agent.return_value
mock_agent_node.run.return_value = 'Agent output'
result = agent(openai_api_key, objective)
assert result == 'Agent output'
|