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

agno/memory/v2/db/postgres.py

@@ -1,214 +0,0 @@
-from typing import Any, Dict, List, Optional
-
-try:
-    from sqlalchemy.dialects import postgresql
-    from sqlalchemy.engine import Engine, create_engine
-    from sqlalchemy.inspection import inspect
-    from sqlalchemy.orm import scoped_session, sessionmaker
-    from sqlalchemy.schema import Column, MetaData, Table
-    from sqlalchemy.sql.expression import delete, select, text
-    from sqlalchemy.types import DateTime, String
-except ImportError:
-    raise ImportError("`sqlalchemy` not installed.  Please install using `pip install sqlalchemy 'psycopg[binary]'`")
-
-from agno.memory.v2.db.base import MemoryDb
-from agno.memory.v2.db.schema import MemoryRow
-from agno.utils.log import log_debug, log_info, logger
-
-
-class PostgresMemoryDb(MemoryDb):
-    def __init__(
-        self,
-        table_name: str,
-        schema: Optional[str] = "ai",
-        db_url: Optional[str] = None,
-        db_engine: Optional[Engine] = None,
-    ):
-        """
-        This class provides a memory store backed by a postgres table.
-
-        The following order is used to determine the database connection:
-            1. Use the db_engine if provided
-            2. Use the db_url to create the engine
-
-        Args:
-            table_name (str): The name of the table to store memory rows.
-            schema (Optional[str]): The schema to store the table in. Defaults to "ai".
-            db_url (Optional[str]): The database URL to connect to. Defaults to None.
-            db_engine (Optional[Engine]): The database engine to use. Defaults to None.
-        """
-        _engine: Optional[Engine] = db_engine
-        if _engine is None and db_url is not None:
-            _engine = create_engine(db_url)
-
-        if _engine is None:
-            raise ValueError("Must provide either db_url or db_engine")
-
-        self.table_name: str = table_name
-        self.schema: Optional[str] = schema
-        self.db_url: Optional[str] = db_url
-        self.db_engine: Engine = _engine
-        self.inspector = inspect(self.db_engine)
-        self.metadata: MetaData = MetaData(schema=self.schema)
-        self.Session: scoped_session = scoped_session(sessionmaker(bind=self.db_engine))
-        self.table: Table = self.get_table()
-
-    def __dict__(self) -> Dict[str, Any]:
-        return {
-            "name": "PostgresMemoryDb",
-            "table_name": self.table_name,
-            "schema": self.schema,
-        }
-
-    def get_table(self) -> Table:
-        return Table(
-            self.table_name,
-            self.metadata,
-            Column("id", String, primary_key=True),
-            Column("user_id", String, index=True),
-            Column("memory", postgresql.JSONB, server_default=text("'{}'::jsonb")),
-            Column("created_at", DateTime(timezone=True), server_default=text("now()")),
-            Column("updated_at", DateTime(timezone=True), onupdate=text("now()")),
-            extend_existing=True,
-        )
-
-    def create(self) -> None:
-        if not self.table_exists():
-            try:
-                with self.Session() as sess, sess.begin():
-                    if self.schema is not None:
-                        log_debug(f"Creating schema: {self.schema}")
-                        sess.execute(text(f"CREATE SCHEMA IF NOT EXISTS {self.schema};"))
-                log_debug(f"Creating table: {self.table_name}")
-                self.table.create(self.db_engine, checkfirst=True)
-            except Exception as e:
-                logger.error(f"Error creating table '{self.table.fullname}': {e}")
-                raise
-
-    def memory_exists(self, memory: MemoryRow) -> bool:
-        columns = [self.table.c.id]
-        with self.Session() as sess, sess.begin():
-            stmt = select(*columns).where(self.table.c.id == memory.id)
-            result = sess.execute(stmt).first()
-            return result is not None
-
-    def read_memories(
-        self, user_id: Optional[str] = None, limit: Optional[int] = None, sort: Optional[str] = None
-    ) -> List[MemoryRow]:
-        memories: List[MemoryRow] = []
-        try:
-            with self.Session() as sess, sess.begin():
-                stmt = select(self.table)
-                if user_id is not None:
-                    stmt = stmt.where(self.table.c.user_id == user_id)
-                if limit is not None:
-                    stmt = stmt.limit(limit)
-
-                if sort == "asc":
-                    stmt = stmt.order_by(self.table.c.created_at.asc())
-                else:
-                    stmt = stmt.order_by(self.table.c.created_at.desc())
-
-                rows = sess.execute(stmt).fetchall()
-                for row in rows:
-                    if row is not None:
-                        memories.append(MemoryRow.model_validate(row))
-        except Exception as e:
-            log_debug(f"Exception reading from table: {e}")
-            log_debug(f"Table does not exist: {self.table.name}")
-            log_debug("Creating table for future transactions")
-            self.create()
-        return memories
-
-    def upsert_memory(self, memory: MemoryRow, create_and_retry: bool = True) -> None:
-        """Create a new memory if it does not exist, otherwise update the existing memory"""
-
-        try:
-            with self.Session() as sess, sess.begin():
-                # Create an insert statement
-                stmt = postgresql.insert(self.table).values(
-                    id=memory.id,
-                    user_id=memory.user_id,
-                    memory=memory.memory,
-                )
-
-                # Define the upsert if the memory already exists
-                # See: https://docs.sqlalchemy.org/en/20/dialects/postgresql.html#postgresql-insert-on-conflict
-                stmt = stmt.on_conflict_do_update(
-                    index_elements=["id"],
-                    set_=dict(
-                        user_id=stmt.excluded.user_id,
-                        memory=stmt.excluded.memory,
-                    ),
-                )
-
-                sess.execute(stmt)
-        except Exception as e:
-            log_debug(f"Exception upserting into table: {e}")
-            log_debug(f"Table does not exist: {self.table.name}")
-            log_debug("Creating table for future transactions")
-            self.create()
-            if create_and_retry:
-                return self.upsert_memory(memory, create_and_retry=False)
-            return None
-
-    def delete_memory(self, memory_id: str) -> None:
-        with self.Session() as sess, sess.begin():
-            stmt = delete(self.table).where(self.table.c.id == memory_id)
-            sess.execute(stmt)
-
-    def drop_table(self) -> None:
-        if self.table_exists():
-            log_debug(f"Deleting table: {self.table_name}")
-            self.table.drop(self.db_engine)
-
-    def table_exists(self) -> bool:
-        log_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 clear(self) -> bool:
-        if self.table_exists():
-            with self.Session() as sess, sess.begin():
-                stmt = delete(self.table)
-                log_info(f"Clearing table: {self.table.name}")
-                sess.execute(stmt)
-                return True
-        return False
-
-    def __deepcopy__(self, memo):
-        """
-        Create a deep copy of the PgMemoryDb instance, handling unpickleable attributes.
-
-        Args:
-            memo (dict): A dictionary of objects already copied during the current copying pass.
-
-        Returns:
-            PostgresMemoryDb: A deep-copied instance of PgMemoryDb.
-        """
-        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=copied_obj.schema)
-        copied_obj.table = copied_obj.get_table()
-
-        return copied_obj

agno/memory/v2/db/redis.py

@@ -1,187 +0,0 @@
-import json
-import time
-from typing import Any, Dict, List, Optional
-
-try:
-    from redis import ConnectionError, Redis
-except ImportError:
-    raise ImportError("`redis` not installed. Please install it using `pip install redis`")
-
-from agno.memory.v2.db.base import MemoryDb
-from agno.memory.v2.db.schema import MemoryRow
-from agno.utils.log import log_debug, log_info, logger
-
-
-class RedisMemoryDb(MemoryDb):
-    def __init__(
-        self,
-        prefix: str = "agno_memory",
-        host: str = "localhost",
-        port: int = 6379,
-        db: int = 0,
-        password: Optional[str] = None,
-        ssl: Optional[bool] = False,
-        expire: Optional[int] = None,
-    ):
-        """
-        Initialize Redis memory store.
-
-        Args:
-            prefix (str): Prefix for Redis keys to namespace the memories
-            host (str): Redis host address
-            port (int): Redis port number
-            db (int): Redis database number
-            password (Optional[str]): Redis password if authentication is required
-            ssl (Optional[bool]): Whether to use SSL for Redis connection
-            expire (Optional[int]): TTL (time to live) in seconds for Redis keys. None means no expiration.
-        """
-        self.prefix = prefix
-        self.expire = expire
-        self.redis_client = Redis(
-            host=host,
-            port=port,
-            db=db,
-            password=password,
-            decode_responses=True,  # Automatically decode responses to str
-            ssl=ssl,
-        )
-        log_debug(f"Created RedisMemoryDb with prefix: '{self.prefix}'")
-
-    def __dict__(self) -> Dict[str, Any]:
-        return {
-            "name": "RedisMemoryDb",
-            "prefix": self.prefix,
-            "expire": self.expire,
-        }
-
-    def _get_key(self, memory_id: str) -> str:
-        """Generate Redis key for a memory."""
-        return f"{self.prefix}:{memory_id}"
-
-    def create(self) -> None:
-        """
-        Test connection to Redis.
-        For Redis, we don't need to create schema as it's schema-less.
-        """
-        try:
-            self.redis_client.ping()
-            log_debug("Redis connection successful")
-        except ConnectionError as e:
-            logger.error(f"Could not connect to Redis: {e}")
-            raise
-
-    def memory_exists(self, memory: MemoryRow) -> bool:
-        """Check if a memory exists"""
-        try:
-            key = self._get_key(memory.id)  # type: ignore
-            return self.redis_client.exists(key) > 0  # type: ignore
-        except Exception as e:
-            logger.error(f"Error checking memory existence: {e}")
-            return False
-
-    def read_memories(
-        self, user_id: Optional[str] = None, limit: Optional[int] = None, sort: Optional[str] = None
-    ) -> List[MemoryRow]:
-        """Read memories from Redis"""
-        memories: List[MemoryRow] = []
-        try:
-            # Get all keys matching the prefix pattern
-            pattern = f"{self.prefix}:*"
-            memory_data = []
-
-            # Collect all matching memories
-            for key in self.redis_client.scan_iter(match=pattern):
-                data_str = self.redis_client.get(key)
-                if data_str:
-                    data = json.loads(data_str)  # type: ignore
-
-                    # Filter by user_id if specified
-                    if user_id is None or data.get("user_id") == user_id:
-                        memory_data.append(data)
-
-            # Sort by created_at timestamp
-            if sort == "asc":
-                memory_data.sort(key=lambda x: x.get("created_at", 0))
-            else:
-                memory_data.sort(key=lambda x: x.get("created_at", 0), reverse=True)
-
-            # Apply limit if specified
-            if limit is not None and limit > 0:
-                memory_data = memory_data[:limit]
-
-            # Convert to MemoryRow objects
-            for data in memory_data:
-                memories.append(MemoryRow.model_validate(data))
-
-        except Exception as e:
-            logger.error(f"Error reading memories: {e}")
-
-        return memories
-
-    def upsert_memory(self, memory: MemoryRow) -> Optional[MemoryRow]:
-        """Upsert a memory in Redis"""
-        try:
-            # Prepare data
-            timestamp = int(time.time())
-
-            # Convert to dict and handle datetime objects
-            memory_data = memory.model_dump(mode="json")
-
-            # Add timestamps if not present
-            if "created_at" not in memory_data:
-                memory_data["created_at"] = timestamp
-
-            memory_data["updated_at"] = timestamp
-
-            # Save to Redis
-            key = self._get_key(memory.id)  # type: ignore
-            if self.expire is not None:
-                self.redis_client.set(key, json.dumps(memory_data), ex=self.expire)
-            else:
-                self.redis_client.set(key, json.dumps(memory_data))
-
-            return memory
-
-        except Exception as e:
-            logger.error(f"Error upserting memory: {e}")
-            return None
-
-    def delete_memory(self, memory_id: str) -> None:
-        """Delete a memory from Redis"""
-        try:
-            key = self._get_key(memory_id)
-            self.redis_client.delete(key)
-            log_debug(f"Deleted memory: {memory_id}")
-        except Exception as e:
-            logger.error(f"Error deleting memory: {e}")
-
-    def drop_table(self) -> None:
-        """Drop all memories from Redis (same as clear)"""
-        self.clear()
-
-    def table_exists(self) -> bool:
-        """Check if any memories exist with our prefix"""
-        try:
-            pattern = f"{self.prefix}:*"
-            # Use scan_iter with count=1 to efficiently check if any keys exist
-            for _ in self.redis_client.scan_iter(match=pattern, count=1):
-                return True
-            return False
-        except Exception as e:
-            logger.error(f"Error checking if table exists: {e}")
-            return False
-
-    def clear(self) -> bool:
-        """Clear all memories with our prefix"""
-        try:
-            pattern = f"{self.prefix}:*"
-            keys_to_delete = list(self.redis_client.scan_iter(match=pattern))
-
-            if keys_to_delete:
-                self.redis_client.delete(*keys_to_delete)
-                log_info(f"Cleared {len(keys_to_delete)} memories with prefix: {self.prefix}")
-
-            return True
-        except Exception as e:
-            logger.error(f"Error clearing memories: {e}")
-            return False

agno/memory/v2/db/schema.py

@@ -1,54 +0,0 @@
-from datetime import datetime
-from typing import Any, Dict, Optional
-
-from pydantic import BaseModel, ConfigDict, model_validator
-
-
-class MemoryRow(BaseModel):
-    """Memory Row that is stored in the database"""
-
-    # id for this memory, auto-generated if not provided
-    id: Optional[str] = None
-    memory: Dict[str, Any]
-    user_id: Optional[str] = None
-    last_updated: Optional[datetime] = None
-
-    model_config = ConfigDict(from_attributes=True, arbitrary_types_allowed=True)
-
-    @model_validator(mode="after")
-    def generate_id(self) -> "MemoryRow":
-        if self.id is None:
-            from uuid import uuid4
-
-            self.id = str(uuid4())
-        return self
-
-    def to_dict(self) -> Dict[str, Any]:
-        _dict = self.model_dump(exclude={"last_updated"})
-        _dict["last_updated"] = self.last_updated.isoformat() if self.last_updated else None
-        return _dict
-
-
-class SummaryRow(BaseModel):
-    """Session Summary Row that is stored in the database"""
-
-    # id for this summary
-    id: Optional[str] = None
-    summary: Dict[str, Any]
-    user_id: Optional[str] = None
-    last_updated: Optional[datetime] = None
-
-    model_config = ConfigDict(from_attributes=True, arbitrary_types_allowed=True)
-
-    @model_validator(mode="after")
-    def generate_id(self) -> "SummaryRow":
-        if self.id is None:
-            from uuid import uuid4
-
-            self.id = str(uuid4())
-        return self
-
-    def to_dict(self) -> Dict[str, Any]:
-        _dict = self.model_dump(exclude={"last_updated"})
-        _dict["last_updated"] = self.last_updated.isoformat() if self.last_updated else None
-        return _dict

agno/memory/v2/db/sqlite.py

@@ -1,209 +0,0 @@
-from pathlib import Path
-from typing import Any, Dict, List, Optional
-
-try:
-    from sqlalchemy import (
-        Column,
-        DateTime,
-        Engine,
-        MetaData,
-        String,
-        Table,
-        create_engine,
-        delete,
-        inspect,
-        select,
-        text,
-    )
-    from sqlalchemy.exc import SQLAlchemyError
-    from sqlalchemy.orm import scoped_session, sessionmaker
-except ImportError:
-    raise ImportError("`sqlalchemy` not installed. Please install it with `pip install sqlalchemy`")
-
-from agno.memory.v2.db.base import MemoryDb
-from agno.memory.v2.db.schema import MemoryRow
-from agno.utils.log import log_debug, log_info, logger
-
-
-class SqliteMemoryDb(MemoryDb):
-    def __init__(
-        self,
-        table_name: str = "memory",
-        db_url: Optional[str] = None,
-        db_file: Optional[str] = None,
-        db_engine: Optional[Engine] = None,
-    ):
-        """
-        This class provides a memory store backed by a SQLite table.
-
-        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 database engine to use.
-        """
-        self.db_file = db_file
-        _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)
-
-        # Database session
-        self.Session = scoped_session(sessionmaker(bind=self.db_engine))
-        # Database table for memories
-        self.table: Table = self.get_table()
-
-    def __dict__(self) -> Dict[str, Any]:
-        return {
-            "name": "SqliteMemoryDb",
-            "table_name": self.table_name,
-            "db_file": self.db_file,
-        }
-
-    def get_table(self) -> Table:
-        return Table(
-            self.table_name,
-            self.metadata,
-            Column("id", String, primary_key=True),
-            Column("user_id", String, index=True),
-            Column("memory", String),
-            Column("created_at", DateTime, server_default=text("CURRENT_TIMESTAMP")),
-            Column(
-                "updated_at", DateTime, server_default=text("CURRENT_TIMESTAMP"), onupdate=text("CURRENT_TIMESTAMP")
-            ),
-            extend_existing=True,
-        )
-
-    def create(self) -> None:
-        if not self.table_exists():
-            try:
-                log_debug(f"Creating table: {self.table_name}")
-                self.table.create(self.db_engine, checkfirst=True)
-            except Exception as e:
-                logger.error(f"Error creating table '{self.table_name}': {e}")
-                raise
-
-    def memory_exists(self, memory: MemoryRow) -> bool:
-        with self.Session() as session:
-            stmt = select(self.table.c.id).where(self.table.c.id == memory.id)
-            result = session.execute(stmt).first()
-            return result is not None
-
-    def read_memories(
-        self, user_id: Optional[str] = None, limit: Optional[int] = None, sort: Optional[str] = None
-    ) -> List[MemoryRow]:
-        memories: List[MemoryRow] = []
-        try:
-            with self.Session() as session:
-                stmt = select(self.table)
-                if user_id is not None:
-                    stmt = stmt.where(self.table.c.user_id == user_id)
-
-                if sort == "asc":
-                    stmt = stmt.order_by(self.table.c.created_at.asc())
-                else:
-                    stmt = stmt.order_by(self.table.c.created_at.desc())
-
-                if limit is not None:
-                    stmt = stmt.limit(limit)
-
-                result = session.execute(stmt)
-                for row in result:
-                    memories.append(
-                        MemoryRow(
-                            id=row.id,
-                            user_id=row.user_id,
-                            memory=eval(row.memory),
-                            last_updated=row.updated_at or row.created_at,
-                        )
-                    )
-        except SQLAlchemyError as e:
-            log_debug(f"Exception reading from table: {e}")
-            log_debug(f"Table does not exist: {self.table_name}")
-            log_debug("Creating table for future transactions")
-            self.create()
-        return memories
-
-    def upsert_memory(self, memory: MemoryRow, create_and_retry: bool = True) -> None:
-        try:
-            with self.Session() as session:
-                # Check if the memory already exists
-                existing = session.execute(select(self.table).where(self.table.c.id == memory.id)).first()
-
-                if existing:
-                    # Update existing memory
-                    stmt = (
-                        self.table.update()
-                        .where(self.table.c.id == memory.id)
-                        .values(user_id=memory.user_id, memory=str(memory.memory), updated_at=text("CURRENT_TIMESTAMP"))
-                    )
-                else:
-                    # Insert new memory
-                    stmt = self.table.insert().values(id=memory.id, user_id=memory.user_id, memory=str(memory.memory))  # type: ignore
-
-                session.execute(stmt)
-                session.commit()
-        except SQLAlchemyError as e:
-            logger.error(f"Exception upserting into table: {e}")
-            if not self.table_exists():
-                log_info(f"Table does not exist: {self.table_name}")
-                log_info("Creating table for future transactions")
-                self.create()
-                if create_and_retry:
-                    return self.upsert_memory(memory, create_and_retry=False)
-            else:
-                raise
-
-    def delete_memory(self, memory_id: str) -> None:
-        with self.Session() as session:
-            stmt = delete(self.table).where(self.table.c.id == memory_id)
-            session.execute(stmt)
-            session.commit()
-
-    def drop_table(self) -> None:
-        if self.table_exists():
-            log_debug(f"Deleting table: {self.table_name}")
-            self.table.drop(self.db_engine)
-
-    def table_exists(self) -> bool:
-        log_debug(f"Checking if table exists: {self.table.name}")
-        try:
-            return self.inspector.has_table(self.table.name)
-        except Exception as e:
-            logger.error(e)
-            return False
-
-    def clear(self) -> bool:
-        with self.Session() as session:
-            if self.table_exists():
-                stmt = delete(self.table)
-                session.execute(stmt)
-                session.commit()
-        return True
-
-    def __del__(self):
-        # self.Session.remove()
-        pass