chainlit 1.0.504__py3-none-any.whl → 1.0.505__py3-none-any.whl

chainlit/data/sql_alchemy.py

@@ -1,29 +1,45 @@
-import uuid
+import json
 import ssl
+import uuid
+from dataclasses import asdict
 from datetime import datetime, timezone
-import json
-from typing import Any, Dict, List, Optional, Union, TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+
 import aiofiles
 import aiohttp
-from dataclasses import asdict
-from sqlalchemy import text
-from sqlalchemy.exc import SQLAlchemyError
-from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, AsyncEngine
-from sqlalchemy.orm import sessionmaker
 from chainlit.context import context
-from chainlit.logger import logger
 from chainlit.data import BaseDataLayer, BaseStorageClient, queue_until_user_message
-from chainlit.user import User, PersistedUser
-from chainlit.types import Feedback, FeedbackDict, Pagination, ThreadDict, ThreadFilter, PageInfo, PaginatedResponse
+from chainlit.element import Avatar, ElementDict
+from chainlit.logger import logger
 from chainlit.step import StepDict
-from chainlit.element import ElementDict, Avatar
+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
 
 if TYPE_CHECKING:
     from chainlit.element import Element, ElementDict
     from chainlit.step import StepDict
 
+
 class SQLAlchemyDataLayer(BaseDataLayer):
-    def __init__(self, conninfo: str, ssl_require: bool = False, storage_provider: Optional[BaseStorageClient] = None, user_thread_limit: Optional[int] = 1000):
+    def __init__(
+        self,
+        conninfo: str,
+        ssl_require: bool = False,
+        storage_provider: Optional[BaseStorageClient] = None,
+        user_thread_limit: Optional[int] = 1000,
+    ):
         self._conninfo = conninfo
         self.user_thread_limit = user_thread_limit
         ssl_args = {}
@@ -32,17 +48,24 @@ class SQLAlchemyDataLayer(BaseDataLayer):
             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)
+            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)  # type: ignore
         if storage_provider:
-            self.storage_provider = storage_provider
+            self.storage_provider: Optional[BaseStorageClient] = storage_provider
             logger.info("SQLAlchemyDataLayer storage client initialized")
         else:
-            logger.warn("SQLAlchemyDataLayer storage client is not initialized and elements will not be persisted!")
+            self.storage_provider = None
+            logger.warn(
+                "SQLAlchemyDataLayer storage client is not initialized and elements will not be persisted!"
+            )
 
     ###### SQL Helpers ######
-    async def execute_sql(self, query: str, parameters: dict) -> Union[List[Dict[str, Any]], int, None]:
+    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:
@@ -66,7 +89,7 @@ class SQLAlchemyDataLayer(BaseDataLayer):
 
     async def get_current_timestamp(self) -> str:
         return datetime.now().isoformat() + "Z"
-    
+
     def clean_result(self, obj):
         """Recursively change UUID -> str and serialize dictionaries"""
         if isinstance(obj, dict):
@@ -76,7 +99,7 @@ class SQLAlchemyDataLayer(BaseDataLayer):
         elif isinstance(obj, uuid.UUID):
             return str(obj)
         return obj
-              
+
     ###### User ######
     async def get_user(self, identifier: str) -> Optional[PersistedUser]:
         logger.info(f"SQLAlchemy: get_user, identifier={identifier}")
@@ -87,26 +110,28 @@ class SQLAlchemyDataLayer(BaseDataLayer):
             user_data = result[0]
             return PersistedUser(**user_data)
         return None
-        
+
     async def create_user(self, user: User) -> Optional[PersistedUser]:
         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]  = {
+        existing_user: Optional["PersistedUser"] = await self.get_user(user.identifier)
+        user_dict: Dict[str, Any] = {
             "identifier": str(user.identifier),
-            "metadata": json.dumps(user.metadata) or {}
-            }
-        if not existing_user: # create the user
+            "metadata": json.dumps(user.metadata) or {},
+        }
+        if not existing_user:  # create the user
             logger.info("SQLAlchemy: create_user, creating the user")
-            user_dict['id'] = str(uuid.uuid4())
-            user_dict['createdAt'] = await self.get_current_timestamp()
+            user_dict["id"] = str(uuid.uuid4())
+            user_dict["createdAt"] = await self.get_current_timestamp()
             query = """INSERT INTO users ("id", "identifier", "createdAt", "metadata") VALUES (:id, :identifier, :createdAt, :metadata)"""
             await self.execute_sql(query=query, parameters=user_dict)
-        else: # update the user
+        else:  # update the user
             logger.info("SQLAlchemy: update user metadata")
             query = """UPDATE users SET "metadata" = :metadata WHERE "identifier" = :identifier"""
-            await self.execute_sql(query=query, parameters=user_dict) # We want to update the metadata
+            await self.execute_sql(
+                query=query, parameters=user_dict
+            )  # We want to update the metadata
         return await self.get_user(user.identifier)
-        
+
     ###### Threads ######
     async def get_thread_author(self, thread_id: str) -> str:
         logger.info(f"SQLAlchemy: get_thread_author, thread_id={thread_id}")
@@ -114,21 +139,30 @@ class SQLAlchemyDataLayer(BaseDataLayer):
         parameters = {"id": thread_id}
         result = await self.execute_sql(query=query, parameters=parameters)
         if isinstance(result, list) and result[0]:
-            author_identifier = result[0].get('userIdentifier')
+            author_identifier = result[0].get("userIdentifier")
             if author_identifier is not None:
-                print(f'Author found: {author_identifier}')
+                print(f"Author found: {author_identifier}")
                 return author_identifier
-        raise ValueError (f"Author not found for thread_id {thread_id}")
-    
+        raise ValueError(f"Author not found for thread_id {thread_id}")
+
     async def get_thread(self, thread_id: str) -> Optional[ThreadDict]:
         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)
+        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):
+    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,
+    ):
         logger.info(f"SQLAlchemy: update_thread, thread_id={thread_id}")
         if context.session.user is not None:
             user_identifier = context.session.user.identifier
@@ -136,17 +170,25 @@ class SQLAlchemyDataLayer(BaseDataLayer):
             raise ValueError("User not found in session context")
         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),
+            "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": tags,
             "metadata": json.dumps(metadata) if metadata else None,
         }
-        parameters = {key: value for key, value in data.items() if value is not None} # Remove keys with None values
-        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')
+        parameters = {
+            key: value for key, value in data.items() if value is not None
+        }  # Remove keys with None values
+        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})
@@ -167,12 +209,18 @@ class SQLAlchemyDataLayer(BaseDataLayer):
         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:
-        logger.info(f"SQLAlchemy: list_threads, pagination={pagination}, filters={filters}")
+
+    async def list_threads(
+        self, pagination: Pagination, filters: ThreadFilter
+    ) -> PaginatedResponse:
+        logger.info(
+            f"SQLAlchemy: list_threads, pagination={pagination}, filters={filters}"
+        )
         if not filters.userId:
             raise ValueError("userId is required")
-        all_user_threads: List[ThreadDict] = await self.get_all_user_threads(user_id=filters.userId) or []
+        all_user_threads: 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
@@ -183,47 +231,68 @@ class SQLAlchemyDataLayer(BaseDataLayer):
             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)
+                    keyword_match = any(
+                        search_keyword in step["output"].lower()
+                        for step in thread["steps"]
+                        if "output" in step
+                    )
                 if feedback_value is not None:
                     feedback_match = False  # Assume no match until found
-                    for step in thread['steps']:
-                        feedback = step.get('feedback')
-                        if feedback and feedback.get('value') == feedback_value:
+                    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, thread in enumerate(filtered_threads):
-                if thread['id'] == pagination.cursor:  # Find the start index using pagination.cursor
+                if (
+                    thread["id"] == pagination.cursor
+                ):  # Find the start index using 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
+        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
+            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'):
+    async def create_step(self, step_dict: "StepDict"):
         logger.info(f"SQLAlchemy: create_step, step_id={step_dict.get('id')}")
-        if not getattr(context.session.user, 'id', None):
+        if not getattr(context.session.user, "id", None):
             raise ValueError("No authenticated user in context")
-        step_dict['showInput'] = str(step_dict.get('showInput', '')).lower() if 'showInput' in step_dict else None
-        parameters = {key: value for key, value in step_dict.items() if value is not None and not (isinstance(value, dict) and not value)}
-        parameters['metadata'] = json.dumps(step_dict.get('metadata', {}))
-        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')
+        step_dict["showInput"] = (
+            str(step_dict.get("showInput", "")).lower()
+            if "showInput" in step_dict
+            else None
+        )
+        parameters = {
+            key: value
+            for key, value in step_dict.items()
+            if value is not None and not (isinstance(value, dict) and not value)
+        }
+        parameters["metadata"] = json.dumps(step_dict.get("metadata", {}))
+        parameters["generation"] = json.dumps(step_dict.get("generation", {}))
+        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})
@@ -231,9 +300,9 @@ class SQLAlchemyDataLayer(BaseDataLayer):
             SET {updates};
         """
         await self.execute_sql(query=query, parameters=parameters)
-    
+
     @queue_until_user_message()
-    async def update_step(self, step_dict: 'StepDict'):
+    async def update_step(self, step_dict: "StepDict"):
         logger.info(f"SQLAlchemy: update_step, step_id={step_dict.get('id')}")
         await self.create_step(step_dict)
 
@@ -248,17 +317,21 @@ class SQLAlchemyDataLayer(BaseDataLayer):
         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:
         logger.info(f"SQLAlchemy: upsert_feedback, feedback_id={feedback.id}")
         feedback.id = feedback.id or str(uuid.uuid4())
         feedback_dict = asdict(feedback)
-        parameters = {key: value for key, value in feedback_dict.items() if value is not None}
+        parameters = {
+            key: value for key, value in feedback_dict.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}" = :{key}' for key in parameters.keys() if key != 'id')
+        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})
@@ -267,24 +340,26 @@ class SQLAlchemyDataLayer(BaseDataLayer):
         """
         await self.execute_sql(query=query, parameters=parameters)
         return feedback.id
-    
+
     async def delete_feedback(self, feedback_id: str) -> bool:
         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 ######
     @queue_until_user_message()
-    async def create_element(self, element: 'Element'):
+    async def create_element(self, element: "Element"):
         logger.info(f"SQLAlchemy: create_element, element_id = {element.id}")
-        if not getattr(context.session.user, 'id', None):
+        if not getattr(context.session.user, "id", None):
             raise ValueError("No authenticated user in context")
-        if isinstance(element, Avatar): # Skip creating elements of type avatar
+        if isinstance(element, Avatar):  # Skip creating elements of type avatar
             return
         if not self.storage_provider:
-            logger.warn(f"SQLAlchemy: create_element error. No blob_storage_client is configured!")
+            logger.warn(
+                f"SQLAlchemy: create_element error. No blob_storage_client is configured!"
+            )
             return
         if not element.for_id:
             return
@@ -310,24 +385,30 @@ class SQLAlchemyDataLayer(BaseDataLayer):
 
         context_user = context.session.user
 
-        user_folder = getattr(context_user, 'id', 'unknown')
-        file_object_key = f"{user_folder}/{element.id}" + (f"/{element.name}" if element.name else "")
+        user_folder = getattr(context_user, "id", "unknown")
+        file_object_key = f"{user_folder}/{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)
+        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("SQLAlchemy Error: create_element, Failed to persist data in storage_provider")
+            raise ValueError(
+                "SQLAlchemy Error: create_element, 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["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())
+        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)
 
@@ -339,9 +420,11 @@ class SQLAlchemyDataLayer(BaseDataLayer):
         await self.execute_sql(query=query, parameters=parameters)
 
     async def delete_user_session(self, id: str) -> bool:
-        return False # Not sure why documentation wants this
+        return False  # Not sure why documentation wants this
 
-    async def get_all_user_threads(self, user_id: Optional[str] = None, thread_id: Optional[str] = None) -> Optional[List[ThreadDict]]:
+    async def get_all_user_threads(
+        self, user_id: Optional[str] = None, thread_id: Optional[str] = None
+    ) -> Optional[List[ThreadDict]]:
         """Fetch all user threads up to self.user_thread_limit, or one thread by id if thread_id is provided."""
         logger.info(f"SQLAlchemy: get_all_user_threads")
         user_threads_query = """
@@ -358,14 +441,25 @@ class SQLAlchemyDataLayer(BaseDataLayer):
             ORDER BY "createdAt" DESC
             LIMIT :limit
         """
-        user_threads = await self.execute_sql(query=user_threads_query, parameters={"user_id": user_id, "limit": self.user_thread_limit, "thread_id": thread_id})
+        user_threads = await self.execute_sql(
+            query=user_threads_query,
+            parameters={
+                "user_id": user_id,
+                "limit": self.user_thread_limit,
+                "thread_id": thread_id,
+            },
+        )
         if not isinstance(user_threads, list):
             return None
         if not user_threads:
             return []
         else:
-            thread_ids = "('" + "','".join(map(str, [thread['thread_id'] for thread in user_threads])) + "')"
-        
+            thread_ids = (
+                "('"
+                + "','".join(map(str, [thread["thread_id"] for thread in user_threads]))
+                + "')"
+            )
+
         steps_feedbacks_query = f"""
             SELECT
                 s."id" AS step_id,
@@ -394,8 +488,10 @@ class SQLAlchemyDataLayer(BaseDataLayer):
             WHERE s."threadId" IN {thread_ids}
             ORDER BY s."createdAt" ASC
         """
-        steps_feedbacks = await self.execute_sql(query=steps_feedbacks_query, parameters={})
-        
+        steps_feedbacks = await self.execute_sql(
+            query=steps_feedbacks_query, parameters={}
+        )
+
         elements_query = f"""
             SELECT
                 e."id" AS element_id,
@@ -418,77 +514,83 @@ class SQLAlchemyDataLayer(BaseDataLayer):
 
         thread_dicts = {}
         for thread in user_threads:
-            thread_id = thread['thread_id']
+            thread_id = thread["thread_id"]
             if thread_id is not None:
                 thread_dicts[thread_id] = ThreadDict(
                     id=thread_id,
-                    createdAt=thread['thread_createdat'],
-                    name=thread['thread_name'],
-                    userId=thread['user_id'],
-                    userIdentifier=thread['user_identifier'],
-                    tags=thread['thread_tags'],
-                    metadata=thread['thread_metadata'],
+                    createdAt=thread["thread_createdat"],
+                    name=thread["thread_name"],
+                    userId=thread["user_id"],
+                    userIdentifier=thread["user_identifier"],
+                    tags=thread["thread_tags"],
+                    metadata=thread["thread_metadata"],
                     steps=[],
-                    elements=[]
+                    elements=[],
                 )
         # Process steps_feedbacks to populate the steps in the corresponding ThreadDict
         if isinstance(steps_feedbacks, list):
             for step_feedback in steps_feedbacks:
-                thread_id = step_feedback['step_threadid']
+                thread_id = step_feedback["step_threadid"]
                 if thread_id is not None:
                     feedback = None
-                    if step_feedback['feedback_value'] is not 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')
+                            forId=step_feedback["step_id"],
+                            id=step_feedback.get("feedback_id"),
+                            value=step_feedback["feedback_value"],
+                            comment=step_feedback.get("feedback_comment"),
                         )
                     step_dict = StepDict(
-                        id=step_feedback['step_id'],
-                        name=step_feedback['step_name'],
-                        type=step_feedback['step_type'],
+                        id=step_feedback["step_id"],
+                        name=step_feedback["step_name"],
+                        type=step_feedback["step_type"],
                         threadId=thread_id,
-                        parentId=step_feedback.get('step_parentid'),
-                        disableFeedback=step_feedback.get('step_disablefeedback', False),
-                        streaming=step_feedback.get('step_streaming', False),
-                        waitForAnswer=step_feedback.get('step_waitforanswer'),
-                        isError=step_feedback.get('step_iserror'),
-                        metadata=step_feedback['step_metadata'] if step_feedback.get('step_metadata') is not None else {},
-                        tags=step_feedback.get('step_tags'),
-                        input=step_feedback.get('step_input', '') if step_feedback['step_showinput'] else '',
-                        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
+                        parentId=step_feedback.get("step_parentid"),
+                        disableFeedback=step_feedback.get(
+                            "step_disablefeedback", False
+                        ),
+                        streaming=step_feedback.get("step_streaming", False),
+                        waitForAnswer=step_feedback.get("step_waitforanswer"),
+                        isError=step_feedback.get("step_iserror"),
+                        metadata=step_feedback["step_metadata"]
+                        if step_feedback.get("step_metadata") is not None
+                        else {},
+                        tags=step_feedback.get("step_tags"),
+                        input=step_feedback.get("step_input", "")
+                        if step_feedback["step_showinput"]
+                        else "",
+                        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,
                     )
                     # Append the step to the steps list of the corresponding ThreadDict
-                    thread_dicts[thread_id]['steps'].append(step_dict)
+                    thread_dicts[thread_id]["steps"].append(step_dict)
 
         if isinstance(elements, list):
             for element in elements:
-                thread_id = element['element_threadid']
+                thread_id = element["element_threadid"]
                 if thread_id is not None:
                     element_dict = ElementDict(
-                        id=element['element_id'],
+                        id=element["element_id"],
                         threadId=thread_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'),
-                        page=element.get('element_page'),
-                        forId=element.get('element_forid'),
-                        mime=element.get('element_mime'),
+                        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"),
+                        page=element.get("element_page"),
+                        forId=element.get("element_forid"),
+                        mime=element.get("element_mime"),
                     )
-                    thread_dicts[thread_id]['elements'].append(element_dict)   # type: ignore
+                    thread_dicts[thread_id]["elements"].append(element_dict)  # type: ignore
 
-        return list(thread_dicts.values()) 
+        return list(thread_dicts.values())