forked from cleveragents/cleveragents-core
712 lines
25 KiB
Python
712 lines
25 KiB
Python
"""
|
|
Unit tests for agents/composite.py
|
|
|
|
Tests the CompositeAgent class.
|
|
"""
|
|
|
|
import pytest
|
|
import asyncio
|
|
from typing import Any, Dict, List, Optional
|
|
from unittest.mock import Mock, AsyncMock
|
|
|
|
from cleveragents.agents.composite import CompositeAgent
|
|
from cleveragents.agents.base import Agent
|
|
from cleveragents.core.exceptions import ConfigurationError, ExecutionError
|
|
from cleveragents.templates.renderer import TemplateRenderer
|
|
|
|
|
|
# Test helper agent
|
|
class MockAgent(Agent):
|
|
"""Mock agent for testing."""
|
|
|
|
def __init__(self, name: str, config: Dict[str, Any], template_renderer: TemplateRenderer):
|
|
super().__init__(name, config, template_renderer)
|
|
self._response = "mock response"
|
|
|
|
async def process_message(self, message: str, context: Optional[Dict[str, Any]] = None) -> str:
|
|
return f"MockAgent {self.name}: {message}"
|
|
|
|
async def process(self, message: str, context: Optional[Dict[str, Any]] = None) -> str:
|
|
return f"MockAgent {self.name}: {message}"
|
|
|
|
def get_capabilities(self) -> List[str]:
|
|
return ["mock"]
|
|
|
|
|
|
class TestCompositeAgent:
|
|
"""Test suite for the CompositeAgent class."""
|
|
|
|
@pytest.fixture
|
|
def template_renderer(self):
|
|
"""Create a basic template renderer."""
|
|
return TemplateRenderer()
|
|
|
|
def _setup_stream_composite(self, template_renderer):
|
|
"""Helper method to set up composite agent with stream routing."""
|
|
from rx.subject import Subject
|
|
|
|
config = {
|
|
"type": "composite",
|
|
"routing": {
|
|
"input": {"type": "stream", "name": "stream1"},
|
|
"output": {"name": "__output__"},
|
|
},
|
|
}
|
|
|
|
output_subject = Subject()
|
|
|
|
mock_router = Mock()
|
|
mock_router.send_message = Mock()
|
|
mock_router.observables = {"__output__": output_subject}
|
|
|
|
composite = CompositeAgent(
|
|
"test_composite",
|
|
config,
|
|
template_renderer,
|
|
stream_router=mock_router
|
|
)
|
|
composite.add_stream("stream1", Mock())
|
|
|
|
return composite, output_subject, mock_router
|
|
|
|
@pytest.fixture
|
|
def simple_config(self):
|
|
"""Create a simple composite configuration."""
|
|
return {
|
|
"type": "composite",
|
|
"components": {},
|
|
"routing": {},
|
|
}
|
|
|
|
@pytest.fixture
|
|
def config_with_agents(self):
|
|
"""Create a composite configuration with agents."""
|
|
return {
|
|
"type": "composite",
|
|
"components": {
|
|
"agents": {
|
|
"agent1": {"type": "mock"},
|
|
"agent2": {"type": "mock"},
|
|
}
|
|
},
|
|
"routing": {
|
|
"input": {"type": "agent", "name": "agent1"},
|
|
"output": {"name": "agent2"},
|
|
},
|
|
}
|
|
|
|
def test_composite_agent_initialization(self, simple_config, template_renderer):
|
|
"""Test CompositeAgent initialization."""
|
|
agent = CompositeAgent("test_composite", simple_config, template_renderer)
|
|
|
|
assert agent.name == "test_composite"
|
|
assert agent.components == {}
|
|
assert agent.routing == {}
|
|
assert agent.expose_params == {}
|
|
assert agent.agents == {}
|
|
assert agent.graphs == {}
|
|
assert agent.streams == {}
|
|
|
|
def test_composite_agent_initialization_with_components(self, config_with_agents, template_renderer):
|
|
"""Test CompositeAgent initialization with components."""
|
|
agent = CompositeAgent("test_composite", config_with_agents, template_renderer)
|
|
|
|
assert "agents" in agent.components
|
|
assert "input" in agent.routing
|
|
|
|
def test_composite_agent_initialization_with_stream_router(self, simple_config, template_renderer):
|
|
"""Test CompositeAgent initialization with stream router."""
|
|
mock_router = Mock()
|
|
agent = CompositeAgent(
|
|
"test_composite",
|
|
simple_config,
|
|
template_renderer,
|
|
stream_router=mock_router
|
|
)
|
|
|
|
assert agent.stream_router == mock_router
|
|
|
|
def test_composite_agent_initialization_with_langgraph_bridge(self, simple_config, template_renderer):
|
|
"""Test CompositeAgent initialization with LangGraph bridge."""
|
|
mock_bridge = Mock()
|
|
agent = CompositeAgent(
|
|
"test_composite",
|
|
simple_config,
|
|
template_renderer,
|
|
langgraph_bridge=mock_bridge
|
|
)
|
|
|
|
assert agent.langgraph_bridge == mock_bridge
|
|
|
|
def test_composite_agent_rejects_legacy_strategy(self, template_renderer):
|
|
"""Test that legacy strategy-based config is rejected."""
|
|
config = {
|
|
"type": "composite",
|
|
"strategy": "sequential",
|
|
}
|
|
|
|
with pytest.raises(ConfigurationError) as exc_info:
|
|
CompositeAgent("test_composite", config, template_renderer)
|
|
|
|
assert "deprecated strategy-based configuration" in str(exc_info.value)
|
|
|
|
def test_add_agent(self, simple_config, template_renderer):
|
|
"""Test adding an agent to composite."""
|
|
composite = CompositeAgent("test_composite", simple_config, template_renderer)
|
|
mock_agent = MockAgent("child", {}, template_renderer)
|
|
|
|
composite.add_agent("child", mock_agent)
|
|
|
|
assert "child" in composite.agents
|
|
assert composite.agents["child"] == mock_agent
|
|
assert "agents" in composite.components
|
|
assert "child" in composite.components["agents"]
|
|
|
|
def test_add_graph(self, simple_config, template_renderer):
|
|
"""Test adding a graph to composite."""
|
|
composite = CompositeAgent("test_composite", simple_config, template_renderer)
|
|
mock_graph = Mock()
|
|
|
|
composite.add_graph("graph1", mock_graph)
|
|
|
|
assert "graph1" in composite.graphs
|
|
assert composite.graphs["graph1"] == mock_graph
|
|
assert "graphs" in composite.components
|
|
|
|
def test_add_stream(self, simple_config, template_renderer):
|
|
"""Test adding a stream to composite."""
|
|
composite = CompositeAgent("test_composite", simple_config, template_renderer)
|
|
mock_stream = Mock()
|
|
|
|
composite.add_stream("stream1", mock_stream)
|
|
|
|
assert "stream1" in composite.streams
|
|
assert composite.streams["stream1"] == mock_stream
|
|
assert "streams" in composite.components
|
|
|
|
def test_set_param(self, simple_config, template_renderer):
|
|
"""Test setting parameters."""
|
|
composite = CompositeAgent("test_composite", simple_config, template_renderer)
|
|
|
|
composite.set_param("temperature", 0.7)
|
|
|
|
assert composite.expose_params["temperature"] == 0.7
|
|
|
|
def test_set_multiple_params(self, simple_config, template_renderer):
|
|
"""Test setting multiple parameters."""
|
|
composite = CompositeAgent("test_composite", simple_config, template_renderer)
|
|
|
|
composite.set_param("temperature", 0.7)
|
|
composite.set_param("max_tokens", 1000)
|
|
|
|
assert composite.expose_params["temperature"] == 0.7
|
|
assert composite.expose_params["max_tokens"] == 1000
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_message_via_agent(self, simple_config, template_renderer):
|
|
"""Test processing message via agent."""
|
|
composite = CompositeAgent("test_composite", simple_config, template_renderer)
|
|
mock_agent = MockAgent("child", {}, template_renderer)
|
|
composite.add_agent("child", mock_agent)
|
|
|
|
result = await composite.process_message("test message")
|
|
|
|
assert "MockAgent child" in result
|
|
assert "test message" in result
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_via_agent_with_routing(self, template_renderer):
|
|
"""Test processing with explicit agent routing."""
|
|
config = {
|
|
"type": "composite",
|
|
"routing": {"input": {"type": "agent", "name": "agent1"}},
|
|
}
|
|
composite = CompositeAgent("test_composite", config, template_renderer)
|
|
mock_agent = MockAgent("agent1", {}, template_renderer)
|
|
composite.add_agent("agent1", mock_agent)
|
|
|
|
result = await composite.process("test")
|
|
|
|
assert "MockAgent agent1" in result
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_via_agent_missing(self, template_renderer):
|
|
"""Test processing with missing agent raises error."""
|
|
config = {
|
|
"type": "composite",
|
|
"routing": {"input": {"type": "agent", "name": "missing"}},
|
|
}
|
|
composite = CompositeAgent("test_composite", config, template_renderer)
|
|
|
|
with pytest.raises(ExecutionError) as exc_info:
|
|
await composite.process("test")
|
|
|
|
assert "missing" in str(exc_info.value)
|
|
assert "not found" in str(exc_info.value)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_no_components(self, simple_config, template_renderer):
|
|
"""Test processing with no components raises error."""
|
|
composite = CompositeAgent("test_composite", simple_config, template_renderer)
|
|
|
|
with pytest.raises(ConfigurationError) as exc_info:
|
|
await composite.process("test")
|
|
|
|
assert "no components" in str(exc_info.value)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_with_context(self, simple_config, template_renderer):
|
|
"""Test processing with context."""
|
|
composite = CompositeAgent("test_composite", simple_config, template_renderer)
|
|
mock_agent = MockAgent("child", {}, template_renderer)
|
|
composite.add_agent("child", mock_agent)
|
|
|
|
result = await composite.process("test", {"key": "value"})
|
|
|
|
assert "MockAgent child" in result
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_merges_expose_params(self, simple_config, template_renderer):
|
|
"""Test that exposed params are merged into context."""
|
|
composite = CompositeAgent("test_composite", simple_config, template_renderer)
|
|
composite.set_param("param1", "value1")
|
|
|
|
mock_agent = MockAgent("child", {}, template_renderer)
|
|
composite.add_agent("child", mock_agent)
|
|
|
|
# Should merge expose_params into context
|
|
result = await composite.process("test", {"param2": "value2"})
|
|
|
|
assert "MockAgent child" in result
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_via_graph_no_bridge(self, template_renderer):
|
|
"""Test processing via graph without bridge raises error."""
|
|
config = {
|
|
"type": "composite",
|
|
"routing": {"input": {"type": "graph", "name": "graph1"}},
|
|
}
|
|
composite = CompositeAgent("test_composite", config, template_renderer)
|
|
|
|
with pytest.raises(ExecutionError) as exc_info:
|
|
await composite.process("test")
|
|
|
|
assert "LangGraph bridge not available" in str(exc_info.value)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_via_graph_with_bridge(self, template_renderer):
|
|
"""Test processing via graph with bridge."""
|
|
config = {
|
|
"type": "composite",
|
|
"routing": {"input": {"type": "graph", "name": "graph1"}},
|
|
}
|
|
|
|
mock_bridge = Mock()
|
|
mock_graph = Mock()
|
|
mock_graph.execute = AsyncMock(return_value=Mock(
|
|
messages=[{"role": "assistant", "content": "graph result"}]
|
|
))
|
|
|
|
composite = CompositeAgent(
|
|
"test_composite",
|
|
config,
|
|
template_renderer,
|
|
langgraph_bridge=mock_bridge
|
|
)
|
|
composite.add_graph("graph1", mock_graph)
|
|
|
|
result = await composite.process("test")
|
|
|
|
assert result == "graph result"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_via_graph_from_bridge(self, template_renderer):
|
|
"""Test processing via graph retrieved from bridge."""
|
|
config = {
|
|
"type": "composite",
|
|
"routing": {"input": {"type": "graph", "name": "graph1"}},
|
|
}
|
|
|
|
mock_graph = Mock()
|
|
mock_graph.execute = AsyncMock(return_value=Mock(
|
|
messages=[{"role": "assistant", "content": "bridge graph result"}]
|
|
))
|
|
|
|
mock_bridge = Mock()
|
|
mock_bridge.get_graph = Mock(return_value=mock_graph)
|
|
|
|
composite = CompositeAgent(
|
|
"test_composite",
|
|
config,
|
|
template_renderer,
|
|
langgraph_bridge=mock_bridge
|
|
)
|
|
|
|
result = await composite.process("test")
|
|
|
|
assert result == "bridge graph result"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_via_stream_no_router(self, template_renderer):
|
|
"""Test processing via stream without router raises error."""
|
|
config = {
|
|
"type": "composite",
|
|
"routing": {"input": {"type": "stream", "name": "stream1"}},
|
|
}
|
|
composite = CompositeAgent("test_composite", config, template_renderer)
|
|
|
|
with pytest.raises(ExecutionError) as exc_info:
|
|
await composite.process("test")
|
|
|
|
assert "Stream router not available" in str(exc_info.value)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_unknown_input_type(self, template_renderer):
|
|
"""Test processing with unknown input type raises error."""
|
|
config = {
|
|
"type": "composite",
|
|
"routing": {"input": {"type": "unknown", "name": "test"}},
|
|
}
|
|
composite = CompositeAgent("test_composite", config, template_renderer)
|
|
|
|
with pytest.raises(ConfigurationError) as exc_info:
|
|
await composite.process("test")
|
|
|
|
assert "Unknown input type" in str(exc_info.value)
|
|
assert "unknown" in str(exc_info.value)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_via_agent_directly(self, simple_config, template_renderer):
|
|
"""Test _process_via_agent method directly."""
|
|
composite = CompositeAgent("test_composite", simple_config, template_renderer)
|
|
mock_agent = MockAgent("agent1", {}, template_renderer)
|
|
composite.add_agent("agent1", mock_agent)
|
|
|
|
result = await composite._process_via_agent("agent1", "test", {})
|
|
|
|
assert "MockAgent agent1" in result
|
|
|
|
def test_get_capabilities_empty(self, simple_config, template_renderer):
|
|
"""Test getting capabilities with no components."""
|
|
composite = CompositeAgent("test_composite", simple_config, template_renderer)
|
|
capabilities = composite.get_capabilities()
|
|
|
|
assert "composite" in capabilities
|
|
|
|
def test_get_capabilities_with_agents(self, simple_config, template_renderer):
|
|
"""Test getting capabilities includes agent capabilities."""
|
|
composite = CompositeAgent("test_composite", simple_config, template_renderer)
|
|
mock_agent = MockAgent("child", {}, template_renderer)
|
|
composite.add_agent("child", mock_agent)
|
|
|
|
capabilities = composite.get_capabilities()
|
|
|
|
assert "composite" in capabilities
|
|
assert "mock" in capabilities
|
|
|
|
def test_get_capabilities_with_graphs(self, simple_config, template_renderer):
|
|
"""Test getting capabilities with graphs."""
|
|
composite = CompositeAgent("test_composite", simple_config, template_renderer)
|
|
composite.add_graph("graph1", Mock())
|
|
|
|
capabilities = composite.get_capabilities()
|
|
|
|
assert "composite" in capabilities
|
|
assert "stateful-workflow" in capabilities
|
|
|
|
def test_get_capabilities_with_streams(self, simple_config, template_renderer):
|
|
"""Test getting capabilities with streams."""
|
|
composite = CompositeAgent("test_composite", simple_config, template_renderer)
|
|
composite.add_stream("stream1", Mock())
|
|
|
|
capabilities = composite.get_capabilities()
|
|
|
|
assert "composite" in capabilities
|
|
assert "reactive-processing" in capabilities
|
|
|
|
def test_get_capabilities_no_duplicates(self, simple_config, template_renderer):
|
|
"""Test that capabilities removes duplicates."""
|
|
composite = CompositeAgent("test_composite", simple_config, template_renderer)
|
|
|
|
# Add multiple agents with same capability
|
|
agent1 = MockAgent("agent1", {}, template_renderer)
|
|
agent2 = MockAgent("agent2", {}, template_renderer)
|
|
composite.add_agent("agent1", agent1)
|
|
composite.add_agent("agent2", agent2)
|
|
|
|
capabilities = composite.get_capabilities()
|
|
|
|
# Should only have one instance of "mock"
|
|
assert capabilities.count("mock") == 1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_message_calls_process(self, simple_config, template_renderer):
|
|
"""Test that process_message delegates to process."""
|
|
composite = CompositeAgent("test_composite", simple_config, template_renderer)
|
|
mock_agent = MockAgent("child", {}, template_renderer)
|
|
composite.add_agent("child", mock_agent)
|
|
|
|
result1 = await composite.process_message("test")
|
|
result2 = await composite.process("test")
|
|
|
|
# Both should produce same result
|
|
assert result1 == result2
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_defaults_to_first_graph(self, template_renderer):
|
|
"""Test that processing defaults to first graph when no routing specified."""
|
|
config = {"type": "composite"}
|
|
|
|
mock_graph = Mock()
|
|
mock_graph.execute = AsyncMock(return_value="graph result")
|
|
|
|
mock_bridge = Mock()
|
|
composite = CompositeAgent(
|
|
"test_composite",
|
|
config,
|
|
template_renderer,
|
|
langgraph_bridge=mock_bridge
|
|
)
|
|
composite.add_graph("graph1", mock_graph)
|
|
|
|
result = await composite.process("test")
|
|
|
|
# Should process via first graph
|
|
assert result == "graph result"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_graph_non_dict_result(self, template_renderer):
|
|
"""Test processing graph with non-dict result."""
|
|
config = {
|
|
"type": "composite",
|
|
"routing": {"input": {"type": "graph", "name": "graph1"}},
|
|
}
|
|
|
|
mock_graph = Mock()
|
|
mock_graph.execute = AsyncMock(return_value="simple string result")
|
|
|
|
composite = CompositeAgent(
|
|
"test_composite",
|
|
config,
|
|
template_renderer,
|
|
langgraph_bridge=Mock()
|
|
)
|
|
composite.add_graph("graph1", mock_graph)
|
|
|
|
result = await composite.process("test")
|
|
|
|
assert result == "simple string result"
|
|
|
|
def test_add_graph_initializes_components_dict(self, template_renderer):
|
|
"""Test that add_graph initializes components['graphs'] if not present."""
|
|
config = {"type": "composite"}
|
|
|
|
composite = CompositeAgent(
|
|
"test_composite",
|
|
config,
|
|
template_renderer,
|
|
langgraph_bridge=Mock()
|
|
)
|
|
|
|
# Initially components dict doesn't have 'graphs'
|
|
assert "graphs" not in composite.components
|
|
|
|
mock_graph = Mock()
|
|
composite.add_graph("test_graph", mock_graph)
|
|
|
|
# Now it should be initialized
|
|
assert "graphs" in composite.components
|
|
assert "test_graph" in composite.components["graphs"]
|
|
assert composite.components["graphs"]["test_graph"] == mock_graph
|
|
|
|
def test_add_stream_initializes_components_dict(self, template_renderer):
|
|
"""Test that add_stream initializes components['streams'] if not present."""
|
|
config = {"type": "composite"}
|
|
|
|
composite = CompositeAgent(
|
|
"test_composite",
|
|
config,
|
|
template_renderer,
|
|
stream_router=Mock()
|
|
)
|
|
|
|
# Initially components dict doesn't have 'streams'
|
|
assert "streams" not in composite.components
|
|
|
|
mock_stream = Mock()
|
|
composite.add_stream("test_stream", mock_stream)
|
|
|
|
# Now it should be initialized
|
|
assert "streams" in composite.components
|
|
assert "test_stream" in composite.components["streams"]
|
|
assert composite.components["streams"]["test_stream"] == mock_stream
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_defaults_to_first_stream(self, template_renderer):
|
|
"""Test that processing defaults to first stream when no routing specified."""
|
|
config = {"type": "composite"}
|
|
|
|
mock_router = Mock()
|
|
mock_router.send_message = Mock()
|
|
mock_router.observables = {"__output__": Mock()}
|
|
|
|
from rx.subject import Subject
|
|
output_subject = Subject()
|
|
mock_router.observables["__output__"] = output_subject
|
|
|
|
composite = CompositeAgent(
|
|
"test_composite",
|
|
config,
|
|
template_renderer,
|
|
stream_router=mock_router
|
|
)
|
|
composite.add_stream("stream1", Mock())
|
|
|
|
# Start the coroutine but don't wait for it to complete
|
|
# since we need to timeout
|
|
try:
|
|
await asyncio.wait_for(
|
|
composite.process("test"),
|
|
timeout=0.1
|
|
)
|
|
except asyncio.TimeoutError:
|
|
# Expected - we're testing routing, not full execution
|
|
pass
|
|
|
|
# Verify send_message was called
|
|
mock_router.send_message.assert_called()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_via_graph_missing_from_bridge(self, template_renderer):
|
|
"""Test processing via graph when graph not in bridge."""
|
|
config = {
|
|
"type": "composite",
|
|
"routing": {"input": {"type": "graph", "name": "missing_graph"}},
|
|
}
|
|
|
|
mock_bridge = Mock()
|
|
mock_bridge.get_graph = Mock(return_value=None)
|
|
|
|
composite = CompositeAgent(
|
|
"test_composite",
|
|
config,
|
|
template_renderer,
|
|
langgraph_bridge=mock_bridge
|
|
)
|
|
|
|
with pytest.raises(ExecutionError) as exc_info:
|
|
await composite.process("test")
|
|
|
|
assert "missing_graph" in str(exc_info.value)
|
|
assert "not found" in str(exc_info.value)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_graph_with_non_string_content(self, template_renderer):
|
|
"""Test processing graph result with non-string content."""
|
|
config = {
|
|
"type": "composite",
|
|
"routing": {"input": {"type": "graph", "name": "graph1"}},
|
|
}
|
|
|
|
mock_graph = Mock()
|
|
# Return result with non-string content
|
|
mock_graph.execute = AsyncMock(return_value=Mock(
|
|
messages=[{"role": "assistant", "content": 12345}] # Integer content
|
|
))
|
|
|
|
composite = CompositeAgent(
|
|
"test_composite",
|
|
config,
|
|
template_renderer,
|
|
langgraph_bridge=Mock()
|
|
)
|
|
composite.add_graph("graph1", mock_graph)
|
|
|
|
result = await composite.process("test")
|
|
|
|
# Should convert to string
|
|
assert result == "12345"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_via_stream_with_timeout(self, template_renderer):
|
|
"""Test processing via stream with timeout."""
|
|
composite, _, _ = self._setup_stream_composite(template_renderer)
|
|
|
|
# Process with timeout - should timeout since we don't send result
|
|
with pytest.raises(asyncio.TimeoutError):
|
|
await asyncio.wait_for(composite.process("test"), timeout=0.1)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_via_stream_completes(self, template_renderer):
|
|
"""Test successful processing via stream."""
|
|
composite, output_subject, _ = self._setup_stream_composite(template_renderer)
|
|
|
|
# Start processing in background
|
|
async def process_and_send():
|
|
await asyncio.sleep(0.01)
|
|
# Send result to output
|
|
output_subject.on_next(Mock(content="stream result"))
|
|
|
|
task = asyncio.create_task(process_and_send())
|
|
|
|
try:
|
|
result = await asyncio.wait_for(composite.process("test"), timeout=0.1)
|
|
assert result == "stream result"
|
|
finally:
|
|
task.cancel()
|
|
try:
|
|
await task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_process_via_stream_with_message_object(self, template_renderer):
|
|
"""Test processing via stream with message that has content attribute."""
|
|
composite, output_subject, _ = self._setup_stream_composite(template_renderer)
|
|
|
|
# Create message with content attribute
|
|
class Message:
|
|
def __init__(self, content):
|
|
self.content = content
|
|
|
|
# Start processing
|
|
async def send_result():
|
|
await asyncio.sleep(0.01)
|
|
output_subject.on_next(Message("message content"))
|
|
|
|
task = asyncio.create_task(send_result())
|
|
|
|
try:
|
|
result = await asyncio.wait_for(composite.process("test"), timeout=0.1)
|
|
assert result == "message content"
|
|
finally:
|
|
task.cancel()
|
|
try:
|
|
await task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
def test_get_capabilities_with_legacy_strategy(self, template_renderer):
|
|
"""Test getting capabilities with legacy strategy attribute."""
|
|
config = {"type": "composite"}
|
|
|
|
composite = CompositeAgent(
|
|
"test_composite",
|
|
config,
|
|
template_renderer
|
|
)
|
|
|
|
# Add legacy strategy attribute
|
|
composite.strategy = "sequential"
|
|
|
|
capabilities = composite.get_capabilities()
|
|
|
|
assert "composite" in capabilities
|
|
assert "sequential" in capabilities
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v"])
|
|
|
|
|