agno 0.1.2__py3-none-any.whl → 2.3.13__py3-none-any.whl

agno/storage/agent/dynamodb.py

@@ -1,350 +0,0 @@
-import time
-from dataclasses import asdict
-from decimal import Decimal
-from typing import Any, Dict, List, Optional
-
-from agno.storage.agent.base import AgentStorage
-from agno.storage.agent.session import AgentSession
-from agno.utils.log import logger
-
-try:
-    import boto3
-    from boto3.dynamodb.conditions import Key
-    from botocore.exceptions import ClientError
-except ImportError:
-    raise ImportError("`boto3` not installed. Please install using `pip install boto3`.")
-
-
-class DynamoDbAgentStorage(AgentStorage):
-    def __init__(
-        self,
-        table_name: str,
-        region_name: Optional[str] = None,
-        aws_access_key_id: Optional[str] = None,
-        aws_secret_access_key: Optional[str] = None,
-        endpoint_url: Optional[str] = None,
-        create_table_if_not_exists: bool = True,
-    ):
-        """
-        Initialize the DynamoDbAgentStorage.
-
-        Args:
-            table_name (str): The name of the DynamoDB table.
-            region_name (Optional[str]): AWS region name.
-            aws_access_key_id (Optional[str]): AWS access key ID.
-            aws_secret_access_key (Optional[str]): AWS secret access key.
-            endpoint_url (Optional[str]): The complete URL to use for the constructed client.
-            create_table_if_not_exists (bool): Whether to create the table if it does not exist.
-        """
-        self.table_name = table_name
-        self.region_name = region_name
-        self.endpoint_url = endpoint_url
-        self.aws_access_key_id = aws_access_key_id
-        self.aws_secret_access_key = aws_secret_access_key
-        self.create_table_if_not_exists = create_table_if_not_exists
-
-        # Initialize DynamoDB resource
-        self.dynamodb = boto3.resource(
-            "dynamodb",
-            region_name=self.region_name,
-            aws_access_key_id=self.aws_access_key_id,
-            aws_secret_access_key=self.aws_secret_access_key,
-            endpoint_url=self.endpoint_url,
-        )
-
-        # Initialize table
-        self.table = self.dynamodb.Table(self.table_name)
-
-        # Optionally create table if it does not exist
-        if self.create_table_if_not_exists:
-            self.create()
-        logger.debug(f"Initialized DynamoDbAgentStorage with table '{self.table_name}'")
-
-    def create(self) -> None:
-        """
-        Create the DynamoDB table if it does not exist.
-        """
-        try:
-            # Check if table exists
-            self.dynamodb.meta.client.describe_table(TableName=self.table_name)
-            logger.debug(f"Table '{self.table_name}' already exists.")
-        except ClientError as e:
-            if e.response["Error"]["Code"] == "ResourceNotFoundException":
-                logger.debug(f"Creating table '{self.table_name}'.")
-                # Create the table
-                self.table = self.dynamodb.create_table(
-                    TableName=self.table_name,
-                    KeySchema=[{"AttributeName": "session_id", "KeyType": "HASH"}],
-                    AttributeDefinitions=[
-                        {"AttributeName": "session_id", "AttributeType": "S"},
-                        {"AttributeName": "user_id", "AttributeType": "S"},
-                        {"AttributeName": "agent_id", "AttributeType": "S"},
-                        {"AttributeName": "created_at", "AttributeType": "N"},
-                    ],
-                    GlobalSecondaryIndexes=[
-                        {
-                            "IndexName": "user_id-index",
-                            "KeySchema": [
-                                {"AttributeName": "user_id", "KeyType": "HASH"},
-                                {"AttributeName": "created_at", "KeyType": "RANGE"},
-                            ],
-                            "Projection": {"ProjectionType": "ALL"},
-                            "ProvisionedThroughput": {
-                                "ReadCapacityUnits": 5,
-                                "WriteCapacityUnits": 5,
-                            },
-                        },
-                        {
-                            "IndexName": "agent_id-index",
-                            "KeySchema": [
-                                {"AttributeName": "agent_id", "KeyType": "HASH"},
-                                {"AttributeName": "created_at", "KeyType": "RANGE"},
-                            ],
-                            "Projection": {"ProjectionType": "ALL"},
-                            "ProvisionedThroughput": {
-                                "ReadCapacityUnits": 5,
-                                "WriteCapacityUnits": 5,
-                            },
-                        },
-                    ],
-                    ProvisionedThroughput={"ReadCapacityUnits": 5, "WriteCapacityUnits": 5},
-                )
-                # Wait until the table exists.
-                self.table.wait_until_exists()
-                logger.debug(f"Table '{self.table_name}' created successfully.")
-            else:
-                logger.error(f"Unable to create table '{self.table_name}': {e.response['Error']['Message']}")
-        except Exception as e:
-            logger.error(f"Exception during table creation: {e}")
-
-    def read(self, session_id: str, user_id: Optional[str] = None) -> Optional[AgentSession]:
-        """
-        Read and return an AgentSession from the database.
-
-        Args:
-            session_id (str): ID of the session to read.
-            user_id (Optional[str]): User ID to filter by. Defaults to None.
-
-        Returns:
-            Optional[AgentSession]: AgentSession object if found, None otherwise.
-        """
-        try:
-            key = {"session_id": session_id}
-            if user_id is not None:
-                key["user_id"] = user_id
-
-            response = self.table.get_item(Key=key)
-            item = response.get("Item", None)
-            if item is not None:
-                # Convert Decimal to int or float
-                item = self._deserialize_item(item)
-                return AgentSession.from_dict(item)
-        except Exception as e:
-            logger.error(f"Error reading session_id '{session_id}' with user_id '{user_id}': {e}")
-        return None
-
-    def get_all_session_ids(self, user_id: Optional[str] = None, agent_id: Optional[str] = None) -> List[str]:
-        """
-        Retrieve all session IDs, optionally filtered by user_id and/or agent_id.
-
-        Args:
-            user_id (Optional[str], optional): User ID to filter by. Defaults to None.
-            agent_id (Optional[str], optional): Agent ID to filter by. Defaults to None.
-
-        Returns:
-            List[str]: List of session IDs matching the criteria.
-        """
-        session_ids: List[str] = []
-        try:
-            if user_id is not None:
-                # Query using user_id index
-                response = self.table.query(
-                    IndexName="user_id-index",
-                    KeyConditionExpression=Key("user_id").eq(user_id),
-                    ProjectionExpression="session_id",
-                )
-                items = response.get("Items", [])
-                session_ids.extend([item["session_id"] for item in items if "session_id" in item])
-            elif agent_id is not None:
-                # Query using agent_id index
-                response = self.table.query(
-                    IndexName="agent_id-index",
-                    KeyConditionExpression=Key("agent_id").eq(agent_id),
-                    ProjectionExpression="session_id",
-                )
-                items = response.get("Items", [])
-                session_ids.extend([item["session_id"] for item in items if "session_id" in item])
-            else:
-                # Scan the whole table
-                response = self.table.scan(ProjectionExpression="session_id")
-                items = response.get("Items", [])
-                session_ids.extend([item["session_id"] for item in items if "session_id" in item])
-        except Exception as e:
-            logger.error(f"Error retrieving session IDs: {e}")
-        return session_ids
-
-    def get_all_sessions(self, user_id: Optional[str] = None, agent_id: Optional[str] = None) -> List[AgentSession]:
-        """
-        Retrieve all sessions, optionally filtered by user_id and/or agent_id.
-
-        Args:
-            user_id (Optional[str], optional): User ID to filter by. Defaults to None.
-            agent_id (Optional[str], optional): Agent ID to filter by. Defaults to None.
-
-        Returns:
-            List[AgentSession]: List of AgentSession objects matching the criteria.
-        """
-        sessions: List[AgentSession] = []
-        try:
-            if user_id is not None:
-                # Query using user_id index
-                response = self.table.query(
-                    IndexName="user_id-index",
-                    KeyConditionExpression=Key("user_id").eq(user_id),
-                    ProjectionExpression="session_id, agent_id, user_id, memory, agent_data, session_data, extra_data, created_at, updated_at",
-                )
-                items = response.get("Items", [])
-                for item in items:
-                    item = self._deserialize_item(item)
-                    _agent_session = AgentSession.from_dict(item)
-                    if _agent_session is not None:
-                        sessions.append(_agent_session)
-            elif agent_id is not None:
-                # Query using agent_id index
-                response = self.table.query(
-                    IndexName="agent_id-index",
-                    KeyConditionExpression=Key("agent_id").eq(agent_id),
-                    ProjectionExpression="session_id, agent_id, user_id, memory, agent_data, session_data, extra_data, created_at, updated_at",
-                )
-                items = response.get("Items", [])
-                for item in items:
-                    item = self._deserialize_item(item)
-                    _agent_session = AgentSession.from_dict(item)
-                    if _agent_session is not None:
-                        sessions.append(_agent_session)
-            else:
-                # Scan the whole table
-                response = self.table.scan(
-                    ProjectionExpression="session_id, agent_id, user_id, memory, agent_data, session_data, extra_data, created_at, updated_at"
-                )
-                items = response.get("Items", [])
-                for item in items:
-                    item = self._deserialize_item(item)
-                    _agent_session = AgentSession.from_dict(item)
-                    if _agent_session is not None:
-                        sessions.append(_agent_session)
-        except Exception as e:
-            logger.error(f"Error retrieving sessions: {e}")
-        return sessions
-
-    def upsert(self, session: AgentSession) -> Optional[AgentSession]:
-        """
-        Create or update an AgentSession in the database.
-
-        Args:
-            session (AgentSession): The session data to upsert.
-
-        Returns:
-            Optional[AgentSession]: The upserted AgentSession, or None if operation failed.
-        """
-        try:
-            item = asdict(session)
-
-            # Add timestamps
-            current_time = int(time.time())
-            if "created_at" not in item or item["created_at"] is None:
-                item["created_at"] = current_time
-            item["updated_at"] = current_time
-
-            # Convert data to DynamoDB compatible format
-            item = self._serialize_item(item)
-
-            # Put item into DynamoDB
-            self.table.put_item(Item=item)
-            return self.read(session.session_id)
-        except Exception as e:
-            logger.error(f"Error upserting session: {e}")
-            return None
-
-    def delete_session(self, session_id: Optional[str] = None):
-        """
-        Delete a session from the database.
-
-        Args:
-            session_id (Optional[str], optional): ID of the session to delete. Defaults to None.
-        """
-        if session_id is None:
-            logger.warning("No session_id provided for deletion.")
-            return
-        try:
-            self.table.delete_item(Key={"session_id": session_id})
-            logger.info(f"Successfully deleted session with session_id: {session_id}")
-        except Exception as e:
-            logger.error(f"Error deleting session: {e}")
-
-    def drop(self) -> None:
-        """
-        Drop the table from the database if it exists.
-        """
-        try:
-            self.table.delete()
-            self.table.wait_until_not_exists()
-            logger.debug(f"Table '{self.table_name}' deleted successfully.")
-        except Exception as e:
-            logger.error(f"Error deleting table '{self.table_name}': {e}")
-
-    def upgrade_schema(self) -> None:
-        """
-        Upgrade the schema to the latest version.
-        This method is currently a placeholder and does not perform any actions.
-        """
-        pass
-
-    def _serialize_item(self, item: Dict[str, Any]) -> Dict[str, Any]:
-        """
-        Serialize item to be compatible with DynamoDB.
-
-        Args:
-            item (Dict[str, Any]): The item to serialize.
-
-        Returns:
-            Dict[str, Any]: The serialized item.
-        """
-
-        def serialize_value(value):
-            if isinstance(value, float):
-                return Decimal(str(value))
-            elif isinstance(value, dict):
-                return {k: serialize_value(v) for k, v in value.items()}
-            elif isinstance(value, list):
-                return [serialize_value(v) for v in value]
-            else:
-                return value
-
-        return {k: serialize_value(v) for k, v in item.items() if v is not None}
-
-    def _deserialize_item(self, item: Dict[str, Any]) -> Dict[str, Any]:
-        """
-        Deserialize item from DynamoDB format.
-
-        Args:
-            item (Dict[str, Any]): The item to deserialize.
-
-        Returns:
-            Dict[str, Any]: The deserialized item.
-        """
-
-        def deserialize_value(value):
-            if isinstance(value, Decimal):
-                if value % 1 == 0:
-                    return int(value)
-                else:
-                    return float(value)
-            elif isinstance(value, dict):
-                return {k: deserialize_value(v) for k, v in value.items()}
-            elif isinstance(value, list):
-                return [deserialize_value(v) for v in value]
-            else:
-                return value
-
-        return {k: deserialize_value(v) for k, v in item.items()}

agno/storage/agent/json.py

@@ -1,92 +0,0 @@
-import json
-import time
-from dataclasses import asdict
-from pathlib import Path
-from typing import List, Optional, Union
-
-from agno.storage.agent.base import AgentStorage
-from agno.storage.agent.session import AgentSession
-from agno.utils.log import logger
-
-
-class JsonAgentStorage(AgentStorage):
-    def __init__(self, dir_path: Union[str, Path]):
-        self.dir_path = Path(dir_path)
-        self.dir_path.mkdir(parents=True, exist_ok=True)
-
-    def serialize(self, data: dict) -> str:
-        return json.dumps(data, ensure_ascii=False, indent=4)
-
-    def deserialize(self, data: str) -> dict:
-        return json.loads(data)
-
-    def create(self) -> None:
-        """Create the storage if it doesn't exist."""
-        if not self.dir_path.exists():
-            self.dir_path.mkdir(parents=True, exist_ok=True)
-
-    def read(self, session_id: str, user_id: Optional[str] = None) -> Optional[AgentSession]:
-        """Read an AgentSession from storage."""
-        try:
-            with open(self.dir_path / f"{session_id}.json", "r", encoding="utf-8") as f:
-                data = self.deserialize(f.read())
-                if user_id and data["user_id"] != user_id:
-                    return None
-                return AgentSession.from_dict(data)
-        except FileNotFoundError:
-            return None
-
-    def get_all_session_ids(self, user_id: Optional[str] = None, agent_id: Optional[str] = None) -> List[str]:
-        """Get all session IDs, optionally filtered by user_id and/or agent_id."""
-        session_ids = []
-        for file in self.dir_path.glob("*.json"):
-            with open(file, "r", encoding="utf-8") as f:
-                data = self.deserialize(f.read())
-                if (not user_id or data["user_id"] == user_id) and (not agent_id or data["agent_id"] == agent_id):
-                    session_ids.append(data["session_id"])
-        return session_ids
-
-    def get_all_sessions(self, user_id: Optional[str] = None, agent_id: Optional[str] = None) -> List[AgentSession]:
-        """Get all sessions, optionally filtered by user_id and/or agent_id."""
-        sessions = []
-        for file in self.dir_path.glob("*.json"):
-            with open(file, "r", encoding="utf-8") as f:
-                data = self.deserialize(f.read())
-                if (not user_id or data["user_id"] == user_id) and (not agent_id or data["agent_id"] == agent_id):
-                    _agent_session = AgentSession.from_dict(data)
-                    if _agent_session is not None:
-                        sessions.append(_agent_session)
-        return sessions
-
-    def upsert(self, session: AgentSession) -> Optional[AgentSession]:
-        """Insert or update an AgentSession in storage."""
-        try:
-            data = asdict(session)
-            data["updated_at"] = int(time.time())
-            if "created_at" not in data:
-                data["created_at"] = data["updated_at"]
-
-            with open(self.dir_path / f"{session.session_id}.json", "w", encoding="utf-8") as f:
-                f.write(self.serialize(data))
-            return session
-        except Exception as e:
-            logger.error(f"Error upserting session: {e}")
-            return None
-
-    def delete_session(self, session_id: Optional[str] = None):
-        """Delete a session from storage."""
-        if session_id is None:
-            return
-        try:
-            (self.dir_path / f"{session_id}.json").unlink(missing_ok=True)
-        except Exception as e:
-            logger.error(f"Error deleting session: {e}")
-
-    def drop(self) -> None:
-        """Drop all sessions from storage."""
-        for file in self.dir_path.glob("*.json"):
-            file.unlink()
-
-    def upgrade_schema(self) -> None:
-        """Upgrade the schema of the storage."""
-        pass

agno/storage/agent/mongodb.py

@@ -1,228 +0,0 @@
-from datetime import datetime, timezone
-from typing import List, Optional
-from uuid import UUID
-
-try:
-    from pymongo import MongoClient
-    from pymongo.collection import Collection
-    from pymongo.database import Database
-    from pymongo.errors import PyMongoError
-except ImportError:
-    raise ImportError("`pymongo` not installed. Please install it with `pip install pymongo`")
-
-from agno.storage.agent.base import AgentStorage
-from agno.storage.agent.session import AgentSession
-from agno.utils.log import logger
-
-
-class MongoDbAgentStorage(AgentStorage):
-    def __init__(
-        self,
-        collection_name: str,
-        db_url: Optional[str] = None,
-        db_name: str = "agno",
-        client: Optional[MongoClient] = None,
-    ):
-        """
-        This class provides agent storage using MongoDB.
-
-        Args:
-            collection_name: Name of the collection to store agent sessions
-            db_url: MongoDB connection URL
-            db_name: Name of the database
-            client: Optional existing MongoDB client
-        """
-        self._client: Optional[MongoClient] = client
-        if self._client is None and db_url is not None:
-            self._client = MongoClient(db_url)
-        elif self._client is None:
-            self._client = MongoClient()
-
-        if self._client is None:
-            raise ValueError("Must provide either db_url or client")
-
-        self.collection_name: str = collection_name
-        self.db_name: str = db_name
-        self.db: Database = self._client[self.db_name]
-        self.collection: Collection = self.db[self.collection_name]
-
-    def create(self) -> None:
-        """Create necessary indexes for the collection"""
-        try:
-            # Create indexes
-            self.collection.create_index("session_id", unique=True)
-            self.collection.create_index("user_id")
-            self.collection.create_index("agent_id")
-            self.collection.create_index("created_at")
-        except PyMongoError as e:
-            logger.error(f"Error creating indexes: {e}")
-            raise
-
-    def read(self, session_id: str, user_id: Optional[str] = None) -> Optional[AgentSession]:
-        """Read an agent session from MongoDB
-        Args:
-            session_id: ID of the session to read
-            user_id: ID of the user to read
-        Returns:
-            AgentSession: The session if found, otherwise None
-        """
-        try:
-            query = {"session_id": session_id}
-            if user_id:
-                query["user_id"] = user_id
-
-            doc = self.collection.find_one(query)
-            if doc:
-                # Remove MongoDB _id before converting to AgentSession
-                doc.pop("_id", None)
-                return AgentSession.from_dict(doc)
-            return None
-        except PyMongoError as e:
-            logger.error(f"Error reading session: {e}")
-            return None
-
-    def get_all_session_ids(self, user_id: Optional[str] = None, agent_id: Optional[str] = None) -> List[str]:
-        """Get all session IDs matching the criteria
-        Args:
-            user_id: ID of the user to read
-            agent_id: ID of the agent to read
-        Returns:
-            List[str]: List of session IDs
-        """
-        try:
-            query = {}
-            if user_id is not None:
-                query["user_id"] = user_id
-            if agent_id is not None:
-                query["agent_id"] = agent_id
-
-            cursor = self.collection.find(query, {"session_id": 1}).sort("created_at", -1)
-
-            return [str(doc["session_id"]) for doc in cursor]
-        except PyMongoError as e:
-            logger.error(f"Error getting session IDs: {e}")
-            return []
-
-    def get_all_sessions(self, user_id: Optional[str] = None, agent_id: Optional[str] = None) -> List[AgentSession]:
-        """Get all sessions matching the criteria
-        Args:
-            user_id: ID of the user to read
-            agent_id: ID of the agent to read
-        Returns:
-            List[AgentSession]: List of sessions
-        """
-        try:
-            query = {}
-            if user_id is not None:
-                query["user_id"] = user_id
-            if agent_id is not None:
-                query["agent_id"] = agent_id
-
-            cursor = self.collection.find(query).sort("created_at", -1)
-            sessions = []
-            for doc in cursor:
-                # Remove MongoDB _id before converting to AgentSession
-                doc.pop("_id", None)
-                _agent_session = AgentSession.from_dict(doc)
-                if _agent_session is not None:
-                    sessions.append(_agent_session)
-            return sessions
-        except PyMongoError as e:
-            logger.error(f"Error getting sessions: {e}")
-            return []
-
-    def upsert(self, session: AgentSession, create_and_retry: bool = True) -> Optional[AgentSession]:
-        """Upsert an agent session
-        Args:
-            session: AgentSession to upsert
-            create_and_retry: Whether to create a new session if the session_id already exists
-        Returns:
-            AgentSession: The session if upserted, otherwise None
-        """
-        try:
-            # Convert session to dict and add timestamps
-            session_dict = session.to_dict()
-            now = datetime.now(timezone.utc)
-            timestamp = int(now.timestamp())
-
-            # Handle UUID serialization
-            if isinstance(session.session_id, UUID):
-                session_dict["session_id"] = str(session.session_id)
-
-            # Add version field for optimistic locking
-            if "_version" not in session_dict:
-                session_dict["_version"] = 1
-            else:
-                session_dict["_version"] += 1
-
-            update_data = {**session_dict, "updated_at": timestamp}
-
-            # For new documents, set created_at
-            query = {"session_id": session_dict["session_id"]}
-
-            doc = self.collection.find_one(query)
-            if not doc:
-                update_data["created_at"] = timestamp
-
-            result = self.collection.update_one(query, {"$set": update_data}, upsert=True)
-
-            if result.acknowledged:
-                return self.read(session_id=session_dict["session_id"])
-            return None
-
-        except PyMongoError as e:
-            logger.error(f"Error upserting session: {e}")
-            return None
-
-    def delete_session(self, session_id: Optional[str] = None) -> None:
-        """Delete an agent session
-        Args:
-            session_id: ID of the session to delete
-        Returns:
-            None
-        """
-        if session_id is None:
-            logger.warning("No session_id provided for deletion")
-            return
-
-        try:
-            result = self.collection.delete_one({"session_id": session_id})
-            if result.deleted_count == 0:
-                logger.debug(f"No session found with session_id: {session_id}")
-            else:
-                logger.debug(f"Successfully deleted session with session_id: {session_id}")
-        except PyMongoError as e:
-            logger.error(f"Error deleting session: {e}")
-
-    def drop(self) -> None:
-        """Drop the collection
-        Returns:
-            None
-        """
-        try:
-            self.collection.drop()
-        except PyMongoError as e:
-            logger.error(f"Error dropping collection: {e}")
-
-    def upgrade_schema(self) -> None:
-        """Placeholder for schema upgrades"""
-        pass
-
-    def __deepcopy__(self, memo):
-        """Create a deep copy of the MongoDbAgentStorage instance"""
-        from copy import deepcopy
-
-        # Create a new instance without calling __init__
-        cls = self.__class__
-        copied_obj = cls.__new__(cls)
-        memo[id(self)] = copied_obj
-
-        # Deep copy attributes
-        for k, v in self.__dict__.items():
-            if k in {"_client", "db", "collection"}:
-                # Reuse MongoDB connections without copying
-                setattr(copied_obj, k, v)
-            else:
-                setattr(copied_obj, k, deepcopy(v, memo))
-
-        return copied_obj