From e041103c042c2665ed77119863a5a8288ffbd489 Mon Sep 17 00:00:00 2001
From: Brandon Hancock <brandon@brandonhancock.io>
Date: Thu, 2 Jan 2025 09:56:31 -0500
Subject: [PATCH] Trying out timeouts

---
 tests/agent_test.py | 57 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 57 insertions(+)

diff --git a/tests/agent_test.py b/tests/agent_test.py
index 679f4832e..9e2e2fa66 100644
--- a/tests/agent_test.py
+++ b/tests/agent_test.py
@@ -20,6 +20,7 @@ from crewai.utilities import RPMController
 from crewai.utilities.events import Emitter
 
 
+@pytest.mark.timeout(60)  # 1 minute timeout
 def test_agent_llm_creation_with_env_vars():
     # Store original environment variables
     original_api_key = os.environ.get("OPENAI_API_KEY")
@@ -63,6 +64,7 @@ def test_agent_llm_creation_with_env_vars():
         os.environ["OPENAI_MODEL_NAME"] = original_model_name
 
 
+@pytest.mark.timeout(60)
 def test_agent_creation():
     agent = Agent(role="test role", goal="test goal", backstory="test backstory")
 
@@ -72,12 +74,14 @@ def test_agent_creation():
     assert agent.tools == []
 
 
+@pytest.mark.timeout(60)
 def test_agent_default_values():
     agent = Agent(role="test role", goal="test goal", backstory="test backstory")
     assert agent.llm.model == "gpt-4o-mini"
     assert agent.allow_delegation is False
 
 
+@pytest.mark.timeout(60)
 def test_custom_llm():
     agent = Agent(
         role="test role", goal="test goal", backstory="test backstory", llm="gpt-4"
@@ -85,6 +89,7 @@ def test_custom_llm():
     assert agent.llm.model == "gpt-4"
 
 
+@pytest.mark.timeout(60)
 def test_custom_llm_with_langchain():
     from langchain_openai import ChatOpenAI
 
@@ -98,6 +103,7 @@ def test_custom_llm_with_langchain():
     assert agent.llm.model == "gpt-4"
 
 
+@pytest.mark.timeout(60)
 def test_custom_llm_temperature_preservation():
     from langchain_openai import ChatOpenAI
 
@@ -115,6 +121,7 @@ def test_custom_llm_temperature_preservation():
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_agent_execute_task():
     from langchain_openai import ChatOpenAI
 
@@ -144,6 +151,7 @@ def test_agent_execute_task():
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_agent_execution():
     agent = Agent(
         role="test role",
@@ -163,6 +171,7 @@ def test_agent_execution():
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_agent_execution_with_tools():
     @tool
     def multiplier(first_number: int, second_number: int) -> float:
@@ -194,6 +203,7 @@ def test_agent_execution_with_tools():
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_logging_tool_usage():
     @tool
     def multiplier(first_number: int, second_number: int) -> float:
@@ -228,6 +238,7 @@ def test_logging_tool_usage():
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_cache_hitting():
     @tool
     def multiplier(first_number: int, second_number: int) -> float:
@@ -300,6 +311,7 @@ def test_cache_hitting():
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_disabling_cache_for_agent():
     @tool
     def multiplier(first_number: int, second_number: int) -> float:
@@ -364,6 +376,7 @@ def test_disabling_cache_for_agent():
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_agent_execution_with_specific_tools():
     @tool
     def multiplier(first_number: int, second_number: int) -> float:
@@ -387,6 +400,7 @@ def test_agent_execution_with_specific_tools():
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_agent_powered_by_new_o_model_family_that_allows_skipping_tool():
     @tool
     def multiplier(first_number: int, second_number: int) -> float:
@@ -413,6 +427,7 @@ def test_agent_powered_by_new_o_model_family_that_allows_skipping_tool():
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_agent_powered_by_new_o_model_family_that_uses_tool():
     @tool
     def comapny_customer_data() -> float:
@@ -439,6 +454,7 @@ def test_agent_powered_by_new_o_model_family_that_uses_tool():
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_agent_custom_max_iterations():
     @tool
     def get_final_answer() -> float:
@@ -469,6 +485,7 @@ def test_agent_custom_max_iterations():
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_agent_repeated_tool_usage(capsys):
     @tool
     def get_final_answer() -> float:
@@ -506,6 +523,7 @@ def test_agent_repeated_tool_usage(capsys):
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_agent_repeated_tool_usage_check_even_with_disabled_cache(capsys):
     @tool
     def get_final_answer(anything: str) -> float:
@@ -542,6 +560,7 @@ def test_agent_repeated_tool_usage_check_even_with_disabled_cache(capsys):
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_agent_moved_on_after_max_iterations():
     @tool
     def get_final_answer() -> float:
@@ -569,6 +588,7 @@ def test_agent_moved_on_after_max_iterations():
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_agent_respect_the_max_rpm_set(capsys):
     @tool
     def get_final_answer() -> float:
@@ -603,6 +623,7 @@ def test_agent_respect_the_max_rpm_set(capsys):
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_agent_respect_the_max_rpm_set_over_crew_rpm(capsys):
     from unittest.mock import patch
 
@@ -641,6 +662,7 @@ def test_agent_respect_the_max_rpm_set_over_crew_rpm(capsys):
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_agent_without_max_rpm_respet_crew_rpm(capsys):
     from unittest.mock import patch
 
@@ -695,6 +717,7 @@ def test_agent_without_max_rpm_respet_crew_rpm(capsys):
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_agent_error_on_parsing_tool(capsys):
     from unittest.mock import patch
 
@@ -738,6 +761,7 @@ def test_agent_error_on_parsing_tool(capsys):
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_agent_remembers_output_format_after_using_tools_too_many_times():
     from unittest.mock import patch
 
@@ -773,6 +797,7 @@ def test_agent_remembers_output_format_after_using_tools_too_many_times():
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_agent_use_specific_tasks_output_as_context(capsys):
     agent1 = Agent(role="test role", goal="test goal", backstory="test backstory")
     agent2 = Agent(role="test role2", goal="test goal2", backstory="test backstory2")
@@ -800,6 +825,7 @@ def test_agent_use_specific_tasks_output_as_context(capsys):
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_agent_step_callback():
     class StepCallback:
         def callback(self, step):
@@ -834,6 +860,7 @@ def test_agent_step_callback():
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_agent_function_calling_llm():
     llm = "gpt-4o"
 
@@ -880,6 +907,7 @@ def test_agent_function_calling_llm():
         mock_original_tool_calling.assert_called()
 
 
+@pytest.mark.timeout(60)
 def test_agent_count_formatting_error():
     from unittest.mock import patch
 
@@ -900,6 +928,7 @@ def test_agent_count_formatting_error():
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_tool_result_as_answer_is_the_final_answer_for_the_agent():
     from crewai.tools import BaseTool
 
@@ -930,6 +959,7 @@ def test_tool_result_as_answer_is_the_final_answer_for_the_agent():
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_tool_usage_information_is_appended_to_agent():
     from crewai.tools import BaseTool
 
@@ -966,6 +996,7 @@ def test_tool_usage_information_is_appended_to_agent():
     ]
 
 
+@pytest.mark.timeout(60)
 def test_agent_definition_based_on_dict():
     config = {
         "role": "test role",
@@ -985,6 +1016,7 @@ def test_agent_definition_based_on_dict():
 
 # test for human input
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_agent_human_input():
     # Agent configuration
     config = {
@@ -1025,6 +1057,7 @@ def test_agent_human_input():
         assert output.strip().lower() == "hello"  # Final output should be 'Hello'
 
 
+@pytest.mark.timeout(60)
 def test_interpolate_inputs():
     agent = Agent(
         role="{topic} specialist",
@@ -1043,6 +1076,7 @@ def test_interpolate_inputs():
     assert agent.backstory == "I am the master of nothing"
 
 
+@pytest.mark.timeout(60)
 def test_not_using_system_prompt():
     agent = Agent(
         role="{topic} specialist",
@@ -1056,6 +1090,7 @@ def test_not_using_system_prompt():
     assert not agent.agent_executor.prompt.get("system")
 
 
+@pytest.mark.timeout(60)
 def test_using_system_prompt():
     agent = Agent(
         role="{topic} specialist",
@@ -1068,6 +1103,7 @@ def test_using_system_prompt():
     assert agent.agent_executor.prompt.get("system")
 
 
+@pytest.mark.timeout(60)
 def test_system_and_prompt_template():
     agent = Agent(
         role="{topic} specialist",
@@ -1121,6 +1157,7 @@ Thought:<|eot_id|>
 
 
 @patch("crewai.agent.CrewTrainingHandler")
+@pytest.mark.timeout(60)
 def test_agent_training_handler(crew_training_handler):
     task_prompt = "What is 1 + 1?"
     agent = Agent(
@@ -1143,6 +1180,7 @@ def test_agent_training_handler(crew_training_handler):
 
 
 @patch("crewai.agent.CrewTrainingHandler")
+@pytest.mark.timeout(60)
 def test_agent_use_trained_data(crew_training_handler):
     task_prompt = "What is 1 + 1?"
     agent = Agent(
@@ -1171,6 +1209,7 @@ def test_agent_use_trained_data(crew_training_handler):
     )
 
 
+@pytest.mark.timeout(60)
 def test_agent_max_retry_limit():
     agent = Agent(
         role="test role",
@@ -1224,6 +1263,7 @@ def test_agent_max_retry_limit():
         )
 
 
+@pytest.mark.timeout(60)
 def test_agent_with_llm():
     agent = Agent(
         role="test role",
@@ -1237,6 +1277,7 @@ def test_agent_with_llm():
     assert agent.llm.temperature == 0.7
 
 
+@pytest.mark.timeout(60)
 def test_agent_with_custom_stop_words():
     stop_words = ["STOP", "END"]
     agent = Agent(
@@ -1252,6 +1293,7 @@ def test_agent_with_custom_stop_words():
     assert "\nObservation:" in agent.llm.stop
 
 
+@pytest.mark.timeout(60)
 def test_agent_with_callbacks():
     def dummy_callback(response):
         pass
@@ -1268,6 +1310,7 @@ def test_agent_with_callbacks():
     assert agent.llm.callbacks[0] == dummy_callback
 
 
+@pytest.mark.timeout(60)
 def test_agent_with_additional_kwargs():
     agent = Agent(
         role="test role",
@@ -1291,6 +1334,7 @@ def test_agent_with_additional_kwargs():
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_llm_call():
     llm = LLM(model="gpt-3.5-turbo")
     messages = [{"role": "user", "content": "Say 'Hello, World!'"}]
@@ -1300,6 +1344,7 @@ def test_llm_call():
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_llm_call_with_error():
     llm = LLM(model="non-existent-model")
     messages = [{"role": "user", "content": "This should fail"}]
@@ -1309,6 +1354,7 @@ def test_llm_call_with_error():
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_handle_context_length_exceeds_limit():
     agent = Agent(
         role="test role",
@@ -1353,6 +1399,7 @@ def test_handle_context_length_exceeds_limit():
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_handle_context_length_exceeds_limit_cli_no():
     agent = Agent(
         role="test role",
@@ -1380,6 +1427,7 @@ def test_handle_context_length_exceeds_limit_cli_no():
             mock_handle_context.assert_not_called()
 
 
+@pytest.mark.timeout(60)
 def test_agent_with_all_llm_attributes():
     agent = Agent(
         role="test role",
@@ -1428,6 +1476,7 @@ def test_agent_with_all_llm_attributes():
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_llm_call_with_all_attributes():
     llm = LLM(
         model="gpt-3.5-turbo",
@@ -1445,6 +1494,7 @@ def test_llm_call_with_all_attributes():
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_agent_with_ollama_llama3():
     agent = Agent(
         role="test role",
@@ -1466,6 +1516,7 @@ def test_agent_with_ollama_llama3():
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_llm_call_with_ollama_llama3():
     llm = LLM(
         model="ollama/llama3.2:3b",
@@ -1483,6 +1534,7 @@ def test_llm_call_with_ollama_llama3():
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_agent_execute_task_basic():
     agent = Agent(
         role="test role",
@@ -1502,6 +1554,7 @@ def test_agent_execute_task_basic():
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_agent_execute_task_with_context():
     agent = Agent(
         role="test role",
@@ -1524,6 +1577,7 @@ def test_agent_execute_task_with_context():
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_agent_execute_task_with_tool():
     @tool
     def dummy_tool(query: str) -> str:
@@ -1549,6 +1603,7 @@ def test_agent_execute_task_with_tool():
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_agent_execute_task_with_custom_llm():
     agent = Agent(
         role="test role",
@@ -1570,6 +1625,7 @@ def test_agent_execute_task_with_custom_llm():
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_agent_execute_task_with_ollama():
     agent = Agent(
         role="test role",
@@ -1590,6 +1646,7 @@ def test_agent_execute_task_with_ollama():
 
 
 @pytest.mark.vcr(filter_headers=["authorization"])
+@pytest.mark.timeout(60)
 def test_agent_with_knowledge_sources():
     # Create a knowledge source with some content
     content = "Brandon's favorite color is red and he likes Mexican food."