mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-23 07:08:14 +00:00
Fix type annotations in guardrail classes to accept BaseLLM
- Update LLMGuardrail to accept BaseLLM instead of LLM in constructor - Update HallucinationGuardrail to accept BaseLLM instead of LLM in constructor - Add BaseLLM imports to both guardrail classes - Fixes type-checker CI failure: 'Argument llm to LLMGuardrail has incompatible type BaseLLM; expected LLM' Co-Authored-By: João <joao@crewai.com>
This commit is contained in:
@@ -9,6 +9,7 @@ Classes:
|
|||||||
from typing import Any, Optional, Tuple
|
from typing import Any, Optional, Tuple
|
||||||
|
|
||||||
from crewai.llm import LLM
|
from crewai.llm import LLM
|
||||||
|
from crewai.llms.base_llm import BaseLLM
|
||||||
from crewai.tasks.task_output import TaskOutput
|
from crewai.tasks.task_output import TaskOutput
|
||||||
from crewai.utilities.logger import Logger
|
from crewai.utilities.logger import Logger
|
||||||
|
|
||||||
@@ -47,7 +48,7 @@ class HallucinationGuardrail:
|
|||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
context: str,
|
context: str,
|
||||||
llm: LLM,
|
llm: BaseLLM,
|
||||||
threshold: Optional[float] = None,
|
threshold: Optional[float] = None,
|
||||||
tool_response: str = "",
|
tool_response: str = "",
|
||||||
):
|
):
|
||||||
@@ -60,7 +61,7 @@ class HallucinationGuardrail:
|
|||||||
tool_response: Optional tool response information that would be used in evaluation.
|
tool_response: Optional tool response information that would be used in evaluation.
|
||||||
"""
|
"""
|
||||||
self.context = context
|
self.context = context
|
||||||
self.llm: LLM = llm
|
self.llm: BaseLLM = llm
|
||||||
self.threshold = threshold
|
self.threshold = threshold
|
||||||
self.tool_response = tool_response
|
self.tool_response = tool_response
|
||||||
self._logger = Logger(verbose=True)
|
self._logger = Logger(verbose=True)
|
||||||
|
|||||||
@@ -4,6 +4,7 @@ from pydantic import BaseModel, Field
|
|||||||
|
|
||||||
from crewai.agent import Agent, LiteAgentOutput
|
from crewai.agent import Agent, LiteAgentOutput
|
||||||
from crewai.llm import LLM
|
from crewai.llm import LLM
|
||||||
|
from crewai.llms.base_llm import BaseLLM
|
||||||
from crewai.task import Task
|
from crewai.task import Task
|
||||||
from crewai.tasks.task_output import TaskOutput
|
from crewai.tasks.task_output import TaskOutput
|
||||||
|
|
||||||
@@ -32,11 +33,11 @@ class LLMGuardrail:
|
|||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
description: str,
|
description: str,
|
||||||
llm: LLM,
|
llm: BaseLLM,
|
||||||
):
|
):
|
||||||
self.description = description
|
self.description = description
|
||||||
|
|
||||||
self.llm: LLM = llm
|
self.llm: BaseLLM = llm
|
||||||
|
|
||||||
def _validate_output(self, task_output: TaskOutput) -> LiteAgentOutput:
|
def _validate_output(self, task_output: TaskOutput) -> LiteAgentOutput:
|
||||||
agent = Agent(
|
agent = Agent(
|
||||||
|
|||||||
Reference in New Issue
Block a user