forked from HAL9000/cleveragents-core
113 lines
3.5 KiB
Python
113 lines
3.5 KiB
Python
import asyncio
|
|
from unittest.mock import MagicMock
|
|
|
|
from behave import given
|
|
from behave import then
|
|
from behave import when
|
|
|
|
from cleveragents.agents.tool import CalculatorTool
|
|
from cleveragents.agents.tool import ToolAgent
|
|
from cleveragents.agents.tool import WebSearchTool
|
|
from cleveragents.templates.renderer import TemplateRenderer
|
|
|
|
|
|
@given("a tool agent with calculator capability")
|
|
def step_impl(context):
|
|
context.template_renderer = TemplateRenderer()
|
|
context.agent = ToolAgent(
|
|
"calculator_agent",
|
|
{"tools": [{"name": "calculator"}]},
|
|
context.template_renderer,
|
|
)
|
|
context.agent.tools = {"calculator": CalculatorTool()}
|
|
|
|
|
|
@when('I send a calculation request "{expression}"')
|
|
def step_impl(context, expression):
|
|
context.expression = expression
|
|
context.result = asyncio.run(context.agent.process(f"calculator: {expression}"))
|
|
|
|
|
|
@then("I should receive the correct calculation result")
|
|
def step_impl(context):
|
|
expected = str(eval(context.expression))
|
|
assert context.result == expected
|
|
|
|
|
|
@when("I send an invalid calculation request")
|
|
def step_impl(context):
|
|
try:
|
|
context.result = asyncio.run(
|
|
context.agent.process("calculator: invalid expression")
|
|
)
|
|
context.error = None
|
|
except Exception as e:
|
|
context.error = e
|
|
|
|
|
|
@then("I should receive an error message")
|
|
def step_impl(context):
|
|
assert context.error is not None
|
|
|
|
|
|
@given("a tool agent with calculator and web search tools")
|
|
def step_impl(context):
|
|
context.template_renderer = TemplateRenderer()
|
|
context.agent = ToolAgent(
|
|
"multi_tool_agent",
|
|
{
|
|
"tools": [
|
|
{"name": "calculator"},
|
|
{"name": "web_search", "api_key": "mock-key"},
|
|
]
|
|
},
|
|
context.template_renderer,
|
|
)
|
|
context.agent.tools = {"calculator": CalculatorTool(), "web_search": MagicMock()}
|
|
|
|
async def mock_execute(input_data, context=None):
|
|
return f"Web search results for '{input_data}'"
|
|
|
|
context.agent.tools["web_search"].execute = mock_execute
|
|
|
|
|
|
@when("I send a tool request in different formats")
|
|
def step_impl(context):
|
|
context.results = {}
|
|
context.results["format1"] = asyncio.run(context.agent.process("calculator: 1 + 1"))
|
|
context.results["format2"] = asyncio.run(
|
|
context.agent.process('{"tool": "web_search", "input": "CleverAgents"}')
|
|
)
|
|
context.results["format3"] = asyncio.run(
|
|
context.agent.process("Please use web_search to find information about AI")
|
|
)
|
|
|
|
|
|
@then("the correct tool should be selected for each format")
|
|
def step_impl(context):
|
|
assert context.results["format1"] == "2"
|
|
assert "Web search results for 'CleverAgents'" in context.results["format2"]
|
|
assert "Web search results for" in context.results["format3"]
|
|
|
|
|
|
from cleveragents.core.exceptions import ConfigurationError
|
|
|
|
|
|
@when("I initialize a WebSearchTool with a missing API key")
|
|
@when("I try to initialize a WebSearchTool with a missing API key")
|
|
def step_initialize_web_search_missing_api(context):
|
|
context.error = None
|
|
try:
|
|
# Config is missing the 'api_key'
|
|
config = {"search_engine": "google", "cx": "test-cx"}
|
|
context.tool = WebSearchTool(config)
|
|
except Exception as e:
|
|
context.error = e
|
|
|
|
|
|
@then('the tool\'s API key should be "{api_key}"')
|
|
def step_check_tool_api_key(context, api_key):
|
|
assert not context.error, f"Tool initialization failed with: {context.error}"
|
|
assert hasattr(context, "tool")
|
|
assert context.tool.api_key == api_key
|