import os import os.path import pickle from typing import Dict, List, Optional import chainlit as cl import chainlit.data as cl_data from chainlit.data.utils import queue_until_user_message from chainlit.element import Element, ElementDict from chainlit.socket import persist_user_session from chainlit.step import StepDict from chainlit.types import ( Feedback, PageInfo, PaginatedResponse, Pagination, ThreadDict, ThreadFilter, ) from chainlit.utils import utc_now os.environ["CHAINLIT_AUTH_SECRET"] = "SUPER_SECRET" # nosec B105 now = utc_now() thread_history = [ { "id": "test1", "name": "thread 1", "createdAt": now, "userId": "user1_id", "userIdentifier": "user1", "steps": [ { "id": "test1", "name": "test", "createdAt": now, "type": "user_message", "output": "Message 1", }, { "id": "test2", "name": "test", "createdAt": now, "type": "assistant_message", "output": "Message 2", }, ], }, { "id": "test2", "createdAt": now, "userId": "user1_id", "userIdentifier": "user1", "name": "thread 2", "steps": [ { "id": "test3", "createdAt": now, "name": "test", "type": "user_message", "output": "Message 3", }, { "id": "test4", "createdAt": now, "name": "test", "type": "assistant_message", "output": "Message 4", }, ], }, ] # type: List[ThreadDict] deleted_thread_ids = [] # type: List[str] ELEMENTS_STORAGE = [] THREAD_HISTORY_PICKLE_PATH = os.path.join( os.path.dirname(__file__), "thread_history.pickle" ) if THREAD_HISTORY_PICKLE_PATH and os.path.exists(THREAD_HISTORY_PICKLE_PATH): with open(THREAD_HISTORY_PICKLE_PATH, "rb") as f: thread_history = pickle.load(f) async def save_thread_history(): # Force saving of thread history for reload when server restarts await persist_user_session( cl.context.session.thread_id, cl.context.session.to_persistable() ) with open(THREAD_HISTORY_PICKLE_PATH, "wb") as out_file: pickle.dump(thread_history, out_file) class TestDataLayer(cl_data.BaseDataLayer): async def get_user(self, identifier: str): if identifier == "user1": return cl.PersistedUser(id="user1_id", createdAt=now, identifier=identifier) elif identifier == "user2": return cl.PersistedUser(id="user2_id", createdAt=now, identifier=identifier) return None async def create_user(self, user: cl.User): if user.identifier == "user1": return cl.PersistedUser( id="user1_id", createdAt=now, identifier=user.identifier ) elif user.identifier == "user2": return cl.PersistedUser( id="user2_id", createdAt=now, identifier=user.identifier ) 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, ): thread = next((t for t in thread_history if t["id"] == thread_id), None) if thread: if name: thread["name"] = name if metadata: thread["metadata"] = metadata if tags: thread["tags"] = tags else: thread_history.append( { "id": thread_id, "name": name, "metadata": metadata, "tags": tags, "createdAt": utc_now(), "userId": user_id, "userIdentifier": "user1" if user_id == "user1_id" else "user2" if user_id == "user2_id" else "unknown", "steps": [], } ) @cl_data.queue_until_user_message() async def create_step(self, step_dict: StepDict): cl.user_session.set( "create_step_counter", cl.user_session.get("create_step_counter") + 1 ) thread = next( (t for t in thread_history if t["id"] == step_dict.get("threadId")), None ) if thread: thread["steps"].append(step_dict) async def get_thread_author(self, thread_id: str): thread = await self.get_thread(thread_id) return thread["userIdentifier"] if thread else None async def list_threads( self, pagination: Pagination, filters: ThreadFilter ) -> PaginatedResponse[ThreadDict]: return PaginatedResponse( data=[t for t in thread_history if t["id"] not in deleted_thread_ids], pageInfo=PageInfo(hasNextPage=False, startCursor=None, endCursor=None), ) async def get_thread(self, thread_id: str): thread = next((t for t in thread_history if t["id"] == thread_id), None) if not thread: return None thread["steps"] = sorted(thread["steps"], key=lambda x: x["createdAt"]) return thread async def delete_thread(self, thread_id: str): deleted_thread_ids.append(thread_id) async def delete_feedback( self, feedback_id: str, ) -> bool: return True async def upsert_feedback( self, feedback: Feedback, ) -> str: return "" @queue_until_user_message() async def create_element(self, element: "Element"): if element.url == "http://example.org/test.txt": element.url = "http://example.com/test.txt" ELEMENTS_STORAGE.append(element.to_dict()) async def get_element( self, thread_id: str, element_id: str ) -> Optional["ElementDict"]: return next((e for e in ELEMENTS_STORAGE if e["id"] == element_id), None) @queue_until_user_message() async def delete_element(self, element_id: str, thread_id: Optional[str] = None): pass @queue_until_user_message() async def update_step(self, step_dict: "StepDict"): pass @queue_until_user_message() async def delete_step(self, step_id: str): pass async def get_favorite_steps(self, user_id: str) -> List["StepDict"]: return [] async def build_debug_url(self) -> str: return "" async def close(self) -> None: pass @cl.data_layer def data_layer(): return TestDataLayer() async def send_count(): create_step_counter = cl.user_session.get("create_step_counter") await cl.Message(f"Create step counter: {create_step_counter}").send() @cl.on_chat_start async def main(): # Add step counter to session so that it is saved in thread metadata cl.user_session.set("create_step_counter", 0) await cl.Message("Hello, send me a message!").send() await send_count() @cl.on_message async def handle_message(): # Wait for queue to be flushed await cl.sleep(2) await send_count() async with cl.Step(type="tool", name="thinking") as step: step.output = "Thinking..." await cl.Message("Ok!").send() await send_count() await save_thread_history() @cl.password_auth_callback def auth_callback(username: str, password: str) -> Optional[cl.User]: if (username, password) == ("user1", "user1"): return cl.User(identifier="user1") elif (username, password) == ("user2", "user2"): return cl.User(identifier="user2") else: return None @cl.on_chat_resume async def on_chat_resume(thread: ThreadDict): await cl.Message(f"Welcome back to {thread['name']}").send() if "metadata" in thread: await cl.Message(thread["metadata"], author="metadata", language="json").send() if "tags" in thread: await cl.Message(thread["tags"], author="tags", language="json").send()