Spaces:
Runtime error
Runtime error
import pytest | |
import logging | |
from unittest.mock import patch | |
from swarms.swarms.swarms import HierarchicalSwarm # replace with your actual module name | |
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 | |
) | |
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 | |
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) | |
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 | |
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 | |
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), {}) | |