Spaces:
Running
Running
File size: 2,898 Bytes
d8d14f1 |
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 |
from unittest.mock import Mock, patch
from transformers import AutoModelForCausalLM, AutoTokenizer
from swarms import ToolAgent
def test_tool_agent_init():
model = Mock(spec=AutoModelForCausalLM)
tokenizer = Mock(spec=AutoTokenizer)
json_schema = {
"type": "object",
"properties": {
"name": {"type": "string"},
"age": {"type": "number"},
"is_student": {"type": "boolean"},
"courses": {"type": "array", "items": {"type": "string"}},
},
}
name = "Test Agent"
description = "This is a test agent"
agent = ToolAgent(
name, description, model, tokenizer, json_schema
)
assert agent.name == name
assert agent.description == description
assert agent.model == model
assert agent.tokenizer == tokenizer
assert agent.json_schema == json_schema
@patch.object(ToolAgent, "run")
def test_tool_agent_run(mock_run):
model = Mock(spec=AutoModelForCausalLM)
tokenizer = Mock(spec=AutoTokenizer)
json_schema = {
"type": "object",
"properties": {
"name": {"type": "string"},
"age": {"type": "number"},
"is_student": {"type": "boolean"},
"courses": {"type": "array", "items": {"type": "string"}},
},
}
name = "Test Agent"
description = "This is a test agent"
task = (
"Generate a person's information based on the following"
" schema:"
)
agent = ToolAgent(
name, description, model, tokenizer, json_schema
)
agent.run(task)
mock_run.assert_called_once_with(task)
def test_tool_agent_init_with_kwargs():
model = Mock(spec=AutoModelForCausalLM)
tokenizer = Mock(spec=AutoTokenizer)
json_schema = {
"type": "object",
"properties": {
"name": {"type": "string"},
"age": {"type": "number"},
"is_student": {"type": "boolean"},
"courses": {"type": "array", "items": {"type": "string"}},
},
}
name = "Test Agent"
description = "This is a test agent"
kwargs = {
"debug": True,
"max_array_length": 20,
"max_number_tokens": 12,
"temperature": 0.5,
"max_string_token_length": 20,
}
agent = ToolAgent(
name, description, model, tokenizer, json_schema, **kwargs
)
assert agent.name == name
assert agent.description == description
assert agent.model == model
assert agent.tokenizer == tokenizer
assert agent.json_schema == json_schema
assert agent.debug == kwargs["debug"]
assert agent.max_array_length == kwargs["max_array_length"]
assert agent.max_number_tokens == kwargs["max_number_tokens"]
assert agent.temperature == kwargs["temperature"]
assert (
agent.max_string_token_length
== kwargs["max_string_token_length"]
)
|