agno 2.0.11__py3-none-any.whl → 2.1.1__py3-none-any.whl

agno/exceptions.py

@@ -1,4 +1,5 @@
-from typing import List, Optional, Union
+from enum import Enum
+from typing import Any, Dict, List, Optional, Union
 
 from agno.models.message import Message
 
@@ -17,6 +18,8 @@ class AgentRunException(Exception):
         self.agent_message = agent_message
         self.messages = messages
         self.stop_execution = stop_execution
+        self.type = "agent_run_error"
+        self.error_id = "agent_run_error"
 
 
 class RetryAgentRun(AgentRunException):
@@ -32,6 +35,7 @@ class RetryAgentRun(AgentRunException):
         super().__init__(
             exc, user_message=user_message, agent_message=agent_message, messages=messages, stop_execution=False
         )
+        self.error_id = "retry_agent_run_error"
 
 
 class StopAgentRun(AgentRunException):
@@ -47,6 +51,7 @@ class StopAgentRun(AgentRunException):
         super().__init__(
             exc, user_message=user_message, agent_message=agent_message, messages=messages, stop_execution=True
         )
+        self.error_id = "stop_agent_run_error"
 
 
 class RunCancelledException(Exception):
@@ -54,6 +59,8 @@ class RunCancelledException(Exception):
 
     def __init__(self, message: str = "Operation cancelled by user"):
         super().__init__(message)
+        self.type = "run_cancelled_error"
+        self.error_id = "run_cancelled_error"
 
 
 class AgnoError(Exception):
@@ -63,6 +70,8 @@ class AgnoError(Exception):
         super().__init__(message)
         self.message = message
         self.status_code = status_code
+        self.type = "agno_error"
+        self.error_id = "agno_error"
 
     def __str__(self) -> str:
         return str(self.message)
@@ -78,6 +87,9 @@ class ModelProviderError(AgnoError):
         self.model_name = model_name
         self.model_id = model_id
 
+        self.type = "model_provider_error"
+        self.error_id = "model_provider_error"
+
 
 class ModelRateLimitError(ModelProviderError):
     """Exception raised when a model provider returns a rate limit error."""
@@ -86,9 +98,58 @@ class ModelRateLimitError(ModelProviderError):
         self, message: str, status_code: int = 429, model_name: Optional[str] = None, model_id: Optional[str] = None
     ):
         super().__init__(message, status_code, model_name, model_id)
+        self.error_id = "model_rate_limit_error"
 
 
 class EvalError(Exception):
     """Exception raised when an evaluation fails."""
 
     pass
+
+
+class CheckTrigger(Enum):
+    """Enum for guardrail triggers."""
+
+    OFF_TOPIC = "off_topic"
+    INPUT_NOT_ALLOWED = "input_not_allowed"
+    OUTPUT_NOT_ALLOWED = "output_not_allowed"
+    VALIDATION_FAILED = "validation_failed"
+
+    PROMPT_INJECTION = "prompt_injection"
+    PII_DETECTED = "pii_detected"
+
+
+class InputCheckError(Exception):
+    """Exception raised when an input check fails."""
+
+    def __init__(
+        self,
+        message: str,
+        check_trigger: CheckTrigger = CheckTrigger.INPUT_NOT_ALLOWED,
+        additional_data: Optional[Dict[str, Any]] = None,
+    ):
+        super().__init__(message)
+        self.type = "input_check_error"
+        self.error_id = check_trigger.value
+
+        self.message = message
+        self.check_trigger = check_trigger
+        self.additional_data = additional_data
+
+
+class OutputCheckError(Exception):
+    """Exception raised when an output check fails."""
+
+    def __init__(
+        self,
+        message: str,
+        check_trigger: CheckTrigger = CheckTrigger.OUTPUT_NOT_ALLOWED,
+        additional_data: Optional[Dict[str, Any]] = None,
+    ):
+        super().__init__(message)
+        self.type = "output_check_error"
+        self.error_id = check_trigger.value
+
+        self.message = message
+        self.check_trigger = check_trigger
+        self.additional_data = additional_data

agno/guardrails/__init__.py

@@ -0,0 +1,6 @@
+from agno.guardrails.base import BaseGuardrail
+from agno.guardrails.openai import OpenAIModerationGuardrail
+from agno.guardrails.pii import PIIDetectionGuardrail
+from agno.guardrails.prompt_injection import PromptInjectionGuardrail
+
+__all__ = ["BaseGuardrail", "OpenAIModerationGuardrail", "PIIDetectionGuardrail", "PromptInjectionGuardrail"]

agno/guardrails/base.py

@@ -0,0 +1,19 @@
+from abc import ABC, abstractmethod
+from typing import Union
+
+from agno.run.agent import RunInput
+from agno.run.team import TeamRunInput
+
+
+class BaseGuardrail(ABC):
+    """Abstract base class for all guardrail implementations."""
+
+    @abstractmethod
+    def check(self, run_input: Union[RunInput, TeamRunInput]) -> None:
+        """Perform synchronous guardrail check."""
+        pass
+
+    @abstractmethod
+    async def async_check(self, run_input: Union[RunInput, TeamRunInput]) -> None:
+        """Perform asynchronous guardrail check."""
+        pass

agno/guardrails/openai.py

@@ -0,0 +1,144 @@
+from os import getenv
+from typing import Any, Dict, List, Literal, Optional, Union
+
+from agno.exceptions import CheckTrigger, InputCheckError
+from agno.guardrails.base import BaseGuardrail
+from agno.run.agent import RunInput
+from agno.run.team import TeamRunInput
+from agno.utils.log import log_debug
+from agno.utils.openai import images_to_message
+
+
+class OpenAIModerationGuardrail(BaseGuardrail):
+    """Guardrail for detecting content that violates OpenAI's content policy.
+
+    Args:
+        moderation_model (str): The model to use for moderation. Defaults to "omni-moderation-latest".
+        raise_for_categories (List[str]): The categories to raise for.
+                                          Options are: "sexual", "sexual/minors", "harassment",
+                                          "harassment/threatening", "hate", "hate/threatening",
+                                          "illicit", "illicit/violent", "self-harm", "self-harm/intent",
+                                          "self-harm/instructions", "violence", "violence/graphic".
+                                          Defaults to include all categories.
+        api_key (str): The API key to use for moderation. Defaults to the OPENAI_API_KEY environment variable.
+    """
+
+    def __init__(
+        self,
+        moderation_model: str = "omni-moderation-latest",
+        raise_for_categories: Optional[
+            List[
+                Literal[
+                    "sexual",
+                    "sexual/minors",
+                    "harassment",
+                    "harassment/threatening",
+                    "hate",
+                    "hate/threatening",
+                    "illicit",
+                    "illicit/violent",
+                    "self-harm",
+                    "self-harm/intent",
+                    "self-harm/instructions",
+                    "violence",
+                    "violence/graphic",
+                ]
+            ]
+        ] = None,
+        api_key: Optional[str] = None,
+    ):
+        self.moderation_model = moderation_model
+        self.api_key = api_key or getenv("OPENAI_API_KEY")
+        self.raise_for_categories = raise_for_categories
+
+    def check(self, run_input: Union[RunInput, TeamRunInput]) -> None:
+        """Check for content that violates OpenAI's content policy."""
+        try:
+            from openai import OpenAI as OpenAIClient
+        except ImportError:
+            raise ImportError("`openai` not installed. Please install using `pip install openai`")
+
+        content = run_input.input_content_string()
+        images = run_input.images
+
+        log_debug(f"Moderating content using {self.moderation_model}")
+        client = OpenAIClient(api_key=self.api_key)
+
+        model_input: Union[str, List[Dict[str, Any]]] = content
+
+        if images is not None:
+            model_input = [{"type": "text", "text": content}, *images_to_message(images=images)]
+
+        # Prepare input based on content type
+        response = client.moderations.create(model=self.moderation_model, input=model_input)  # type: ignore
+
+        result = response.results[0]
+
+        if result.flagged:
+            moderation_result = {
+                "categories": result.categories.model_dump(),
+                "category_scores": result.category_scores.model_dump(),
+            }
+
+            trigger_validation = False
+
+            if self.raise_for_categories is not None:
+                for category in self.raise_for_categories:
+                    if moderation_result["categories"][category]:
+                        trigger_validation = True
+            else:
+                # Since at least one category is flagged, we need to raise the check
+                trigger_validation = True
+
+            if trigger_validation:
+                raise InputCheckError(
+                    "OpenAI moderation violation detected.",
+                    additional_data=moderation_result,
+                    check_trigger=CheckTrigger.INPUT_NOT_ALLOWED,
+                )
+
+    async def async_check(self, run_input: Union[RunInput, TeamRunInput]) -> None:
+        """Check for content that violates OpenAI's content policy."""
+        try:
+            from openai import AsyncOpenAI as OpenAIClient
+        except ImportError:
+            raise ImportError("`openai` not installed. Please install using `pip install openai`")
+
+        content = run_input.input_content_string()
+        images = run_input.images
+
+        log_debug(f"Moderating content using {self.moderation_model}")
+        client = OpenAIClient(api_key=self.api_key)
+
+        model_input: Union[str, List[Dict[str, Any]]] = content
+
+        if images is not None:
+            model_input = [{"type": "text", "text": content}, *images_to_message(images=images)]
+
+        # Prepare input based on content type
+        response = await client.moderations.create(model=self.moderation_model, input=model_input)  # type: ignore
+
+        result = response.results[0]
+
+        if result.flagged:
+            moderation_result = {
+                "categories": result.categories.model_dump(),
+                "category_scores": result.category_scores.model_dump(),
+            }
+
+            trigger_validation = False
+
+            if self.raise_for_categories is not None:
+                for category in self.raise_for_categories:
+                    if moderation_result["categories"][category]:
+                        trigger_validation = True
+            else:
+                # Since at least one category is flagged, we need to raise the check
+                trigger_validation = True
+
+            if trigger_validation:
+                raise InputCheckError(
+                    "OpenAI moderation violation detected.",
+                    additional_data=moderation_result,
+                    check_trigger=CheckTrigger.INPUT_NOT_ALLOWED,
+                )

agno/guardrails/pii.py

@@ -0,0 +1,94 @@
+from re import Pattern
+from typing import Dict, Optional, Union
+
+from agno.exceptions import CheckTrigger, InputCheckError
+from agno.guardrails.base import BaseGuardrail
+from agno.run.agent import RunInput
+from agno.run.team import TeamRunInput
+
+
+class PIIDetectionGuardrail(BaseGuardrail):
+    """Guardrail for detecting Personally Identifiable Information (PII).
+
+    Args:
+        mask_pii: Whether to mask the PII in the input, rather than raising an error.
+        enable_ssn_check: Whether to check for Social Security Numbers. True by default.
+        enable_credit_card_check: Whether to check for credit cards. True by default.
+        enable_email_check: Whether to check for emails. True by default.
+        enable_phone_check: Whether to check for phone numbers. True by default.
+        custom_patterns: A dictionary of custom PII patterns to detect. This is added to the default patterns.
+    """
+
+    def __init__(
+        self,
+        mask_pii: bool = False,
+        enable_ssn_check: bool = True,
+        enable_credit_card_check: bool = True,
+        enable_email_check: bool = True,
+        enable_phone_check: bool = True,
+        custom_patterns: Optional[Dict[str, Pattern[str]]] = None,
+    ):
+        import re
+
+        self.mask_pii = mask_pii
+        self.pii_patterns = {}
+
+        if enable_ssn_check:
+            self.pii_patterns["SSN"] = re.compile(r"\b\d{3}-\d{2}-\d{4}\b")
+        if enable_credit_card_check:
+            self.pii_patterns["Credit Card"] = re.compile(r"\b\d{4}[\s-]?\d{4}[\s-]?\d{4}[\s-]?\d{4}\b")
+        if enable_email_check:
+            self.pii_patterns["Email"] = re.compile(r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b")
+        if enable_phone_check:
+            self.pii_patterns["Phone"] = re.compile(r"\b\d{3}[\s.-]?\d{3}[\s.-]?\d{4}\b")
+
+        if custom_patterns:
+            self.pii_patterns.update(custom_patterns)
+
+    def check(self, run_input: Union[RunInput, TeamRunInput]) -> None:
+        """Check for PII patterns in the input."""
+        content = run_input.input_content_string()
+        detected_pii = []
+        for pii_type, pattern in self.pii_patterns.items():
+            if pattern.search(content):
+                detected_pii.append(pii_type)
+        if detected_pii:
+            if self.mask_pii:
+                for pii_type in detected_pii:
+
+                    def mask_match(match):
+                        return "*" * len(match.group(0))
+
+                    content = self.pii_patterns[pii_type].sub(mask_match, content)
+                run_input.input_content = content
+                return
+            else:
+                raise InputCheckError(
+                    "Potential PII detected in input",
+                    additional_data={"detected_pii": detected_pii},
+                    check_trigger=CheckTrigger.PII_DETECTED,
+                )
+
+    async def async_check(self, run_input: Union[RunInput, TeamRunInput]) -> None:
+        """Asynchronously check for PII patterns in the input."""
+        content = run_input.input_content_string()
+        detected_pii = []
+        for pii_type, pattern in self.pii_patterns.items():
+            if pattern.search(content):
+                detected_pii.append(pii_type)
+        if detected_pii:
+            if self.mask_pii:
+                for pii_type in detected_pii:
+
+                    def mask_match(match):
+                        return "*" * len(match.group(0))
+
+                    content = self.pii_patterns[pii_type].sub(mask_match, content)
+                run_input.input_content = content
+                return
+            else:
+                raise InputCheckError(
+                    "Potential PII detected in input",
+                    additional_data={"detected_pii": detected_pii},
+                    check_trigger=CheckTrigger.PII_DETECTED,
+                )

agno/guardrails/prompt_injection.py

@@ -0,0 +1,51 @@
+from typing import List, Optional, Union
+
+from agno.exceptions import CheckTrigger, InputCheckError
+from agno.guardrails.base import BaseGuardrail
+from agno.run.agent import RunInput
+from agno.run.team import TeamRunInput
+
+
+class PromptInjectionGuardrail(BaseGuardrail):
+    """Guardrail for detecting prompt injection attempts.
+
+    Args:
+        injection_patterns (Optional[List[str]]): A list of patterns to check for. Defaults to a list of common prompt injection patterns.
+    """
+
+    def __init__(self, injection_patterns: Optional[List[str]] = None):
+        self.injection_patterns = injection_patterns or [
+            "ignore previous instructions",
+            "ignore your instructions",
+            "you are now a",
+            "forget everything above",
+            "developer mode",
+            "override safety",
+            "disregard guidelines",
+            "system prompt",
+            "jailbreak",
+            "act as if",
+            "pretend you are",
+            "roleplay as",
+            "simulate being",
+            "bypass restrictions",
+            "ignore safeguards",
+            "admin override",
+            "root access",
+        ]
+
+    def check(self, run_input: Union[RunInput, TeamRunInput]) -> None:
+        """Check for prompt injection patterns in the input."""
+        if any(keyword in run_input.input_content_string().lower() for keyword in self.injection_patterns):
+            raise InputCheckError(
+                "Potential jailbreaking or prompt injection detected.",
+                check_trigger=CheckTrigger.PROMPT_INJECTION,
+            )
+
+    async def async_check(self, run_input: Union[RunInput, TeamRunInput]) -> None:
+        """Asynchronously check for prompt injection patterns in the input."""
+        if any(keyword in run_input.input_content_string().lower() for keyword in self.injection_patterns):
+            raise InputCheckError(
+                "Potential jailbreaking or prompt injection detected.",
+                check_trigger=CheckTrigger.PROMPT_INJECTION,
+            )

agno/knowledge/embedder/aws_bedrock.py

@@ -5,7 +5,7 @@ from typing import Any, Dict, List, Optional, Tuple
 
 from agno.exceptions import AgnoError, ModelProviderError
 from agno.knowledge.embedder.base import Embedder
-from agno.utils.log import log_error, logger
+from agno.utils.log import log_error, log_warning
 
 try:
     from boto3 import client as AwsClient
@@ -69,6 +69,11 @@ class AwsBedrockEmbedder(Embedder):
     client_params: Optional[Dict[str, Any]] = None
     client: Optional[AwsClient] = None
 
+    def __post_init__(self):
+        if self.enable_batch:
+            log_warning("AwsBedrockEmbedder does not support batch embeddings, setting enable_batch to False")
+            self.enable_batch = False
+
     def get_client(self) -> AwsClient:
         """
         Returns an AWS Bedrock client.
@@ -220,10 +225,10 @@ class AwsBedrockEmbedder(Embedder):
                     # Fallback to the first available embedding type
                     for embedding_type in response["embeddings"]:
                         return response["embeddings"][embedding_type][0]
-            logger.warning("No embeddings found in response")
+            log_warning("No embeddings found in response")
             return []
         except Exception as e:
-            logger.warning(f"Error extracting embeddings: {e}")
+            log_warning(f"Error extracting embeddings: {e}")
             return []
 
     def get_embedding_and_usage(self, text: str) -> Tuple[List[float], Optional[Dict[str, Any]]]:
@@ -286,7 +291,7 @@ class AwsBedrockEmbedder(Embedder):
                         # Fallback to the first available embedding type
                         for embedding_type in response_body["embeddings"]:
                             return response_body["embeddings"][embedding_type][0]
-                logger.warning("No embeddings found in response")
+                log_warning("No embeddings found in response")
                 return []
         except ClientError as e:
             log_error(f"Unexpected error calling Bedrock API: {str(e)}")

agno/knowledge/embedder/azure_openai.py

@@ -154,3 +154,57 @@ class AzureOpenAIEmbedder(Embedder):
         embedding = response.data[0].embedding
         usage = response.usage
         return embedding, usage.model_dump()
+
+    async def async_get_embeddings_batch_and_usage(
+        self, texts: List[str]
+    ) -> Tuple[List[List[float]], List[Optional[Dict]]]:
+        """
+        Get embeddings and usage for multiple texts in batches.
+
+        Args:
+            texts: List of text strings to embed
+
+        Returns:
+            Tuple of (List of embedding vectors, List of usage dictionaries)
+        """
+        all_embeddings = []
+        all_usage = []
+        logger.info(f"Getting embeddings and usage for {len(texts)} texts in batches of {self.batch_size}")
+
+        for i in range(0, len(texts), self.batch_size):
+            batch_texts = texts[i : i + self.batch_size]
+
+            req: Dict[str, Any] = {
+                "input": batch_texts,
+                "model": self.id,
+                "encoding_format": self.encoding_format,
+            }
+            if self.user is not None:
+                req["user"] = self.user
+            if self.id.startswith("text-embedding-3"):
+                req["dimensions"] = self.dimensions
+            if self.request_params:
+                req.update(self.request_params)
+
+            try:
+                response: CreateEmbeddingResponse = await self.aclient.embeddings.create(**req)
+                batch_embeddings = [data.embedding for data in response.data]
+                all_embeddings.extend(batch_embeddings)
+
+                # For each embedding in the batch, add the same usage information
+                usage_dict = response.usage.model_dump() if response.usage else None
+                all_usage.extend([usage_dict] * len(batch_embeddings))
+            except Exception as e:
+                logger.warning(f"Error in async batch embedding: {e}")
+                # Fallback to individual calls for this batch
+                for text in batch_texts:
+                    try:
+                        embedding, usage = await self.async_get_embedding_and_usage(text)
+                        all_embeddings.append(embedding)
+                        all_usage.append(usage)
+                    except Exception as e2:
+                        logger.warning(f"Error in individual async embedding fallback: {e2}")
+                        all_embeddings.append([])
+                        all_usage.append(None)
+
+        return all_embeddings, all_usage

agno/knowledge/embedder/base.py

@@ -7,6 +7,8 @@ class Embedder:
     """Base class for managing embedders"""
 
     dimensions: Optional[int] = 1536
+    enable_batch: bool = False
+    batch_size: int = 100  # Number of texts to process in each API call
 
     def get_embedding(self, text: str) -> List[float]:
         raise NotImplementedError