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

agno/storage/postgres.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
 
@@ -79,12 +80,12 @@ class PostgresStorage(Storage):
         log_debug(f"Created PostgresStorage: '{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(PostgresStorage, type(self)).mode.fset(self, value)  # type: ignore
         if value is not None:
@@ -127,6 +128,11 @@ class PostgresStorage(Storage):
                 Column("workflow_id", String, index=True),
                 Column("workflow_data", postgresql.JSONB),
             ]
+        elif self.mode == "workflow_v2":
+            specific_columns = [
+                Column("workflow_id", String, index=True),
+                Column("workflow_data", postgresql.JSONB),
+            ]
 
         # Create table with all columns
         table = Table(
@@ -263,6 +269,8 @@ class PostgresStorage(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 "does not exist" in str(e):
                 log_debug(f"Table does not exist: {self.table.name}")
@@ -296,7 +304,8 @@ class PostgresStorage(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
@@ -331,7 +340,9 @@ class PostgresStorage(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())
@@ -384,7 +395,8 @@ class PostgresStorage(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:
@@ -402,7 +414,8 @@ class PostgresStorage(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
@@ -479,7 +492,7 @@ class PostgresStorage(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,
@@ -492,7 +505,7 @@ class PostgresStorage(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 +518,7 @@ class PostgresStorage(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,
@@ -518,19 +531,19 @@ class PostgresStorage(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()),
                         ),  # The updated value for each column
                     )
-                else:
+                elif self.mode == "workflow":
                     stmt = postgresql.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,
@@ -542,13 +555,42 @@ class PostgresStorage(Storage):
                         set_=dict(
                             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()),
                         ),  # The updated value for each column
                     )
+                elif self.mode == "workflow_v2":
+                    # Convert session to dict to ensure proper serialization
+                    session_dict = session.to_dict()
+
+                    stmt = postgresql.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/postgresql.html#postgresql-insert-on-conflict
+                    stmt = stmt.on_conflict_do_update(
+                        index_elements=["session_id"],
+                        set_=dict(
+                            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,
+                            updated_at=int(time.time()),
+                        ),
+                    )
 
                 sess.execute(stmt)
         except Exception as e:

agno/storage/redis.py

@@ -8,6 +8,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
 
@@ -32,7 +33,7 @@ class RedisStorage(Storage):
         port: int = 6379,
         db: int = 0,
         password: Optional[str] = None,
-        mode: Optional[Literal["agent", "team", "workflow"]] = "agent",
+        mode: Optional[Literal["agent", "team", "workflow", "workflow_v2"]] = "agent",
         ssl: Optional[bool] = False,
         expire: Optional[int] = None,
     ):
@@ -45,7 +46,7 @@ class RedisStorage(Storage):
             port (int): Redis port number
             db (int): Redis database number
             password (Optional[str]): Redis password if authentication is required
-            mode (Optional[Literal["agent", "team", "workflow"]]): Storage mode
+            mode (Optional[Literal["agent", "team", "workflow", "workflow_v2"]]): Storage mode
             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.
         """
@@ -104,6 +105,8 @@ class RedisStorage(Storage):
                 return TeamSession.from_dict(session_data)
             elif self.mode == "workflow":
                 return WorkflowSession.from_dict(session_data)
+            elif self.mode == "workflow_v2":
+                return WorkflowSessionV2.from_dict(session_data)
 
         except Exception as e:
             logger.error(f"Error reading session: {e}")
@@ -128,6 +131,8 @@ class RedisStorage(Storage):
                             self.mode == "workflow" and data["workflow_id"] == entity_id and data["user_id"] == user_id
                         ):
                             session_ids.append(data["session_id"])
+                        elif self.mode == "workflow_v2" and data["workflow_id"] == entity_id:
+                            session_ids.append(data["session_id"])
                     elif user_id and data["user_id"] == user_id:
                         session_ids.append(data["session_id"])
                     elif entity_id:
@@ -137,6 +142,8 @@ class RedisStorage(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"])
@@ -166,6 +173,8 @@ class RedisStorage(Storage):
                             self.mode == "workflow" and data["workflow_id"] == entity_id and data["user_id"] == user_id
                         ):
                             _session = WorkflowSession.from_dict(data)
+                        elif self.mode == "workflow_v2" and data["workflow_id"] == entity_id:
+                            _session = WorkflowSessionV2.from_dict(data)
                     elif user_id and data["user_id"] == user_id:
                         if self.mode == "agent":
                             _session = AgentSession.from_dict(data)
@@ -173,6 +182,8 @@ class RedisStorage(Storage):
                             _session = TeamSession.from_dict(data)
                         elif self.mode == "workflow":
                             _session = WorkflowSession.from_dict(data)
+                        elif self.mode == "workflow_v2" and data["workflow_id"] == entity_id:
+                            _session = WorkflowSessionV2.from_dict(data)
                     elif entity_id:
                         if self.mode == "agent" and data["agent_id"] == entity_id:
                             _session = AgentSession.from_dict(data)
@@ -180,6 +191,8 @@ class RedisStorage(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)
@@ -192,6 +205,8 @@ class RedisStorage(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)
 
@@ -237,7 +252,8 @@ class RedisStorage(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))
@@ -260,7 +276,8 @@ class RedisStorage(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)
 
@@ -272,7 +289,10 @@ class RedisStorage(Storage):
     def upsert(self, session: Session) -> Optional[Session]:
         """Insert or update a Session in Redis."""
         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/session/__init__.py

@@ -2,13 +2,15 @@ from typing import Union
 
 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
 
-Session = Union[AgentSession, TeamSession, WorkflowSession]
+Session = Union[AgentSession, TeamSession, WorkflowSession, WorkflowSessionV2]
 
 __all__ = [
     "AgentSession",
     "TeamSession",
     "WorkflowSession",
+    "WorkflowSessionV2",
     "Session",
 ]

agno/storage/session/agent.py

@@ -16,6 +16,8 @@ class AgentSession:
     user_id: Optional[str] = None
     # ID of the team session this agent session is associated with
     team_session_id: Optional[str] = None
+    # ID of the workflow session this agent session is associated with
+    workflow_session_id: Optional[str] = None
     # Agent Memory
     memory: Optional[Dict[str, Any]] = None
     # Session Data: session_name, session_state, images, videos, audio
@@ -51,6 +53,7 @@ class AgentSession:
             session_id=data.get("session_id"),  # type: ignore
             agent_id=data.get("agent_id"),
             team_session_id=data.get("team_session_id"),
+            workflow_session_id=data.get("workflow_session_id"),
             user_id=data.get("user_id"),
             memory=data.get("memory"),
             agent_data=data.get("agent_data"),

agno/storage/session/team.py

@@ -14,6 +14,8 @@ class TeamSession:
     session_id: str
     # ID of the team session this team session is associated with (so for sub-teams)
     team_session_id: Optional[str] = None
+    # ID of the workflow session this team session is associated with
+    workflow_session_id: Optional[str] = None
     # ID of the team that this session is associated with
     team_id: Optional[str] = None
     # ID of the user interacting with this team
@@ -50,6 +52,7 @@ class TeamSession:
             session_id=data.get("session_id"),  # type: ignore
             team_id=data.get("team_id"),
             team_session_id=data.get("team_session_id"),
+            workflow_session_id=data.get("workflow_session_id"),
             user_id=data.get("user_id"),
             memory=data.get("memory"),
             team_data=data.get("team_data"),

agno/storage/session/v2/__init__.py

@@ -0,0 +1,5 @@
+from agno.storage.session.workflow import WorkflowSession
+
+__all__ = [
+    "WorkflowSession",
+]

agno/storage/session/v2/workflow.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Mapping, Optional
+
+from agno.run.v2.workflow import WorkflowRunResponse
+from agno.utils.log import logger
+
+
+@dataclass
+class WorkflowSession:
+    """Workflow Session V2 for pipeline-based workflows"""
+
+    # Session UUID - this is the workflow_session_id that gets set on agents/teams
+    session_id: str
+    # ID of the user interacting with this workflow
+    user_id: Optional[str] = None
+
+    # ID of the workflow that this session is associated with
+    workflow_id: Optional[str] = None
+    # Workflow name
+    workflow_name: Optional[str] = None
+
+    # Workflow runs - stores WorkflowRunResponse objects in memory
+    runs: Optional[List[WorkflowRunResponse]] = None
+
+    # Session Data: session_name, session_state, images, videos, audio
+    session_data: Optional[Dict[str, Any]] = None
+    # Workflow configuration and metadata
+    workflow_data: Optional[Dict[str, Any]] = None
+    # Extra Data stored with this workflow session
+    extra_data: Optional[Dict[str, Any]] = None
+
+    # The unix timestamp when this session was created
+    created_at: Optional[int] = None
+    # The unix timestamp when this session was last updated
+    updated_at: Optional[int] = None
+
+    def __post_init__(self):
+        if self.runs is None:
+            self.runs = []
+
+    def add_run(self, run_response: WorkflowRunResponse) -> None:
+        """Add a workflow run response to this session"""
+        if self.runs is None:
+            self.runs = []
+        # Store the actual WorkflowRunResponse object
+        self.runs.append(run_response)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for storage, serializing runs to dicts"""
+        return {
+            "session_id": self.session_id,
+            "user_id": self.user_id,
+            "workflow_id": self.workflow_id,
+            "workflow_name": self.workflow_name,
+            "runs": [run.to_dict() for run in self.runs] if self.runs else None,
+            "session_data": self.session_data,
+            "workflow_data": self.workflow_data,
+            "extra_data": self.extra_data,
+            "created_at": self.created_at,
+            "updated_at": self.updated_at,
+        }
+
+    @classmethod
+    def from_dict(cls, data: Mapping[str, Any]) -> Optional[WorkflowSession]:
+        """Create WorkflowSession from dictionary, deserializing runs from dicts"""
+        if data is None or data.get("session_id") is None:
+            logger.warning("WorkflowSession is missing session_id")
+            return None
+
+        # Deserialize runs from dictionaries back to WorkflowRunResponse objects
+        runs_data = data.get("runs")
+        runs: Optional[List[WorkflowRunResponse]] = None
+        if runs_data is not None:
+            runs = [WorkflowRunResponse.from_dict(run_dict) for run_dict in runs_data]
+
+        return cls(
+            session_id=data.get("session_id"),  # type: ignore
+            user_id=data.get("user_id"),
+            workflow_id=data.get("workflow_id"),
+            workflow_name=data.get("workflow_name"),
+            runs=runs,
+            session_data=data.get("session_data"),
+            workflow_data=data.get("workflow_data"),
+            extra_data=data.get("extra_data"),
+            created_at=data.get("created_at"),
+            updated_at=data.get("updated_at"),
+        )

agno/storage/singlestore.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
 
@@ -30,7 +31,7 @@ class SingleStoreStorage(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 singlestore table.
@@ -46,7 +47,7 @@ class SingleStoreStorage(Storage):
             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.
-            mode (Optional[Literal["agent", "team", "workflow"]], optional): The mode of the storage. Defaults to "agent".
+            mode (Optional[Literal["agent", "team", "workflow", "workflow_v2"]], optional): The mode of the storage. Defaults to "agent".
         """
         super().__init__(mode)
         _engine: Optional[Engine] = db_engine
@@ -75,12 +76,12 @@ class SingleStoreStorage(Storage):
         self.table: Table = self.get_table()
 
     @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(SingleStoreStorage, type(self)).mode.fset(self, value)  # type: ignore
         if value is not None:
@@ -115,6 +116,11 @@ class SingleStoreStorage(Storage):
                 Column("workflow_id", mysql.TEXT),
                 Column("workflow_data", mysql.JSON),
             ]
+        elif self.mode == "workflow_v2":
+            specific_columns = [
+                Column("workflow_id", mysql.TEXT),
+                Column("workflow_data", mysql.JSON),
+            ]
 
         # Create table with all columns
         table = Table(
@@ -171,6 +177,8 @@ class SingleStoreStorage(Storage):
                     return TeamSession.from_dict(existing_row._mapping)  # type: ignore
                 elif self.mode == "workflow":
                     return WorkflowSession.from_dict(existing_row._mapping)  # type: ignore
+                elif self.mode == "workflow_v2":
+                    return WorkflowSessionV2.from_dict(existing_row._mapping)  # type: ignore
             return None
 
     def get_all_session_ids(self, user_id: Optional[str] = None, entity_id: Optional[str] = None) -> List[str]:
@@ -188,6 +196,8 @@ class SingleStoreStorage(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
@@ -214,6 +224,8 @@ class SingleStoreStorage(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
@@ -232,6 +244,10 @@ class SingleStoreStorage(Storage):
                             _workflow_session = WorkflowSession.from_dict(row._mapping)  # type: ignore
                             if _workflow_session is not None:
                                 sessions.append(_workflow_session)
+                        elif self.mode == "workflow_v2":
+                            _workflow_session = WorkflowSessionV2.from_dict(row._mapping)  # type: ignore
+                            if _workflow_session is not None:
+                                sessions.append(_workflow_session)
         except Exception:
             log_debug(f"Table does not exist: {self.table.name}")
         return sessions
@@ -268,7 +284,8 @@ class SingleStoreStorage(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:
@@ -290,7 +307,10 @@ class SingleStoreStorage(Storage):
                             session = WorkflowSession.from_dict(row._mapping)  # type: ignore
                             if session is not None:
                                 sessions.append(session)
-
+                        elif self.mode == "workflow_v2":
+                            session = WorkflowSessionV2.from_dict(row._mapping)  # type: ignore
+                            if session is not None:
+                                sessions.append(session)
         except Exception as e:
             if "doesn't exist" in str(e):
                 log_debug(f"Table does not exist: {self.table.name}")
@@ -401,6 +421,25 @@ class SingleStoreStorage(Storage):
                         updated_at = UNIX_TIMESTAMP();
                     """
                 )
+            elif self.mode == "workflow_v2":
+                # Convert session to dict to ensure proper serialization
+                upsert_sql = text(
+                    f"""
+                    INSERT INTO {self.schema}.{self.table_name}
+                    (session_id, workflow_id, user_id, workflow_name, runs, workflow_data, session_data, extra_data, created_at, updated_at)
+                    VALUES      
+                    (:session_id, :workflow_id, :user_id, :workflow_name, :runs, :workflow_data, :session_data, :extra_data, UNIX_TIMESTAMP(), NULL)
+                    ON DUPLICATE KEY UPDATE
+                        workflow_id = VALUES(workflow_id),
+                        user_id = VALUES(user_id),
+                        workflow_name = VALUES(workflow_name),
+                        runs = VALUES(runs),
+                        workflow_data = VALUES(workflow_data),  
+                        session_data = VALUES(session_data),
+                        extra_data = VALUES(extra_data),
+                        updated_at = UNIX_TIMESTAMP();
+                    """
+                )
 
             try:
                 if self.mode == "agent":
@@ -411,8 +450,8 @@ class SingleStoreStorage(Storage):
                             "agent_id": session.agent_id,  # type: ignore
                             "team_session_id": session.team_session_id,  # type: ignore
                             "user_id": session.user_id,
-                            "memory": json.dumps(session.memory, ensure_ascii=False)
-                            if session.memory is not None
+                            "memory": json.dumps(getattr(session, "memory", None), ensure_ascii=False)
+                            if getattr(session, "memory", None) is not None
                             else None,
                             "agent_data": json.dumps(session.agent_data, ensure_ascii=False)  # type: ignore
                             if session.agent_data is not None  # type: ignore
@@ -433,8 +472,8 @@ class SingleStoreStorage(Storage):
                             "team_id": session.team_id,  # type: ignore
                             "user_id": session.user_id,
                             "team_session_id": session.team_session_id,  # type: ignore
-                            "memory": json.dumps(session.memory, ensure_ascii=False)
-                            if session.memory is not None
+                            "memory": json.dumps(getattr(session, "memory", None), ensure_ascii=False)
+                            if getattr(session, "memory", None) is not None
                             else None,
                             "team_data": json.dumps(session.team_data, ensure_ascii=False)  # type: ignore
                             if session.team_data is not None  # type: ignore
@@ -454,9 +493,32 @@ class SingleStoreStorage(Storage):
                             "session_id": session.session_id,
                             "workflow_id": session.workflow_id,  # type: ignore
                             "user_id": session.user_id,
-                            "memory": json.dumps(session.memory, ensure_ascii=False)
-                            if session.memory is not None
+                            "memory": json.dumps(getattr(session, "memory", None), ensure_ascii=False)
+                            if getattr(session, "memory", None) is not None
+                            else None,
+                            "workflow_data": json.dumps(session.workflow_data, ensure_ascii=False)  # type: ignore
+                            if session.workflow_data is not None  # type: ignore
                             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,
+                        },
+                    )
+                elif self.mode == "workflow_v2":
+                    # Convert session to dict to ensure proper serialization
+                    session_dict = session.to_dict()
+
+                    sess.execute(
+                        upsert_sql,
+                        {
+                            "session_id": session.session_id,
+                            "workflow_id": session.workflow_id,  # type: ignore
+                            "user_id": session.user_id,
+                            "workflow_name": session.workflow_name,  # type: ignore
+                            "runs": session_dict.get("runs"),
                             "workflow_data": json.dumps(session.workflow_data, ensure_ascii=False)  # type: ignore
                             if session.workflow_data is not None  # type: ignore
                             else None,