forked from HAL9000/cleveragents-core
102 lines
3.3 KiB
Python
102 lines
3.3 KiB
Python
import asyncio
|
|
import json
|
|
from typing import Any
|
|
from typing import Dict
|
|
from typing import List
|
|
from typing import Optional
|
|
|
|
import yaml
|
|
from behave import given
|
|
from behave import then
|
|
from behave import when
|
|
|
|
from cleveragents.agents.base import Agent
|
|
from cleveragents.agents.composite import CompositeAgent
|
|
from cleveragents.agents.factory import AgentFactory
|
|
from cleveragents.core.config import ConfigurationManager
|
|
from cleveragents.templates.renderer import TemplateRenderer
|
|
|
|
|
|
class MockAgent(Agent):
|
|
"""A mock agent that returns a configured response."""
|
|
|
|
def __init__(
|
|
self, name: str, config: Dict[str, Any], template_renderer: TemplateRenderer
|
|
):
|
|
super().__init__(name, config, template_renderer)
|
|
self.response = self.config.get("mock_response", f"Response from {name}")
|
|
|
|
async def process(
|
|
self, message: str, context: Optional[Dict[str, Any]] = None
|
|
) -> str:
|
|
return self.response
|
|
|
|
def get_capabilities(self) -> List[str]:
|
|
return ["mock"]
|
|
|
|
|
|
def setup_context_with_config(context, agent_configs: Dict[str, Any]):
|
|
"""Set up the test context with configuration and agent factory."""
|
|
context.config_manager = ConfigurationManager()
|
|
full_config = {"agents": agent_configs}
|
|
context.config_manager.config = full_config
|
|
context.template_renderer = TemplateRenderer()
|
|
context.agent_factory = AgentFactory(
|
|
config=context.config_manager.config,
|
|
template_renderer=context.template_renderer,
|
|
)
|
|
context.agent_factory.register_agent_type("mock", MockAgent)
|
|
|
|
|
|
@given(
|
|
'a CompositeAgent configuration with parallel strategy and format "{output_format}"'
|
|
)
|
|
def step_impl(context, output_format: str):
|
|
"""Set up a CompositeAgent with parallel strategy and specified output format."""
|
|
config_yaml = f"""
|
|
agents:
|
|
AgentA:
|
|
type: mock
|
|
config:
|
|
mock_response: "Hello from A"
|
|
AgentB:
|
|
type: mock
|
|
config:
|
|
mock_response: "Hello from B"
|
|
ParallelAgent:
|
|
type: composite
|
|
config:
|
|
strategy: parallel
|
|
output_format: {output_format}
|
|
steps:
|
|
- AgentA
|
|
- AgentB
|
|
"""
|
|
agent_configs = yaml.safe_load(config_yaml)["agents"]
|
|
setup_context_with_config(context, agent_configs)
|
|
|
|
|
|
@when('I process the message "{message}" with the composite agent')
|
|
def step_impl(context, message: str):
|
|
"""Process a message using the created composite agent."""
|
|
context.response = asyncio.run(context.agent.process(message))
|
|
|
|
|
|
@then("the response should be the plain text combination of child agent outputs")
|
|
def step_impl(context):
|
|
"""Verify the plain text output from the parallel strategy."""
|
|
expected_response = "AgentA: Hello from A\n\nAgentB: Hello from B"
|
|
assert (
|
|
context.response == expected_response
|
|
), f"Expected:\n{expected_response}\nGot:\n{context.response}"
|
|
|
|
|
|
@then("the response should be a JSON object with child agent outputs")
|
|
def step_impl(context):
|
|
"""Verify the JSON output from the parallel strategy."""
|
|
response_dict = json.loads(context.response)
|
|
expected_dict = {"AgentA": "Hello from A", "AgentB": "Hello from B"}
|
|
assert (
|
|
response_dict == expected_dict
|
|
), f"Expected:\n{expected_dict}\nGot:\n{response_dict}"
|