mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-05-03 08:12:39 +00:00
Gl/feat/a2a refactor (#3793)
* feat: agent metaclass, refactor a2a to wrappers * feat: a2a schemas and utils * chore: move agent class, update imports * refactor: organize imports to avoid circularity, add a2a to console * feat: pass response_model through call chain * feat: add standard openapi spec serialization to tools and structured output * feat: a2a events * chore: add a2a to pyproject * docs: minimal base for learn docs * fix: adjust a2a conversation flow, allow llm to decide exit until max_retries * fix: inject agent skills into initial prompt * fix: format agent card as json in prompt * refactor: simplify A2A agent prompt formatting and improve skill display * chore: wide cleanup * chore: cleanup logic, add auth cache, use json for messages in prompt * chore: update docs * fix: doc snippets formatting * feat: optimize A2A agent card fetching and improve error reporting * chore: move imports to top of file * chore: refactor hasattr check * chore: add httpx-auth, update lockfile * feat: create base public api * chore: cleanup modules, add docstrings, types * fix: exclude extra fields in prompt * chore: update docs * tests: update to correct import * chore: lint for ruff, add missing import * fix: tweak openai streaming logic for response model * tests: add reimport for test * tests: add reimport for test * fix: don't set a2a attr if not set * fix: don't set a2a attr if not set * chore: update cassettes * tests: fix tests * fix: use instructor and dont pass response_format for litellm * chore: consolidate event listeners, add typing * fix: address race condition in test, update cassettes * tests: add correct mocks, rerun cassette for json * tests: update cassette * chore: regenerate cassette after new run * fix: make token manager access-safe * fix: make token manager access-safe * merge * chore: update test and cassete for output pydantic * fix: tweak to disallow deadlock * chore: linter * fix: adjust event ordering for threading * fix: use conditional for batch check * tests: tweak for emission * tests: simplify api + event check * fix: ensure non-function calling llms see json formatted string * tests: tweak message comparison * fix: use internal instructor for litellm structure responses --------- Co-authored-by: Mike Plachta <mike@crewai.com>
This commit is contained in:
Greyson LaLonde
@@ -57,6 +57,7 @@ class TestTraceListenerSetup:
|
||||
if hasattr(TraceCollectionListener, "_instance"):
|
||||
TraceCollectionListener._instance = None
|
||||
TraceCollectionListener._initialized = False
|
||||
TraceCollectionListener._listeners_setup = False
|
||||
|
||||
# Reset EventListener singleton
|
||||
if hasattr(EventListener, "_instance"):
|
||||
@@ -74,6 +75,7 @@ class TestTraceListenerSetup:
|
||||
if hasattr(TraceCollectionListener, "_instance"):
|
||||
TraceCollectionListener._instance = None
|
||||
TraceCollectionListener._initialized = False
|
||||
TraceCollectionListener._listeners_setup = False
|
||||
|
||||
if hasattr(EventListener, "_instance"):
|
||||
EventListener._instance = None
|
||||
@@ -131,25 +133,16 @@ class TestTraceListenerSetup:
|
||||
)
|
||||
crew = Crew(agents=[agent], tasks=[task], verbose=True)
|
||||
|
||||
from crewai.events.listeners.tracing.trace_listener import TraceCollectionListener
|
||||
trace_listener = TraceCollectionListener()
|
||||
from crewai.events.event_bus import crewai_event_bus
|
||||
|
||||
trace_listener.setup_listeners(crewai_event_bus)
|
||||
crew.kickoff()
|
||||
|
||||
with patch.object(
|
||||
trace_listener.batch_manager,
|
||||
"initialize_batch",
|
||||
return_value=None,
|
||||
) as initialize_mock:
|
||||
crew.kickoff()
|
||||
initialized = trace_listener.batch_manager.wait_for_batch_initialization(timeout=5.0)
|
||||
|
||||
assert initialize_mock.call_count >= 1
|
||||
|
||||
call_args = initialize_mock.call_args_list[0]
|
||||
assert len(call_args[0]) == 2 # user_context, execution_metadata
|
||||
_, execution_metadata = call_args[0]
|
||||
assert isinstance(execution_metadata, dict)
|
||||
assert "crew_name" in execution_metadata
|
||||
assert initialized, "Batch should have been initialized"
|
||||
assert trace_listener.batch_manager.is_batch_initialized()
|
||||
assert trace_listener.batch_manager.current_batch is not None
|
||||
|
||||
@pytest.mark.vcr(filter_headers=["authorization"])
|
||||
def test_batch_manager_finalizes_batch_clears_buffer(self):
|
||||
@@ -364,24 +357,21 @@ class TestTraceListenerSetup:
|
||||
)
|
||||
crew = Crew(agents=[agent], tasks=[task], tracing=True)
|
||||
|
||||
with patch.object(TraceBatchManager, "initialize_batch") as mock_initialize:
|
||||
crew.kickoff()
|
||||
from crewai.events.listeners.tracing.trace_listener import TraceCollectionListener
|
||||
trace_listener = TraceCollectionListener()
|
||||
|
||||
assert mock_initialize.call_count >= 1
|
||||
assert mock_initialize.call_args_list[0][1]["use_ephemeral"] is True
|
||||
crew.kickoff()
|
||||
|
||||
wait_for_event_handlers()
|
||||
|
||||
assert trace_listener.batch_manager.is_batch_initialized(), (
|
||||
"Batch should have been initialized for unauthenticated user"
|
||||
)
|
||||
|
||||
@pytest.mark.vcr(filter_headers=["authorization"])
|
||||
def test_trace_listener_with_authenticated_user(self):
|
||||
"""Test that trace listener properly handles authenticated batches"""
|
||||
with (
|
||||
patch.dict(os.environ, {"CREWAI_TRACING_ENABLED": "true"}),
|
||||
patch(
|
||||
"crewai.events.listeners.tracing.trace_batch_manager.PlusAPI"
|
||||
) as mock_plus_api_class,
|
||||
):
|
||||
mock_plus_api_instance = MagicMock()
|
||||
mock_plus_api_class.return_value = mock_plus_api_instance
|
||||
|
||||
with patch.dict(os.environ, {"CREWAI_TRACING_ENABLED": "true"}):
|
||||
agent = Agent(
|
||||
role="Test Agent",
|
||||
goal="Test goal",
|
||||
@@ -394,21 +384,17 @@ class TestTraceListenerSetup:
|
||||
agent=agent,
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object(TraceBatchManager, "initialize_batch") as mock_initialize,
|
||||
patch.object(
|
||||
TraceBatchManager, "finalize_batch"
|
||||
) as mock_finalize_backend_batch,
|
||||
):
|
||||
crew = Crew(agents=[agent], tasks=[task], tracing=True)
|
||||
crew.kickoff()
|
||||
wait_for_event_handlers()
|
||||
from crewai.events.listeners.tracing.trace_listener import TraceCollectionListener
|
||||
trace_listener = TraceCollectionListener()
|
||||
|
||||
mock_plus_api_class.assert_called_with(api_key="mock_token_12345")
|
||||
crew = Crew(agents=[agent], tasks=[task], tracing=True)
|
||||
crew.kickoff()
|
||||
|
||||
assert mock_initialize.call_count >= 1
|
||||
mock_finalize_backend_batch.assert_called_with()
|
||||
assert mock_finalize_backend_batch.call_count >= 1
|
||||
wait_for_event_handlers()
|
||||
|
||||
assert trace_listener.batch_manager.is_batch_initialized(), (
|
||||
"Batch should have been initialized for authenticated user"
|
||||
)
|
||||
|
||||
# Helper method to ensure cleanup
|
||||
def teardown_method(self):
|
||||
@@ -489,30 +475,19 @@ class TestTraceListenerSetup:
|
||||
assert trace_listener.first_time_handler.is_first_time is True
|
||||
assert trace_listener.first_time_handler.collected_events is False
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
trace_listener.first_time_handler,
|
||||
"handle_execution_completion",
|
||||
wraps=trace_listener.first_time_handler.handle_execution_completion,
|
||||
) as mock_handle_completion,
|
||||
patch.object(
|
||||
trace_listener.batch_manager,
|
||||
"add_event",
|
||||
wraps=trace_listener.batch_manager.add_event,
|
||||
) as mock_add_event,
|
||||
):
|
||||
result = crew.kickoff()
|
||||
wait_for_event_handlers()
|
||||
assert result is not None
|
||||
trace_listener.batch_manager.batch_owner_type = "crew"
|
||||
|
||||
assert mock_handle_completion.call_count >= 1
|
||||
assert mock_add_event.call_count >= 1
|
||||
result = crew.kickoff()
|
||||
wait_for_event_handlers()
|
||||
assert result is not None
|
||||
|
||||
assert trace_listener.first_time_handler.collected_events is True
|
||||
assert trace_listener.first_time_handler.collected_events is True, (
|
||||
"Events should have been collected"
|
||||
)
|
||||
|
||||
mock_prompt.assert_called_once()
|
||||
mock_prompt.assert_called_once()
|
||||
|
||||
mock_mark_completed.assert_called_once()
|
||||
mock_mark_completed.assert_called_once()
|
||||
|
||||
@pytest.mark.vcr(filter_headers=["authorization"])
|
||||
def test_first_time_user_trace_collection_user_accepts(self, mock_plus_api_calls):
|
||||
@@ -556,9 +531,10 @@ class TestTraceListenerSetup:
|
||||
from crewai.events.event_bus import crewai_event_bus
|
||||
|
||||
trace_listener = TraceCollectionListener()
|
||||
trace_listener.setup_listeners(crewai_event_bus)
|
||||
|
||||
assert trace_listener.first_time_handler.is_first_time is True
|
||||
trace_listener.batch_manager.ephemeral_trace_url = (
|
||||
"https://crewai.com/trace/mock-id"
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
@@ -569,26 +545,17 @@ class TestTraceListenerSetup:
|
||||
patch.object(
|
||||
trace_listener.first_time_handler, "_display_ephemeral_trace_link"
|
||||
) as mock_display_link,
|
||||
patch.object(
|
||||
trace_listener.first_time_handler,
|
||||
"handle_execution_completion",
|
||||
wraps=trace_listener.first_time_handler.handle_execution_completion,
|
||||
) as mock_handle_completion,
|
||||
):
|
||||
trace_listener.batch_manager.ephemeral_trace_url = (
|
||||
"https://crewai.com/trace/mock-id"
|
||||
)
|
||||
trace_listener.setup_listeners(crewai_event_bus)
|
||||
|
||||
assert trace_listener.first_time_handler.is_first_time is True
|
||||
|
||||
trace_listener.first_time_handler.collected_events = True
|
||||
|
||||
crew.kickoff()
|
||||
wait_for_event_handlers()
|
||||
|
||||
assert mock_handle_completion.call_count >= 1, (
|
||||
"handle_execution_completion should be called"
|
||||
)
|
||||
|
||||
assert trace_listener.first_time_handler.collected_events is True, (
|
||||
"Events should be marked as collected"
|
||||
)
|
||||
trace_listener.first_time_handler.handle_execution_completion()
|
||||
|
||||
mock_init_backend.assert_called_once()
|
||||
|
||||
@@ -636,15 +603,14 @@ class TestTraceListenerSetup:
|
||||
)
|
||||
crew = Crew(agents=[agent], tasks=[task])
|
||||
|
||||
with patch.object(TraceBatchManager, "initialize_batch") as mock_initialize:
|
||||
result = crew.kickoff()
|
||||
result = crew.kickoff()
|
||||
|
||||
assert trace_listener.batch_manager.wait_for_pending_events(timeout=5.0), (
|
||||
"Timeout waiting for trace event handlers to complete"
|
||||
)
|
||||
assert mock_initialize.call_count >= 1
|
||||
assert mock_initialize.call_args_list[0][1]["use_ephemeral"] is True
|
||||
assert result is not None
|
||||
wait_for_event_handlers()
|
||||
|
||||
assert trace_listener.batch_manager.is_batch_initialized(), (
|
||||
"Batch should have been initialized for first-time user"
|
||||
)
|
||||
assert result is not None
|
||||
|
||||
def test_first_time_handler_timeout_behavior(self):
|
||||
"""Test the timeout behavior of the first-time trace prompt"""
|
||||
@@ -699,60 +665,43 @@ class TestTraceListenerSetup:
|
||||
|
||||
mock_mark_completed.assert_called_once()
|
||||
|
||||
@pytest.mark.vcr(filter_headers=["authorization"])
|
||||
def test_trace_batch_marked_as_failed_on_finalize_error(self, mock_plus_api_calls):
|
||||
def test_trace_batch_marked_as_failed_on_finalize_error(self):
|
||||
"""Test that trace batch is marked as failed when finalization returns non-200 status"""
|
||||
# Test the error handling logic directly in TraceBatchManager
|
||||
batch_manager = TraceBatchManager()
|
||||
|
||||
with patch.dict(os.environ, {"CREWAI_TRACING_ENABLED": "true"}):
|
||||
agent = Agent(
|
||||
role="Test Agent",
|
||||
goal="Test goal",
|
||||
backstory="Test backstory",
|
||||
llm="gpt-4o-mini",
|
||||
# Initialize a batch
|
||||
batch_manager.current_batch = batch_manager.initialize_batch(
|
||||
user_context={"privacy_level": "standard"},
|
||||
execution_metadata={
|
||||
"execution_type": "crew",
|
||||
"crew_name": "test_crew",
|
||||
},
|
||||
)
|
||||
batch_manager.trace_batch_id = "test_batch_id_12345"
|
||||
batch_manager.backend_initialized = True
|
||||
|
||||
# Mock the API responses
|
||||
with (
|
||||
patch.object(
|
||||
batch_manager.plus_api,
|
||||
"send_trace_events",
|
||||
return_value=MagicMock(status_code=200),
|
||||
),
|
||||
patch.object(
|
||||
batch_manager.plus_api,
|
||||
"finalize_trace_batch",
|
||||
return_value=MagicMock(status_code=500, text="Internal Server Error"),
|
||||
),
|
||||
patch.object(
|
||||
batch_manager.plus_api,
|
||||
"mark_trace_batch_as_failed",
|
||||
) as mock_mark_failed,
|
||||
):
|
||||
# Call finalize_batch directly
|
||||
batch_manager.finalize_batch()
|
||||
|
||||
# Verify that mark_trace_batch_as_failed was called with the error message
|
||||
mock_mark_failed.assert_called_once_with(
|
||||
"test_batch_id_12345", "Internal Server Error"
|
||||
)
|
||||
task = Task(
|
||||
description="Say hello to the world",
|
||||
expected_output="hello world",
|
||||
agent=agent,
|
||||
)
|
||||
crew = Crew(agents=[agent], tasks=[task], verbose=True)
|
||||
|
||||
trace_listener = TraceCollectionListener()
|
||||
from crewai.events.event_bus import crewai_event_bus
|
||||
|
||||
trace_listener.setup_listeners(crewai_event_bus)
|
||||
|
||||
mock_init_response = MagicMock()
|
||||
mock_init_response.status_code = 200
|
||||
mock_init_response.json.return_value = {"trace_id": "test_batch_id_12345"}
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
trace_listener.batch_manager.plus_api,
|
||||
"initialize_trace_batch",
|
||||
return_value=mock_init_response,
|
||||
),
|
||||
patch.object(
|
||||
trace_listener.batch_manager.plus_api,
|
||||
"send_trace_events",
|
||||
return_value=MagicMock(status_code=200),
|
||||
),
|
||||
patch.object(
|
||||
trace_listener.batch_manager.plus_api,
|
||||
"finalize_trace_batch",
|
||||
return_value=MagicMock(
|
||||
status_code=500, text="Internal Server Error"
|
||||
),
|
||||
),
|
||||
patch.object(
|
||||
trace_listener.batch_manager.plus_api,
|
||||
"mark_trace_batch_as_failed",
|
||||
wraps=mock_plus_api_calls["mark_trace_batch_as_failed"],
|
||||
) as mock_mark_failed,
|
||||
):
|
||||
crew.kickoff()
|
||||
wait_for_event_handlers()
|
||||
|
||||
mock_mark_failed.assert_called_once()
|
||||
call_args = mock_mark_failed.call_args_list[0]
|
||||
assert call_args[0][1] == "Internal Server Error"
|
||||
|
||||
Reference in New Issue
Block a user