mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-09 08:08:32 +00:00
84 lines
2.7 KiB
Python
84 lines
2.7 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}'
|
|
|
|
def test_flow_tool_result_as_answer():
|
|
"""Test that tools with result_as_answer=True are properly processed in Flow mode."""
|
|
# Create a test tool
|
|
test_tool = TestTool()
|
|
|
|
# Create a test agent with the tool
|
|
agent = Agent(
|
|
role='Tester',
|
|
goal='Test tools',
|
|
backstory='Testing tools in Flow vs Crew',
|
|
tools=[test_tool]
|
|
)
|
|
|
|
# Create a task with the tool
|
|
task = Task(
|
|
description='Test task using the tool',
|
|
expected_output='Test output',
|
|
agent=agent,
|
|
tools=[test_tool]
|
|
)
|
|
|
|
# Create a simple Flow with direct access to the agent
|
|
class SimpleFlow(Flow):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.test_agent = agent
|
|
|
|
@start()
|
|
def start_task(self):
|
|
return 'Task started'
|
|
|
|
@listen('start_task')
|
|
@Flow.with_agent(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=[agent], tasks=[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"
|