Comparing BaseLLM class instead of LLM (#3120)

* Compaing BaseLLM class instead of LLM

* Fixed test cases

* Fixed Linting Issues

* removed last line

---------

Co-authored-by: Lucas Gomide <lucaslg200@gmail.com>

src/crewai/lite_agent.py

@@ -28,7 +28,7 @@ from pydantic import (
     InstanceOf,
     PrivateAttr,
     model_validator,
-    field_validator,
+    field_validator
 )
 
 from crewai.agents.agent_builder.base_agent import BaseAgent
@@ -40,7 +40,7 @@ from crewai.agents.parser import (
     OutputParserException,
 )
 from crewai.flow.flow_trackable import FlowTrackable
-from crewai.llm import LLM
+from crewai.llm import LLM, BaseLLM
 from crewai.tools.base_tool import BaseTool
 from crewai.tools.structured_tool import CrewStructuredTool
 from crewai.utilities import I18N
@@ -135,7 +135,7 @@ class LiteAgent(FlowTrackable, BaseModel):
     role: str = Field(description="Role of the agent")
     goal: str = Field(description="Goal of the agent")
     backstory: str = Field(description="Backstory of the agent")
-    llm: Optional[Union[str, InstanceOf[LLM], Any]] = Field(
+    llm: Optional[Union[str, InstanceOf[BaseLLM], Any]] = Field(
         default=None, description="Language model that will run the agent"
     )
     tools: List[BaseTool] = Field(
@@ -209,8 +209,8 @@ class LiteAgent(FlowTrackable, BaseModel):
     def setup_llm(self):
         """Set up the LLM and other components after initialization."""
         self.llm = create_llm(self.llm)
-        if not isinstance(self.llm, LLM):
+        if not isinstance(self.llm, BaseLLM):
-            raise ValueError("Unable to create LLM instance")
+            raise ValueError(f"Expected LLM instance of type BaseLLM, got {type(self.llm).__name__}")
 
         # Initialize callbacks
         token_callback = TokenCalcHandler(token_cost_process=self._token_process)
@@ -232,7 +232,8 @@ class LiteAgent(FlowTrackable, BaseModel):
         elif isinstance(self.guardrail, str):
             from crewai.tasks.llm_guardrail import LLMGuardrail
 
-            assert isinstance(self.llm, LLM)
+            if not isinstance(self.llm, BaseLLM):
+                raise TypeError(f"Guardrail requires LLM instance of type BaseLLM, got {type(self.llm).__name__}")
 
             self._guardrail = LLMGuardrail(description=self.guardrail, llm=self.llm)
 
@@ -620,4 +621,4 @@ class LiteAgent(FlowTrackable, BaseModel):
 
     def _append_message(self, text: str, role: str = "assistant") -> None:
         """Append a message to the message list with the given role."""
-        self._messages.append(format_message_for_llm(text, role=role))
+        self._messages.append(format_message_for_llm(text, role=role))

src/crewai/tasks/llm_guardrail.py

@@ -1,10 +1,9 @@
-from typing import Any, Optional, Tuple
+from typing import Any, Tuple
 
 from pydantic import BaseModel, Field
 
 from crewai.agent import Agent, LiteAgentOutput
-from crewai.llm import LLM
+from crewai.llm import BaseLLM
-from crewai.task import Task
 from crewai.tasks.task_output import TaskOutput
 
 
@@ -32,11 +31,11 @@ class LLMGuardrail:
     def __init__(
         self,
         description: str,
-        llm: LLM,
+        llm: BaseLLM,
     ):
         self.description = description
 
-        self.llm: LLM = llm
+        self.llm: BaseLLM = llm
 
     def _validate_output(self, task_output: TaskOutput) -> LiteAgentOutput:
         agent = Agent(

tests/test_lite_agent.py

@@ -12,6 +12,8 @@ from crewai.tools import BaseTool
 from crewai.utilities.events import crewai_event_bus
 from crewai.utilities.events.agent_events import LiteAgentExecutionStartedEvent
 from crewai.utilities.events.tool_usage_events import ToolUsageStartedEvent
+from crewai.llms.base_llm import BaseLLM
+from unittest.mock import patch
 
 
 # A simple test tool
@@ -418,3 +420,76 @@ def test_agent_output_when_guardrail_returns_base_model():
     result = agent.kickoff(messages="Top 10 best players in the world?")
 
     assert result.pydantic == Player(name="Lionel Messi", country="Argentina")
+
+def test_lite_agent_with_custom_llm_and_guardrails():
+    """Test that CustomLLM (inheriting from BaseLLM) works with guardrails."""
+    class CustomLLM(BaseLLM):
+        def __init__(self, response: str = "Custom response"):
+            super().__init__(model="custom-model")
+            self.response = response
+            self.call_count = 0
+
+        def call(self, messages, tools=None, callbacks=None, available_functions=None, from_task=None, from_agent=None) -> str:
+            self.call_count += 1
+
+            if "valid" in str(messages) and "feedback" in str(messages):
+                return '{"valid": true, "feedback": null}'
+
+            if "Thought:" in str(messages):
+                return f"Thought: I will analyze soccer players\nFinal Answer: {self.response}"
+
+            return self.response
+
+        def supports_function_calling(self) -> bool:
+            return False
+
+        def supports_stop_words(self) -> bool:
+            return False
+
+        def get_context_window_size(self) -> int:
+            return 4096
+
+    custom_llm = CustomLLM(response="Brazilian soccer players are the best!")
+
+    agent = LiteAgent(
+        role="Sports Analyst",
+        goal="Analyze soccer players",
+        backstory="You analyze soccer players and their performance.",
+        llm=custom_llm,
+        guardrail="Only include Brazilian players"
+    )
+
+    result = agent.kickoff("Tell me about the best soccer players")
+
+    assert custom_llm.call_count > 0
+    assert "Brazilian" in result.raw
+
+    custom_llm2 = CustomLLM(response="Original response")
+
+    def test_guardrail(output):
+        return (True, "Modified by guardrail")
+
+    agent2 = LiteAgent(
+        role="Test Agent",
+        goal="Test goal",
+        backstory="Test backstory",
+        llm=custom_llm2,
+        guardrail=test_guardrail
+    )
+
+    result2 = agent2.kickoff("Test message")
+    assert result2.raw == "Modified by guardrail"
+
+
+@pytest.mark.vcr(filter_headers=["authorization"])
+def test_lite_agent_with_invalid_llm():
+    """Test that LiteAgent raises proper error when create_llm returns None."""
+    with patch('crewai.lite_agent.create_llm', return_value=None):
+        with pytest.raises(ValueError) as exc_info:
+            LiteAgent(
+                role="Test Agent",
+                goal="Test goal", 
+                backstory="Test backstory",
+                llm="invalid-model"
+            )
+        assert "Expected LLM instance of type BaseLLM" in str(exc_info.value)