forked from HAL9000/cleveragents-core
123 lines
3.7 KiB
Python
123 lines
3.7 KiB
Python
import asyncio
|
|
from typing import Any
|
|
from typing import Dict
|
|
from typing import List
|
|
from typing import Optional
|
|
from unittest.mock import AsyncMock
|
|
from unittest.mock import MagicMock
|
|
from unittest.mock import patch
|
|
|
|
from behave import given
|
|
from behave import then
|
|
from behave import when
|
|
|
|
from cleveragents.agents.chain import ChainAgent
|
|
from cleveragents.templates.renderer import TemplateRenderer
|
|
|
|
|
|
@given("a chain agent with multiple processing steps")
|
|
def step_impl(context):
|
|
# Create a template renderer
|
|
context.template_renderer = TemplateRenderer()
|
|
|
|
# Create a chain agent with multiple steps
|
|
context.agent = ChainAgent(
|
|
"test_chain_agent",
|
|
{"steps": ["parse_input", "process_data", "format_output"]},
|
|
context.template_renderer,
|
|
)
|
|
|
|
|
|
@when("I send a message to the chain agent")
|
|
def step_impl(context):
|
|
# Process a message
|
|
context.response = asyncio.run(context.agent.process("Test message"))
|
|
|
|
|
|
@then("the message should be processed through all steps")
|
|
def step_impl(context):
|
|
# Check that the response includes all steps
|
|
assert "ChainAgent processed:" in context.response
|
|
assert "parse_input" in context.response
|
|
assert "process_data" in context.response
|
|
assert "format_output" in context.response
|
|
|
|
|
|
@given("a chain agent with custom step implementations")
|
|
def step_impl(context):
|
|
# Create a template renderer
|
|
context.template_renderer = TemplateRenderer()
|
|
|
|
# Create a chain agent with custom step implementations
|
|
context.agent = ChainAgent(
|
|
"custom_chain_agent",
|
|
{
|
|
"steps": ["custom_step_1", "custom_step_2"],
|
|
"implementations": {
|
|
"custom_step_1": "def process(message, context): return f'Step 1: {message}'",
|
|
"custom_step_2": "def process(message, context): return f'Step 2: {message}'",
|
|
},
|
|
},
|
|
context.template_renderer,
|
|
)
|
|
|
|
# Store the original steps for verification
|
|
context.steps = context.agent.steps
|
|
|
|
# Mock the process method to simulate custom implementations
|
|
async def mock_process(message, ctx=None):
|
|
if ctx is None:
|
|
ctx = {}
|
|
|
|
result = message
|
|
# Use the stored steps from the behave context
|
|
for step in context.steps:
|
|
result = f"{result} -> {step}"
|
|
|
|
return f"Custom ChainAgent processed: {result}"
|
|
|
|
# Apply the mock
|
|
context.agent.process = mock_process
|
|
|
|
|
|
@when("I send a message requiring custom processing")
|
|
def step_impl(context):
|
|
# Process a message
|
|
context.response = asyncio.run(context.agent.process("Custom test message"))
|
|
|
|
|
|
@then("the message should be processed using the custom implementations")
|
|
def step_impl(context):
|
|
# Check that the response includes custom processing
|
|
assert "Custom ChainAgent processed:" in context.response
|
|
assert "custom_step_1" in context.response
|
|
assert "custom_step_2" in context.response
|
|
|
|
|
|
@given("a chain agent with various capabilities")
|
|
def step_impl(context):
|
|
# Create a template renderer
|
|
context.template_renderer = TemplateRenderer()
|
|
|
|
# Create a chain agent
|
|
context.agent = ChainAgent(
|
|
"capability_chain_agent",
|
|
{
|
|
"steps": ["step1", "step2"],
|
|
"capabilities": ["custom-capability-1", "custom-capability-2"],
|
|
},
|
|
context.template_renderer,
|
|
)
|
|
|
|
|
|
@when("I request the agent's capabilities")
|
|
def step_impl(context):
|
|
# Get the capabilities
|
|
context.capabilities = context.agent.get_capabilities()
|
|
|
|
|
|
@then("the capabilities should include chain-processing")
|
|
def step_impl(context):
|
|
# Check that the capabilities include chain-processing
|
|
assert "chain-processing" in context.capabilities
|