agno 1.8.1__py3-none-any.whl → 2.0.0a1__py3-none-any.whl

agno/storage/firestore.py

@@ -1,297 +0,0 @@
-from datetime import datetime, timezone
-from typing import Any, Dict, List, Literal, Optional, Union
-from uuid import UUID
-
-from agno.storage.base import Storage
-from agno.storage.session import Session
-from agno.storage.session.agent import AgentSession
-from agno.storage.session.team import TeamSession
-from agno.storage.session.v2.workflow import WorkflowSession as WorkflowSessionV2
-from agno.storage.session.workflow import WorkflowSession
-from agno.utils.log import log_debug, logger
-
-try:
-    from google.api_core import exceptions as google_exceptions
-    from google.cloud.firestore_v1 import (
-        Client,
-        CollectionReference,
-        DocumentReference,
-        Query,
-    )
-    from google.cloud.firestore_v1.base_query import FieldFilter
-except ImportError:
-    raise ImportError("`firestore` not installed. Please install it with `pip install google-cloud-firestore`")
-
-
-class FirestoreStorage(Storage):
-    def __init__(
-        self,
-        collection_name: str,
-        db_name: Optional[str] = "(default)",
-        project_id: Optional[str] = None,
-        client: Optional[Client] = None,
-        mode: Optional[Literal["agent", "team", "workflow", "workflow_v2"]] = "agent",
-    ):
-        super().__init__(mode)
-        self.collection_name = collection_name
-        self.db_name = db_name
-        self.project_id = project_id
-
-        # Initialize Firestore client
-        self._client = client
-        if self._client is None:
-            self._client = self._initialize_client()
-
-        # Get collection reference
-        self.collection: CollectionReference = self._client.collection(self.collection_name)
-        log_debug(f"Created FirestoreStorage with collection: '{self.collection_name}'")
-
-    def _initialize_client(self) -> Client:
-        """Initialize and return a Firestore client with proper error handling."""
-        try:
-            client = Client(database=self.db_name, project=self.project_id)
-            log_debug(f"Firestore client initialized with database: '{self.db_name}'")
-            return client
-        except google_exceptions.Unauthenticated as e:
-            raise ImportError(
-                "Failed to authenticate with Google Cloud. Please set up authentication:\n"
-                "1. Run: gcloud auth application-default login\n"
-                "2. Or set GOOGLE_APPLICATION_CREDENTIALS environment variable"
-            ) from e
-        except google_exceptions.PermissionDenied as e:
-            raise ImportError(
-                "Permission denied when accessing Firestore. Please ensure:\n"
-                "1. Your service account has the 'Cloud Datastore User' role\n"
-                "2. The Firestore API is enabled for your project"
-            ) from e
-        except Exception as e:
-            raise ImportError(f"Failed to initialize Firestore client: {e}") from e
-
-    def _delete_document(self, document: DocumentReference) -> None:
-        """Recursively delete a document and all its subcollections."""
-        log_debug(f"Deleting document: {document.path}")
-        for collection in document.collections():
-            self._delete_collection(collection)
-        document.delete()
-
-    def _delete_collection(self, collection: CollectionReference) -> None:
-        """Recursively delete all documents in a collection."""
-        for document in collection.list_documents():
-            self._delete_document(document)
-
-    def _build_query(
-        self, base_query: CollectionReference, user_id: Optional[str] = None, entity_id: Optional[str] = None
-    ) -> Union[Query, CollectionReference]:
-        """Build a Firestore query with optional filters."""
-        query: Union[Query, CollectionReference] = base_query
-
-        if user_id:
-            query = query.where(filter=FieldFilter("user_id", "==", user_id))
-
-        if entity_id:
-            if self.mode == "agent":
-                query = query.where(filter=FieldFilter("agent_id", "==", entity_id))
-            elif self.mode == "team":
-                query = query.where(filter=FieldFilter("team_id", "==", entity_id))
-            elif self.mode == "workflow":
-                query = query.where(filter=FieldFilter("workflow_id", "==", entity_id))
-            elif self.mode == "workflow_v2":
-                query = query.where(filter=FieldFilter("workflow_id", "==", entity_id))
-
-        return query
-
-    def _parse_session(self, doc_data: Optional[Dict[str, Any]]) -> Optional[Session]:
-        """Parse document data into appropriate Session type."""
-        if not doc_data:
-            return None
-
-        try:
-            if self.mode == "agent":
-                return AgentSession.from_dict(doc_data)
-            elif self.mode == "team":
-                return TeamSession.from_dict(doc_data)
-            elif self.mode == "workflow":
-                return WorkflowSession.from_dict(doc_data)
-            elif self.mode == "workflow_v2":
-                return WorkflowSessionV2.from_dict(doc_data)
-        except Exception as e:
-            logger.error(f"Error parsing session data: {e}")
-            return None
-
-    def create(self) -> None:
-        """
-        Create storage if it doesn't exist.
-        For Firestore, this is a no-operation as collections are created automatically.
-        """
-        try:
-            if self._client:
-                list(self._client.collections())
-            log_debug("Firestore connection successful")
-        except Exception as e:
-            logger.error(f"Could not connect to Firestore: {e}")
-            raise
-
-    def read(self, session_id: str, user_id: Optional[str] = None) -> Optional[Session]:
-        """Read a session from Firestore."""
-        try:
-            query = self.collection.where(filter=FieldFilter("session_id", "==", session_id))
-
-            if user_id:
-                query = query.where(filter=FieldFilter("user_id", "==", user_id))
-
-            docs = query.get()
-            for doc in docs:
-                doc_dict = doc.to_dict()
-                if doc_dict:
-                    return self._parse_session(doc_dict)
-
-            return None
-        except Exception as e:
-            logger.error(f"Error reading session: {e}")
-            return None
-
-    def get_all_session_ids(self, user_id: Optional[str] = None, entity_id: Optional[str] = None) -> List[str]:
-        """Get all session IDs, optionally filtered by user_id and/or entity_id."""
-        try:
-            query = self._build_query(self.collection, user_id, entity_id)
-            docs = query.get()
-
-            # Sort by created_at descending and extract session IDs
-            session_ids: List[str] = []
-            doc_list = []
-            for doc in docs:
-                doc_data = doc.to_dict()
-                if doc_data:
-                    doc_list.append(doc_data)
-
-            sorted_docs = sorted(doc_list, key=lambda x: x.get("created_at", 0), reverse=True)
-
-            for doc_data in sorted_docs:
-                session_id = doc_data.get("session_id")
-                if session_id:
-                    session_ids.append(session_id)
-
-            return session_ids
-        except Exception as e:
-            logger.error(f"Error getting session IDs: {e}")
-            return []
-
-    def get_all_sessions(self, user_id: Optional[str] = None, entity_id: Optional[str] = None) -> List[Session]:
-        """Get all sessions, optionally filtered by user_id and/or entity_id."""
-        sessions: List[Session] = []
-        try:
-            query = self._build_query(self.collection, user_id, entity_id)
-            docs = query.get()
-
-            # Sort by created_at descending
-            doc_list = []
-            for doc in docs:
-                doc_data = doc.to_dict()
-                if doc_data:
-                    doc_list.append(doc_data)
-
-            sorted_docs = sorted(doc_list, key=lambda x: x.get("created_at", 0), reverse=True)
-
-            for doc_data in sorted_docs:
-                session = self._parse_session(doc_data)
-                if session:
-                    sessions.append(session)
-
-            return sessions
-        except Exception as e:
-            logger.error(f"Error getting all sessions: {e}")
-            return []
-
-    def get_recent_sessions(
-        self,
-        user_id: Optional[str] = None,
-        entity_id: Optional[str] = None,
-        limit: Optional[int] = 2,
-    ) -> List[Session]:
-        """Get the last N sessions, ordered by created_at descending."""
-        sessions: List[Session] = []
-        try:
-            query = self._build_query(self.collection, user_id, entity_id)
-            docs = query.get()
-
-            # Sort by created_at descending
-            doc_list = []
-            for doc in docs:
-                doc_data = doc.to_dict()
-                if doc_data:
-                    doc_list.append(doc_data)
-
-            sorted_docs = sorted(doc_list, key=lambda x: x.get("created_at", 0), reverse=True)
-
-            # Apply limit
-            if limit is not None:
-                sorted_docs = sorted_docs[:limit]
-
-            for doc_data in sorted_docs:
-                session = self._parse_session(doc_data)
-                if session:
-                    sessions.append(session)
-
-            return sessions
-        except Exception as e:
-            logger.error(f"Error getting recent sessions: {e}")
-            return []
-
-    def upsert(self, session: Session, create_and_retry: bool = True) -> Optional[Session]:
-        """Insert or update a session in Firestore."""
-        try:
-            # Prepare session data
-            session_dict = session.to_dict()
-            now = datetime.now(timezone.utc)
-            timestamp = int(now.timestamp())
-
-            if isinstance(session.session_id, UUID):
-                session_dict["session_id"] = str(session.session_id)
-
-            # Add timestamps
-            session_dict["updated_at"] = timestamp
-
-            # Get document reference
-            doc_ref = self.collection.document(session_dict["session_id"])
-            doc = doc_ref.get()
-
-            # Add created_at for new documents
-            if not doc.exists:
-                session_dict["created_at"] = timestamp
-
-            # Save to Firestore
-            doc_ref.set(session_dict)
-
-            # Return the updated session
-            return self.read(session_id=session_dict["session_id"])
-
-        except Exception as e:
-            logger.error(f"Error upserting session: {e}")
-            return None
-
-    def delete_session(self, session_id: Optional[str] = None) -> None:
-        """Delete a session from Firestore."""
-        if session_id is None:
-            logger.warning("No session_id provided for deletion")
-            return
-
-        try:
-            self.collection.document(session_id).delete()
-            log_debug(f"Deleted session: {session_id}")
-        except Exception as e:
-            logger.error(f"Error deleting session: {e}")
-
-    def drop(self) -> None:
-        """Drop all sessions from storage."""
-        try:
-            self._delete_collection(self.collection)
-            log_debug(f"Dropped all sessions in collection: {self.collection_name}")
-        except Exception as e:
-            logger.error(f"Error dropping collection: {e}")
-
-    def upgrade_schema(self) -> None:
-        """
-        Upgrade the schema of the storage.
-        For Firestore, this is a no-op as it's schema-less.
-        """
-        pass

agno/storage/gcs_json.py

@@ -1,261 +0,0 @@
-import json
-import time
-from typing import Any, List, Literal, Optional
-
-from agno.storage.json import JsonStorage, Storage
-from agno.storage.session import Session
-from agno.storage.session.agent import AgentSession
-from agno.storage.session.team import TeamSession
-from agno.storage.session.v2.workflow import WorkflowSession as WorkflowSessionV2
-from agno.storage.session.workflow import WorkflowSession
-from agno.utils.log import logger
-
-try:
-    from google.cloud import storage as gcs
-except ImportError:
-    raise ImportError("`google-cloud-storage` not installed. Please install it with `pip install google-cloud-storage`")
-
-
-class GCSJsonStorage(JsonStorage):
-    """
-    A Cloud-based JSON storage for agent sessions that stores session (memory) data
-    in a GCS bucket. This class derives from JsonStorage and replaces local
-    file system operations with Cloud Storage operations. The GCS client and bucket
-    are initialized once in the constructor and then reused for all subsequent operations.
-
-    Parameters:
-      - bucket_name: The GCS bucket name (must be provided).
-      - prefix: The GCS folder path prefix). See (Flat Namespace)[https://cloud.google.com/storage/docs/objects#flat-namespace] in GCS docs for details.
-      - mode: One of "agent", "team", or "workflow". Defaults to "agent".
-      - project: Optional; the GCP project ID. Defaults to current Google Cloud's project (set with `gcloud init`).
-      - location: Optional; the GCP location for the bucket. Default's to current project's location.
-      - credentials: Optional credentials object; if not provided, defaults will be used.
-    """
-
-    def __init__(
-        self,
-        bucket_name: str,
-        prefix: Optional[str] = "",
-        mode: Optional[Literal["agent", "team", "workflow", "workflow_v2"]] = "agent",
-        project: Optional[str] = None,
-        location: Optional[str] = None,
-        credentials: Optional[Any] = None,
-    ):
-        # Call Storage's __init__ directly to bypass the folder creation logic in JsonStorage.
-        Storage.__init__(self, mode=mode)
-        self.bucket_name = bucket_name
-        if prefix is not None and prefix != "" and not prefix.endswith("/"):
-            prefix += "/"
-        self.prefix = prefix
-        self.project = project
-        self.location = location
-
-        # Initialize the GCS client once; if STORAGE_EMULATOR_HOST is set, it will be used automatically.
-        self.client = gcs.Client(project=self.project, credentials=credentials)
-        self.bucket = self.client.bucket(self.bucket_name)
-
-    def _get_blob_path(self, session_id: str) -> str:
-        """Returns the blob path for a given session."""
-        return f"{self.prefix}{session_id}.json"
-
-    def create(self) -> None:
-        """
-        Creates the bucket if it doesn't exist
-        The client and bucket are already stored in self.
-        """
-        try:
-            self.bucket = self.client.create_bucket(self.bucket_name, self.location, self.project)
-            logger.info(f"Bucket {self.bucket_name} created successfully.")
-        except Exception as e:
-            # If the bucket already exists, check for conflict (HTTP 409) and continue.
-            if hasattr(e, "code") and e.code == 409:
-                logger.info(f"Bucket {self.bucket_name} already exists.")
-            else:
-                logger.error(f"Failed to create bucket {self.bucket_name}: {e}")
-                raise
-
-    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 read(self, session_id: str, user_id: Optional[str] = None) -> Optional[Session]:
-        """
-        Reads a session JSON blob from the GCS bucket and returns a Session object.
-        If the blob is not found, returns None.
-        """
-        blob = self.bucket.blob(self._get_blob_path(session_id))
-        try:
-            data_str = blob.download_as_bytes().decode("utf-8")
-            data = self.deserialize(data_str)
-        except Exception as e:
-            # If the error indicates that the blob was not found (404), return None.
-            if "404" in str(e):
-                return None
-            logger.error(f"Error reading session {session_id} from GCS: {e}")
-            return None
-
-        if user_id and data.get("user_id") != user_id:
-            return None
-
-        if self.mode == "agent":
-            return AgentSession.from_dict(data)
-        elif self.mode == "team":
-            return TeamSession.from_dict(data)
-        elif self.mode == "workflow":
-            return WorkflowSession.from_dict(data)
-        elif self.mode == "workflow_v2":
-            return WorkflowSessionV2.from_dict(data)
-        return None
-
-    def get_all_session_ids(self, user_id: Optional[str] = None, entity_id: Optional[str] = None) -> List[str]:
-        """
-        Lists all session IDs stored in the bucket.
-        """
-        session_ids = []
-        for blob in self.client.list_blobs(self.bucket, prefix=self.prefix):
-            if blob.name.endswith(".json"):
-                session_ids.append(blob.name.replace(".json", ""))
-        return session_ids
-
-    def get_all_sessions(self, user_id: Optional[str] = None, entity_id: Optional[str] = None) -> List[Session]:
-        """
-        Retrieves all sessions stored in the bucket.
-        """
-        sessions: List[Session] = []
-        for blob in self.client.list_blobs(self.bucket, prefix=self.prefix):
-            if blob.name.endswith(".json"):
-                try:
-                    data_str = blob.download_as_bytes().decode("utf-8")
-                    data = self.deserialize(data_str)
-
-                    if user_id and data.get("user_id") != user_id:
-                        continue
-                    session: Optional[Session] = None
-                    if self.mode == "agent":
-                        session = AgentSession.from_dict(data)
-                    elif self.mode == "team":
-                        session = TeamSession.from_dict(data)
-                    elif self.mode == "workflow":
-                        session = WorkflowSession.from_dict(data)
-                    elif self.mode == "workflow_v2":
-                        session = WorkflowSessionV2.from_dict(data)
-                    if session is not None:
-                        sessions.append(session)
-                except Exception as e:
-                    logger.error(f"Error reading session from blob {blob.name}: {e}")
-                    continue
-        return sessions
-
-    def get_recent_sessions(
-        self,
-        user_id: Optional[str] = None,
-        entity_id: Optional[str] = None,
-        limit: Optional[int] = 2,
-    ) -> List[Session]:
-        """Get the last N sessions, ordered by created_at descending.
-
-        Args:
-            num_history_sessions: Number of most recent sessions to return
-            user_id: Filter by user ID
-            entity_id: Filter by entity ID (agent_id, team_id, or workflow_id)
-
-        Returns:
-            List[Session]: List of most recent sessions
-        """
-        sessions: List[Session] = []
-        # List of (created_at, data) tuples for sorting
-        session_data: List[tuple[int, dict]] = []
-
-        try:
-            # Get all blobs with the specified prefix
-            for blob in self.client.list_blobs(self.bucket, prefix=self.prefix):
-                if not blob.name.endswith(".json"):
-                    continue
-
-                try:
-                    data_str = blob.download_as_bytes().decode("utf-8")
-                    data = self.deserialize(data_str)
-
-                    # Apply filters
-                    if user_id and data.get("user_id") != user_id:
-                        continue
-
-                    # Store with created_at for sorting
-                    created_at = data.get("created_at", 0)
-                    session_data.append((created_at, data))
-
-                except Exception as e:
-                    logger.error(f"Error reading session from blob {blob.name}: {e}")
-                    continue
-
-            # Sort by created_at descending and take only num_history_sessions
-            session_data.sort(key=lambda x: x[0], reverse=True)
-            if limit is not None:
-                session_data = session_data[:limit]
-
-            # Convert filtered and sorted data to Session objects
-            for _, data in session_data:
-                session: Optional[Session] = None
-                if self.mode == "agent":
-                    session = AgentSession.from_dict(data)
-                elif self.mode == "team":
-                    session = TeamSession.from_dict(data)
-                elif self.mode == "workflow":
-                    session = WorkflowSession.from_dict(data)
-                elif self.mode == "workflow_v2":
-                    session = WorkflowSessionV2.from_dict(data)
-                if session is not None:
-                    sessions.append(session)
-
-        except Exception as e:
-            logger.error(f"Error getting last {limit} sessions: {e}")
-
-        return sessions
-
-    def upsert(self, session: Session) -> Optional[Session]:
-        """
-        Inserts or updates a session JSON blob in the GCS bucket.
-        """
-        blob = self.bucket.blob(self._get_blob_path(session.session_id))
-        try:
-            data = session.to_dict()
-            data["updated_at"] = int(time.time())
-            if "created_at" not in data or data["created_at"] is None:
-                data["created_at"] = data["updated_at"]
-            json_data = self.serialize(data)
-            blob.upload_from_string(json_data, content_type="application/json")
-            return session
-        except Exception as e:
-            logger.error(f"Error upserting session {session.session_id}: {e}")
-            return None
-
-    def delete_session(self, session_id: Optional[str] = None):
-        """
-        Deletes a session JSON blob from the GCS bucket.
-        """
-        if session_id is None:
-            return
-        blob = self.bucket.blob(self._get_blob_path(session_id))
-        try:
-            blob.delete()
-        except Exception as e:
-            logger.error(f"Error deleting session {session_id}: {e}")
-
-    def drop(self) -> None:
-        """
-        Deletes all session JSON blobs from the bucket.
-        """
-        prefix = ""
-        for blob in self.client.list_blobs(self.bucket, prefix=prefix):
-            try:
-                blob.delete()
-            except Exception as e:
-                logger.error(f"Error deleting blob {blob.name}: {e}")
-
-    def upgrade_schema(self) -> None:
-        """
-        Schema upgrade is not implemented.
-        """
-        pass