From 63ce0c91f9e306ca6b10dbb59d718bc335b3add3 Mon Sep 17 00:00:00 2001
From: Brandon Hancock <brandon@brandonhancock.io>
Date: Fri, 24 Jan 2025 14:58:04 -0500
Subject: [PATCH] Fix error

---
 src/crewai/agent.py                      |  3 +--
 src/crewai/agents/crew_agent_executor.py | 17 +++++++++++------
 tests/agent_test.py                      | 23 ++++++++++++++++-------
 3 files changed, 28 insertions(+), 15 deletions(-)

diff --git a/src/crewai/agent.py b/src/crewai/agent.py
index 94b8ba065..3a4d083d4 100644
--- a/src/crewai/agent.py
+++ b/src/crewai/agent.py
@@ -2,7 +2,6 @@ import shutil
 import subprocess
 from typing import Any, Dict, List, Literal, Optional, Union
 
-from litellm.llms.base_llm.chat.transformation import BaseLLMException
 from pydantic import Field, InstanceOf, PrivateAttr, model_validator
 
 from crewai.agents import CacheHandler
@@ -260,7 +259,7 @@ class Agent(BaseAgent):
                 }
             )["output"]
         except Exception as e:
-            if isinstance(e, BaseLLMException):
+            if e.__class__.__module__.startswith("litellm"):
                 # Do not retry on litellm errors
                 raise e
             self._times_executed += 1
diff --git a/src/crewai/agents/crew_agent_executor.py b/src/crewai/agents/crew_agent_executor.py
index dd5252ea7..b9797193c 100644
--- a/src/crewai/agents/crew_agent_executor.py
+++ b/src/crewai/agents/crew_agent_executor.py
@@ -3,9 +3,6 @@ import re
 from dataclasses import dataclass
 from typing import Any, Callable, Dict, List, Optional, Union
 
-from litellm.exceptions import AuthenticationError as LiteLLMAuthenticationError
-from litellm.llms.base_llm.chat.transformation import BaseLLMException
-
 from crewai.agents.agent_builder.base_agent import BaseAgent
 from crewai.agents.agent_builder.base_agent_executor_mixin import CrewAgentExecutorMixin
 from crewai.agents.parser import (
@@ -104,7 +101,12 @@ class CrewAgentExecutor(CrewAgentExecutorMixin):
         try:
             formatted_answer = self._invoke_loop()
         except Exception as e:
-            raise e
+            if e.__class__.__module__.startswith("litellm"):
+                # Do not retry on litellm errors
+                raise e
+            else:
+                self._handle_unknown_error(e)
+                raise e
 
         if self.ask_for_human_input:
             formatted_answer = self._handle_human_feedback(formatted_answer)
@@ -143,9 +145,12 @@ class CrewAgentExecutor(CrewAgentExecutorMixin):
                 self._invoke_step_callback(formatted_answer)
                 self._append_message(formatted_answer.text, role="assistant")
 
+            except OutputParserException as e:
+                formatted_answer = self._handle_output_parser_exception(e)
+
             except Exception as e:
-                if isinstance(e, BaseLLMException):
-                    # Stop execution on litellm errors
+                if e.__class__.__module__.startswith("litellm"):
+                    # Do not retry on litellm errors
                     raise e
                 if self._is_context_length_exceeded(e):
                     self._handle_context_length()
diff --git a/tests/agent_test.py b/tests/agent_test.py
index db909ef2f..21ce76347 100644
--- a/tests/agent_test.py
+++ b/tests/agent_test.py
@@ -1676,7 +1676,7 @@ def test_crew_agent_executor_litellm_auth_error():
     with (
         patch.object(LLM, "call") as mock_llm_call,
         patch.object(Printer, "print") as mock_printer,
-        pytest.raises(AuthenticationError, match="Invalid API key"),
+        pytest.raises(AuthenticationError) as exc_info,
     ):
         mock_llm_call.side_effect = AuthenticationError(
             message="Invalid API key", llm_provider="openai", model="gpt-4"
@@ -1689,18 +1689,27 @@ def test_crew_agent_executor_litellm_auth_error():
             }
         )
 
-    # Verify error handling
+    # Verify error handling messages
+    error_message = f"Error during LLM call: {str(mock_llm_call.side_effect)}"
     mock_printer.assert_any_call(
-        content="An unknown error occurred. Please check the details below.",
-        color="red",
-    )
-    mock_printer.assert_any_call(
-        content="Error details: litellm.AuthenticationError: Invalid API key",
+        content=error_message,
         color="red",
     )
+
     # Verify the call was only made once (no retries)
     mock_llm_call.assert_called_once()
 
+    # Assert that the exception was raised and has the expected attributes
+    assert exc_info.type is AuthenticationError
+    assert "Invalid API key".lower() in exc_info.value.message.lower()
+    assert exc_info.value.llm_provider == "openai"
+    assert exc_info.value.model == "gpt-4"
+
+    # Optionally, assert that the exception is an instance of BaseLLMException
+    from litellm.llms.base_llm.chat.transformation import BaseLLMException
+
+    assert isinstance(exc_info.value, BaseLLMException)
+
 
 def test_litellm_anthropic_error_handling():
     """Test that AnthropicError from LiteLLM is handled correctly and not retried."""