PraisonAI 3.0.0__py3-none-any.whl

praisonai/ui/sql_alchemy.py

@@ -0,0 +1,710 @@
+import json
+import ssl
+import uuid
+from dataclasses import asdict
+from datetime import datetime
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+import os
+
+import aiofiles
+import aiohttp
+
+from chainlit.data.base import BaseDataLayer
+from chainlit.data.utils import queue_until_user_message
+from praisonai.ui.chainlit_compat import EXPIRY_TIME, BaseStorageClient
+from chainlit.element import ElementDict
+from chainlit.logger import logger
+from chainlit.step import StepDict
+from chainlit.types import (
+    Feedback,
+    FeedbackDict,
+    PageInfo,
+    PaginatedResponse,
+    Pagination,
+    ThreadDict,
+    ThreadFilter,
+)
+from chainlit.user import PersistedUser, User
+from sqlalchemy import text
+from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
+from sqlalchemy.orm import sessionmaker
+from database_config import get_database_config_for_sqlalchemy
+
+if TYPE_CHECKING:
+    from chainlit.element import Element
+    from chainlit.step import StepDict
+
+# Check FORCE_SQLITE flag to bypass external database detection
+DATABASE_URL, SUPABASE_DATABASE_URL = get_database_config_for_sqlalchemy()
+
+class SQLAlchemyDataLayer(BaseDataLayer):
+    def __init__(
+        self,
+        conninfo: str,
+        ssl_require: bool = False,
+        storage_provider: Optional[BaseStorageClient] = None,
+        user_thread_limit: Optional[int] = 1000,
+        show_logger: Optional[bool] = False,
+    ):
+        self._conninfo = conninfo
+        self.user_thread_limit = user_thread_limit
+        self.show_logger = show_logger
+        ssl_args = {}
+        if ssl_require:
+            ssl_context = ssl.create_default_context()
+            ssl_context.check_hostname = False
+            ssl_context.verify_mode = ssl.CERT_NONE
+            ssl_args["ssl"] = ssl_context
+        self.engine: AsyncEngine = create_async_engine(
+            self._conninfo, connect_args=ssl_args
+        )
+        self.async_session = sessionmaker(
+            bind=self.engine, expire_on_commit=False, class_=AsyncSession
+        )
+        if storage_provider:
+            self.storage_provider: Optional[BaseStorageClient] = storage_provider
+            if self.show_logger:
+                logger.info("SQLAlchemyDataLayer storage client initialized")
+        else:
+            self.storage_provider = None
+            logger.warning(
+                "SQLAlchemyDataLayer storage client is not initialized and elements will not be persisted!"
+            )
+
+    async def build_debug_url(self) -> str:
+        return ""
+
+    ###### SQL Helpers ######
+    async def execute_sql(
+        self, query: str, parameters: dict
+    ) -> Union[List[Dict[str, Any]], int, None]:
+        parameterized_query = text(query)
+        async with self.async_session() as session:
+            try:
+                await session.begin()
+                result = await session.execute(parameterized_query, parameters)
+                await session.commit()
+                if result.returns_rows:
+                    json_result = [dict(row._mapping) for row in result.fetchall()]
+                    clean_json_result = self.clean_result(json_result)
+                    assert isinstance(clean_json_result, list) or isinstance(
+                        clean_json_result, int
+                    )
+                    return clean_json_result
+                else:
+                    return result.rowcount
+            except SQLAlchemyError as e:
+                await session.rollback()
+                logger.warning(f"An error occurred: {e}")
+                return None
+            except Exception as e:
+                await session.rollback()
+                logger.warning(f"An unexpected error occurred: {e}")
+                return None
+
+    async def get_current_timestamp(self) -> str:
+        return datetime.now().isoformat() + "Z"
+
+    def clean_result(self, obj):
+        if isinstance(obj, dict):
+            return {k: self.clean_result(v) for k, v in obj.items()}
+        elif isinstance(obj, list):
+            return [self.clean_result(item) for item in obj]
+        elif isinstance(obj, uuid.UUID):
+            return str(obj)
+        return obj
+
+    ###### User ######
+    async def get_user(self, identifier: str) -> Optional[PersistedUser]:
+        if self.show_logger:
+            logger.info(f"SQLAlchemy: get_user, identifier={identifier}")
+        query = 'SELECT * FROM users WHERE "identifier" = :identifier'
+        parameters = {"identifier": identifier}
+        result = await self.execute_sql(query=query, parameters=parameters)
+        if result and isinstance(result, list):
+            user_data = result[0]
+
+            meta = user_data.get("meta", "{}")
+            if isinstance(meta, str):
+                meta = json.loads(meta)
+
+            return PersistedUser(
+                id=user_data["id"],
+                identifier=user_data["identifier"],
+                createdAt=user_data["createdAt"],
+                metadata=meta,
+            )
+        return None
+
+    async def _get_user_identifer_by_id(self, user_id: str) -> str:
+        if self.show_logger:
+            logger.info(f"SQLAlchemy: _get_user_identifer_by_id, user_id={user_id}")
+        query = 'SELECT "identifier" FROM users WHERE "id" = :user_id'
+        parameters = {"user_id": user_id}
+        result = await self.execute_sql(query=query, parameters=parameters)
+        assert result
+        assert isinstance(result, list)
+        return result[0]["identifier"]
+
+    async def _get_user_id_by_thread(self, thread_id: str) -> Optional[str]:
+        if self.show_logger:
+            logger.info(f"SQLAlchemy: _get_user_id_by_thread, thread_id={thread_id}")
+        query = 'SELECT "userId" FROM threads WHERE "id" = :thread_id'
+        parameters = {"thread_id": thread_id}
+        result = await self.execute_sql(query=query, parameters=parameters)
+        if result and isinstance(result, list):
+            return result[0]["userId"]
+        return None
+
+    async def create_user(self, user: User) -> Optional[PersistedUser]:
+        if self.show_logger:
+            logger.info(f"SQLAlchemy: create_user, user_identifier={user.identifier}")
+        existing_user: Optional["PersistedUser"] = await self.get_user(user.identifier)
+        user_dict: Dict[str, Any] = {
+            "identifier": str(user.identifier),
+            "meta": json.dumps(user.metadata) or "{}",
+        }
+        if not existing_user:
+            user_dict["id"] = str(uuid.uuid4())
+            user_dict["createdAt"] = await self.get_current_timestamp()
+            query = 'INSERT INTO users ("id", "identifier", "createdAt", "meta") VALUES (:id, :identifier, :createdAt, :meta)'
+            await self.execute_sql(query=query, parameters=user_dict)
+        else:
+            query = 'UPDATE users SET "meta" = :meta WHERE "identifier" = :identifier'
+            await self.execute_sql(query=query, parameters=user_dict)
+        return await self.get_user(user.identifier)
+
+    ###### Threads ######
+    async def get_thread_author(self, thread_id: str) -> str:
+        if self.show_logger:
+            logger.info(f"SQLAlchemy: get_thread_author, thread_id={thread_id}")
+        query = 'SELECT "userIdentifier" FROM threads WHERE "id" = :id'
+        parameters = {"id": thread_id}
+        result = await self.execute_sql(query=query, parameters=parameters)
+        if isinstance(result, list) and result:
+            author_identifier = result[0].get("userIdentifier")
+            if author_identifier is not None:
+                return author_identifier
+        raise ValueError(f"Author not found for thread_id {thread_id}")
+
+    async def get_thread(self, thread_id: str) -> Optional[ThreadDict]:
+        if self.show_logger:
+            logger.info(f"SQLAlchemy: get_thread, thread_id={thread_id}")
+        user_threads: Optional[List[ThreadDict]] = await self.get_all_user_threads(
+            thread_id=thread_id
+        )
+        if user_threads:
+            return user_threads[0]
+        else:
+            return None
+
+    async def update_thread(
+        self,
+        thread_id: str,
+        name: Optional[str] = None,
+        user_id: Optional[str] = None,
+        metadata: Optional[Dict] = None,
+        tags: Optional[List[str]] = None,
+    ):
+        if self.show_logger:
+            logger.info(f"SQLAlchemy: update_thread, thread_id={thread_id}")
+
+        user_identifier = None
+        if user_id:
+            user_identifier = await self._get_user_identifer_by_id(user_id)
+
+        data = {
+            "id": thread_id,
+            "createdAt": (
+                await self.get_current_timestamp() if metadata is None else None
+            ),
+            "name": (
+                name
+                if name is not None
+                else (metadata.get("name") if metadata and "name" in metadata else None)
+            ),
+            "userId": user_id,
+            "userIdentifier": user_identifier,
+            "tags": json.dumps(tags) if tags else None,
+            "meta": json.dumps(metadata) if metadata else None,
+        }
+        parameters = {key: value for key, value in data.items() if value is not None}
+        columns = ", ".join(f'"{key}"' for key in parameters.keys())
+        values = ", ".join(f':{key}' for key in parameters.keys())
+        updates = ", ".join(
+            f'"{key}" = EXCLUDED."{key}"' for key in parameters.keys() if key != "id"
+        )
+        query = f"""
+            INSERT INTO threads ({columns})
+            VALUES ({values})
+            ON CONFLICT ("id") DO UPDATE
+            SET {updates};
+        """
+        await self.execute_sql(query=query, parameters=parameters)
+
+    async def delete_thread(self, thread_id: str):
+        if self.show_logger:
+            logger.info(f"SQLAlchemy: delete_thread, thread_id={thread_id}")
+        feedbacks_query = 'DELETE FROM feedbacks WHERE "forId" IN (SELECT "id" FROM steps WHERE "threadId" = :id)'
+        elements_query = 'DELETE FROM elements WHERE "threadId" = :id'
+        steps_query = 'DELETE FROM steps WHERE "threadId" = :id'
+        thread_query = 'DELETE FROM threads WHERE "id" = :id'
+        parameters = {"id": thread_id}
+        await self.execute_sql(query=feedbacks_query, parameters=parameters)
+        await self.execute_sql(query=elements_query, parameters=parameters)
+        await self.execute_sql(query=steps_query, parameters=parameters)
+        await self.execute_sql(query=thread_query, parameters=parameters)
+
+    async def list_threads(
+        self, pagination: Pagination, filters: ThreadFilter
+    ) -> PaginatedResponse:
+        if self.show_logger:
+            logger.info(
+                f"SQLAlchemy: list_threads, pagination={pagination}, filters={filters}"
+            )
+        if not filters.userId:
+            raise ValueError("userId is required")
+        all_user_threads: Optional[List[ThreadDict]] = (
+            await self.get_all_user_threads(user_id=filters.userId) or []
+        )
+
+        search_keyword = filters.search.lower() if filters.search else None
+        feedback_value = int(filters.feedback) if filters.feedback else None
+
+        filtered_threads = []
+        for thread in all_user_threads:
+            keyword_match = True
+            feedback_match = True
+            if search_keyword or feedback_value is not None:
+                if search_keyword:
+                    keyword_match = any(
+                        search_keyword in step["output"].lower()
+                        for step in thread["steps"]
+                        if "output" in step and isinstance(step["output"], str)
+                    )
+                if feedback_value is not None:
+                    feedback_match = False
+                    for step in thread["steps"]:
+                        feedback = step.get("feedback")
+                        if feedback and feedback.get("value") == feedback_value:
+                            feedback_match = True
+                            break
+            if keyword_match and feedback_match:
+                filtered_threads.append(thread)
+
+        start = 0
+        if pagination.cursor:
+            for i, thr in enumerate(filtered_threads):
+                if thr["id"] == pagination.cursor:
+                    start = i + 1
+                    break
+        end = start + pagination.first
+        paginated_threads = filtered_threads[start:end] or []
+
+        has_next_page = len(filtered_threads) > end
+        start_cursor = paginated_threads[0]["id"] if paginated_threads else None
+        end_cursor = paginated_threads[-1]["id"] if paginated_threads else None
+
+        return PaginatedResponse(
+            pageInfo=PageInfo(
+                hasNextPage=has_next_page,
+                startCursor=start_cursor,
+                endCursor=end_cursor,
+            ),
+            data=paginated_threads,
+        )
+
+    ###### Steps ######
+    @queue_until_user_message()
+    async def create_step(self, step_dict: "StepDict"):
+        if self.show_logger:
+            logger.info(f"SQLAlchemy: create_step, step_id={step_dict.get('id')}")
+
+        step_dict["showInput"] = (
+            str(step_dict.get("showInput", "")).lower()
+            if "showInput" in step_dict
+            else None
+        )
+
+        tags = step_dict.get("tags")
+        if not tags:
+            tags = []
+        meta = json.dumps(step_dict.get("metadata", {}))
+        generation = json.dumps(step_dict.get("generation", {}))
+        parameters = {
+            "id": step_dict.get("id"),
+            "name": step_dict.get("name"),
+            "type": step_dict.get("type"),
+            "threadId": step_dict.get("threadId"),
+            "parentId": step_dict.get("parentId"),
+            "disableFeedback": step_dict.get("disableFeedback", False),
+            "streaming": step_dict.get("streaming", False),
+            "waitForAnswer": step_dict.get("waitForAnswer", False),
+            "isError": step_dict.get("isError", False),
+            "meta": meta,
+            "tags": json.dumps(tags),
+            "input": step_dict.get("input"),
+            "output": step_dict.get("output"),
+            "createdAt": step_dict.get("createdAt"),
+            "startTime": step_dict.get("start"),
+            "endTime": step_dict.get("end"),
+            "generation": generation,
+            "showInput": step_dict.get("showInput"),
+            "language": step_dict.get("language"),
+            "indent": step_dict.get("indent"),
+        }
+        parameters = {k: v for k, v in parameters.items() if v is not None}
+        columns = ", ".join(f'"{key}"' for key in parameters.keys())
+        values = ", ".join(f':{key}' for key in parameters.keys())
+        updates = ", ".join(
+            f'"{key}" = :{key}' for key in parameters.keys() if key != "id"
+        )
+        query = f"""
+            INSERT INTO steps ({columns})
+            VALUES ({values})
+            ON CONFLICT ("id") DO UPDATE
+            SET {updates};
+        """
+        await self.execute_sql(query=query, parameters=parameters)
+
+    @queue_until_user_message()
+    async def update_step(self, step_dict: "StepDict"):
+        if self.show_logger:
+            logger.info(f"SQLAlchemy: update_step, step_id={step_dict.get('id')}")
+        await self.create_step(step_dict)
+
+    @queue_until_user_message()
+    async def delete_step(self, step_id: str):
+        if self.show_logger:
+            logger.info(f"SQLAlchemy: delete_step, step_id={step_id}")
+        feedbacks_query = 'DELETE FROM feedbacks WHERE "forId" = :id'
+        elements_query = 'DELETE FROM elements WHERE "forId" = :id'
+        steps_query = 'DELETE FROM steps WHERE "id" = :id'
+        parameters = {"id": step_id}
+        await self.execute_sql(query=feedbacks_query, parameters=parameters)
+        await self.execute_sql(query=elements_query, parameters=parameters)
+        await self.execute_sql(query=steps_query, parameters=parameters)
+
+    ###### Feedback ######
+    async def upsert_feedback(self, feedback: Feedback) -> str:
+        if self.show_logger:
+            logger.info(f"SQLAlchemy: upsert_feedback, feedback_id={feedback.id}")
+        feedback.id = feedback.id or str(uuid.uuid4())
+        feedback_dict = asdict(feedback)
+        parameters = {k: v for k, v in feedback_dict.items() if v is not None}
+        columns = ", ".join(f'"{key}"' for key in parameters.keys())
+        values = ", ".join(f':{key}' for key in parameters.keys())
+        updates = ", ".join(
+            f'"{key}" = :{key}' for key in parameters.keys() if key != "id"
+        )
+        query = f"""
+            INSERT INTO feedbacks ({columns})
+            VALUES ({values})
+            ON CONFLICT ("id") DO UPDATE
+            SET {updates};
+        """
+        await self.execute_sql(query=query, parameters=parameters)
+        return feedback.id
+
+    async def delete_feedback(self, feedback_id: str) -> bool:
+        if self.show_logger:
+            logger.info(f"SQLAlchemy: delete_feedback, feedback_id={feedback_id}")
+        query = 'DELETE FROM feedbacks WHERE "id" = :feedback_id'
+        parameters = {"feedback_id": feedback_id}
+        await self.execute_sql(query=query, parameters=parameters)
+        return True
+
+    ###### Elements ######
+    async def get_element(
+        self, thread_id: str, element_id: str
+    ) -> Optional["ElementDict"]:
+        if self.show_logger:
+            logger.info(
+                f"SQLAlchemy: get_element, thread_id={thread_id}, element_id={element_id}"
+            )
+        query = 'SELECT * FROM elements WHERE "threadId" = :thread_id AND "id" = :element_id'
+        parameters = {"thread_id": thread_id, "element_id": element_id}
+        element = await self.execute_sql(query=query, parameters=parameters)
+        if isinstance(element, list) and element:
+            element_dict = element[0]
+            return ElementDict(
+                id=element_dict["id"],
+                threadId=element_dict.get("threadId"),
+                type=element_dict.get("type"),
+                chainlitKey=element_dict.get("chainlitKey"),
+                url=element_dict.get("url"),
+                objectKey=element_dict.get("objectKey"),
+                name=element_dict["name"],
+                display=element_dict["display"],
+                size=element_dict.get("size"),
+                language=element_dict.get("language"),
+                page=element_dict.get("page"),
+                autoPlay=element_dict.get("autoPlay"),
+                playerConfig=element_dict.get("playerConfig"),
+                forId=element_dict.get("forId"),
+                mime=element_dict.get("mime"),
+            )
+        else:
+            return None
+
+    @queue_until_user_message()
+    async def create_element(self, element: "Element"):
+        if self.show_logger:
+            logger.info(f"SQLAlchemy: create_element, element_id = {element.id}")
+
+        if not self.storage_provider:
+            logger.warning("SQLAlchemy: create_element error. No storage client!")
+            return
+        if not element.for_id:
+            return
+
+        content: Optional[Union[bytes, str]] = None
+
+        if element.path:
+            async with aiofiles.open(element.path, "rb") as f:
+                content = await f.read()
+        elif element.url:
+            async with aiohttp.ClientSession() as session:
+                async with session.get(element.url) as response:
+                    if response.status == 200:
+                        content = await response.read()
+                    else:
+                        content = None
+        elif element.content:
+            content = element.content
+        else:
+            raise ValueError("Element url, path or content must be provided")
+        if content is None:
+            raise ValueError("Content is None, cannot upload file")
+
+        user_id: str = await self._get_user_id_by_thread(element.thread_id) or "unknown"
+        file_object_key = f"{user_id}/{element.id}" + (
+            f"/{element.name}" if element.name else ""
+        )
+
+        if not element.mime:
+            element.mime = "application/octet-stream"
+
+        uploaded_file = await self.storage_provider.upload_file(
+            object_key=file_object_key, data=content, mime=element.mime, overwrite=True
+        )
+        if not uploaded_file:
+            raise ValueError("Failed to persist data in storage_provider")
+
+        element_dict: ElementDict = element.to_dict()
+        element_dict["url"] = uploaded_file.get("url")
+        element_dict["objectKey"] = uploaded_file.get("object_key")
+        element_dict_cleaned = {k: v for k, v in element_dict.items() if v is not None}
+
+        columns = ", ".join(f'"{column}"' for column in element_dict_cleaned.keys())
+        placeholders = ", ".join(f':{column}' for column in element_dict_cleaned.keys())
+        query = f"INSERT INTO elements ({columns}) VALUES ({placeholders})"
+        await self.execute_sql(query=query, parameters=element_dict_cleaned)
+
+    @queue_until_user_message()
+    async def delete_element(self, element_id: str, thread_id: Optional[str] = None):
+        if self.show_logger:
+            logger.info(f"SQLAlchemy: delete_element, element_id={element_id}")
+        query = 'DELETE FROM elements WHERE "id" = :id'
+        parameters = {"id": element_id}
+        await self.execute_sql(query=query, parameters=parameters)
+
+    async def get_all_user_threads(
+        self, user_id: Optional[str] = None, thread_id: Optional[str] = None
+    ) -> Optional[List[ThreadDict]]:
+        if self.show_logger:
+            logger.info("SQLAlchemy: get_all_user_threads")
+        user_threads_query = """
+            SELECT
+                "id" AS thread_id,
+                "createdAt" AS thread_createdat,
+                "name" AS thread_name,
+                "userId" AS user_id,
+                "userIdentifier" AS user_identifier,
+                "tags" AS thread_tags,
+                "meta" AS thread_meta
+            FROM threads
+            WHERE ("userId" = :user_id OR :user_id IS NULL)
+              AND ("id" = :thread_id OR :thread_id IS NULL)
+            ORDER BY "createdAt" DESC
+            LIMIT :limit
+        """
+        params = {
+            "user_id": user_id,
+            "thread_id": thread_id,
+            "limit": self.user_thread_limit,
+        }
+        user_threads = await self.execute_sql(
+            query=user_threads_query,
+            parameters=params,
+        )
+        if not isinstance(user_threads, list):
+            return None
+        if not user_threads:
+            return []
+        else:
+            thread_ids = (
+                "('" + "','".join([t["thread_id"] for t in user_threads]) + "')"
+            )
+
+        steps_feedbacks_query = f"""
+            SELECT
+                s."id" AS step_id,
+                s."name" AS step_name,
+                s."type" AS step_type,
+                s."threadId" AS step_threadid,
+                s."parentId" AS step_parentid,
+                s."streaming" AS step_streaming,
+                s."waitForAnswer" AS step_waitforanswer,
+                s."isError" AS step_iserror,
+                s."meta" AS step_meta,
+                s."tags" AS step_tags,
+                s."input" AS step_input,
+                s."output" AS step_output,
+                s."createdAt" AS step_createdat,
+                s."startTime" AS step_start,
+                s."endTime" AS step_end,
+                s."generation" AS step_generation,
+                s."showInput" AS step_showinput,
+                s."language" AS step_language,
+                s."indent" AS step_indent,
+                f."value" AS feedback_value,
+                f."comment" AS feedback_comment,
+                f."id" AS feedback_id
+            FROM steps s LEFT JOIN feedbacks f ON s."id" = f."forId"
+            WHERE s."threadId" IN {thread_ids}
+            ORDER BY s."createdAt" ASC
+        """
+        steps_feedbacks = await self.execute_sql(
+            query=steps_feedbacks_query, parameters={}
+        )
+
+        elements_query = f"""
+            SELECT
+                e."id" AS element_id,
+                e."threadId" as element_threadid,
+                e."type" AS element_type,
+                e."chainlitKey" AS element_chainlitkey,
+                e."url" AS element_url,
+                e."objectKey" as element_objectkey,
+                e."name" AS element_name,
+                e."display" AS element_display,
+                e."size" AS element_size,
+                e."language" AS element_language,
+                e."page" AS element_page,
+                e."forId" AS element_forid,
+                e."mime" AS element_mime
+            FROM elements e
+            WHERE e."threadId" IN {thread_ids}
+        """
+        elements = await self.execute_sql(query=elements_query, parameters={})
+
+        thread_dicts = {}
+        for thread in user_threads:
+            t_id = thread["thread_id"]
+            meta = thread["thread_meta"]
+            if isinstance(meta, str):
+                try:
+                    meta = json.loads(meta)
+                except:
+                    meta = {}
+            tags = thread["thread_tags"]
+            if isinstance(tags, str):
+                try:
+                    tags = json.loads(tags)
+                except:
+                    tags = []
+            thread_dicts[t_id] = ThreadDict(
+                id=t_id,
+                createdAt=thread["thread_createdat"],
+                name=thread["thread_name"],
+                userId=thread["user_id"],
+                userIdentifier=thread["user_identifier"],
+                tags=tags,
+                metadata=meta,
+                steps=[],
+                elements=[],
+            )
+
+        if isinstance(steps_feedbacks, list):
+            for step_feedback in steps_feedbacks:
+                t_id = step_feedback["step_threadid"]
+                if t_id in thread_dicts:
+                    meta = step_feedback["step_meta"]
+                    if isinstance(meta, str):
+                        try:
+                            meta = json.loads(meta)
+                        except:
+                            meta = {}
+                    tags = step_feedback["step_tags"]
+                    if isinstance(tags, str):
+                        try:
+                            tags = json.loads(tags)
+                        except:
+                            tags = []
+                    feedback = None
+                    if step_feedback["feedback_value"] is not None:
+                        feedback = FeedbackDict(
+                            forId=step_feedback["step_id"],
+                            id=step_feedback.get("feedback_id"),
+                            value=step_feedback["feedback_value"],
+                            comment=step_feedback.get("feedback_comment"),
+                        )
+                    input_val = step_feedback.get("step_input", "")
+                    show_input = step_feedback.get("step_showinput", "false")
+                    if show_input == "false":
+                        input_val = ""
+                    step_dict = StepDict(
+                        id=step_feedback["step_id"],
+                        name=step_feedback["step_name"],
+                        type=step_feedback["step_type"],
+                        threadId=t_id,
+                        parentId=step_feedback.get("step_parentid"),
+                        streaming=step_feedback.get("step_streaming", False),
+                        waitForAnswer=step_feedback.get("step_waitforanswer"),
+                        isError=step_feedback.get("step_iserror"),
+                        metadata=meta,
+                        tags=tags,
+                        input=input_val,
+                        output=step_feedback.get("step_output", ""),
+                        createdAt=step_feedback.get("step_createdat"),
+                        start=step_feedback.get("step_start"),
+                        end=step_feedback.get("step_end"),
+                        generation=step_feedback.get("step_generation"),
+                        showInput=step_feedback.get("step_showinput"),
+                        language=step_feedback.get("step_language"),
+                        indent=step_feedback.get("step_indent"),
+                        feedback=feedback,
+                    )
+                    thread_dicts[t_id]["steps"].append(step_dict)
+
+        if isinstance(elements, list):
+            for element in elements:
+                t_id = element["element_threadid"]
+                if t_id in thread_dicts:
+                    element_dict = ElementDict(
+                        id=element["element_id"],
+                        threadId=t_id,
+                        type=element["element_type"],
+                        chainlitKey=element.get("element_chainlitkey"),
+                        url=element.get("element_url"),
+                        objectKey=element.get("element_objectkey"),
+                        name=element["element_name"],
+                        display=element["element_display"],
+                        size=element.get("element_size"),
+                        language=element.get("element_language"),
+                        autoPlay=element.get("element_autoPlay"),
+                        playerConfig=element.get("element_playerconfig"),
+                        page=element.get("element_page"),
+                        forId=element.get("element_forid"),
+                        mime=element.get("element_mime"),
+                    )
+                    thread_dicts[t_id]["elements"].append(element_dict)
+
+        return list(thread_dicts.values())
+
+    async def close(self) -> None:
+        """Close the database connection and cleanup resources."""
+        if hasattr(self, 'engine') and self.engine:
+            await self.engine.dispose()