forked from cleveragents/cleveragents-core
215 lines
7.8 KiB
Python
215 lines
7.8 KiB
Python
"""
|
|
Unit tests for agents/chain.py
|
|
|
|
Tests the ChainAgent class.
|
|
"""
|
|
|
|
import pytest
|
|
from unittest.mock import Mock
|
|
|
|
from cleveragents.agents.chain import ChainAgent
|
|
from cleveragents.core.exceptions import ConfigurationError
|
|
from cleveragents.templates.renderer import TemplateRenderer
|
|
|
|
|
|
class TestChainAgent:
|
|
"""Test suite for the ChainAgent class."""
|
|
|
|
@pytest.fixture
|
|
def template_renderer(self):
|
|
"""Create a basic template renderer."""
|
|
return TemplateRenderer()
|
|
|
|
@pytest.fixture
|
|
def simple_config(self):
|
|
"""Create a simple chain configuration."""
|
|
return {
|
|
"type": "chain",
|
|
"steps": ["step1", "step2", "step3"],
|
|
}
|
|
|
|
@pytest.fixture
|
|
def config_with_prompt(self):
|
|
"""Create a chain configuration with prompt template."""
|
|
return {
|
|
"type": "chain",
|
|
"prompt": "Process this: {{ message }}",
|
|
"steps": ["analyze", "summarize"],
|
|
}
|
|
|
|
def test_chain_agent_initialization(self, simple_config, template_renderer):
|
|
"""Test ChainAgent initialization."""
|
|
agent = ChainAgent("test_chain", simple_config, template_renderer)
|
|
|
|
assert agent.name == "test_chain"
|
|
assert agent.steps == ["step1", "step2", "step3"]
|
|
assert agent.prompt_template is None
|
|
|
|
def test_chain_agent_initialization_with_prompt(self, config_with_prompt, template_renderer):
|
|
"""Test ChainAgent initialization with prompt template."""
|
|
agent = ChainAgent("test_chain", config_with_prompt, template_renderer)
|
|
|
|
assert agent.prompt_template == "Process this: {{ message }}"
|
|
assert agent.steps == ["analyze", "summarize"]
|
|
|
|
def test_chain_agent_initialization_with_prompt_reference(self, template_renderer):
|
|
"""Test ChainAgent initialization with prompt reference."""
|
|
config = {
|
|
"type": "chain",
|
|
"prompt_reference": "my_template",
|
|
"steps": ["step1"],
|
|
}
|
|
|
|
# Mock the template renderer
|
|
template_renderer.get_template = Mock(return_value="Referenced template")
|
|
|
|
agent = ChainAgent("test_chain", config, template_renderer)
|
|
|
|
assert agent.prompt_template == "Referenced template"
|
|
template_renderer.get_template.assert_called_once_with("my_template")
|
|
|
|
def test_chain_agent_prompt_and_reference_conflict(self, template_renderer):
|
|
"""Test that having both prompt and prompt_reference raises error."""
|
|
config = {
|
|
"type": "chain",
|
|
"prompt": "Direct prompt",
|
|
"prompt_reference": "my_template",
|
|
"steps": ["step1"],
|
|
}
|
|
|
|
with pytest.raises(ConfigurationError) as exc_info:
|
|
ChainAgent("test_chain", config, template_renderer)
|
|
|
|
assert "cannot contain both" in str(exc_info.value)
|
|
assert "prompt" in str(exc_info.value)
|
|
assert "prompt_reference" in str(exc_info.value)
|
|
|
|
def test_chain_agent_empty_steps(self, template_renderer):
|
|
"""Test ChainAgent with empty steps list."""
|
|
config = {"type": "chain"}
|
|
agent = ChainAgent("test_chain", config, template_renderer)
|
|
|
|
assert agent.steps == []
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_message_simple(self, simple_config, template_renderer):
|
|
"""Test processing a simple message without prompt."""
|
|
agent = ChainAgent("test_chain", simple_config, template_renderer)
|
|
result = await agent.process_message("Hello")
|
|
|
|
assert "ChainAgent processed:" in result
|
|
assert "step1" in result
|
|
assert "step2" in result
|
|
assert "step3" in result
|
|
assert "Hello" in result
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_message_with_prompt(self, template_renderer):
|
|
"""Test processing a message with prompt template."""
|
|
config = {
|
|
"type": "chain",
|
|
"prompt": "Transform: {{ message }}",
|
|
"steps": ["step1", "step2"],
|
|
}
|
|
|
|
agent = ChainAgent("test_chain", config, template_renderer)
|
|
result = await agent.process_message("test input")
|
|
|
|
assert "ChainAgent processed:" in result
|
|
assert "Transform: test input" in result
|
|
assert "step1" in result
|
|
assert "step2" in result
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_message_with_context(self, template_renderer):
|
|
"""Test processing a message with context variables."""
|
|
config = {
|
|
"type": "chain",
|
|
"prompt": "User: {{ user_name }}, Message: {{ message }}",
|
|
"steps": ["process"],
|
|
}
|
|
|
|
agent = ChainAgent("test_chain", config, template_renderer)
|
|
context = {"user_name": "Alice"}
|
|
result = await agent.process_message("Hello", context)
|
|
|
|
assert "User: Alice" in result
|
|
assert "Message: Hello" in result
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_message_no_steps(self, template_renderer):
|
|
"""Test processing with no steps defined."""
|
|
config = {"type": "chain", "steps": []}
|
|
agent = ChainAgent("test_chain", config, template_renderer)
|
|
result = await agent.process_message("test")
|
|
|
|
assert "ChainAgent processed:" in result
|
|
assert "test" in result
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_message_chain_sequence(self, simple_config, template_renderer):
|
|
"""Test that steps are applied in sequence."""
|
|
agent = ChainAgent("test_chain", simple_config, template_renderer)
|
|
result = await agent.process_message("input")
|
|
|
|
# Verify steps appear in order
|
|
assert result.index("step1") < result.index("step2")
|
|
assert result.index("step2") < result.index("step3")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_legacy_method(self, simple_config, template_renderer):
|
|
"""Test the legacy process method."""
|
|
agent = ChainAgent("test_chain", simple_config, template_renderer)
|
|
result = await agent.process("test message")
|
|
|
|
assert "ChainAgent processed:" in result
|
|
assert "test message" in result
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_legacy_method_with_context(self, template_renderer):
|
|
"""Test legacy process method with context."""
|
|
config = {
|
|
"type": "chain",
|
|
"prompt": "{{ prefix }}: {{ message }}",
|
|
"steps": [],
|
|
}
|
|
agent = ChainAgent("test_chain", config, template_renderer)
|
|
result = await agent.process("msg", {"prefix": "LOG"})
|
|
|
|
assert "LOG: msg" in result
|
|
|
|
def test_get_capabilities(self, simple_config, template_renderer):
|
|
"""Test getting chain agent capabilities."""
|
|
agent = ChainAgent("test_chain", simple_config, template_renderer)
|
|
capabilities = agent.get_capabilities()
|
|
|
|
assert isinstance(capabilities, list)
|
|
assert "chain-processing" in capabilities
|
|
|
|
def test_chain_agent_multiple_instances(self, template_renderer):
|
|
"""Test creating multiple chain agents with different configs."""
|
|
config1 = {"type": "chain", "steps": ["a", "b"]}
|
|
config2 = {"type": "chain", "steps": ["x", "y", "z"]}
|
|
|
|
agent1 = ChainAgent("chain1", config1, template_renderer)
|
|
agent2 = ChainAgent("chain2", config2, template_renderer)
|
|
|
|
assert agent1.steps == ["a", "b"]
|
|
assert agent2.steps == ["x", "y", "z"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_message_context_without_prompt(self, simple_config, template_renderer):
|
|
"""Test that context is ignored when there's no prompt template."""
|
|
agent = ChainAgent("test_chain", simple_config, template_renderer)
|
|
result = await agent.process_message("test", {"extra": "data"})
|
|
|
|
# Should process normally, ignoring the context
|
|
assert "ChainAgent processed:" in result
|
|
assert "test" in result
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v"])
|
|
|
|
|