Former-commit-id: e5dce12e5ba03be619ed8f2da9d69d6287a546cfclean-history
Kye
1 year ago
parent
97b09c37f1
commit
f6c230e6cc
@ -1,133 +0,0 @@
|
||||
import pytest
|
||||
from unittest.mock import Mock, patch
|
||||
from swarms.tools.agent_tools import *
|
||||
from swarms.boss.boss_node import BossNodeInitializer, BossNode
|
||||
|
||||
|
||||
# For initializing BossNodeInitializer in multiple tests
|
||||
@pytest.fixture
|
||||
def mock_boss_node_initializer():
|
||||
llm = Mock()
|
||||
vectorstore = Mock()
|
||||
agent_executor = Mock()
|
||||
max_iterations = 5
|
||||
|
||||
boss_node_initializer = BossNodeInitializer(
|
||||
llm, vectorstore, agent_executor, max_iterations
|
||||
)
|
||||
|
||||
return boss_node_initializer
|
||||
|
||||
|
||||
# Test BossNodeInitializer class __init__ method
|
||||
def test_boss_node_initializer_init(mock_boss_node_initializer):
|
||||
with patch("swarms.tools.agent_tools.BabyAGI.from_llm") as mock_from_llm:
|
||||
assert isinstance(mock_boss_node_initializer, BossNodeInitializer)
|
||||
mock_from_llm.assert_called_once()
|
||||
|
||||
|
||||
# Test initialize_vectorstore method of BossNodeInitializer class
|
||||
def test_boss_node_initializer_initialize_vectorstore(mock_boss_node_initializer):
|
||||
with patch("swarms.tools.agent_tools.OpenAIEmbeddings") as mock_embeddings, patch(
|
||||
"swarms.tools.agent_tools.FAISS"
|
||||
) as mock_faiss:
|
||||
result = mock_boss_node_initializer.initialize_vectorstore()
|
||||
mock_embeddings.assert_called_once()
|
||||
mock_faiss.assert_called_once()
|
||||
assert result is not None
|
||||
|
||||
|
||||
# Test initialize_llm method of BossNodeInitializer class
|
||||
def test_boss_node_initializer_initialize_llm(mock_boss_node_initializer):
|
||||
with patch("swarms.tools.agent_tools.OpenAI") as mock_llm:
|
||||
result = mock_boss_node_initializer.initialize_llm(mock_llm)
|
||||
mock_llm.assert_called_once()
|
||||
assert result is not None
|
||||
|
||||
|
||||
# Test create_task method of BossNodeInitializer class
|
||||
@pytest.mark.parametrize("objective", ["valid objective", ""])
|
||||
def test_boss_node_initializer_create_task(objective, mock_boss_node_initializer):
|
||||
if objective == "":
|
||||
with pytest.raises(ValueError):
|
||||
mock_boss_node_initializer.create_task(objective)
|
||||
else:
|
||||
assert mock_boss_node_initializer.create_task(objective) == {
|
||||
"objective": objective
|
||||
}
|
||||
|
||||
|
||||
# Test run method of BossNodeInitializer class
|
||||
@pytest.mark.parametrize("task", ["valid task", ""])
|
||||
def test_boss_node_initializer_run(task, mock_boss_node_initializer):
|
||||
with patch.object(mock_boss_node_initializer, "baby_agi"):
|
||||
if task == "":
|
||||
with pytest.raises(ValueError):
|
||||
mock_boss_node_initializer.run(task)
|
||||
else:
|
||||
try:
|
||||
mock_boss_node_initializer.run(task)
|
||||
mock_boss_node_initializer.baby_agi.assert_called_once_with(task)
|
||||
except Exception:
|
||||
pytest.fail("Unexpected Error!")
|
||||
|
||||
|
||||
# Test BossNode function
|
||||
@pytest.mark.parametrize(
|
||||
"api_key, objective, llm_class, max_iterations",
|
||||
[
|
||||
("valid_key", "valid_objective", OpenAI, 5),
|
||||
("", "valid_objective", OpenAI, 5),
|
||||
("valid_key", "", OpenAI, 5),
|
||||
("valid_key", "valid_objective", "", 5),
|
||||
("valid_key", "valid_objective", OpenAI, 0),
|
||||
],
|
||||
)
|
||||
def test_boss_node(api_key, objective, llm_class, max_iterations):
|
||||
with patch("os.getenv") as mock_getenv, patch(
|
||||
"swarms.tools.agent_tools.PromptTemplate.from_template"
|
||||
) as mock_from_template, patch(
|
||||
"swarms.tools.agent_tools.LLMChain"
|
||||
) as mock_llm_chain, patch(
|
||||
"swarms.tools.agent_tools.ZeroShotAgent.create_prompt"
|
||||
) as mock_create_prompt, patch(
|
||||
"swarms.tools.agent_tools.ZeroShotAgent"
|
||||
) as mock_zero_shot_agent, patch(
|
||||
"swarms.tools.agent_tools.AgentExecutor.from_agent_and_tools"
|
||||
) as mock_from_agent_and_tools, patch(
|
||||
"swarms.tools.agent_tools.BossNodeInitializer"
|
||||
) as mock_boss_node_initializer, patch.object(
|
||||
mock_boss_node_initializer, "create_task"
|
||||
) as mock_create_task, patch.object(
|
||||
mock_boss_node_initializer, "run"
|
||||
) as mock_run:
|
||||
if api_key == "" or objective == "" or llm_class == "" or max_iterations <= 0:
|
||||
with pytest.raises(ValueError):
|
||||
BossNode(
|
||||
objective,
|
||||
api_key,
|
||||
vectorstore=None,
|
||||
worker_node=None,
|
||||
llm_class=llm_class,
|
||||
max_iterations=max_iterations,
|
||||
verbose=False,
|
||||
)
|
||||
else:
|
||||
mock_getenv.return_value = "valid_key"
|
||||
BossNode(
|
||||
objective,
|
||||
api_key,
|
||||
vectorstore=None,
|
||||
worker_node=None,
|
||||
llm_class=llm_class,
|
||||
max_iterations=max_iterations,
|
||||
verbose=False,
|
||||
)
|
||||
mock_from_template.assert_called_once()
|
||||
mock_llm_chain.assert_called_once()
|
||||
mock_create_prompt.assert_called_once()
|
||||
mock_zero_shot_agent.assert_called_once()
|
||||
mock_from_agent_and_tools.assert_called_once()
|
||||
mock_boss_node_initializer.assert_called_once()
|
||||
mock_create_task.assert_called_once()
|
||||
mock_run.assert_called_once()
|
||||
@ -1,83 +0,0 @@
|
||||
import pytest
|
||||
from swarms.chunkers.base import (
|
||||
BaseChunker,
|
||||
TextArtifact,
|
||||
ChunkSeparator,
|
||||
OpenAITokenizer,
|
||||
) # adjust the import paths accordingly
|
||||
|
||||
|
||||
# 1. Test Initialization
|
||||
def test_chunker_initialization():
|
||||
chunker = BaseChunker()
|
||||
assert isinstance(chunker, BaseChunker)
|
||||
assert chunker.max_tokens == chunker.tokenizer.max_tokens
|
||||
|
||||
|
||||
def test_default_separators():
|
||||
chunker = BaseChunker()
|
||||
assert chunker.separators == BaseChunker.DEFAULT_SEPARATORS
|
||||
|
||||
|
||||
def test_default_tokenizer():
|
||||
chunker = BaseChunker()
|
||||
assert isinstance(chunker.tokenizer, OpenAITokenizer)
|
||||
|
||||
|
||||
# 2. Test Basic Chunking
|
||||
@pytest.mark.parametrize(
|
||||
"input_text, expected_output",
|
||||
[
|
||||
("This is a test.", [TextArtifact("This is a test.")]),
|
||||
("Hello World!", [TextArtifact("Hello World!")]),
|
||||
# Add more simple cases
|
||||
],
|
||||
)
|
||||
def test_basic_chunk(input_text, expected_output):
|
||||
chunker = BaseChunker()
|
||||
result = chunker.chunk(input_text)
|
||||
assert result == expected_output
|
||||
|
||||
|
||||
# 3. Test Chunking with Different Separators
|
||||
def test_custom_separators():
|
||||
custom_separator = ChunkSeparator(";")
|
||||
chunker = BaseChunker(separators=[custom_separator])
|
||||
input_text = "Hello;World!"
|
||||
expected_output = [TextArtifact("Hello;"), TextArtifact("World!")]
|
||||
result = chunker.chunk(input_text)
|
||||
assert result == expected_output
|
||||
|
||||
|
||||
# 4. Test Recursive Chunking
|
||||
def test_recursive_chunking():
|
||||
chunker = BaseChunker(max_tokens=5)
|
||||
input_text = "This is a more complex text."
|
||||
expected_output = [
|
||||
TextArtifact("This"),
|
||||
TextArtifact("is a"),
|
||||
TextArtifact("more"),
|
||||
TextArtifact("complex"),
|
||||
TextArtifact("text."),
|
||||
]
|
||||
result = chunker.chunk(input_text)
|
||||
assert result == expected_output
|
||||
|
||||
|
||||
# 5. Test Edge Cases and Special Scenarios
|
||||
def test_empty_text():
|
||||
chunker = BaseChunker()
|
||||
result = chunker.chunk("")
|
||||
assert result == []
|
||||
|
||||
|
||||
def test_whitespace_text():
|
||||
chunker = BaseChunker()
|
||||
result = chunker.chunk(" ")
|
||||
assert result == [TextArtifact(" ")]
|
||||
|
||||
|
||||
def test_single_word():
|
||||
chunker = BaseChunker()
|
||||
result = chunker.chunk("Hello")
|
||||
assert result == [TextArtifact("Hello")]
|
||||
Loading…
Reference in new issue