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

agno/storage/agent/singlestore.py

@@ -1,303 +0,0 @@
-import json
-from typing import Any, List, Optional
-
-try:
-    from sqlalchemy.dialects import mysql
-    from sqlalchemy.engine import Engine, create_engine
-    from sqlalchemy.engine.row import Row
-    from sqlalchemy.inspection import inspect
-    from sqlalchemy.orm import Session, sessionmaker
-    from sqlalchemy.schema import Column, MetaData, Table
-    from sqlalchemy.sql.expression import select, text
-except ImportError:
-    raise ImportError("`sqlalchemy` not installed")
-
-from agno.storage.agent.base import AgentStorage
-from agno.storage.agent.session import AgentSession
-from agno.utils.log import logger
-
-
-class SingleStoreAgentStorage(AgentStorage):
-    def __init__(
-        self,
-        table_name: str,
-        schema: Optional[str] = "ai",
-        db_url: Optional[str] = None,
-        db_engine: Optional[Engine] = None,
-        schema_version: int = 1,
-        auto_upgrade_schema: bool = False,
-    ):
-        """
-        This class provides Agent storage using a singlestore table.
-
-        The following order is used to determine the database connection:
-            1. Use the db_engine if provided
-            2. Use the db_url if provided
-
-        Args:
-            table_name (str): The name of the table to store the agent data.
-            schema (Optional[str], optional): The schema of the table. Defaults to "ai".
-            db_url (Optional[str], optional): The database URL. Defaults to None.
-            db_engine (Optional[Engine], optional): The database engine. Defaults to None.
-            schema_version (int, optional): The schema version. Defaults to 1.
-            auto_upgrade_schema (bool, optional): Automatically upgrade the schema. Defaults to False.
-        """
-        _engine: Optional[Engine] = db_engine
-        if _engine is None and db_url is not None:
-            _engine = create_engine(db_url, connect_args={"charset": "utf8mb4"})
-
-        if _engine is None:
-            raise ValueError("Must provide either db_url or db_engine")
-
-        # Database attributes
-        self.table_name: str = table_name
-        self.schema: Optional[str] = schema
-        self.db_url: Optional[str] = db_url
-        self.db_engine: Engine = _engine
-        self.metadata: MetaData = MetaData(schema=self.schema)
-
-        # Table schema version
-        self.schema_version: int = schema_version
-        # Automatically upgrade schema if True
-        self.auto_upgrade_schema: bool = auto_upgrade_schema
-
-        # Database session
-        self.Session: sessionmaker[Session] = sessionmaker(bind=self.db_engine)
-        # Database table for storage
-        self.table: Table = self.get_table()
-
-    def get_table_v1(self) -> Table:
-        return Table(
-            self.table_name,
-            self.metadata,
-            # Session UUID: Primary Key
-            Column("session_id", mysql.TEXT, primary_key=True),
-            # ID of the agent that this session is associated with
-            Column("agent_id", mysql.TEXT),
-            # ID of the user interacting with this agent
-            Column("user_id", mysql.TEXT),
-            # Agent memory
-            Column("memory", mysql.JSON),
-            # Agent Data
-            Column("agent_data", mysql.JSON),
-            # Session Data
-            Column("session_data", mysql.JSON),
-            # Extra Data stored with this agent
-            Column("extra_data", mysql.JSON),
-            # The Unix timestamp of when this session was created.
-            Column("created_at", mysql.BIGINT),
-            # The Unix timestamp of when this session was last updated.
-            Column("updated_at", mysql.BIGINT),
-            extend_existing=True,
-        )
-
-    def get_table(self) -> Table:
-        if self.schema_version == 1:
-            return self.get_table_v1()
-        else:
-            raise ValueError(f"Unsupported schema version: {self.schema_version}")
-
-    def table_exists(self) -> bool:
-        logger.debug(f"Checking if table exists: {self.table.name}")
-        try:
-            return inspect(self.db_engine).has_table(self.table.name, schema=self.schema)
-        except Exception as e:
-            logger.error(e)
-            return False
-
-    def create(self) -> None:
-        if not self.table_exists():
-            logger.info(f"\nCreating table: {self.table_name}\n")
-            self.table.create(self.db_engine)
-
-    def _read(self, session: Session, session_id: str, user_id: Optional[str] = None) -> Optional[Row[Any]]:
-        stmt = select(self.table).where(self.table.c.session_id == session_id)
-        if user_id is not None:
-            stmt = stmt.where(self.table.c.user_id == user_id)
-        try:
-            return session.execute(stmt).first()
-        except Exception as e:
-            logger.debug(f"Exception reading from table: {e}")
-            logger.debug(f"Table does not exist: {self.table.name}")
-            logger.debug(f"Creating table: {self.table_name}")
-            self.create()
-        return None
-
-    def read(self, session_id: str, user_id: Optional[str] = None) -> Optional[AgentSession]:
-        with self.Session.begin() as sess:
-            existing_row: Optional[Row[Any]] = self._read(session=sess, session_id=session_id, user_id=user_id)
-            return AgentSession.from_dict(existing_row._mapping) if existing_row is not None else None  # type: ignore
-
-    def get_all_session_ids(self, user_id: Optional[str] = None, agent_id: Optional[str] = None) -> List[str]:
-        session_ids: List[str] = []
-        try:
-            with self.Session.begin() as sess:
-                # get all session_ids for this user
-                stmt = select(self.table)
-                if user_id is not None:
-                    stmt = stmt.where(self.table.c.user_id == user_id)
-                if agent_id is not None:
-                    stmt = stmt.where(self.table.c.agent_id == agent_id)
-                # order by created_at desc
-                stmt = stmt.order_by(self.table.c.created_at.desc())
-                # execute query
-                rows = sess.execute(stmt).fetchall()
-                for row in rows:
-                    if row is not None and row.session_id is not None:
-                        session_ids.append(row.session_id)
-        except Exception as e:
-            logger.error(f"An unexpected error occurred: {str(e)}")
-        return session_ids
-
-    def get_all_sessions(self, user_id: Optional[str] = None, agent_id: Optional[str] = None) -> List[AgentSession]:
-        sessions: List[AgentSession] = []
-        try:
-            with self.Session.begin() as sess:
-                # get all sessions for this user
-                stmt = select(self.table)
-                if user_id is not None:
-                    stmt = stmt.where(self.table.c.user_id == user_id)
-                if agent_id is not None:
-                    stmt = stmt.where(self.table.c.agent_id == agent_id)
-                # order by created_at desc
-                stmt = stmt.order_by(self.table.c.created_at.desc())
-                # execute query
-                rows = sess.execute(stmt).fetchall()
-                for row in rows:
-                    if row.session_id is not None:
-                        _agent_session = AgentSession.from_dict(row._mapping)  # type: ignore
-                        if _agent_session is not None:
-                            sessions.append(_agent_session)
-        except Exception:
-            logger.debug(f"Table does not exist: {self.table.name}")
-        return sessions
-
-    def upsert(self, session: AgentSession) -> Optional[AgentSession]:
-        """
-        Create a new session if it does not exist, otherwise update the existing session.
-        """
-
-        with self.Session.begin() as sess:
-            # Create an insert statement using MySQL's ON DUPLICATE KEY UPDATE syntax
-            upsert_sql = text(
-                f"""
-                INSERT INTO {self.schema}.{self.table_name}
-                (session_id, agent_id, user_id, memory, agent_data, session_data, extra_data, created_at, updated_at)
-                VALUES
-                (:session_id, :agent_id, :user_id, :memory, :agent_data, :session_data, :extra_data, UNIX_TIMESTAMP(), NULL)
-                ON DUPLICATE KEY UPDATE
-                    agent_id = VALUES(agent_id),
-                    user_id = VALUES(user_id),
-                    memory = VALUES(memory),
-                    agent_data = VALUES(agent_data),
-                    session_data = VALUES(session_data),
-                    extra_data = VALUES(extra_data),
-                    updated_at = UNIX_TIMESTAMP();
-                """
-            )
-
-            try:
-                sess.execute(
-                    upsert_sql,
-                    {
-                        "session_id": session.session_id,
-                        "agent_id": session.agent_id,
-                        "user_id": session.user_id,
-                        "memory": json.dumps(session.memory, ensure_ascii=False)
-                        if session.memory is not None
-                        else None,
-                        "agent_data": json.dumps(session.agent_data, ensure_ascii=False)
-                        if session.agent_data is not None
-                        else None,
-                        "session_data": json.dumps(session.session_data, ensure_ascii=False)
-                        if session.session_data is not None
-                        else None,
-                        "extra_data": json.dumps(session.extra_data, ensure_ascii=False)
-                        if session.extra_data is not None
-                        else None,
-                    },
-                )
-            except Exception:
-                # Create table and try again
-                self.create()
-                sess.execute(
-                    upsert_sql,
-                    {
-                        "session_id": session.session_id,
-                        "agent_id": session.agent_id,
-                        "user_id": session.user_id,
-                        "memory": json.dumps(session.memory, ensure_ascii=False)
-                        if session.memory is not None
-                        else None,
-                        "agent_data": json.dumps(session.agent_data, ensure_ascii=False)
-                        if session.agent_data is not None
-                        else None,
-                        "session_data": json.dumps(session.session_data, ensure_ascii=False)
-                        if session.session_data is not None
-                        else None,
-                        "extra_data": json.dumps(session.extra_data, ensure_ascii=False)
-                        if session.extra_data is not None
-                        else None,
-                    },
-                )
-        return self.read(session_id=session.session_id)
-
-    def delete_session(self, session_id: Optional[str] = None):
-        if session_id is None:
-            logger.warning("No session_id provided for deletion.")
-            return
-
-        with self.Session() as sess, sess.begin():
-            try:
-                # Delete the session with the given session_id
-                delete_stmt = self.table.delete().where(self.table.c.session_id == session_id)
-                result = sess.execute(delete_stmt)
-
-                if result.rowcount == 0:
-                    logger.warning(f"No session found with session_id: {session_id}")
-                else:
-                    logger.info(f"Successfully deleted session with session_id: {session_id}")
-            except Exception as e:
-                logger.error(f"Error deleting session: {e}")
-                raise
-
-    def drop(self) -> None:
-        if self.table_exists():
-            logger.info(f"Deleting table: {self.table_name}")
-            self.table.drop(self.db_engine)
-
-    def upgrade_schema(self) -> None:
-        pass
-
-    def __deepcopy__(self, memo):
-        """
-        Create a deep copy of the SingleStoreAgentStorage instance, handling unpickleable attributes.
-
-        Args:
-            memo (dict): A dictionary of objects already copied during the current copying pass.
-
-        Returns:
-            SingleStoreAgentStorage: A deep-copied instance of SingleStoreAgentStorage.
-        """
-        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 {"metadata", "table"}:
-                continue
-            # Reuse db_engine and Session without copying
-            elif k in {"db_engine", "Session"}:
-                setattr(copied_obj, k, v)
-            else:
-                setattr(copied_obj, k, deepcopy(v, memo))
-
-        # Recreate metadata and table for the copied instance
-        copied_obj.metadata = MetaData(schema=self.schema)
-        copied_obj.table = copied_obj.get_table()
-
-        return copied_obj

agno/storage/agent/sqlite.py

@@ -1,357 +0,0 @@
-import time
-from pathlib import Path
-from typing import List, Optional
-
-try:
-    from sqlalchemy.dialects import sqlite
-    from sqlalchemy.engine import Engine, create_engine
-    from sqlalchemy.inspection import inspect
-    from sqlalchemy.orm import Session, sessionmaker
-    from sqlalchemy.schema import Column, MetaData, Table
-    from sqlalchemy.sql.expression import select
-    from sqlalchemy.types import String
-except ImportError:
-    raise ImportError("`sqlalchemy` not installed. Please install it using `pip install sqlalchemy`")
-
-from agno.storage.agent.base import AgentStorage
-from agno.storage.agent.session import AgentSession
-from agno.utils.log import logger
-
-
-class SqliteAgentStorage(AgentStorage):
-    def __init__(
-        self,
-        table_name: str,
-        db_url: Optional[str] = None,
-        db_file: Optional[str] = None,
-        db_engine: Optional[Engine] = None,
-        schema_version: int = 1,
-        auto_upgrade_schema: bool = False,
-    ):
-        """
-        This class provides agent storage using a sqlite database.
-
-        The following order is used to determine the database connection:
-            1. Use the db_engine if provided
-            2. Use the db_url
-            3. Use the db_file
-            4. Create a new in-memory database
-
-        Args:
-            table_name: The name of the table to store Agent sessions.
-            db_url: The database URL to connect to.
-            db_file: The database file to connect to.
-            db_engine: The SQLAlchemy database engine to use.
-        """
-        _engine: Optional[Engine] = db_engine
-        if _engine is None and db_url is not None:
-            _engine = create_engine(db_url)
-        elif _engine is None and db_file is not None:
-            # Use the db_file to create the engine
-            db_path = Path(db_file).resolve()
-            # Ensure the directory exists
-            db_path.parent.mkdir(parents=True, exist_ok=True)
-            _engine = create_engine(f"sqlite:///{db_path}")
-        else:
-            _engine = create_engine("sqlite://")
-
-        if _engine is None:
-            raise ValueError("Must provide either db_url, db_file or db_engine")
-
-        # Database attributes
-        self.table_name: str = table_name
-        self.db_url: Optional[str] = db_url
-        self.db_engine: Engine = _engine
-        self.metadata: MetaData = MetaData()
-        self.inspector = inspect(self.db_engine)
-
-        # Table schema version
-        self.schema_version: int = schema_version
-        # Automatically upgrade schema if True
-        self.auto_upgrade_schema: bool = auto_upgrade_schema
-
-        # Database session
-        self.Session: sessionmaker[Session] = sessionmaker(bind=self.db_engine)
-        # Database table for storage
-        self.table: Table = self.get_table()
-
-    def get_table_v1(self) -> Table:
-        """
-        Define the table schema for version 1.
-
-        Returns:
-            Table: SQLAlchemy Table object representing the schema.
-        """
-        return Table(
-            self.table_name,
-            self.metadata,
-            # Session UUID: Primary Key
-            Column("session_id", String, primary_key=True),
-            # ID of the agent that this session is associated with
-            Column("agent_id", String),
-            # ID of the user interacting with this agent
-            Column("user_id", String),
-            # Agent Memory
-            Column("memory", sqlite.JSON),
-            # Agent Data
-            Column("agent_data", sqlite.JSON),
-            # Session Data
-            Column("session_data", sqlite.JSON),
-            # Extra Data stored with this agent
-            Column("extra_data", sqlite.JSON),
-            # The Unix timestamp of when this session was created.
-            Column("created_at", sqlite.INTEGER, default=lambda: int(time.time())),
-            # The Unix timestamp of when this session was last updated.
-            Column("updated_at", sqlite.INTEGER, onupdate=lambda: int(time.time())),
-            extend_existing=True,
-            sqlite_autoincrement=True,
-        )
-
-    def get_table(self) -> Table:
-        """
-        Get the table schema based on the schema version.
-
-        Returns:
-            Table: SQLAlchemy Table object for the current schema version.
-
-        Raises:
-            ValueError: If an unsupported schema version is specified.
-        """
-        if self.schema_version == 1:
-            return self.get_table_v1()
-        else:
-            raise ValueError(f"Unsupported schema version: {self.schema_version}")
-
-    def table_exists(self) -> bool:
-        """
-        Check if the table exists in the database.
-
-        Returns:
-            bool: True if the table exists, False otherwise.
-        """
-        logger.debug(f"Checking if table exists: {self.table.name}")
-        try:
-            return self.inspector.has_table(self.table.name)
-        except Exception as e:
-            logger.error(f"Error checking if table exists: {e}")
-            return False
-
-    def create(self) -> None:
-        """
-        Create the table if it doesn't exist.
-        """
-        if not self.table_exists():
-            logger.debug(f"Creating table: {self.table.name}")
-            self.table.create(self.db_engine, checkfirst=True)
-
-    def read(self, session_id: str, user_id: Optional[str] = None) -> Optional[AgentSession]:
-        """
-        Read 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:
-            with self.Session() as sess:
-                stmt = select(self.table).where(self.table.c.session_id == session_id)
-                if user_id:
-                    stmt = stmt.where(self.table.c.user_id == user_id)
-                result = sess.execute(stmt).fetchone()
-                return AgentSession.from_dict(result._mapping) if result is not None else None  # type: ignore
-        except Exception as e:
-            logger.debug(f"Exception reading from table: {e}")
-            logger.debug(f"Table does not exist: {self.table.name}")
-            logger.debug("Creating table for future transactions")
-            self.create()
-        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.
-
-        Args:
-            user_id (Optional[str]): The ID of the user to filter by.
-            agent_id (Optional[str]): The ID of the agent to filter by.
-
-        Returns:
-            List[str]: List of session IDs matching the criteria.
-        """
-        try:
-            with self.Session() as sess, sess.begin():
-                # get all session_ids
-                stmt = select(self.table.c.session_id)
-                if user_id is not None:
-                    stmt = stmt.where(self.table.c.user_id == user_id)
-                if agent_id is not None:
-                    stmt = stmt.where(self.table.c.agent_id == agent_id)
-                # order by created_at desc
-                stmt = stmt.order_by(self.table.c.created_at.desc())
-                # execute query
-                rows = sess.execute(stmt).fetchall()
-                return [row[0] for row in rows] if rows is not None else []
-        except Exception as e:
-            logger.debug(f"Exception reading from table: {e}")
-            logger.debug(f"Table does not exist: {self.table.name}")
-            logger.debug("Creating table for future transactions")
-            self.create()
-        return []
-
-    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.
-
-        Args:
-            user_id (Optional[str]): The ID of the user to filter by.
-            agent_id (Optional[str]): The ID of the agent to filter by.
-
-        Returns:
-            List[AgentSession]: List of AgentSession objects matching the criteria.
-        """
-        try:
-            with self.Session() as sess, sess.begin():
-                # get all sessions
-                stmt = select(self.table)
-                if user_id is not None:
-                    stmt = stmt.where(self.table.c.user_id == user_id)
-                if agent_id is not None:
-                    stmt = stmt.where(self.table.c.agent_id == agent_id)
-                # order by created_at desc
-                stmt = stmt.order_by(self.table.c.created_at.desc())
-                # execute query
-                rows = sess.execute(stmt).fetchall()
-                return [AgentSession.from_dict(row._mapping) for row in rows] if rows is not None else []  # type: ignore
-        except Exception as e:
-            logger.debug(f"Exception reading from table: {e}")
-            logger.debug(f"Table does not exist: {self.table.name}")
-            logger.debug("Creating table for future transactions")
-            self.create()
-        return []
-
-    def upsert(self, session: AgentSession, create_and_retry: bool = True) -> Optional[AgentSession]:
-        """
-        Insert or update an AgentSession in the database.
-
-        Args:
-            session (AgentSession): The session data to upsert.
-            create_and_retry (bool): Retry upsert if table does not exist.
-
-        Returns:
-            Optional[AgentSession]: The upserted AgentSession, or None if operation failed.
-        """
-        try:
-            with self.Session() as sess, sess.begin():
-                # Create an insert statement
-                stmt = sqlite.insert(self.table).values(
-                    session_id=session.session_id,
-                    agent_id=session.agent_id,
-                    user_id=session.user_id,
-                    memory=session.memory,
-                    agent_data=session.agent_data,
-                    session_data=session.session_data,
-                    extra_data=session.extra_data,
-                )
-
-                # Define the upsert if the session_id already exists
-                # See: https://docs.sqlalchemy.org/en/20/dialects/sqlite.html#insert-on-conflict-upsert
-                stmt = stmt.on_conflict_do_update(
-                    index_elements=["session_id"],
-                    set_=dict(
-                        agent_id=session.agent_id,
-                        user_id=session.user_id,
-                        memory=session.memory,
-                        agent_data=session.agent_data,
-                        session_data=session.session_data,
-                        extra_data=session.extra_data,
-                        updated_at=int(time.time()),
-                    ),  # The updated value for each column
-                )
-
-                sess.execute(stmt)
-        except Exception as e:
-            logger.debug(f"Exception upserting into table: {e}")
-            if create_and_retry and not self.table_exists():
-                logger.debug(f"Table does not exist: {self.table.name}")
-                logger.debug("Creating table and retrying upsert")
-                self.create()
-                return self.upsert(session, create_and_retry=False)
-            return None
-        return self.read(session_id=session.session_id)
-
-    def delete_session(self, session_id: Optional[str] = None):
-        """
-        Delete a workflow session from the database.
-
-        Args:
-            session_id (Optional[str]): The ID of the session to delete.
-
-        Raises:
-            ValueError: If session_id is not provided.
-        """
-        if session_id is None:
-            logger.warning("No session_id provided for deletion.")
-            return
-
-        try:
-            with self.Session() as sess, sess.begin():
-                # Delete the session with the given session_id
-                delete_stmt = self.table.delete().where(self.table.c.session_id == session_id)
-                result = sess.execute(delete_stmt)
-                if result.rowcount == 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 Exception as e:
-            logger.error(f"Error deleting session: {e}")
-
-    def drop(self) -> None:
-        """
-        Drop the table from the database if it exists.
-        """
-        if self.table_exists():
-            logger.debug(f"Deleting table: {self.table_name}")
-            self.table.drop(self.db_engine)
-
-    def upgrade_schema(self) -> None:
-        """
-        Upgrade the schema of the workflow storage table.
-        This method is currently a placeholder and does not perform any actions.
-        """
-        pass
-
-    def __deepcopy__(self, memo):
-        """
-        Create a deep copy of the SqliteAgentStorage instance, handling unpickleable attributes.
-
-        Args:
-            memo (dict): A dictionary of objects already copied during the current copying pass.
-
-        Returns:
-            SqliteAgentStorage: A deep-copied instance of SqliteAgentStorage.
-        """
-        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 {"metadata", "table", "inspector"}:
-                continue
-            # Reuse db_engine and Session without copying
-            elif k in {"db_engine", "Session"}:
-                setattr(copied_obj, k, v)
-            else:
-                setattr(copied_obj, k, deepcopy(v, memo))
-
-        # Recreate metadata and table for the copied instance
-        copied_obj.metadata = MetaData()
-        copied_obj.inspector = inspect(copied_obj.db_engine)
-        copied_obj.table = copied_obj.get_table()
-
-        return copied_obj