chainlit 1.0.500__py3-none-any.whl → 1.0.501__py3-none-any.whl

chainlit/data/__init__.py

@@ -10,10 +10,8 @@ from chainlit.context import context
 from chainlit.logger import logger
 from chainlit.session import WebsocketSession
 from chainlit.types import Feedback, Pagination, ThreadDict, ThreadFilter
-from chainlit.user import PersistedUser, User, UserDict
-from literalai import Attachment, PageInfo, PaginatedResponse
-from literalai import Score as LiteralScore
-from literalai import Step as LiteralStep
+from chainlit.user import PersistedUser, User
+from literalai import Attachment, PageInfo, PaginatedResponse, Score as LiteralScore, Step as LiteralStep
 from literalai.filter import threads_filters as LiteralThreadsFilters
 from literalai.step import StepDict as LiteralStepDict
 
@@ -128,9 +126,9 @@ _data_layer: Optional[BaseDataLayer] = None
 
 class ChainlitDataLayer(BaseDataLayer):
     def __init__(self, api_key: str, server: Optional[str]):
-        from literalai import LiteralClient
+        from literalai import AsyncLiteralClient
 
-        self.client = LiteralClient(api_key=api_key, url=server)
+        self.client = AsyncLiteralClient(api_key=api_key, url=server)
         logger.info("Chainlit data layer initialized")
 
     def attachment_to_element_dict(self, attachment: Attachment) -> "ElementDict":
@@ -327,13 +325,16 @@ class ChainlitDataLayer(BaseDataLayer):
 
     @queue_until_user_message()
     async def create_step(self, step_dict: "StepDict"):
-        metadata = {
-            "disableFeedback": step_dict.get("disableFeedback"),
-            "isError": step_dict.get("isError"),
-            "waitForAnswer": step_dict.get("waitForAnswer"),
-            "language": step_dict.get("language"),
-            "showInput": step_dict.get("showInput"),
-        }
+        metadata = dict(
+            step_dict.get("metadata", {}),
+            **{
+                "disableFeedback": step_dict.get("disableFeedback"),
+                "isError": step_dict.get("isError"),
+                "waitForAnswer": step_dict.get("waitForAnswer"),
+                "language": step_dict.get("language"),
+                "showInput": step_dict.get("showInput"),
+            },
+        )
 
         step: LiteralStepDict = {
             "createdAt": step_dict.get("createdAt"),
@@ -345,6 +346,7 @@ class ChainlitDataLayer(BaseDataLayer):
             "name": step_dict.get("name"),
             "threadId": step_dict.get("threadId"),
             "type": step_dict.get("type"),
+            "tags": step_dict.get("tags"),
             "metadata": metadata,
         }
         if step_dict.get("input"):
@@ -453,7 +455,7 @@ class ChainlitDataLayer(BaseDataLayer):
         tags: Optional[List[str]] = None,
     ):
         await self.client.api.upsert_thread(
-            thread_id=thread_id,
+            id=thread_id,
             name=name,
             participant_id=user_id,
             metadata=metadata,
@@ -462,7 +464,8 @@ class ChainlitDataLayer(BaseDataLayer):
 
 
 if api_key := os.environ.get("LITERAL_API_KEY"):
-    server = os.environ.get("LITERAL_SERVER")
+    # support legacy LITERAL_SERVER variable as fallback
+    server = os.environ.get("LITERAL_API_URL", os.environ.get("LITERAL_SERVER"))
     _data_layer = ChainlitDataLayer(api_key=api_key, server=server)