Spaces:
Sleeping
Sleeping
| import time | |
| from pathlib import Path | |
| from typing import Optional, List | |
| try: | |
| from sqlalchemy.dialects import sqlite | |
| from sqlalchemy.engine import create_engine, Engine | |
| from sqlalchemy.inspection import inspect | |
| from sqlalchemy.orm import Session, sessionmaker | |
| from sqlalchemy.schema import MetaData, Table, Column | |
| from sqlalchemy.sql.expression import select | |
| from sqlalchemy.types import String | |
| except ImportError: | |
| raise ImportError("`sqlalchemy` not installed. Please install it using `pip install sqlalchemy`") | |
| from phi.workflow import WorkflowSession | |
| from phi.storage.workflow.base import WorkflowStorage | |
| from phi.utils.log import logger | |
| class SqlWorkflowStorage(WorkflowStorage): | |
| def __init__( | |
| self, | |
| table_name: str, | |
| db_url: Optional[str] = None, | |
| db_file: Optional[str] = None, | |
| db_engine: Optional[Engine] = None, | |
| schema_version: int = 1, | |
| auto_upgrade_schema: bool = False, | |
| ): | |
| """ | |
| This class provides workflow storage using a sqlite database. | |
| The following order is used to determine the database connection: | |
| 1. Use the db_engine if provided | |
| 2. Use the db_url | |
| 3. Use the db_file | |
| 4. Create a new in-memory database | |
| Args: | |
| table_name: The name of the table to store Workflow sessions. | |
| db_url: The database URL to connect to. | |
| db_file: The database file to connect to. | |
| db_engine: The SQLAlchemy database engine to use. | |
| """ | |
| _engine: Optional[Engine] = db_engine | |
| if _engine is None and db_url is not None: | |
| _engine = create_engine(db_url) | |
| elif _engine is None and db_file is not None: | |
| # Use the db_file to create the engine | |
| db_path = Path(db_file).resolve() | |
| # Ensure the directory exists | |
| db_path.parent.mkdir(parents=True, exist_ok=True) | |
| _engine = create_engine(f"sqlite:///{db_path}") | |
| else: | |
| _engine = create_engine("sqlite://") | |
| if _engine is None: | |
| raise ValueError("Must provide either db_url, db_file or db_engine") | |
| # Database attributes | |
| self.table_name: str = table_name | |
| self.db_url: Optional[str] = db_url | |
| self.db_engine: Engine = _engine | |
| self.metadata: MetaData = MetaData() | |
| self.inspector = inspect(self.db_engine) | |
| # Table schema version | |
| self.schema_version: int = schema_version | |
| # Automatically upgrade schema if True | |
| self.auto_upgrade_schema: bool = auto_upgrade_schema | |
| # Database session | |
| self.Session: sessionmaker[Session] = sessionmaker(bind=self.db_engine) | |
| # Database table for storage | |
| self.table: Table = self.get_table() | |
| def get_table_v1(self) -> Table: | |
| """ | |
| Define the table schema for version 1. | |
| Returns: | |
| Table: SQLAlchemy Table object representing the schema. | |
| """ | |
| return Table( | |
| self.table_name, | |
| self.metadata, | |
| # Session UUID: Primary Key | |
| Column("session_id", String, primary_key=True), | |
| # ID of the workflow that this session is associated with | |
| Column("workflow_id", String), | |
| # ID of the user interacting with this workflow | |
| Column("user_id", String), | |
| # Workflow Memory | |
| Column("memory", sqlite.JSON), | |
| # Workflow Metadata | |
| Column("workflow_data", sqlite.JSON), | |
| # User Metadata | |
| Column("user_data", sqlite.JSON), | |
| # Session Metadata | |
| Column("session_data", sqlite.JSON), | |
| # The Unix timestamp of when this session was created. | |
| Column("created_at", sqlite.INTEGER, default=lambda: int(time.time())), | |
| # The Unix timestamp of when this session was last updated. | |
| Column("updated_at", sqlite.INTEGER, onupdate=lambda: int(time.time())), | |
| extend_existing=True, | |
| sqlite_autoincrement=True, | |
| ) | |
| def get_table(self) -> Table: | |
| """ | |
| Get the table schema based on the schema version. | |
| Returns: | |
| Table: SQLAlchemy Table object for the current schema version. | |
| Raises: | |
| ValueError: If an unsupported schema version is specified. | |
| """ | |
| if self.schema_version == 1: | |
| return self.get_table_v1() | |
| else: | |
| raise ValueError(f"Unsupported schema version: {self.schema_version}") | |
| def table_exists(self) -> bool: | |
| """ | |
| Check if the table exists in the database. | |
| Returns: | |
| bool: True if the table exists, False otherwise. | |
| """ | |
| logger.debug(f"Checking if table exists: {self.table.name}") | |
| try: | |
| return self.inspector.has_table(self.table.name) | |
| except Exception as e: | |
| logger.error(f"Error checking if table exists: {e}") | |
| return False | |
| def create(self) -> None: | |
| """ | |
| Create the table if it doesn't exist. | |
| """ | |
| if not self.table_exists(): | |
| logger.debug(f"Creating table: {self.table.name}") | |
| self.table.create(self.db_engine, checkfirst=True) | |
| def read(self, session_id: str, user_id: Optional[str] = None) -> Optional[WorkflowSession]: | |
| """ | |
| Read a WorkflowSession from the database. | |
| Args: | |
| session_id (str): The ID of the session to read. | |
| user_id (Optional[str]): The ID of the user associated with the session. | |
| Returns: | |
| Optional[WorkflowSession]: The WorkflowSession object if found, None otherwise. | |
| """ | |
| try: | |
| with self.Session() as sess: | |
| stmt = select(self.table).where(self.table.c.session_id == session_id) | |
| if user_id: | |
| stmt = stmt.where(self.table.c.user_id == user_id) | |
| result = sess.execute(stmt).fetchone() | |
| return WorkflowSession.model_validate(result) if result is not None else None | |
| except Exception as e: | |
| logger.debug(f"Exception reading from table: {e}") | |
| logger.debug(f"Table does not exist: {self.table.name}") | |
| logger.debug("Creating table for future transactions") | |
| self.create() | |
| return None | |
| def get_all_session_ids(self, user_id: Optional[str] = None, workflow_id: Optional[str] = None) -> List[str]: | |
| """ | |
| Get all session IDs, optionally filtered by user_id and/or workflow_id. | |
| Args: | |
| user_id (Optional[str]): The ID of the user to filter by. | |
| workflow_id (Optional[str]): The ID of the workflow to filter by. | |
| Returns: | |
| List[str]: List of session IDs matching the criteria. | |
| """ | |
| try: | |
| with self.Session() as sess, sess.begin(): | |
| # get all session_ids | |
| stmt = select(self.table.c.session_id) | |
| if user_id is not None and user_id != "": | |
| stmt = stmt.where(self.table.c.user_id == user_id) | |
| if workflow_id is not None: | |
| stmt = stmt.where(self.table.c.workflow_id == workflow_id) | |
| # order by created_at desc | |
| stmt = stmt.order_by(self.table.c.created_at.desc()) | |
| # execute query | |
| rows = sess.execute(stmt).fetchall() | |
| return [row[0] for row in rows] if rows is not None else [] | |
| except Exception as e: | |
| logger.debug(f"Exception reading from table: {e}") | |
| logger.debug(f"Table does not exist: {self.table.name}") | |
| logger.debug("Creating table for future transactions") | |
| self.create() | |
| return [] | |
| def get_all_sessions( | |
| self, user_id: Optional[str] = None, workflow_id: Optional[str] = None | |
| ) -> List[WorkflowSession]: | |
| """ | |
| Get all sessions, optionally filtered by user_id and/or workflow_id. | |
| Args: | |
| user_id (Optional[str]): The ID of the user to filter by. | |
| workflow_id (Optional[str]): The ID of the workflow to filter by. | |
| Returns: | |
| List[WorkflowSession]: List of AgentSession objects matching the criteria. | |
| """ | |
| try: | |
| with self.Session() as sess, sess.begin(): | |
| # get all sessions | |
| stmt = select(self.table) | |
| if user_id is not None and user_id != "": | |
| stmt = stmt.where(self.table.c.user_id == user_id) | |
| if workflow_id is not None: | |
| stmt = stmt.where(self.table.c.workflow_id == workflow_id) | |
| # order by created_at desc | |
| stmt = stmt.order_by(self.table.c.created_at.desc()) | |
| # execute query | |
| rows = sess.execute(stmt).fetchall() | |
| return [WorkflowSession.model_validate(row) for row in rows] if rows is not None else [] | |
| except Exception as e: | |
| logger.debug(f"Exception reading from table: {e}") | |
| logger.debug(f"Table does not exist: {self.table.name}") | |
| logger.debug("Creating table for future transactions") | |
| self.create() | |
| return [] | |
| def upsert(self, session: WorkflowSession, create_and_retry: bool = True) -> Optional[WorkflowSession]: | |
| """ | |
| Insert or update a WorkflowSession in the database. | |
| Args: | |
| session (WorkflowSession): The WorkflowSession object to upsert. | |
| create_and_retry (bool): Retry upsert if table does not exist. | |
| Returns: | |
| Optional[WorkflowSession]: The upserted WorkflowSession object. | |
| """ | |
| try: | |
| with self.Session() as sess, sess.begin(): | |
| # Create an insert statement | |
| stmt = sqlite.insert(self.table).values( | |
| session_id=session.session_id, | |
| workflow_id=session.workflow_id, | |
| user_id=session.user_id, | |
| memory=session.memory, | |
| workflow_data=session.workflow_data, | |
| user_data=session.user_data, | |
| session_data=session.session_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( | |
| workflow_id=session.workflow_id, | |
| user_id=session.user_id, | |
| memory=session.memory, | |
| workflow_data=session.workflow_data, | |
| user_data=session.user_data, | |
| session_data=session.session_data, | |
| updated_at=int(time.time()), | |
| ), # The updated value for each column | |
| ) | |
| sess.execute(stmt) | |
| except Exception as e: | |
| logger.debug(f"Exception upserting into table: {e}") | |
| if create_and_retry and not self.table_exists(): | |
| logger.debug(f"Table does not exist: {self.table.name}") | |
| logger.debug("Creating table and retrying upsert") | |
| self.create() | |
| return self.upsert(session, create_and_retry=False) | |
| return None | |
| return self.read(session_id=session.session_id) | |
| def delete_session(self, session_id: Optional[str] = None): | |
| """ | |
| Delete a workflow session from the database. | |
| Args: | |
| session_id (Optional[str]): The ID of the session to delete. | |
| Raises: | |
| ValueError: If session_id is not provided. | |
| """ | |
| if session_id is None: | |
| logger.warning("No session_id provided for deletion.") | |
| return | |
| try: | |
| with self.Session() as sess, sess.begin(): | |
| # Delete the session with the given session_id | |
| delete_stmt = self.table.delete().where(self.table.c.session_id == session_id) | |
| result = sess.execute(delete_stmt) | |
| if result.rowcount == 0: | |
| logger.debug(f"No session found with session_id: {session_id}") | |
| else: | |
| logger.debug(f"Successfully deleted session with session_id: {session_id}") | |
| except Exception as e: | |
| logger.error(f"Error deleting session: {e}") | |
| def drop(self) -> None: | |
| """ | |
| Drop the table from the database if it exists. | |
| """ | |
| if self.table_exists(): | |
| logger.debug(f"Deleting table: {self.table_name}") | |
| self.table.drop(self.db_engine) | |
| def upgrade_schema(self) -> None: | |
| """ | |
| Upgrade the schema of the workflow storage table. | |
| This method is currently a placeholder and does not perform any actions. | |
| """ | |
| pass | |
| def __deepcopy__(self, memo): | |
| """ | |
| Create a deep copy of the SqlWorkflowStorage instance, handling unpickleable attributes. | |
| Args: | |
| memo (dict): A dictionary of objects already copied during the current copying pass. | |
| Returns: | |
| SqlWorkflowStorage: A deep-copied instance of SqlWorkflowStorage. | |
| """ | |
| from copy import deepcopy | |
| # Create a new instance without calling __init__ | |
| cls = self.__class__ | |
| copied_obj = cls.__new__(cls) | |
| memo[id(self)] = copied_obj | |
| # Deep copy attributes | |
| for k, v in self.__dict__.items(): | |
| if k in {"metadata", "table", "inspector"}: | |
| continue | |
| # Reuse db_engine and Session without copying | |
| elif k in {"db_engine", "Session"}: | |
| setattr(copied_obj, k, v) | |
| else: | |
| setattr(copied_obj, k, deepcopy(v, memo)) | |
| # Recreate metadata and table for the copied instance | |
| copied_obj.metadata = MetaData() | |
| copied_obj.inspector = inspect(copied_obj.db_engine) | |
| copied_obj.table = copied_obj.get_table() | |
| return copied_obj | |