feat: use force_final_answer prompt on timeout

Co-Authored-By: Joe Moura <joao@crewai.com>

pyproject.toml

@@ -13,6 +13,7 @@ dependencies = [
     "openai>=1.13.3",
     "litellm==1.60.2",
     "instructor>=1.3.3",
+    "timeout-decorator>=0.5.0",
     # Text Processing
     "pdfplumber>=0.11.4",
     "regex>=2024.9.11",

src/crewai/agent.py

@@ -1,9 +1,11 @@
 import re
 import shutil
 import subprocess
+import threading
 from typing import Any, Dict, List, Literal, Optional, Sequence, Union
 
 from pydantic import Field, InstanceOf, PrivateAttr, model_validator
+import timeout_decorator
 
 from crewai.agents import CacheHandler
 from crewai.agents.agent_builder.base_agent import BaseAgent
@@ -64,6 +66,7 @@ class Agent(BaseAgent):
     """
 
     _times_executed: int = PrivateAttr(default=0)
+    _have_forced_answer: bool = PrivateAttr(default=False)
     max_execution_time: Optional[int] = Field(
         default=None,
         description="Maximum execution time for an agent to execute a task",
@@ -159,6 +162,77 @@ class Agent(BaseAgent):
         except (TypeError, ValueError) as e:
             raise ValueError(f"Invalid Knowledge Configuration: {str(e)}")
 
+    def _execute_with_timeout(
+        self,
+        task: Task,
+        context: Optional[str],
+        tools: Optional[List[BaseTool]],
+        timeout: int
+    ) -> str:
+        """Execute task with timeout using thread-based timeout.
+        
+        Args:
+            task: The task to execute
+            context: Optional context for the task
+            tools: Optional list of tools to use
+            timeout: Maximum execution time in seconds (must be > 0)
+            
+        Returns:
+            The result of the task execution, with force_final_answer prompt appended on timeout
+            
+        Raises:
+            ValueError: If timeout is not a positive integer
+            Exception: Any error that occurs during execution
+        """
+        # Validate timeout before creating any resources
+        if not isinstance(timeout, int) or timeout <= 0:
+            raise ValueError("Timeout must be a positive integer greater than zero")
+
+        completion_event: threading.Event = threading.Event()
+        result_container: List[Optional[str]] = [None]
+        error_container: List[Optional[Exception]] = [None]
+        
+        def target() -> None:
+            try:
+                result_container[0] = self._execute_task_without_timeout(task, context, tools)
+            except Exception as e:
+                error_container[0] = e
+            finally:
+                completion_event.set()
+        
+        thread: threading.Thread = threading.Thread(target=target)
+        thread.daemon = True  # Ensures thread doesn't prevent program exit
+        thread.start()
+        
+        # Wait for either completion or timeout
+        completed: bool = completion_event.wait(timeout=timeout)
+        
+        if not completed:
+            self._logger.log("warning", f"Task execution timed out after {timeout} seconds")
+            thread.join(timeout=0.1) 
+            
+            # Clean up resources
+            if hasattr(self, 'agent_executor') and self.agent_executor:
+                self.agent_executor.llm = None  # Release LLM resources
+                if hasattr(self.agent_executor, 'close'):
+                    self.agent_executor.close()
+
+            # Force final answer using the prompt
+            self._have_forced_answer = True
+            forced_answer = self.i18n.errors("force_final_answer")
+            return f"{result_container[0] if result_container[0] else ''}\n{forced_answer}"
+        
+        if error_container[0]:
+            error = error_container[0]
+            self._logger.log("error", f"Task execution failed: {str(error)}")
+            raise error
+        
+        if result_container[0] is None:
+            self._logger.log("warning", "Task execution completed but returned no result")
+            raise timeout_decorator.TimeoutError("Task execution completed but returned no result")  # This is a different kind of failure than timeout
+        
+        return result_container[0]
+
     def execute_task(
         self,
         task: Task,
@@ -174,7 +248,42 @@ class Agent(BaseAgent):
 
         Returns:
             Output of the agent
+
+        Raises:
+            TimeoutError: If the task execution exceeds max_execution_time (if set)
+            Exception: For other execution errors
         """
+        if self.max_execution_time is None:
+            return self._execute_task_without_timeout(task, context, tools)
+
+        original_llm_timeout = getattr(self.llm, 'timeout', None)
+        try:
+            if hasattr(self.llm, 'timeout'):
+                self.llm.timeout = self.max_execution_time
+            
+            result = self._execute_with_timeout(task, context, tools, self.max_execution_time)
+            if self._have_forced_answer:
+                self._logger.log("warning", f"Task '{task.description}' execution timed out after {self.max_execution_time} seconds. Using forced answer.")
+            return result
+        except timeout_decorator.TimeoutError:
+            # This is a different kind of failure (e.g., no result at all)
+            error_msg = (
+                f"Task '{task.description}' execution timed out after {self.max_execution_time} seconds "
+                f"and produced no result. Consider increasing max_execution_time or optimizing the task."
+            )
+            self._logger.log("error", error_msg)
+            raise TimeoutError(error_msg)
+        finally:
+            if original_llm_timeout is not None and hasattr(self.llm, 'timeout'):
+                self.llm.timeout = original_llm_timeout
+
+    def _execute_task_without_timeout(
+        self,
+        task: Task,
+        context: Optional[str] = None,
+        tools: Optional[List[BaseTool]] = None,
+    ) -> str:
+        """Execute task without timeout - contains the original execute_task logic."""
         if self.tools_handler:
             self.tools_handler.last_used_tool = {}  # type: ignore # Incompatible types in assignment (expression has type "dict[Never, Never]", variable has type "ToolCalling")
 

src/crewai/memory/storage/rag_storage.py

@@ -149,17 +149,10 @@ class RAGStorage(BaseRAGStorage):
         )
 
     def reset(self) -> None:
-        """Reset the storage by removing the database files and reinitializing."""
         try:
             if self.app:
                 self.app.reset()
-                # Clean up ChromaDB files
-                storage_path = os.path.join(db_storage_path(), self.type)
-                if os.path.exists(storage_path):
-                    shutil.rmtree(storage_path)
-                # Clean up temporary directory
-                if os.path.exists(self.path):
-                    shutil.rmtree(self.path)
+                shutil.rmtree(f"{db_storage_path()}/{self.type}")
                 self.app = None
                 self.collection = None
         except Exception as e:
@@ -170,3 +163,12 @@ class RAGStorage(BaseRAGStorage):
                 raise Exception(
                     f"An error occurred while resetting the {self.type} memory: {e}"
                 )
+
+    def _create_default_embedding_function(self):
+        from chromadb.utils.embedding_functions.openai_embedding_function import (
+            OpenAIEmbeddingFunction,
+        )
+
+        return OpenAIEmbeddingFunction(
+            api_key=os.getenv("OPENAI_API_KEY"), model_name="text-embedding-3-small"
+        )

src/crewai/utilities/embedding_configurator.py

@@ -25,34 +25,7 @@ class EmbeddingConfigurator:
         self,
         embedder_config: Optional[Dict[str, Any]] = None,
     ) -> EmbeddingFunction:
-        """Configure and return an embedding function based on the provided config.
-        
-        Args:
-            embedder_config: Optional configuration dictionary containing:
-                - provider: Name of the embedding provider or EmbeddingFunction instance
-                - config: Provider-specific configuration dictionary with options like:
-                    - api_key: API key for the provider
-                    - model: Model name to use for embeddings
-                    - url: API endpoint URL (for some providers)
-                    - session: Session object (for some providers)
-        
-        Returns:
-            EmbeddingFunction: Configured embedding function for the specified provider
-            
-        Raises:
-            ValueError: If custom embedding function is invalid
-            Exception: If provider is not supported or configuration is invalid
-            
-        Examples:
-            >>> config = {
-            ...     "provider": "openai",
-            ...     "config": {
-            ...         "api_key": "your-api-key",
-            ...         "model": "text-embedding-3-small"
-            ...     }
-            ... }
-            >>> embedder = EmbeddingConfigurator().configure_embedder(config)
-        """
+        """Configures and returns an embedding function based on the provided config."""
         if embedder_config is None:
             return self._create_default_embedding_function()
 
@@ -66,29 +39,20 @@ class EmbeddingConfigurator:
             )
 
         embedding_function = self.embedding_functions[provider]
-        if provider == "custom":
-            return embedding_function(config)
-        return embedding_function(config, model_name)
+        return (
+            embedding_function(config)
+            if provider == "custom"
+            else embedding_function(config, model_name)
+        )
 
     @staticmethod
     def _create_default_embedding_function():
-        """Create a default embedding function based on environment variables.
-        
-        Environment Variables:
-            CREWAI_EMBEDDING_PROVIDER: The embedding provider to use (default: "openai")
-            CREWAI_EMBEDDING_MODEL: The model to use for embeddings
-            OPENAI_API_KEY: API key for OpenAI (required if using OpenAI provider)
-        
-        Returns:
-            EmbeddingFunction: Configured embedding function
-        """
-        provider = os.getenv("CREWAI_EMBEDDING_PROVIDER", "openai")
-        config = {
-            "api_key": os.getenv("OPENAI_API_KEY"),
-            "model": os.getenv("CREWAI_EMBEDDING_MODEL", "text-embedding-3-small")
-        }
-        return EmbeddingConfigurator().configure_embedder(
-            {"provider": provider, "config": config}
+        from chromadb.utils.embedding_functions.openai_embedding_function import (
+            OpenAIEmbeddingFunction,
+        )
+
+        return OpenAIEmbeddingFunction(
+            api_key=os.getenv("OPENAI_API_KEY"), model_name="text-embedding-3-small"
         )
 
     @staticmethod
@@ -207,31 +171,6 @@ class EmbeddingConfigurator:
             url=config.get("api_url"),
         )
 
-    @staticmethod
-    def _configure_custom(config, model_name=None):
-        """Configure a custom embedding function.
-        
-        Args:
-            config: Configuration dictionary containing:
-                - embedder: Custom EmbeddingFunction instance
-            model_name: Not used for custom embedders, defaults to None
-            
-        Returns:
-            EmbeddingFunction: The validated custom embedding function
-            
-        Raises:
-            ValueError: If embedder is missing or invalid
-        """
-        embedder = config.get("embedder")
-        if not embedder or not isinstance(embedder, EmbeddingFunction):
-            raise ValueError("Custom provider requires a valid EmbeddingFunction instance")
-        
-        try:
-            validate_embedding_function(embedder)
-            return embedder
-        except Exception as e:
-            raise ValueError(f"Invalid custom embedding function: {str(e)}")
-
     @staticmethod
     def _configure_watson(config, model_name):
         try:
@@ -270,3 +209,28 @@ class EmbeddingConfigurator:
                     raise e
 
         return WatsonEmbeddingFunction()
+
+    @staticmethod
+    def _configure_custom(config):
+        custom_embedder = config.get("embedder")
+        if isinstance(custom_embedder, EmbeddingFunction):
+            try:
+                validate_embedding_function(custom_embedder)
+                return custom_embedder
+            except Exception as e:
+                raise ValueError(f"Invalid custom embedding function: {str(e)}")
+        elif callable(custom_embedder):
+            try:
+                instance = custom_embedder()
+                if isinstance(instance, EmbeddingFunction):
+                    validate_embedding_function(instance)
+                    return instance
+                raise ValueError(
+                    "Custom embedder does not create an EmbeddingFunction instance"
+                )
+            except Exception as e:
+                raise ValueError(f"Error instantiating custom embedder: {str(e)}")
+        else:
+            raise ValueError(
+                "Custom embedder must be an instance of `EmbeddingFunction` or a callable that creates one"
+            )

src/crewai/utilities/exceptions/__init__.py

@@ -1,5 +0,0 @@
-"""Exceptions module for crewAI utilities."""
-
-from .embedding_exceptions import EmbeddingConfigurationError, EmbeddingProviderError
-
-__all__ = ["EmbeddingConfigurationError", "EmbeddingProviderError"]

src/crewai/utilities/exceptions/embedding_exceptions.py

@@ -1,9 +0,0 @@
-"""Exceptions related to embedding functionality."""
-
-class EmbeddingConfigurationError(Exception):
-    """Raised when there is an error in the embedding configuration."""
-    pass
-
-class EmbeddingProviderError(Exception):
-    """Raised when there is an error with the embedding provider."""
-    pass

tests/memory/test_memory_reset.py

@@ -1,152 +0,0 @@
-import os
-import tempfile
-from pathlib import Path
-from typing import Generator
-
-import pytest
-from chromadb import Documents, EmbeddingFunction, Embeddings
-
-from crewai.memory import EntityMemory, LongTermMemory, ShortTermMemory
-from crewai.utilities import EmbeddingConfigurator
-from crewai.utilities.exceptions.embedding_exceptions import (
-    EmbeddingConfigurationError,
-    EmbeddingProviderError,
-)
-
-
-@pytest.fixture
-def temp_db_dir() -> Generator[Path, None, None]:
-    """Create a temporary directory for test databases."""
-    with tempfile.TemporaryDirectory() as tmpdir:
-        path = Path(tmpdir)
-        # Ensure directory exists and is writable
-        path.mkdir(parents=True, exist_ok=True)
-        # Set ChromaDB to use in-memory mode for tests
-        os.environ["CHROMA_IN_MEMORY"] = "true"
-        try:
-            yield path
-        finally:
-            # Clean up ChromaDB environment variable
-            if "CHROMA_IN_MEMORY" in os.environ:
-                del os.environ["CHROMA_IN_MEMORY"]
-
-
-def test_memory_reset_with_openai(temp_db_dir):
-    """Test memory reset with default OpenAI provider."""
-    original_key = os.environ.get("OPENAI_API_KEY")
-    original_provider = os.environ.get("CREWAI_EMBEDDING_PROVIDER")
-    try:
-        os.environ["OPENAI_API_KEY"] = "test-key"
-        if "CREWAI_EMBEDDING_PROVIDER" in os.environ:
-            del os.environ["CREWAI_EMBEDDING_PROVIDER"]
-        memory = ShortTermMemory(path=str(temp_db_dir))
-        memory.reset()  # Should work with OpenAI as default
-    finally:
-        if original_key:
-            os.environ["OPENAI_API_KEY"] = original_key
-        else:
-            del os.environ["OPENAI_API_KEY"]
-        if original_provider:
-            os.environ["CREWAI_EMBEDDING_PROVIDER"] = original_provider
-
-
-def test_memory_reset_with_ollama(temp_db_dir):
-    """Test memory reset with Ollama provider."""
-    original_provider = os.environ.get("CREWAI_EMBEDDING_PROVIDER")
-    try:
-        os.environ["CREWAI_EMBEDDING_PROVIDER"] = "ollama"
-        memory = ShortTermMemory(path=str(temp_db_dir))
-        memory.reset()  # Should not raise any OpenAI-related errors
-    finally:
-        if original_provider:
-            os.environ["CREWAI_EMBEDDING_PROVIDER"] = original_provider
-        elif "CREWAI_EMBEDDING_PROVIDER" in os.environ:
-            del os.environ["CREWAI_EMBEDDING_PROVIDER"]
-
-
-def test_memory_reset_with_custom_provider(temp_db_dir):
-    """Test memory reset with custom embedding provider."""
-
-    class CustomEmbedder(EmbeddingFunction):
-        def __call__(self, input: Documents) -> Embeddings:
-            if isinstance(input, str):
-                input = [input]
-            return [[0.5] * 10] * len(input)
-
-    memory = ShortTermMemory(
-        path=str(temp_db_dir),
-        embedder_config={
-            "provider": "custom",
-            "config": {"embedder": CustomEmbedder()}
-        }
-    )
-    memory.reset()  # Should work with custom embedder
-
-
-def test_memory_reset_with_invalid_provider(temp_db_dir):
-    """Test memory reset with invalid provider raises appropriate error."""
-    original_provider = os.environ.get("CREWAI_EMBEDDING_PROVIDER")
-    original_key = os.environ.get("OPENAI_API_KEY")
-    try:
-        os.environ["CREWAI_EMBEDDING_PROVIDER"] = "invalid_provider"
-        with pytest.raises(Exception) as exc_info:
-            memory = ShortTermMemory(path=str(temp_db_dir))
-            memory.reset()
-        assert "Unsupported embedding provider" in str(exc_info.value)
-    finally:
-        if original_provider:
-            os.environ["CREWAI_EMBEDDING_PROVIDER"] = original_provider
-        elif "CREWAI_EMBEDDING_PROVIDER" in os.environ:
-            del os.environ["CREWAI_EMBEDDING_PROVIDER"]
-        if original_key:
-            os.environ["OPENAI_API_KEY"] = original_key
-
-
-def test_memory_reset_with_missing_api_key(temp_db_dir):
-    """Test memory reset with missing API key raises appropriate error."""
-    original_key = os.environ.get("OPENAI_API_KEY")
-    original_provider = os.environ.get("CREWAI_EMBEDDING_PROVIDER")
-    try:
-        if "OPENAI_API_KEY" in os.environ:
-            del os.environ["OPENAI_API_KEY"]
-        os.environ["CREWAI_EMBEDDING_PROVIDER"] = "openai"
-        with pytest.raises(ValueError) as exc_info:
-            memory = ShortTermMemory(path=str(temp_db_dir))
-            memory.reset()
-        assert "openai api key" in str(exc_info.value).lower()
-    finally:
-        if original_key:
-            os.environ["OPENAI_API_KEY"] = original_key
-        if original_provider:
-            os.environ["CREWAI_EMBEDDING_PROVIDER"] = original_provider
-        elif "CREWAI_EMBEDDING_PROVIDER" in os.environ:
-            del os.environ["CREWAI_EMBEDDING_PROVIDER"]
-
-
-def test_memory_reset_cleans_up_files(temp_db_dir):
-    """Test that memory reset properly cleans up database files."""
-    original_provider = os.environ.get("CREWAI_EMBEDDING_PROVIDER")
-    try:
-        class TestEmbedder(EmbeddingFunction):
-            def __call__(self, input: Documents) -> Embeddings:
-                if isinstance(input, str):
-                    input = [input]
-                return [[0.5] * 10] * len(input)
-
-        if "CREWAI_EMBEDDING_PROVIDER" in os.environ:
-            del os.environ["CREWAI_EMBEDDING_PROVIDER"]
-        memory = ShortTermMemory(
-            path=str(temp_db_dir),
-            embedder_config={
-                "provider": "custom",
-                "config": {"embedder": TestEmbedder()}
-            }
-        )
-        memory.save("test memory", {"test": "metadata"})
-        assert any(temp_db_dir.iterdir())  # Directory should have files
-        memory.reset()
-        # After reset, directory should either not exist or be empty
-        assert not os.path.exists(temp_db_dir) or not any(temp_db_dir.iterdir())
-    finally:
-        if original_provider:
-            os.environ["CREWAI_EMBEDDING_PROVIDER"] = original_provider

tests/task_test.py

@@ -13,6 +13,7 @@ from crewai import Agent, Crew, Process, Task
 from crewai.tasks.conditional_task import ConditionalTask
 from crewai.tasks.task_output import TaskOutput
 from crewai.utilities.converter import Converter
+from concurrent.futures import TimeoutError as FuturesTimeoutError
 
 
 def test_task_tool_reflect_agent_tools():
@@ -1283,3 +1284,139 @@ def test_interpolate_valid_types():
     assert parsed["optional"] is None
     assert parsed["nested"]["flag"] is True
     assert parsed["nested"]["empty"] is None
+
+
+def test_task_completes_within_max_execution_time():
+    """Test task completes successfully within specified timeout"""
+    researcher = Agent(
+        role="Researcher",
+        goal="Test goal",
+        backstory="Test backstory",
+        max_execution_time=100  # Ample time for completion
+    )
+
+    task = Task(
+        description="Test task",
+        expected_output="Test output",
+        agent=researcher
+    )
+
+    with patch.object(Agent, '_execute_task_without_timeout', return_value="Success") as mock_execute:
+        result = task.execute_sync(agent=researcher)
+        assert result.raw == "Success"
+        mock_execute.assert_called_once()
+
+
+def test_task_exceeds_max_execution_time():
+    """Test task raises TimeoutError when exceeding max execution time"""
+    researcher = Agent(
+        role="Researcher",
+        goal="Test goal",
+        backstory="Test backstory",
+        max_execution_time=1  # Very short timeout
+    )
+
+    task = Task(
+        description="Test task",
+        expected_output="Test output",
+        agent=researcher
+    )
+
+    # Mock the OpenAI API call to avoid authentication
+    with patch('litellm.completion', side_effect=TimeoutError("Request timed out")):
+        with pytest.raises(TimeoutError) as excinfo:
+            task.execute_sync(agent=researcher)
+        assert "timed out" in str(excinfo.value)
+
+
+def test_task_no_max_execution_time():
+    """Test task executes normally without timeout setting"""
+    researcher = Agent(
+        role="Researcher",
+        goal="Test goal",
+        backstory="Test backstory",
+        max_execution_time=None  # No timeout
+    )
+
+    task = Task(
+        description="Test task",
+        expected_output="Test output",
+        agent=researcher
+    )
+
+    with patch.object(Agent, 'execute_task', return_value="Success") as mock_execute:
+        result = task.execute_sync(agent=researcher)
+        assert result.raw == "Success"
+        mock_execute.assert_called_once()
+
+
+def test_task_max_execution_time_zero():
+    """Test immediate timeout with max_execution_time=0"""
+    researcher = Agent(
+        role="Researcher",
+        goal="Test goal",
+        backstory="Test backstory",
+        max_execution_time=1  # Set to minimum valid value
+    )
+
+    task = Task(
+        description="Test task",
+        expected_output="Test output",
+        agent=researcher
+    )
+
+    # Simulate immediate timeout using FuturesTimeoutError
+    with patch('concurrent.futures.ThreadPoolExecutor') as mock_executor:
+        mock_future = MagicMock()
+        mock_future.result.side_effect = FuturesTimeoutError()
+        mock_executor.return_value.submit.return_value = mock_future
+
+        with pytest.raises(TimeoutError) as excinfo:
+            task.execute_sync(agent=researcher)
+        assert "timed out after 1 seconds" in str(excinfo.value)
+
+
+def test_task_force_final_answer_on_timeout():
+    """Test that force_final_answer is used when task times out"""
+    researcher = Agent(
+        role="Researcher",
+        goal="Test goal",
+        backstory="Test backstory",
+        max_execution_time=1  # Very short timeout
+    )
+
+    task = Task(
+        description="Test task",
+        expected_output="Test output",
+        agent=researcher
+    )
+
+    # Mock the task execution to simulate a partial result before timeout
+    mock_i18n = MagicMock()
+    mock_i18n.errors.return_value = "MUST give your absolute best final answer"
+    researcher.i18n = mock_i18n
+
+    class MockThread:
+        def __init__(self, target, *args, **kwargs):
+            self.target = target
+            self.daemon = kwargs.get('daemon', False)
+            self.args = args
+            self.kwargs = kwargs
+
+        def start(self):
+            # Execute the target function to set the result
+            self.target()
+
+        def join(self, timeout=None):
+            pass
+
+    def mock_thread(*args, **kwargs):
+        return MockThread(*args, **kwargs)
+
+    with patch('threading.Thread', side_effect=mock_thread), \
+         patch('threading.Event.wait', return_value=False), \
+         patch('litellm.completion'), \
+         patch.object(Agent, '_execute_task_without_timeout', return_value="Partial result"):
+        result = task.execute_sync(agent=researcher)
+        assert "MUST give your absolute best final answer" in result.raw
+        assert "Partial result" in result.raw  # Should include partial result