agno 2.0.10__py3-none-any.whl → 2.1.0__py3-none-any.whl

agno/db/in_memory/in_memory_db.py

@@ -1,4 +1,5 @@
 import time
+from copy import deepcopy
 from datetime import date, datetime, timedelta, timezone
 from typing import Any, Dict, List, Optional, Tuple, Union
 from uuid import uuid4
@@ -107,15 +108,17 @@ class InMemoryDb(BaseDb):
                     if session_data.get("session_type") != session_type_value:
                         continue
 
+                    session_data_copy = deepcopy(session_data)
+
                     if not deserialize:
-                        return session_data
+                        return session_data_copy
 
                     if session_type == SessionType.AGENT:
-                        return AgentSession.from_dict(session_data)
+                        return AgentSession.from_dict(session_data_copy)
                     elif session_type == SessionType.TEAM:
-                        return TeamSession.from_dict(session_data)
+                        return TeamSession.from_dict(session_data_copy)
                     else:
-                        return WorkflowSession.from_dict(session_data)
+                        return WorkflowSession.from_dict(session_data_copy)
 
             return None
 
@@ -188,7 +191,7 @@ class InMemoryDb(BaseDb):
                 if session_data.get("session_type") != session_type_value:
                     continue
 
-                filtered_sessions.append(session_data)
+                filtered_sessions.append(deepcopy(session_data))
 
             total_count = len(filtered_sessions)
 
@@ -233,15 +236,16 @@ class InMemoryDb(BaseDb):
 
                     log_debug(f"Renamed session with id '{session_id}' to '{session_name}'")
 
+                    session_copy = deepcopy(session)
                     if not deserialize:
-                        return session
+                        return session_copy
 
                     if session_type == SessionType.AGENT:
-                        return AgentSession.from_dict(session)
+                        return AgentSession.from_dict(session_copy)
                     elif session_type == SessionType.TEAM:
-                        return TeamSession.from_dict(session)
+                        return TeamSession.from_dict(session_copy)
                     else:
-                        return WorkflowSession.from_dict(session)
+                        return WorkflowSession.from_dict(session_copy)
 
             return None
 
@@ -269,22 +273,26 @@ class InMemoryDb(BaseDb):
                 if existing_session.get("session_id") == session_dict.get("session_id") and self._matches_session_key(
                     existing_session, session
                 ):
-                    # Update existing session
                     session_dict["updated_at"] = int(time.time())
-                    self._sessions[i] = session_dict
+                    self._sessions[i] = deepcopy(session_dict)
                     session_updated = True
                     break
 
             if not session_updated:
-                # Add new session
                 session_dict["created_at"] = session_dict.get("created_at", int(time.time()))
                 session_dict["updated_at"] = session_dict.get("created_at")
-                self._sessions.append(session_dict)
+                self._sessions.append(deepcopy(session_dict))
 
+            session_dict_copy = deepcopy(session_dict)
             if not deserialize:
-                return session_dict
+                return session_dict_copy
 
-            return session
+            if session_dict_copy["session_type"] == SessionType.AGENT:
+                return AgentSession.from_dict(session_dict_copy)
+            elif session_dict_copy["session_type"] == SessionType.TEAM:
+                return TeamSession.from_dict(session_dict_copy)
+            else:
+                return WorkflowSession.from_dict(session_dict_copy)
 
         except Exception as e:
             log_error(f"Exception upserting session: {e}")
@@ -378,9 +386,10 @@ class InMemoryDb(BaseDb):
         try:
             for memory_data in self._memories:
                 if memory_data.get("memory_id") == memory_id:
+                    memory_data_copy = deepcopy(memory_data)
                     if not deserialize:
-                        return memory_data
-                    return UserMemory.from_dict(memory_data)
+                        return memory_data_copy
+                    return UserMemory.from_dict(memory_data_copy)
 
             return None
 
@@ -420,7 +429,7 @@ class InMemoryDb(BaseDb):
                     if search_content.lower() not in memory_content.lower():
                         continue
 
-                filtered_memories.append(memory_data)
+                filtered_memories.append(deepcopy(memory_data))
 
             total_count = len(filtered_memories)
 
@@ -499,9 +508,11 @@ class InMemoryDb(BaseDb):
             if not memory_updated:
                 self._memories.append(memory_dict)
 
+            memory_dict_copy = deepcopy(memory_dict)
             if not deserialize:
-                return memory_dict
-            return UserMemory.from_dict(memory_dict)
+                return memory_dict_copy
+
+            return UserMemory.from_dict(memory_dict_copy)
 
         except Exception as e:
             log_warning(f"Exception upserting user memory: {e}")
@@ -657,8 +668,8 @@ class InMemoryDb(BaseDb):
                 # Only include necessary fields for metrics
                 filtered_session = {
                     "user_id": session.get("user_id"),
-                    "session_data": session.get("session_data"),
-                    "runs": session.get("runs"),
+                    "session_data": deepcopy(session.get("session_data")),
+                    "runs": deepcopy(session.get("runs")),
                     "created_at": session.get("created_at"),
                     "session_type": session.get("session_type"),
                 }
@@ -688,7 +699,7 @@ class InMemoryDb(BaseDb):
                 if ending_date and metric_date > ending_date:
                     continue
 
-                filtered_metrics.append(metric)
+                filtered_metrics.append(deepcopy(metric))
 
                 updated_at = metric.get("updated_at")
                 if updated_at and (latest_updated_at is None or updated_at > latest_updated_at):
@@ -763,7 +774,7 @@ class InMemoryDb(BaseDb):
             Exception: If an error occurs during retrieval.
         """
         try:
-            knowledge_items = self._knowledge.copy()
+            knowledge_items = [deepcopy(item) for item in self._knowledge]
 
             total_count = len(knowledge_items)
 
@@ -858,9 +869,10 @@ class InMemoryDb(BaseDb):
         try:
             for run_data in self._eval_runs:
                 if run_data.get("run_id") == eval_run_id:
+                    run_data_copy = deepcopy(run_data)
                     if not deserialize:
-                        return run_data
-                    return EvalRunRecord.model_validate(run_data)
+                        return run_data_copy
+                    return EvalRunRecord.model_validate(run_data_copy)
 
             return None
 
@@ -906,7 +918,7 @@ class InMemoryDb(BaseDb):
                     elif filter_type == EvalFilterType.WORKFLOW and run_data.get("workflow_id") is None:
                         continue
 
-                filtered_runs.append(run_data)
+                filtered_runs.append(deepcopy(run_data))
 
             total_count = len(filtered_runs)
 
@@ -945,10 +957,11 @@ class InMemoryDb(BaseDb):
 
                     log_debug(f"Renamed eval run with id '{eval_run_id}' to '{name}'")
 
+                    run_data_copy = deepcopy(run_data)
                     if not deserialize:
-                        return run_data
+                        return run_data_copy
 
-                    return EvalRunRecord.model_validate(run_data)
+                    return EvalRunRecord.model_validate(run_data_copy)
 
             return None
 

agno/db/postgres/postgres.py

@@ -756,7 +756,7 @@ class PostgresDb(BaseDb):
                     )
 
                 with self.Session() as sess, sess.begin():
-                    stmt = postgresql.insert(table)
+                    stmt: Any = postgresql.insert(table)
                     update_columns = {
                         col.name: stmt.excluded[col.name]
                         for col in table.columns
@@ -1263,13 +1263,15 @@ class PostgresDb(BaseDb):
             results: List[Union[UserMemory, Dict[str, Any]]] = []
 
             with self.Session() as sess, sess.begin():
-                stmt = postgresql.insert(table)
+                insert_stmt = postgresql.insert(table)
                 update_columns = {
-                    col.name: stmt.excluded[col.name]
+                    col.name: insert_stmt.excluded[col.name]
                     for col in table.columns
                     if col.name not in ["memory_id"]  # Don't update primary key
                 }
-                stmt = stmt.on_conflict_do_update(index_elements=["memory_id"], set_=update_columns).returning(table)
+                stmt = insert_stmt.on_conflict_do_update(index_elements=["memory_id"], set_=update_columns).returning(
+                    table
+                )
 
                 result = sess.execute(stmt, memory_records)
                 for row in result.fetchall():

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,
+            )