Spaces:
Runtime error
Runtime error
File size: 2,620 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 |
import pytest
import logging
from unittest.mock import patch
from swarms.swarms.swarms import HierarchicalSwarm # replace with your actual module name
@pytest.fixture
def swarm():
return HierarchicalSwarm(
model_id='gpt-4',
openai_api_key='some_api_key',
use_vectorstore=True,
embedding_size=1024,
use_async=False,
human_in_the_loop=True,
model_type='openai',
boss_prompt='boss',
worker_prompt='worker',
temperature=0.5,
max_iterations=100,
logging_enabled=True
)
@pytest.fixture
def swarm_no_logging():
return HierarchicalSwarm(logging_enabled=False)
def test_swarm_init(swarm):
assert swarm.model_id == 'gpt-4'
assert swarm.openai_api_key == 'some_api_key'
assert swarm.use_vectorstore
assert swarm.embedding_size == 1024
assert not swarm.use_async
assert swarm.human_in_the_loop
assert swarm.model_type == 'openai'
assert swarm.boss_prompt == 'boss'
assert swarm.worker_prompt == 'worker'
assert swarm.temperature == 0.5
assert swarm.max_iterations == 100
assert swarm.logging_enabled
assert isinstance(swarm.logger, logging.Logger)
def test_swarm_no_logging_init(swarm_no_logging):
assert not swarm_no_logging.logging_enabled
assert swarm_no_logging.logger.disabled
@patch('your_module.OpenAI')
@patch('your_module.HuggingFaceLLM')
def test_initialize_llm(mock_huggingface, mock_openai, swarm):
swarm.initialize_llm('openai')
mock_openai.assert_called_once_with(openai_api_key='some_api_key', temperature=0.5)
swarm.initialize_llm('huggingface')
mock_huggingface.assert_called_once_with(model_id='gpt-4', temperature=0.5)
@patch('your_module.HierarchicalSwarm.initialize_llm')
def test_initialize_tools(mock_llm, swarm):
mock_llm.return_value = 'mock_llm_class'
tools = swarm.initialize_tools('openai')
assert 'mock_llm_class' in tools
@patch('your_module.HierarchicalSwarm.initialize_llm')
def test_initialize_tools_with_extra_tools(mock_llm, swarm):
mock_llm.return_value = 'mock_llm_class'
tools = swarm.initialize_tools('openai', extra_tools=['tool1', 'tool2'])
assert 'tool1' in tools
assert 'tool2' in tools
@patch('your_module.OpenAIEmbeddings')
@patch('your_module.FAISS')
def test_initialize_vectorstore(mock_faiss, mock_openai_embeddings, swarm):
mock_openai_embeddings.return_value.embed_query = 'embed_query'
swarm.initialize_vectorstore()
mock_faiss.assert_called_once_with('embed_query', instance_of(faiss.IndexFlatL2), instance_of(InMemoryDocstore), {})
|