agno 1.7.2__py3-none-any.whl → 1.7.4__py3-none-any.whl

agno/storage/dynamodb.py

@@ -7,6 +7,7 @@ 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, log_info, logger
 
@@ -28,7 +29,7 @@ class DynamoDbStorage(Storage):
         aws_secret_access_key: Optional[str] = None,
         endpoint_url: Optional[str] = None,
         create_table_if_not_exists: bool = True,
-        mode: Optional[Literal["agent", "team", "workflow"]] = "agent",
+        mode: Optional[Literal["agent", "team", "workflow", "workflow_v2"]] = "agent",
     ):
         """
         Initialize the DynamoDbStorage.
@@ -41,7 +42,7 @@ class DynamoDbStorage(Storage):
             aws_secret_access_key (Optional[str]): AWS secret access key.
             endpoint_url (Optional[str]): The complete URL to use for the constructed client.
             create_table_if_not_exists (bool): Whether to create the table if it does not exist.
-            mode (Optional[Literal["agent", "team", "workflow"]]): The mode of the storage.
+            mode (Optional[Literal["agent", "team", "workflow", "workflow_v2"]]): The mode of the storage.
         """
         super().__init__(mode)
         self.table_name = table_name
@@ -79,12 +80,12 @@ class DynamoDbStorage(Storage):
         log_debug(f"Initialized DynamoDbStorage with table '{self.table_name}'")
 
     @property
-    def mode(self) -> Literal["agent", "team", "workflow"]:
+    def mode(self) -> Literal["agent", "team", "workflow", "workflow_v2"]:
         """Get the mode of the storage."""
         return super().mode
 
     @mode.setter
-    def mode(self, value: Optional[Literal["agent", "team", "workflow"]]) -> None:
+    def mode(self, value: Optional[Literal["agent", "team", "workflow", "workflow_v2"]]) -> None:
         """Set the mode and refresh the table if mode changes."""
         super(DynamoDbStorage, type(self)).mode.fset(self, value)  # type: ignore
         if value is not None:
@@ -125,7 +126,13 @@ class DynamoDbStorage(Storage):
                         {"AttributeName": "workflow_id", "AttributeType": "S"},
                         {"AttributeName": "created_at", "AttributeType": "N"},
                     ]
-
+                elif self.mode == "workflow_v2":
+                    attribute_definitions = [
+                        {"AttributeName": "session_id", "AttributeType": "S"},
+                        {"AttributeName": "user_id", "AttributeType": "S"},
+                        {"AttributeName": "workflow_id", "AttributeType": "S"},
+                        {"AttributeName": "created_at", "AttributeType": "N"},
+                    ]
                 secondary_indexes = [
                     {
                         "IndexName": "user_id-index",
@@ -185,7 +192,21 @@ class DynamoDbStorage(Storage):
                             },
                         }
                     )
-
+                elif self.mode == "workflow_v2":
+                    secondary_indexes.append(
+                        {
+                            "IndexName": "workflow_id-index",
+                            "KeySchema": [
+                                {"AttributeName": "workflow_id", "KeyType": "HASH"},
+                                {"AttributeName": "created_at", "KeyType": "RANGE"},
+                            ],
+                            "Projection": {"ProjectionType": "ALL"},
+                            "ProvisionedThroughput": {
+                                "ReadCapacityUnits": 5,
+                                "WriteCapacityUnits": 5,
+                            },
+                        }
+                    )
                 # Create the table
                 self.table = self.dynamodb.create_table(
                     TableName=self.table_name,
@@ -229,6 +250,8 @@ class DynamoDbStorage(Storage):
                     return TeamSession.from_dict(item)
                 elif self.mode == "workflow":
                     return WorkflowSession.from_dict(item)
+                elif self.mode == "workflow_v2":
+                    return WorkflowSessionV2.from_dict(item)
         except Exception as e:
             logger.error(f"Error reading session_id '{session_id}' with user_id '{user_id}': {e}")
         return None
@@ -323,7 +346,13 @@ class DynamoDbStorage(Storage):
                         KeyConditionExpression=Key("user_id").eq(user_id),
                         ProjectionExpression="session_id, workflow_id, user_id, memory, workflow_data, session_data, extra_data, created_at, updated_at",
                     )
-
+                elif self.mode == "workflow_v2":
+                    # Query using user_id index
+                    response = self.table.query(
+                        IndexName="user_id-index",
+                        KeyConditionExpression=Key("user_id").eq(user_id),
+                        ProjectionExpression="session_id, workflow_id, user_id, workflow_name, runs, workflow_data, session_data, extra_data, created_at, updated_at",
+                    )
                 items = response.get("Items", [])  # type: ignore
                 for item in items:
                     item = self._deserialize_item(item)
@@ -356,6 +385,13 @@ class DynamoDbStorage(Storage):
                         KeyConditionExpression=Key("workflow_id").eq(entity_id),
                         ProjectionExpression="session_id, workflow_id, user_id, memory, workflow_data, session_data, extra_data, created_at, updated_at",
                     )
+                elif self.mode == "workflow_v2":
+                    # Query using workflow_id index
+                    response = self.table.query(
+                        IndexName="workflow_id-index",
+                        KeyConditionExpression=Key("workflow_id").eq(entity_id),
+                        ProjectionExpression="session_id, workflow_id, user_id, workflow_name, runs, workflow_data, session_data, extra_data, created_at, updated_at",
+                    )
                 items = response.get("Items", [])  # type: ignore
                 for item in items:
                     item = self._deserialize_item(item)
@@ -379,6 +415,10 @@ class DynamoDbStorage(Storage):
                     response = self.table.scan(
                         ProjectionExpression="session_id, workflow_id, user_id, memory, workflow_data, session_data, extra_data, created_at, updated_at"
                     )
+                elif self.mode == "workflow_v2":
+                    response = self.table.scan(
+                        ProjectionExpression="session_id, workflow_id, user_id, workflow_name, runs, workflow_data, session_data, extra_data, created_at, updated_at"
+                    )
                 items = response.get("Items", [])
                 for item in items:
                     item = self._deserialize_item(item)
@@ -437,6 +477,14 @@ class DynamoDbStorage(Storage):
                         ScanIndexForward=False,
                         Limit=limit if limit is not None else None,
                     )
+                elif self.mode == "workflow_v2":
+                    response = self.table.query(
+                        IndexName="user_id-index",
+                        KeyConditionExpression=Key("user_id").eq(user_id),
+                        ProjectionExpression="session_id, workflow_id, user_id, workflow_name, runs, workflow_data, session_data, extra_data, created_at, updated_at",
+                        ScanIndexForward=False,
+                        Limit=limit if limit is not None else None,
+                    )
             elif entity_id is not None:
                 if self.mode == "agent":
                     response = self.table.query(
@@ -462,6 +510,14 @@ class DynamoDbStorage(Storage):
                         ScanIndexForward=False,
                         Limit=limit if limit is not None else None,
                     )
+                elif self.mode == "workflow_v2":
+                    response = self.table.query(
+                        IndexName="workflow_id-index",
+                        KeyConditionExpression=Key("workflow_id").eq(entity_id),
+                        ProjectionExpression="session_id, workflow_id, user_id, workflow_name, runs, workflow_data, session_data, extra_data, created_at, updated_at",
+                        ScanIndexForward=False,
+                        Limit=limit if limit is not None else None,
+                    )
             else:
                 # If no filters, scan the table and sort by created_at
                 if self.mode == "agent":
@@ -479,7 +535,11 @@ class DynamoDbStorage(Storage):
                         ProjectionExpression="session_id, workflow_id, user_id, memory, workflow_data, session_data, extra_data, created_at, updated_at",
                         Limit=limit if limit is not None else None,
                     )
-
+                elif self.mode == "workflow_v2":
+                    response = self.table.scan(
+                        ProjectionExpression="session_id, workflow_id, user_id, workflow_name, runs, workflow_data, session_data, extra_data, created_at, updated_at",
+                        Limit=limit if limit is not None else None,
+                    )
             items = response.get("Items", [])
             for item in items:
                 item = self._deserialize_item(item)
@@ -491,7 +551,8 @@ class DynamoDbStorage(Storage):
                     session = TeamSession.from_dict(item)
                 elif self.mode == "workflow":
                     session = WorkflowSession.from_dict(item)
-
+                elif self.mode == "workflow_v2":
+                    session = WorkflowSessionV2.from_dict(item)
                 if session is not None:
                     sessions.append(session)
 
@@ -511,7 +572,10 @@ class DynamoDbStorage(Storage):
             Optional[Session]: The upserted Session, or None if operation failed.
         """
         try:
-            item = asdict(session)
+            if self.mode == "workflow_v2":
+                item = session.to_dict()
+            else:
+                item = asdict(session)
 
             # Add timestamps
             current_time = int(time.time())

agno/storage/firestore.py

@@ -6,6 +6,7 @@ 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
 
@@ -29,7 +30,7 @@ class FirestoreStorage(Storage):
         db_name: Optional[str] = "(default)",
         project_id: Optional[str] = None,
         client: Optional[Client] = None,
-        mode: Optional[Literal["agent", "team", "workflow"]] = "agent",
+        mode: Optional[Literal["agent", "team", "workflow", "workflow_v2"]] = "agent",
     ):
         super().__init__(mode)
         self.collection_name = collection_name
@@ -94,6 +95,8 @@ class FirestoreStorage(Storage):
                 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
 
@@ -109,6 +112,8 @@ class FirestoreStorage(Storage):
                 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

agno/storage/gcs_json.py

@@ -6,6 +6,7 @@ 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
 
@@ -35,7 +36,7 @@ class GCSJsonStorage(JsonStorage):
         self,
         bucket_name: str,
         prefix: Optional[str] = "",
-        mode: Optional[Literal["agent", "team", "workflow"]] = "agent",
+        mode: Optional[Literal["agent", "team", "workflow", "workflow_v2"]] = "agent",
         project: Optional[str] = None,
         location: Optional[str] = None,
         credentials: Optional[Any] = None,
@@ -104,6 +105,8 @@ class GCSJsonStorage(JsonStorage):
             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]:
@@ -136,6 +139,8 @@ class GCSJsonStorage(JsonStorage):
                         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:
@@ -199,7 +204,8 @@ class GCSJsonStorage(JsonStorage):
                     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)
 

agno/storage/json.py

@@ -8,12 +8,15 @@ 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 logger
 
 
 class JsonStorage(Storage):
-    def __init__(self, dir_path: Union[str, Path], mode: Optional[Literal["agent", "team", "workflow"]] = "agent"):
+    def __init__(
+        self, dir_path: Union[str, Path], mode: Optional[Literal["agent", "team", "workflow", "workflow_v2"]] = "agent"
+    ):
         super().__init__(mode)
         self.dir_path = Path(dir_path)
         self.dir_path.mkdir(parents=True, exist_ok=True)
@@ -42,6 +45,8 @@ class JsonStorage(Storage):
                     return TeamSession.from_dict(data)
                 elif self.mode == "workflow":
                     return WorkflowSession.from_dict(data)
+                elif self.mode == "workflow_v2":
+                    return WorkflowSessionV2.from_dict(data)
         except FileNotFoundError:
             return None
 
@@ -70,6 +75,8 @@ class JsonStorage(Storage):
                             session_ids.append(data["session_id"])
                         elif self.mode == "workflow" and data["workflow_id"] == entity_id:
                             session_ids.append(data["session_id"])
+                        elif self.mode == "workflow_v2" and data["workflow_id"] == entity_id:
+                            session_ids.append(data["session_id"])
                 else:
                     # No filters applied, add all session_ids
                     session_ids.append(data["session_id"])
@@ -107,7 +114,8 @@ class JsonStorage(Storage):
                             _session = TeamSession.from_dict(data)
                         elif self.mode == "workflow" and data["workflow_id"] == entity_id:
                             _session = WorkflowSession.from_dict(data)
-
+                        elif self.mode == "workflow_v2" and data["workflow_id"] == entity_id:
+                            _session = WorkflowSessionV2.from_dict(data)
                     if _session:
                         sessions.append(_session)
                 else:
@@ -118,6 +126,8 @@ class JsonStorage(Storage):
                         _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:
                         sessions.append(_session)
         return sessions
@@ -158,7 +168,8 @@ class JsonStorage(Storage):
                             continue
                         elif self.mode == "workflow" and data["workflow_id"] != entity_id:
                             continue
-
+                        elif self.mode == "workflow_v2" and data["workflow_id"] != entity_id:
+                            continue
                     # Store with created_at for sorting
                     created_at = data.get("created_at", 0)
                     session_data.append((created_at, data))
@@ -181,7 +192,8 @@ class JsonStorage(Storage):
                 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)
 
@@ -190,7 +202,10 @@ class JsonStorage(Storage):
     def upsert(self, session: Session) -> Optional[Session]:
         """Insert or update a Session in storage."""
         try:
-            data = asdict(session)
+            if self.mode == "workflow_v2":
+                data = session.to_dict()
+            else:
+                data = asdict(session)
             data["updated_at"] = int(time.time())
             if "created_at" not in data:
                 data["created_at"] = data["updated_at"]

agno/storage/mongodb.py

@@ -6,6 +6,7 @@ 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
 
@@ -25,7 +26,7 @@ class MongoDbStorage(Storage):
         db_url: Optional[str] = None,
         db_name: str = "agno",
         client: Optional[MongoClient] = None,
-        mode: Optional[Literal["agent", "team", "workflow"]] = "agent",
+        mode: Optional[Literal["agent", "team", "workflow", "workflow_v2"]] = "agent",
     ):
         """
         This class provides agent storage using MongoDB.
@@ -64,6 +65,8 @@ class MongoDbStorage(Storage):
                 self.collection.create_index("team_id")
             elif self.mode == "workflow":
                 self.collection.create_index("workflow_id")
+            elif self.mode == "workflow_v2":
+                self.collection.create_index("workflow_id")
         except PyMongoError as e:
             logger.error(f"Error creating indexes: {e}")
             raise
@@ -91,6 +94,8 @@ class MongoDbStorage(Storage):
                     return TeamSession.from_dict(doc)
                 elif self.mode == "workflow":
                     return WorkflowSession.from_dict(doc)
+                elif self.mode == "workflow_v2":
+                    return WorkflowSessionV2.from_dict(doc)
             return None
         except PyMongoError as e:
             logger.error(f"Error reading session: {e}")
@@ -115,7 +120,8 @@ class MongoDbStorage(Storage):
                     query["team_id"] = entity_id
                 elif self.mode == "workflow":
                     query["workflow_id"] = entity_id
-
+                elif self.mode == "workflow_v2":
+                    query["workflow_id"] = entity_id
             cursor = self.collection.find(query, {"session_id": 1}).sort("created_at", -1)
 
             return [str(doc["session_id"]) for doc in cursor]
@@ -142,7 +148,8 @@ class MongoDbStorage(Storage):
                     query["team_id"] = entity_id
                 elif self.mode == "workflow":
                     query["workflow_id"] = entity_id
-
+                elif self.mode == "workflow_v2":
+                    query["workflow_id"] = entity_id
             cursor = self.collection.find(query).sort("created_at", -1)
             sessions: List[Session] = []
             for doc in cursor:
@@ -193,7 +200,8 @@ class MongoDbStorage(Storage):
                     query["team_id"] = entity_id
                 elif self.mode == "workflow":
                     query["workflow_id"] = entity_id
-
+                elif self.mode == "workflow_v2":
+                    query["workflow_id"] = entity_id
             # Execute query with sort and limit
             cursor = self.collection.find(query)
             cursor = cursor.sort("created_at", -1)  # Sort by created_at descending
@@ -212,7 +220,8 @@ class MongoDbStorage(Storage):
                     session = TeamSession.from_dict(doc)
                 elif self.mode == "workflow":
                     session = WorkflowSession.from_dict(doc)
-
+                elif self.mode == "workflow_v2":
+                    session = WorkflowSessionV2.from_dict(doc)
                 if session is not None:
                     sessions.append(session)
 

agno/storage/mysql.py

@@ -5,6 +5,7 @@ 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, log_info, log_warning, logger
 
@@ -29,7 +30,7 @@ class MySQLStorage(Storage):
         db_engine: Optional[Engine] = None,
         schema_version: int = 1,
         auto_upgrade_schema: bool = False,
-        mode: Optional[Literal["agent", "team", "workflow"]] = "agent",
+        mode: Optional[Literal["agent", "team", "workflow", "workflow_v2"]] = "agent",
     ):
         """
         This class provides agent storage using a MySQL table.
@@ -46,7 +47,7 @@ class MySQLStorage(Storage):
             db_engine (Optional[Engine]): The SQLAlchemy database engine to use.
             schema_version (int): Version of the schema. Defaults to 1.
             auto_upgrade_schema (bool): Whether to automatically upgrade the schema.
-            mode (Optional[Literal["agent", "team", "workflow"]]): The mode of the storage.
+            mode (Optional[Literal["agent", "team", "workflow", "workflow_v2"]]): The mode of the storage.
         Raises:
             ValueError: If neither db_url nor db_engine is provided.
         """
@@ -79,12 +80,12 @@ class MySQLStorage(Storage):
         log_debug(f"Created MySQLStorage: '{self.schema}.{self.table_name}'")
 
     @property
-    def mode(self) -> Literal["agent", "team", "workflow"]:
+    def mode(self) -> Literal["agent", "team", "workflow", "workflow_v2"]:
         """Get the mode of the storage."""
         return super().mode
 
     @mode.setter
-    def mode(self, value: Optional[Literal["agent", "team", "workflow"]]) -> None:
+    def mode(self, value: Optional[Literal["agent", "team", "workflow", "workflow_v2"]]) -> None:
         """Set the mode and refresh the table if mode changes."""
         super(MySQLStorage, type(self)).mode.fset(self, value)  # type: ignore
         if value is not None:
@@ -127,7 +128,11 @@ class MySQLStorage(Storage):
                 Column("workflow_id", String(255), index=True),
                 Column("workflow_data", JSON),
             ]
-
+        elif self.mode == "workflow_v2":
+            specific_columns = [
+                Column("workflow_id", String(255), index=True),
+                Column("workflow_data", JSON),
+            ]
         # Create table with all columns
         table = Table(
             self.table_name,
@@ -278,6 +283,8 @@ class MySQLStorage(Storage):
                     return TeamSession.from_dict(result._mapping) if result is not None else None
                 elif self.mode == "workflow":
                     return WorkflowSession.from_dict(result._mapping) if result is not None else None
+                elif self.mode == "workflow_v2":
+                    return WorkflowSessionV2.from_dict(result._mapping) if result is not None else None
         except Exception as e:
             if "doesn't exist" in str(e) or "doesn't exist" in str(e):
                 log_debug(f"Table does not exist: {self.table.name}")
@@ -311,7 +318,8 @@ class MySQLStorage(Storage):
                         stmt = stmt.where(self.table.c.team_id == entity_id)
                     elif self.mode == "workflow":
                         stmt = stmt.where(self.table.c.workflow_id == entity_id)
-
+                    elif self.mode == "workflow_v2":
+                        stmt = stmt.where(self.table.c.workflow_id == entity_id)
                 # order by created_at desc
                 stmt = stmt.order_by(self.table.c.created_at.desc())
                 # execute query
@@ -346,7 +354,9 @@ class MySQLStorage(Storage):
                         stmt = stmt.where(self.table.c.agent_id == entity_id)
                     elif self.mode == "team":
                         stmt = stmt.where(self.table.c.team_id == entity_id)
-                    else:
+                    elif self.mode == "workflow":
+                        stmt = stmt.where(self.table.c.workflow_id == entity_id)
+                    elif self.mode == "workflow_v2":
                         stmt = stmt.where(self.table.c.workflow_id == entity_id)
                 # order by created_at desc
                 stmt = stmt.order_by(self.table.c.created_at.desc())
@@ -357,8 +367,10 @@ class MySQLStorage(Storage):
                         return [AgentSession.from_dict(row._mapping) for row in rows]  # type: ignore
                     elif self.mode == "team":
                         return [TeamSession.from_dict(row._mapping) for row in rows]  # type: ignore
-                    else:
+                    elif self.mode == "workflow":
                         return [WorkflowSession.from_dict(row._mapping) for row in rows]  # type: ignore
+                    elif self.mode == "workflow_v2":
+                        return [WorkflowSessionV2.from_dict(row._mapping) for row in rows]  # type: ignore[misc]
                 else:
                     return []
         except Exception as e:
@@ -399,7 +411,8 @@ class MySQLStorage(Storage):
                         stmt = stmt.where(self.table.c.team_id == entity_id)
                     elif self.mode == "workflow":
                         stmt = stmt.where(self.table.c.workflow_id == entity_id)
-
+                    elif self.mode == "workflow_v2":
+                        stmt = stmt.where(self.table.c.workflow_id == entity_id)
                 # Order by created_at desc and limit results
                 stmt = stmt.order_by(self.table.c.created_at.desc())
                 if limit is not None:
@@ -417,7 +430,8 @@ class MySQLStorage(Storage):
                             session = TeamSession.from_dict(row._mapping)  # type: ignore
                         elif self.mode == "workflow":
                             session = WorkflowSession.from_dict(row._mapping)  # type: ignore
-
+                        elif self.mode == "workflow_v2":
+                            session = WorkflowSessionV2.from_dict(row._mapping)  # type: ignore
                         if session is not None:
                             sessions.append(session)
                     return sessions
@@ -494,7 +508,7 @@ class MySQLStorage(Storage):
                         agent_id=session.agent_id,  # type: ignore
                         team_session_id=session.team_session_id,  # type: ignore
                         user_id=session.user_id,
-                        memory=session.memory,
+                        memory=getattr(session, "memory", None),
                         agent_data=session.agent_data,  # type: ignore
                         session_data=session.session_data,
                         extra_data=session.extra_data,
@@ -505,7 +519,7 @@ class MySQLStorage(Storage):
                         agent_id=session.agent_id,  # type: ignore
                         team_session_id=session.team_session_id,  # type: ignore
                         user_id=session.user_id,
-                        memory=session.memory,
+                        memory=getattr(session, "memory", None),
                         agent_data=session.agent_data,  # type: ignore
                         session_data=session.session_data,
                         extra_data=session.extra_data,
@@ -517,7 +531,7 @@ class MySQLStorage(Storage):
                         team_id=session.team_id,  # type: ignore
                         user_id=session.user_id,
                         team_session_id=session.team_session_id,  # type: ignore
-                        memory=session.memory,
+                        memory=getattr(session, "memory", None),
                         team_data=session.team_data,  # type: ignore
                         session_data=session.session_data,
                         extra_data=session.extra_data,
@@ -528,18 +542,18 @@ class MySQLStorage(Storage):
                         team_id=session.team_id,  # type: ignore
                         user_id=session.user_id,
                         team_session_id=session.team_session_id,  # type: ignore
-                        memory=session.memory,
+                        memory=getattr(session, "memory", None),
                         team_data=session.team_data,  # type: ignore
                         session_data=session.session_data,
                         extra_data=session.extra_data,
                         updated_at=int(time.time()),
                     )
-                else:
+                elif self.mode == "workflow":
                     stmt = mysql.insert(self.table).values(
                         session_id=session.session_id,
                         workflow_id=session.workflow_id,  # type: ignore
                         user_id=session.user_id,
-                        memory=session.memory,
+                        memory=getattr(session, "memory", None),
                         workflow_data=session.workflow_data,  # type: ignore
                         session_data=session.session_data,
                         extra_data=session.extra_data,
@@ -549,13 +563,38 @@ class MySQLStorage(Storage):
                     stmt = stmt.on_duplicate_key_update(
                         workflow_id=session.workflow_id,  # type: ignore
                         user_id=session.user_id,
-                        memory=session.memory,
+                        memory=getattr(session, "memory", None),
                         workflow_data=session.workflow_data,  # type: ignore
                         session_data=session.session_data,
                         extra_data=session.extra_data,
                         updated_at=int(time.time()),
                     )
+                elif self.mode == "workflow_v2":
+                    # Convert session to dict to ensure proper serialization
+                    session_dict = session.to_dict()
 
+                    stmt = mysql.insert(self.table).values(
+                        session_id=session.session_id,
+                        workflow_id=session.workflow_id,  # type: ignore
+                        workflow_name=session.workflow_name,  # type: ignore
+                        user_id=session.user_id,
+                        runs=session_dict.get("runs"),
+                        workflow_data=session.workflow_data,  # type: ignore
+                        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/mysql.html#insert-on-duplicate-key-update
+                    stmt = stmt.on_duplicate_key_update(
+                        workflow_id=session.workflow_id,  # type: ignore
+                        workflow_name=session.workflow_name,  # type: ignore
+                        user_id=session.user_id,
+                        runs=session_dict.get("runs"),  # type: ignore
+                        workflow_data=session.workflow_data,  # type: ignore
+                        session_data=session.session_data,
+                        extra_data=session.extra_data,
+                        updated_at=int(time.time()),
+                    )
                 sess.execute(stmt)
         except Exception as e:
             if create_and_retry and not self.table_exists():