diff --git a/src/crewai/lite_agent.py b/src/crewai/lite_agent.py index 53b6ab4c9..b768b9985 100644 --- a/src/crewai/lite_agent.py +++ b/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): - raise ValueError("Unable to create LLM instance") + if not isinstance(self.llm, BaseLLM): + 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)) \ No newline at end of file diff --git a/src/crewai/tasks/llm_guardrail.py b/src/crewai/tasks/llm_guardrail.py index 2bb948075..65aea7cec 100644 --- a/src/crewai/tasks/llm_guardrail.py +++ b/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.task import Task +from crewai.llm import BaseLLM 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( diff --git a/tests/test_lite_agent.py b/tests/test_lite_agent.py index 6b0ca4b70..b495cb800 100644 --- a/tests/test_lite_agent.py +++ b/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) \ No newline at end of file