agno 2.3.1__py3-none-any.whl → 2.3.3__py3-none-any.whl

agno/db/mongo/async_mongo.py

@@ -1,6 +1,7 @@
+import asyncio
 import time
 from datetime import date, datetime, timedelta, timezone
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 from uuid import uuid4
 
 from agno.db.base import AsyncBaseDb, SessionType
@@ -25,11 +26,26 @@ from agno.utils.log import log_debug, log_error, log_info
 from agno.utils.string import generate_id
 
 try:
-    import asyncio
-
     from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorCollection, AsyncIOMotorDatabase  # type: ignore
+
+    MOTOR_AVAILABLE = True
+except ImportError:
+    MOTOR_AVAILABLE = False
+    AsyncIOMotorClient = None  # type: ignore
+    AsyncIOMotorCollection = None  # type: ignore
+    AsyncIOMotorDatabase = None  # type: ignore
+
+try:
+    from pymongo import AsyncMongoClient  # type: ignore
+    from pymongo.collection import AsyncCollection  # type: ignore
+    from pymongo.database import AsyncDatabase  # type: ignore
+
+    PYMONGO_ASYNC_AVAILABLE = True
 except ImportError:
-    raise ImportError("`motor` not installed. Please install it using `pip install -U motor`")
+    PYMONGO_ASYNC_AVAILABLE = False
+    AsyncMongoClient = None  # type: ignore
+    AsyncDatabase = None  # type: ignore
+    AsyncCollection = None  # type: ignore
 
 try:
     from pymongo import ReturnDocument
@@ -37,11 +53,89 @@ try:
 except ImportError:
     raise ImportError("`pymongo` not installed. Please install it using `pip install -U pymongo`")
 
+# Ensure at least one async library is available
+if not MOTOR_AVAILABLE and not PYMONGO_ASYNC_AVAILABLE:
+    raise ImportError(
+        "Neither `motor` nor PyMongo async is installed. "
+        "Please install one of them using:\n"
+        "  - `pip install -U 'pymongo>=4.9'` (recommended)"
+        "  - `pip install -U motor` (legacy, deprecated)\n"
+    )
+
+# Create union types for client, database, and collection
+if TYPE_CHECKING:
+    if MOTOR_AVAILABLE and PYMONGO_ASYNC_AVAILABLE:
+        AsyncMongoClientType = Union[AsyncIOMotorClient, AsyncMongoClient]  # type: ignore
+        AsyncMongoDatabaseType = Union[AsyncIOMotorDatabase, AsyncDatabase]  # type: ignore
+        AsyncMongoCollectionType = Union[AsyncIOMotorCollection, AsyncCollection]  # type: ignore
+    elif MOTOR_AVAILABLE:
+        AsyncMongoClientType = AsyncIOMotorClient  # type: ignore
+        AsyncMongoDatabaseType = AsyncIOMotorDatabase  # type: ignore
+        AsyncMongoCollectionType = AsyncIOMotorCollection  # type: ignore
+    else:
+        AsyncMongoClientType = AsyncMongoClient  # type: ignore
+        AsyncMongoDatabaseType = AsyncDatabase  # type: ignore
+        AsyncMongoCollectionType = AsyncCollection  # type: ignore
+else:
+    # Runtime type - use Any to avoid import issues
+    AsyncMongoClientType = Any
+    AsyncMongoDatabaseType = Any
+    AsyncMongoCollectionType = Any
+
+
+# Client type constants (defined before class to allow use in _detect_client_type)
+_CLIENT_TYPE_MOTOR = "motor"
+_CLIENT_TYPE_PYMONGO_ASYNC = "pymongo_async"
+_CLIENT_TYPE_UNKNOWN = "unknown"
+
+
+def _detect_client_type(client: Any) -> str:
+    """Detect whether a client is Motor or PyMongo async."""
+    if client is None:
+        return _CLIENT_TYPE_UNKNOWN
+
+    # Check PyMongo async
+    if PYMONGO_ASYNC_AVAILABLE and AsyncMongoClient is not None:
+        try:
+            if isinstance(client, AsyncMongoClient):
+                return _CLIENT_TYPE_PYMONGO_ASYNC
+        except (TypeError, AttributeError):
+            pass  # Fall through to next check
+
+    if MOTOR_AVAILABLE and AsyncIOMotorClient is not None:
+        try:
+            if isinstance(client, AsyncIOMotorClient):
+                return _CLIENT_TYPE_MOTOR
+        except (TypeError, AttributeError):
+            pass  # Fall through to fallback
+
+    # Fallback to string matching only if isinstance fails
+    # (should rarely happen, but useful for edge cases)
+    client_type_name = type(client).__name__
+    if "Motor" in client_type_name or "AsyncIOMotor" in client_type_name:
+        return _CLIENT_TYPE_MOTOR
+    elif "AsyncMongo" in client_type_name:
+        return _CLIENT_TYPE_PYMONGO_ASYNC
+
+    # Last resort: check module name
+    module_name = type(client).__module__
+    if "motor" in module_name:
+        return _CLIENT_TYPE_MOTOR
+    elif "pymongo" in module_name:
+        return _CLIENT_TYPE_PYMONGO_ASYNC
+
+    return _CLIENT_TYPE_UNKNOWN
+
 
 class AsyncMongoDb(AsyncBaseDb):
+    # Client type constants (class-level access to module constants)
+    CLIENT_TYPE_MOTOR = _CLIENT_TYPE_MOTOR
+    CLIENT_TYPE_PYMONGO_ASYNC = _CLIENT_TYPE_PYMONGO_ASYNC
+    CLIENT_TYPE_UNKNOWN = _CLIENT_TYPE_UNKNOWN
+
     def __init__(
         self,
-        db_client: Optional[AsyncIOMotorClient] = None,
+        db_client: Optional[Union["AsyncIOMotorClient", "AsyncMongoClient"]] = None,
         db_name: Optional[str] = None,
         db_url: Optional[str] = None,
         session_collection: Optional[str] = None,
@@ -53,10 +147,16 @@ class AsyncMongoDb(AsyncBaseDb):
         id: Optional[str] = None,
     ):
         """
-        Async interface for interacting with a MongoDB database using Motor.
+        Async interface for interacting with a MongoDB database.
+
+        Supports both Motor (legacy) and PyMongo async (recommended) clients.
+        When both libraries are available, PyMongo async is preferred.
 
         Args:
-            db_client (Optional[AsyncIOMotorClient]): The MongoDB async client to use.
+            db_client (Optional[Union[AsyncIOMotorClient, AsyncMongoClient]]):
+                The MongoDB async client to use. Can be either Motor's AsyncIOMotorClient
+                or PyMongo's AsyncMongoClient. If not provided, a client will be created
+                from db_url using the preferred available library.
             db_name (Optional[str]): The name of the database to use.
             db_url (Optional[str]): The database URL to connect to.
             session_collection (Optional[str]): Name of the collection to store sessions.
@@ -68,7 +168,8 @@ class AsyncMongoDb(AsyncBaseDb):
             id (Optional[str]): ID of the database.
 
         Raises:
-            ValueError: If neither db_url nor db_client is provided.
+            ValueError: If neither db_url nor db_client is provided, or if db_client type is unsupported.
+            ImportError: If neither motor nor pymongo async is installed.
         """
         if id is None:
             base_seed = db_url or str(db_client)
@@ -86,8 +187,21 @@ class AsyncMongoDb(AsyncBaseDb):
             culture_table=culture_collection,
         )
 
+        # Detect client type if provided
+        if db_client is not None:
+            self._client_type = _detect_client_type(db_client)
+            if self._client_type == self.CLIENT_TYPE_UNKNOWN:
+                raise ValueError(
+                    f"Unsupported MongoDB client type: {type(db_client).__name__}. "
+                    "Only Motor (AsyncIOMotorClient) or PyMongo async (AsyncMongoClient) are supported."
+                )
+        else:
+            # Auto-select preferred library when creating from URL
+            # Prefer PyMongo async if available, fallback to Motor
+            self._client_type = self.CLIENT_TYPE_PYMONGO_ASYNC if PYMONGO_ASYNC_AVAILABLE else self.CLIENT_TYPE_MOTOR
+
         # Store configuration for lazy initialization
-        self._provided_client: Optional[AsyncIOMotorClient] = db_client
+        self._provided_client: Optional[AsyncMongoClientType] = db_client
         self.db_url: Optional[str] = db_url
         self.db_name: str = db_name if db_name is not None else "agno"
 
@@ -95,8 +209,8 @@ class AsyncMongoDb(AsyncBaseDb):
             raise ValueError("One of db_url or db_client must be provided")
 
         # Client and database will be lazily initialized per event loop
-        self._client: Optional[AsyncIOMotorClient] = None
-        self._database: Optional[AsyncIOMotorDatabase] = None
+        self._client: Optional[AsyncMongoClientType] = None
+        self._database: Optional[AsyncMongoDatabaseType] = None
         self._event_loop: Optional[asyncio.AbstractEventLoop] = None
 
     async def table_exists(self, table_name: str) -> bool:
@@ -126,15 +240,16 @@ class AsyncMongoDb(AsyncBaseDb):
             if collection_name and not await self.table_exists(collection_name):
                 await self._get_collection(collection_type, create_collection_if_not_found=True)
 
-    def _ensure_client(self) -> AsyncIOMotorClient:
+    def _ensure_client(self) -> AsyncMongoClientType:
         """
-        Ensure the Motor client is valid for the current event loop.
+        Ensure the MongoDB async client is valid for the current event loop.
 
-        Motor's AsyncIOMotorClient is tied to the event loop it was created in.
-        If we detect a new event loop, we need to refresh the client.
+        Both Motor's AsyncIOMotorClient and PyMongo's AsyncMongoClient are tied to
+        the event loop they were created in. If we detect a new event loop, we need
+        to refresh the client.
 
         Returns:
-            AsyncIOMotorClient: A valid client for the current event loop.
+            Union[AsyncIOMotorClient, AsyncMongoClient]: A valid client for the current event loop.
         """
         try:
             current_loop = asyncio.get_running_loop()
@@ -144,8 +259,13 @@ class AsyncMongoDb(AsyncBaseDb):
                 if self._provided_client is not None:
                     self._client = self._provided_client
                 elif self.db_url is not None:
-                    self._client = AsyncIOMotorClient(self.db_url)
-                    log_debug("Created AsyncIOMotorClient outside event loop")
+                    # Create client based on detected type
+                    if self._client_type == self.CLIENT_TYPE_PYMONGO_ASYNC and PYMONGO_ASYNC_AVAILABLE:
+                        self._client = AsyncMongoClient(self.db_url)  # type: ignore
+                    elif self._client_type == self.CLIENT_TYPE_MOTOR and MOTOR_AVAILABLE:
+                        self._client = AsyncIOMotorClient(self.db_url)  # type: ignore
+                    else:
+                        raise RuntimeError(f"Client type '{self._client_type}' not available")
             return self._client  # type: ignore
 
         # Check if we're in a different event loop
@@ -153,17 +273,21 @@ class AsyncMongoDb(AsyncBaseDb):
             # New event loop detected, create new client
             if self._provided_client is not None:
                 # User provided a client, use it but warn them
+                client_type_name = (
+                    "AsyncMongoClient" if self._client_type == self.CLIENT_TYPE_PYMONGO_ASYNC else "AsyncIOMotorClient"
+                )
                 log_debug(
-                    "New event loop detected. Using provided AsyncIOMotorClient, "
+                    f"New event loop detected. Using provided {client_type_name}, "
                     "which may cause issues if it was created in a different event loop."
                 )
                 self._client = self._provided_client
             elif self.db_url is not None:
-                # Create a new client for this event loop
-                old_loop_id = id(self._event_loop) if self._event_loop else "None"
-                new_loop_id = id(current_loop)
-                log_debug(f"Event loop changed from {old_loop_id} to {new_loop_id}, creating new AsyncIOMotorClient")
-                self._client = AsyncIOMotorClient(self.db_url)
+                if self._client_type == self.CLIENT_TYPE_PYMONGO_ASYNC and PYMONGO_ASYNC_AVAILABLE:
+                    self._client = AsyncMongoClient(self.db_url)  # type: ignore
+                elif self._client_type == self.CLIENT_TYPE_MOTOR and MOTOR_AVAILABLE:
+                    self._client = AsyncIOMotorClient(self.db_url)  # type: ignore
+                else:
+                    raise RuntimeError(f"Client type '{self._client_type}' not available")
 
             self._event_loop = current_loop
             self._database = None  # Reset database reference
@@ -175,21 +299,21 @@ class AsyncMongoDb(AsyncBaseDb):
         return self._client  # type: ignore
 
     @property
-    def db_client(self) -> AsyncIOMotorClient:
+    def db_client(self) -> AsyncMongoClientType:
         """Get the MongoDB client, ensuring it's valid for the current event loop."""
         return self._ensure_client()
 
     @property
-    def database(self) -> AsyncIOMotorDatabase:
+    def database(self) -> AsyncMongoDatabaseType:
         """Get the MongoDB database, ensuring it's valid for the current event loop."""
         try:
             current_loop = asyncio.get_running_loop()
             if self._database is None or self._event_loop != current_loop:
-                self._database = self.db_client[self.db_name]
+                self._database = self.db_client[self.db_name]  # type: ignore
         except RuntimeError:
             # No running loop - fallback to existing database or create new one
             if self._database is None:
-                self._database = self.db_client[self.db_name]
+                self._database = self.db_client[self.db_name]  # type: ignore
         return self._database
 
     # -- DB methods --
@@ -204,7 +328,7 @@ class AsyncMongoDb(AsyncBaseDb):
 
     async def _get_collection(
         self, table_type: str, create_collection_if_not_found: Optional[bool] = True
-    ) -> Optional[AsyncIOMotorCollection]:
+    ) -> Optional[AsyncMongoCollectionType]:
         """Get or create a collection based on table type.
 
         Args:
@@ -212,7 +336,7 @@ class AsyncMongoDb(AsyncBaseDb):
             create_collection_if_not_found (Optional[bool]): Whether to create the collection if it doesn't exist.
 
         Returns:
-            AsyncIOMotorCollection: The collection object.
+            Union[AsyncIOMotorCollection, AsyncCollection]: The collection object.
         """
         # Ensure client is valid for current event loop before accessing collections
         _ = self.db_client  # This triggers _ensure_client()
@@ -290,7 +414,7 @@ class AsyncMongoDb(AsyncBaseDb):
 
     async def _get_or_create_collection(
         self, collection_name: str, collection_type: str, create_collection_if_not_found: Optional[bool] = True
-    ) -> Optional[AsyncIOMotorCollection]:
+    ) -> Optional[AsyncMongoCollectionType]:
         """Get or create a collection with proper indexes.
 
         Args:
@@ -299,7 +423,7 @@ class AsyncMongoDb(AsyncBaseDb):
             create_collection_if_not_found (Optional[bool]): Whether to create the collection if it doesn't exist.
 
         Returns:
-            Optional[AsyncIOMotorCollection]: The collection object.
+            Union[AsyncIOMotorCollection, AsyncCollection]: The collection object.
         """
         try:
             collection = self.database[collection_name]
@@ -307,7 +431,7 @@ class AsyncMongoDb(AsyncBaseDb):
             if not hasattr(self, f"_{collection_name}_initialized"):
                 if not create_collection_if_not_found:
                     return None
-                # Create indexes asynchronously for Motor collections
+                # Create indexes asynchronously for async MongoDB collections
                 await create_collection_indexes_async(collection, collection_type)
                 setattr(self, f"_{collection_name}_initialized", True)
                 log_debug(f"Initialized collection '{collection_name}'")
@@ -1543,7 +1667,7 @@ class AsyncMongoDb(AsyncBaseDb):
             log_error(f"Exception reading from sessions collection: {e}")
             return []
 
-    async def _get_metrics_calculation_starting_date(self, collection: AsyncIOMotorCollection) -> Optional[date]:
+    async def _get_metrics_calculation_starting_date(self, collection: AsyncMongoCollectionType) -> Optional[date]:
         """Get the first date for which metrics calculation is needed."""
         try:
             result = await collection.find_one({}, sort=[("date", -1)], limit=1)

agno/db/mongo/mongo.py

@@ -984,12 +984,14 @@ class MongoDb(BaseDb):
         self,
         limit: Optional[int] = None,
         page: Optional[int] = None,
+        user_id: Optional[str] = None,
     ) -> Tuple[List[Dict[str, Any]], int]:
         """Get user memories stats.
 
         Args:
             limit (Optional[int]): The limit of the memories to get.
             page (Optional[int]): The page number to get.
+            user_id (Optional[str]): User ID for filtering.
 
         Returns:
             Tuple[List[Dict[str, Any]], int]: A tuple containing the memories stats and the total count.
@@ -1002,9 +1004,11 @@ class MongoDb(BaseDb):
             if collection is None:
                 return [], 0
 
-            match_stage = {"user_id": {"$ne": None}}
+            match_stage: Dict[str, Any] = {"user_id": {"$ne": None}}
+            if user_id is not None:
+                match_stage["user_id"] = user_id
 
-            pipeline = [
+            pipeline: List[Dict[str, Any]] = [
                 {"$match": match_stage},
                 {
                     "$group": {

agno/db/mysql/mysql.py

@@ -105,7 +105,7 @@ class MySQLDb(BaseDb):
         self.db_url: Optional[str] = db_url
         self.db_engine: Engine = _engine
         self.db_schema: str = db_schema if db_schema is not None else "ai"
-        self.metadata: MetaData = MetaData()
+        self.metadata: MetaData = MetaData(schema=self.db_schema)
 
         # Initialize database session
         self.Session: scoped_session = scoped_session(sessionmaker(bind=self.db_engine))
@@ -123,14 +123,13 @@ class MySQLDb(BaseDb):
         with self.Session() as sess:
             return is_table_available(session=sess, table_name=table_name, db_schema=self.db_schema)
 
-    def _create_table(self, table_name: str, table_type: str, db_schema: str) -> Table:
+    def _create_table(self, table_name: str, table_type: str) -> Table:
         """
         Create a table with the appropriate schema based on the table type.
 
         Args:
             table_name (str): Name of the table to create
             table_type (str): Type of table (used to get schema definition)
-            db_schema (str): Database schema name
 
         Returns:
             Table: SQLAlchemy Table object
@@ -138,8 +137,6 @@ class MySQLDb(BaseDb):
         try:
             table_schema = get_table_schema_definition(table_type)
 
-            log_debug(f"Creating table {table_name}")
-
             columns: List[Column] = []
             indexes: List[str] = []
             unique_constraints: List[str] = []
@@ -161,8 +158,7 @@ class MySQLDb(BaseDb):
                 columns.append(Column(*column_args, **column_kwargs))  # type: ignore
 
             # Create the table object
-            table_metadata = MetaData(schema=db_schema)
-            table = Table(table_name, table_metadata, *columns, schema=db_schema)
+            table = Table(table_name, self.metadata, *columns, schema=self.db_schema)
 
             # Add multi-column unique constraints with table-specific names
             for constraint in schema_unique_constraints:
@@ -176,16 +172,20 @@ class MySQLDb(BaseDb):
                 table.append_constraint(Index(idx_name, idx_col))
 
             with self.Session() as sess, sess.begin():
-                create_schema(session=sess, db_schema=db_schema)
+                create_schema(session=sess, db_schema=self.db_schema)
 
             # Create table
-            table.create(self.db_engine, checkfirst=True)
+            table_created = False
+            if not self.table_exists(table_name):
+                table.create(self.db_engine, checkfirst=True)
+                log_debug(f"Successfully created table '{table_name}'")
+                table_created = True
+            else:
+                log_debug(f"Table {self.db_schema}.{table_name} already exists, skipping creation")
 
             # Create indexes
             for idx in table.indexes:
                 try:
-                    log_debug(f"Creating index: {idx.name}")
-
                     # Check if index already exists
                     with self.Session() as sess:
                         exists_query = text(
@@ -194,24 +194,35 @@ class MySQLDb(BaseDb):
                         )
                         exists = (
                             sess.execute(
-                                exists_query, {"schema": db_schema, "table_name": table_name, "index_name": idx.name}
+                                exists_query,
+                                {"schema": self.db_schema, "table_name": table_name, "index_name": idx.name},
                             ).scalar()
                             is not None
                         )
                         if exists:
-                            log_debug(f"Index {idx.name} already exists in {db_schema}.{table_name}, skipping creation")
+                            log_debug(
+                                f"Index {idx.name} already exists in {self.db_schema}.{table_name}, skipping creation"
+                            )
                             continue
 
                     idx.create(self.db_engine)
 
+                    log_debug(f"Created index: {idx.name} for table {self.db_schema}.{table_name}")
                 except Exception as e:
                     log_error(f"Error creating index {idx.name}: {e}")
 
-            log_debug(f"Successfully created table {db_schema}.{table_name}")
+            # Store the schema version for the created table
+            if table_name != self.versions_table_name and table_created:
+                latest_schema_version = MigrationManager(self).latest_schema_version
+                self.upsert_schema_version(table_name=table_name, version=latest_schema_version.public)
+                log_info(
+                    f"Successfully stored version {latest_schema_version.public} in database for table {table_name}"
+                )
+
             return table
 
         except Exception as e:
-            log_error(f"Could not create table {db_schema}.{table_name}: {e}")
+            log_error(f"Could not create table {self.db_schema}.{table_name}: {e}")
             raise
 
     def _create_all_tables(self):
@@ -226,19 +237,13 @@ class MySQLDb(BaseDb):
         ]
 
         for table_name, table_type in tables_to_create:
-            if table_name != self.versions_table_name:
-                # Also store the schema version for the created table
-                latest_schema_version = MigrationManager(self).latest_schema_version
-                self.upsert_schema_version(table_name=table_name, version=latest_schema_version.public)
-
-            self._create_table(table_name=table_name, table_type=table_type, db_schema=self.db_schema)
+            self._get_or_create_table(table_name=table_name, table_type=table_type, create_table_if_not_found=True)
 
     def _get_table(self, table_type: str, create_table_if_not_found: Optional[bool] = False) -> Optional[Table]:
         if table_type == "sessions":
             self.session_table = self._get_or_create_table(
                 table_name=self.session_table_name,
                 table_type="sessions",
-                db_schema=self.db_schema,
                 create_table_if_not_found=create_table_if_not_found,
             )
             return self.session_table
@@ -247,7 +252,6 @@ class MySQLDb(BaseDb):
             self.memory_table = self._get_or_create_table(
                 table_name=self.memory_table_name,
                 table_type="memories",
-                db_schema=self.db_schema,
                 create_table_if_not_found=create_table_if_not_found,
             )
             return self.memory_table
@@ -256,7 +260,6 @@ class MySQLDb(BaseDb):
             self.metrics_table = self._get_or_create_table(
                 table_name=self.metrics_table_name,
                 table_type="metrics",
-                db_schema=self.db_schema,
                 create_table_if_not_found=create_table_if_not_found,
             )
             return self.metrics_table
@@ -265,7 +268,6 @@ class MySQLDb(BaseDb):
             self.eval_table = self._get_or_create_table(
                 table_name=self.eval_table_name,
                 table_type="evals",
-                db_schema=self.db_schema,
                 create_table_if_not_found=create_table_if_not_found,
             )
             return self.eval_table
@@ -274,7 +276,6 @@ class MySQLDb(BaseDb):
             self.knowledge_table = self._get_or_create_table(
                 table_name=self.knowledge_table_name,
                 table_type="knowledge",
-                db_schema=self.db_schema,
                 create_table_if_not_found=create_table_if_not_found,
             )
             return self.knowledge_table
@@ -283,7 +284,6 @@ class MySQLDb(BaseDb):
             self.culture_table = self._get_or_create_table(
                 table_name=self.culture_table_name,
                 table_type="culture",
-                db_schema=self.db_schema,
                 create_table_if_not_found=create_table_if_not_found,
             )
             return self.culture_table
@@ -292,7 +292,6 @@ class MySQLDb(BaseDb):
             self.versions_table = self._get_or_create_table(
                 table_name=self.versions_table_name,
                 table_type="versions",
-                db_schema=self.db_schema,
                 create_table_if_not_found=create_table_if_not_found,
             )
             return self.versions_table
@@ -300,7 +299,7 @@ class MySQLDb(BaseDb):
         raise ValueError(f"Unknown table type: {table_type}")
 
     def _get_or_create_table(
-        self, table_name: str, table_type: str, db_schema: str, create_table_if_not_found: Optional[bool] = False
+        self, table_name: str, table_type: str, create_table_if_not_found: Optional[bool] = False
     ) -> Optional[Table]:
         """
         Check if the table exists and is valid, else create it.
@@ -308,25 +307,19 @@ class MySQLDb(BaseDb):
         Args:
             table_name (str): Name of the table to get or create
             table_type (str): Type of table (used to get schema definition)
-            db_schema (str): Database schema name
 
         Returns:
             Table: SQLAlchemy Table object representing the schema.
         """
 
         with self.Session() as sess, sess.begin():
-            table_is_available = is_table_available(session=sess, table_name=table_name, db_schema=db_schema)
+            table_is_available = is_table_available(session=sess, table_name=table_name, db_schema=self.db_schema)
 
         if not table_is_available:
             if not create_table_if_not_found:
                 return None
 
-            created_table = self._create_table(table_name=table_name, table_type=table_type, db_schema=db_schema)
-
-            if table_name != self.versions_table_name:
-                # Also store the schema version for the created table
-                latest_schema_version = MigrationManager(self).latest_schema_version
-                self.upsert_schema_version(table_name=table_name, version=latest_schema_version.public)
+            created_table = self._create_table(table_name=table_name, table_type=table_type)
 
             return created_table
 
@@ -334,17 +327,16 @@ class MySQLDb(BaseDb):
             db_engine=self.db_engine,
             table_name=table_name,
             table_type=table_type,
-            db_schema=db_schema,
+            db_schema=self.db_schema,
         ):
-            raise ValueError(f"Table {db_schema}.{table_name} has an invalid schema")
+            raise ValueError(f"Table {self.db_schema}.{table_name} has an invalid schema")
 
         try:
-            table = Table(table_name, self.metadata, schema=db_schema, autoload_with=self.db_engine)
-            log_debug(f"Loaded existing table {db_schema}.{table_name}")
+            table = Table(table_name, self.metadata, schema=self.db_schema, autoload_with=self.db_engine)
             return table
 
         except Exception as e:
-            log_error(f"Error loading existing table {db_schema}.{table_name}: {e}")
+            log_error(f"Error loading existing table {self.db_schema}.{table_name}: {e}")
             raise
 
     def get_latest_schema_version(self, table_name: str) -> str:
@@ -513,7 +505,7 @@ class MySQLDb(BaseDb):
         Args:
             session_type (Optional[SessionType]): The type of sessions to get.
             user_id (Optional[str]): The ID of the user to filter by.
-            entity_id (Optional[str]): The ID of the agent / workflow to filter by.
+            component_id (Optional[str]): The ID of the agent / workflow to filter by.
             start_timestamp (Optional[int]): The start timestamp to filter by.
             end_timestamp (Optional[int]): The end timestamp to filter by.
             session_name (Optional[str]): The name of the session to filter by.
@@ -522,7 +514,6 @@ class MySQLDb(BaseDb):
             sort_by (Optional[str]): The field to sort by. Defaults to None.
             sort_order (Optional[str]): The sort order. Defaults to None.
             deserialize (Optional[bool]): Whether to serialize the sessions. Defaults to True.
-            create_table_if_not_found (Optional[bool]): Whether to create the table if it doesn't exist.
 
         Returns:
             Union[List[Session], Tuple[List[Dict], int]]:
@@ -1254,7 +1245,7 @@ class MySQLDb(BaseDb):
             log_error(f"Exception clearing user memories: {e}")
 
     def get_user_memory_stats(
-        self, limit: Optional[int] = None, page: Optional[int] = None
+        self, limit: Optional[int] = None, page: Optional[int] = None, user_id: Optional[str] = None
     ) -> Tuple[List[Dict[str, Any]], int]:
         """Get user memories stats.
 
@@ -1283,17 +1274,20 @@ class MySQLDb(BaseDb):
                 return [], 0
 
             with self.Session() as sess, sess.begin():
-                stmt = (
-                    select(
-                        table.c.user_id,
-                        func.count(table.c.memory_id).label("total_memories"),
-                        func.max(table.c.updated_at).label("last_memory_updated_at"),
-                    )
-                    .where(table.c.user_id.is_not(None))
-                    .group_by(table.c.user_id)
-                    .order_by(func.max(table.c.updated_at).desc())
+                stmt = select(
+                    table.c.user_id,
+                    func.count(table.c.memory_id).label("total_memories"),
+                    func.max(table.c.updated_at).label("last_memory_updated_at"),
                 )
 
+                if user_id is not None:
+                    stmt = stmt.where(table.c.user_id == user_id)
+                else:
+                    stmt = stmt.where(table.c.user_id.is_not(None))
+
+                stmt = stmt.group_by(table.c.user_id)
+                stmt = stmt.order_by(func.max(table.c.updated_at).desc())
+
                 count_stmt = select(func.count()).select_from(stmt.alias())
                 total_count = sess.execute(count_stmt).scalar()