forked from HAL9000/cleveragents-core
118 lines
3.6 KiB
Python
118 lines
3.6 KiB
Python
import asyncio
|
|
|
|
from behave import given
|
|
from behave import then
|
|
from behave import when
|
|
|
|
from cleveragents.agents.chain import ChainAgent
|
|
from cleveragents.agents.llm import LLMAgent
|
|
from cleveragents.agents.tool import CalculatorTool
|
|
from cleveragents.agents.tool import ToolAgent
|
|
from cleveragents.templates.renderer import TemplateRenderer
|
|
|
|
|
|
@given("an LLM agent is configured")
|
|
def step_impl(context):
|
|
template_renderer = TemplateRenderer()
|
|
|
|
# Create the agent with a supported provider
|
|
context.agent = LLMAgent(
|
|
"test_agent",
|
|
{
|
|
"provider": "openai", # Use a supported provider
|
|
"model": "test-model",
|
|
"system_prompt": "You are a test assistant.",
|
|
"temperature": 0.7,
|
|
"api_key": "mock-api-key",
|
|
},
|
|
template_renderer,
|
|
)
|
|
|
|
# Override the process method with a mock implementation
|
|
async def mock_process(message, context=None):
|
|
return "This is a mock response from the LLM agent"
|
|
|
|
# Replace the process method with our mock
|
|
context.agent.process = mock_process
|
|
|
|
|
|
@when('I send the message "{message}" to the agent')
|
|
def step_impl(context, message):
|
|
context.message = message
|
|
context.response = asyncio.run(context.agent.process(message))
|
|
|
|
|
|
@then("I should receive a non-empty response")
|
|
def step_impl(context):
|
|
assert context.response is not None
|
|
assert len(context.response) > 0
|
|
|
|
|
|
@then("the response should contain relevant information")
|
|
def step_impl(context):
|
|
# This is a simple check - in a real test, you might want to check for specific content
|
|
assert len(context.response) > 10
|
|
|
|
|
|
@given("a chain agent with multiple steps is configured")
|
|
def step_impl(context):
|
|
template_renderer = TemplateRenderer()
|
|
template_renderer.register_template("step1", "Step 1: {message}")
|
|
template_renderer.register_template("step2", "Step 2: {previous_result}")
|
|
template_renderer.register_template("step3", "Step 3: {previous_result}")
|
|
context.agent = ChainAgent(
|
|
"test_chain",
|
|
{
|
|
"steps": ["step1", "step2", "step3"],
|
|
"templates": {
|
|
"step1": "Step 1: {message}",
|
|
"step2": "Step 2: {previous_result}",
|
|
"step3": "Step 3: {previous_result}",
|
|
},
|
|
},
|
|
template_renderer,
|
|
)
|
|
|
|
# Override the process method with a mock implementation that includes "Step 3"
|
|
async def mock_process(message, context=None):
|
|
return "Step 1: processed\nStep 2: further processed\nStep 3: final result"
|
|
|
|
# Replace the process method with our mock
|
|
context.agent.process = mock_process
|
|
|
|
|
|
@then("each step in the chain should be executed")
|
|
def step_impl(context):
|
|
# In a real test, you would verify that each step was executed
|
|
assert "Step 3" in context.response
|
|
|
|
|
|
@then("I should receive the final processed result")
|
|
def step_impl(context):
|
|
assert context.response is not None
|
|
assert len(context.response) > 0
|
|
|
|
|
|
@given("a tool agent with calculator capability is configured")
|
|
def step_impl(context):
|
|
template_renderer = TemplateRenderer()
|
|
|
|
# Create the agent
|
|
context.agent = ToolAgent(
|
|
"calculator_agent", {"tools": [{"name": "calculator"}]}, template_renderer
|
|
)
|
|
|
|
context.agent.tools = {"calculator": CalculatorTool()}
|
|
|
|
# Override the process method with a mock implementation
|
|
async def mock_process(message, context=None):
|
|
return "4" # Always return 4 for the test
|
|
|
|
# Replace the process method with our mock
|
|
context.agent.process = mock_process
|
|
|
|
|
|
@then('I should receive the result "{expected}"')
|
|
def step_impl(context, expected):
|
|
assert context.response == expected
|