mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-08 23:58:34 +00:00
134 lines
4.8 KiB
Python
134 lines
4.8 KiB
Python
from typing import Type
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
from pydantic import BaseModel, Field
|
|
|
|
from crewai import Agent, Crew, Flow, Task
|
|
from crewai.flow import listen, start
|
|
from crewai.tools import BaseTool
|
|
|
|
|
|
class TestToolInput(BaseModel):
|
|
query: str = Field(..., description='Query to process')
|
|
|
|
|
|
class TestTool(BaseTool):
|
|
name: str = 'Test Tool'
|
|
description: str = 'A test tool to demonstrate the issue'
|
|
args_schema: Type[BaseModel] = TestToolInput
|
|
result_as_answer: bool = True
|
|
|
|
def _run(self, query: str) -> str:
|
|
return f'Result for query: {query}'
|
|
|
|
|
|
class TestFlowToolResult:
|
|
"""Test suite for Flow tool result handling."""
|
|
|
|
@pytest.fixture
|
|
def test_tool(self):
|
|
"""Create a test tool with result_as_answer=True."""
|
|
return TestTool()
|
|
|
|
@pytest.fixture
|
|
def test_agent(self, test_tool):
|
|
"""Create a test agent with the tool."""
|
|
return Agent(
|
|
role='Tester',
|
|
goal='Test tools',
|
|
backstory='Testing tools in Flow vs Crew',
|
|
tools=[test_tool]
|
|
)
|
|
|
|
@pytest.fixture
|
|
def test_task(self, test_agent, test_tool):
|
|
"""Create a task with the tool."""
|
|
return Task(
|
|
description='Test task using the tool',
|
|
expected_output='Test output',
|
|
agent=test_agent,
|
|
tools=[test_tool]
|
|
)
|
|
|
|
def test_flow_tool_result_as_answer(self, test_agent, test_task):
|
|
"""Test that tools with result_as_answer=True are properly processed in Flow mode."""
|
|
# Create a simple Flow with direct access to the agent
|
|
class SimpleFlow(Flow):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.test_agent = test_agent
|
|
|
|
@start()
|
|
def start_task(self):
|
|
return 'Task started'
|
|
|
|
@listen('start_task')
|
|
@Flow.with_agent(test_agent) # Associate the agent with this method
|
|
def execute_task(self):
|
|
# Simulate tool execution and setting tools_results
|
|
self.test_agent.tools_results = [{
|
|
"name": "Test Tool",
|
|
"input": {"query": "test"},
|
|
"result": "Result for query: test",
|
|
"result_as_answer": True
|
|
}]
|
|
return "Agent task execution result"
|
|
|
|
# Create a mock for Crew to return the same result
|
|
with patch('crewai.crew.Crew.kickoff') as mock_crew_kickoff:
|
|
mock_crew_kickoff.return_value = "Result for query: test"
|
|
|
|
# Test Flow
|
|
flow = SimpleFlow()
|
|
flow_result = flow.kickoff()
|
|
|
|
# Verify that Flow returns the tool result with our fix
|
|
assert flow_result == "Result for query: test", "Flow should return tool result with result_as_answer=True"
|
|
|
|
# Test Crew
|
|
crew = Crew(agents=[test_agent], tasks=[test_task])
|
|
crew_result = crew.kickoff()
|
|
|
|
# Verify that Crew returns the tool result
|
|
assert crew_result == "Result for query: test", "Crew should return tool result with result_as_answer=True"
|
|
|
|
def test_flow_multiple_tool_results(self, test_agent, test_task):
|
|
"""Test handling of multiple tool results with result_as_answer=True."""
|
|
# Create a simple Flow with direct access to the agent
|
|
class SimpleFlow(Flow):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.test_agent = test_agent
|
|
|
|
@start()
|
|
def start_task(self):
|
|
return 'Task started'
|
|
|
|
@listen('start_task')
|
|
@Flow.with_agent(test_agent)
|
|
def execute_task(self):
|
|
# Simulate multiple tool executions with result_as_answer=True
|
|
self.test_agent.tools_results = [
|
|
{
|
|
"name": "Test Tool 1",
|
|
"input": {"query": "test1"},
|
|
"result": "Result for query: test1",
|
|
"result_as_answer": True
|
|
},
|
|
{
|
|
"name": "Test Tool 2",
|
|
"input": {"query": "test2"},
|
|
"result": "Result for query: test2",
|
|
"result_as_answer": True
|
|
}
|
|
]
|
|
return "Agent task execution result"
|
|
|
|
# Test Flow with multiple tool results
|
|
flow = SimpleFlow()
|
|
flow_result = flow.kickoff()
|
|
|
|
# Verify that Flow returns the first tool result with result_as_answer=True
|
|
assert flow_result == "Result for query: test1", "Flow should return the first tool result with result_as_answer=True"
|