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

agno/storage/sqlite.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, log_info, log_warning, logger
 
@@ -32,7 +33,7 @@ class SqliteStorage(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 sqlite database.
@@ -84,12 +85,12 @@ class SqliteStorage(Storage):
         self.table: Table = self.get_table()
 
     @property
-    def mode(self) -> Optional[Literal["agent", "team", "workflow"]]:
+    def mode(self) -> Optional[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(SqliteStorage, type(self)).mode.fset(self, value)  # type: ignore
         if value is not None:
@@ -131,6 +132,14 @@ class SqliteStorage(Storage):
                 Column("workflow_id", String, index=True),
                 Column("workflow_data", sqlite.JSON),
             ]
+        elif self.mode == "workflow_v2":
+            specific_columns = [
+                Column("workflow_id", String, index=True),
+                Column("workflow_name", String, index=True),
+                Column("workflow_data", sqlite.JSON),
+                # Store WorkflowRunResponse objects
+                Column("runs", sqlite.JSON),
+            ]
 
         # Create table with all columns
         table = Table(
@@ -241,6 +250,8 @@ class SqliteStorage(Storage):
                     return TeamSession.from_dict(result._mapping) if result is not None else None  # type: ignore
                 elif self.mode == "workflow":
                     return WorkflowSession.from_dict(result._mapping) if result is not None else None  # type: ignore
+                elif self.mode == "workflow_v2":
+                    return WorkflowSessionV2.from_dict(result._mapping) if result is not None else None  # type: ignore
         except Exception as e:
             if "no such table" in str(e):
                 log_debug(f"Table does not exist: {self.table.name}")
@@ -273,6 +284,8 @@ class SqliteStorage(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
@@ -308,11 +321,11 @@ class SqliteStorage(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)
-                    elif self.mode == "workflow":
+                    elif self.mode in ["workflow", "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
+
                 rows = sess.execute(stmt).fetchall()
                 if rows is not None:
                     if self.mode == "agent":
@@ -321,6 +334,8 @@ class SqliteStorage(Storage):
                         return [TeamSession.from_dict(row._mapping) for row in rows]  # type: ignore
                     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
                 else:
                     return []
         except Exception as e:
@@ -341,7 +356,7 @@ class SqliteStorage(Storage):
         Get the last N sessions, ordered by created_at descending.
 
         Args:
-            num_history_sessions: Number of most recent sessions to return
+            limit: Number of most recent sessions to return
             user_id: Filter by user ID
             entity_id: Filter by entity ID (agent_id, team_id, or workflow_id)
 
@@ -359,7 +374,7 @@ class SqliteStorage(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)
-                    elif self.mode == "workflow":
+                    elif self.mode in ["workflow", "workflow_v2"]:
                         stmt = stmt.where(self.table.c.workflow_id == entity_id)
 
                 # Order by created_at desc and limit to num_history_sessions
@@ -370,12 +385,14 @@ class SqliteStorage(Storage):
                 # Execute query
                 rows = sess.execute(stmt).fetchall()
                 if rows is not None:
-                    if self.mode == "agent":
+                    if self.mode == "agent":  # type: ignore
                         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
                     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
                 return []
         except Exception as e:
             if "no such table" in str(e):
@@ -437,26 +454,25 @@ class SqliteStorage(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,
                     )
 
                     # Define the upsert if the session_id already exists
-                    # See: https://docs.sqlalchemy.org/en/20/dialects/sqlite.html#insert-on-conflict-upsert
                     stmt = stmt.on_conflict_do_update(
                         index_elements=["session_id"],
                         set_=dict(
                             agent_id=session.agent_id,  # 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,
                             updated_at=int(time.time()),
-                        ),  # The updated value for each column
+                        ),
                     )
                 elif self.mode == "team":
                     # Create an insert statement
@@ -465,7 +481,7 @@ class SqliteStorage(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,
@@ -479,12 +495,12 @@ class SqliteStorage(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
+                        ),
                     )
                 elif self.mode == "workflow":
                     # Create an insert statement
@@ -492,7 +508,7 @@ class SqliteStorage(Storage):
                         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,
@@ -505,12 +521,42 @@ class SqliteStorage(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()),
+                        ),
+                    )
+                elif self.mode == "workflow_v2":
+                    # Convert session to dict to ensure proper serialization
+                    session_dict = session.to_dict()
+
+                    # Create an insert statement for WorkflowSessionV2
+                    stmt = sqlite.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
+                    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()),
-                        ),  # The updated value for each column
+                        ),
                     )
 
                 sess.execute(stmt)

agno/storage/yaml.py

@@ -9,12 +9,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 YamlStorage(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)
@@ -43,6 +46,8 @@ class YamlStorage(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
 
@@ -62,6 +67,8 @@ class YamlStorage(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:
@@ -71,6 +78,8 @@ class YamlStorage(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"])
@@ -93,6 +102,8 @@ class YamlStorage(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)
@@ -100,6 +111,8 @@ class YamlStorage(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)
                     elif entity_id:
                         if self.mode == "agent" and data["agent_id"] == entity_id:
                             _session = AgentSession.from_dict(data)
@@ -107,7 +120,8 @@ class YamlStorage(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 +132,8 @@ class YamlStorage(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 +174,8 @@ class YamlStorage(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 +198,8 @@ class YamlStorage(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,11 +208,13 @@ class YamlStorage(Storage):
     def upsert(self, session: Session) -> Optional[Session]:
         """Insert or update an 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"]
-
             with open(self.dir_path / f"{session.session_id}.yaml", "w", encoding="utf-8") as f:
                 f.write(self.serialize(data))
             return session