chainlit 0.7.604rc2__py3-none-any.whl → 1.0.0rc0__py3-none-any.whl

chainlit/socket.py

@@ -6,7 +6,7 @@ from chainlit.action import Action
 from chainlit.auth import get_current_user, require_login
 from chainlit.config import config
 from chainlit.context import init_ws_context
-from chainlit.data import chainlit_client
+from chainlit.data import get_data_layer
 from chainlit.logger import logger
 from chainlit.message import ErrorMessage, Message
 from chainlit.server import socket
@@ -27,39 +27,33 @@ def restore_existing_session(sid, session_id, emit_fn, ask_user_fn):
     return False
 
 
-async def persist_user_session(conversation_id: str, metadata: Dict):
-    if not chainlit_client:
-        return
-
-    await chainlit_client.update_conversation_metadata(
-        conversation_id=conversation_id, metadata=metadata
-    )
+async def persist_user_session(thread_id: str, metadata: Dict):
+    if data_layer := get_data_layer():
+        await data_layer.update_thread(thread_id=thread_id, metadata=metadata)
 
 
-async def resume_conversation(session: WebsocketSession):
-    if not chainlit_client or not session.user or not session.conversation_id:
+async def resume_thread(session: WebsocketSession):
+    data_layer = get_data_layer()
+    if not data_layer or not session.user or not session.thread_id_to_resume:
+        return
+    thread = await data_layer.get_thread(thread_id=session.thread_id_to_resume)
+    if not thread:
         return
 
-    conversation = await chainlit_client.get_conversation(
-        conversation_id=session.conversation_id
-    )
+    author = thread.get("user").get("identifier") if thread["user"] else None
+    user_is_author = author == session.user.identifier
 
-    author = (
-        conversation["appUser"].get("username") if conversation["appUser"] else None
-    )
-    user_is_author = author == session.user.username
-
-    if conversation and user_is_author:
-        metadata = conversation["metadata"] or {}
-        user_sessions[session.id] = metadata
+    if user_is_author:
+        metadata = thread["metadata"] or {}
+        user_sessions[session.id] = metadata.copy()
         if chat_profile := metadata.get("chat_profile"):
             session.chat_profile = chat_profile
         if chat_settings := metadata.get("chat_settings"):
             session.chat_settings = chat_settings
 
-        trace_event("conversation_resumed")
+        trace_event("thread_resumed")
 
-        return conversation
+        return thread
 
 
 def load_user_env(user_env):
@@ -128,9 +122,8 @@ async def connect(sid, environ, auth):
         user=user,
         token=token,
         chat_profile=environ.get("HTTP_X_CHAINLIT_CHAT_PROFILE"),
-        conversation_id=environ.get("HTTP_X_CHAINLIT_CONVERSATION_ID"),
+        thread_id=environ.get("HTTP_X_CHAINLIT_THREAD_ID"),
     )
-
     trace_event("connection_successful")
     return True
 
@@ -142,22 +135,27 @@ async def connection_successful(sid):
     if context.session.restored:
         return
 
-    if context.session.conversation_id and config.code.on_chat_resume:
-        conversation = await resume_conversation(context.session)
-        if conversation:
+    if context.session.thread_id_to_resume and config.code.on_chat_resume:
+        thread = await resume_thread(context.session)
+        if thread:
             context.session.has_user_message = True
-            await config.code.on_chat_resume(conversation)
-            await context.emitter.resume_conversation(conversation)
+            await context.emitter.clear_ask()
+            await context.emitter.resume_thread(thread)
+            await config.code.on_chat_resume(thread)
             return
 
     if config.code.on_chat_start:
         """Call the on_chat_start function provided by the developer."""
+        await context.emitter.clear_ask()
         await config.code.on_chat_start()
 
 
 @socket.on("clear_session")
 async def clean_session(sid):
     if session := WebsocketSession.get(sid):
+        if config.code.on_chat_end:
+            init_ws_context(session)
+            await config.code.on_chat_end()
         # Clean up the user session
         if session.id in user_sessions:
             user_sessions.pop(session.id)
@@ -169,13 +167,14 @@ async def clean_session(sid):
 @socket.on("disconnect")
 async def disconnect(sid):
     session = WebsocketSession.get(sid)
-    if config.code.on_chat_end and session:
+    if session:
         init_ws_context(session)
-        """Call the on_chat_end function provided by the developer."""
+
+    if config.code.on_chat_end and session:
         await config.code.on_chat_end()
 
-    if session and session.conversation_id:
-        await persist_user_session(session.conversation_id, session.to_persistable())
+    if session and session.thread_id and session.has_user_message:
+        await persist_user_session(session.thread_id, session.to_persistable())
 
     async def disconnect_on_timeout(sid):
         await asyncio.sleep(config.project.session_timeout)
@@ -195,7 +194,9 @@ async def stop(sid):
         trace_event("stop_task")
 
         init_ws_context(session)
-        await Message(author="System", content="Task stopped by the user.").send()
+        await Message(
+            author="System", content="Task stopped by the user.", disable_feedback=True
+        ).send()
 
         session.should_stop = True
 
@@ -235,7 +236,8 @@ async def message(sid, payload: UIMessagePayload):
 async def process_action(action: Action):
     callback = config.code.action_callbacks.get(action.name)
     if callback:
-        await callback(action)
+        res = await callback(action)
+        return res
     else:
         logger.warning("No callback found for action %s", action.name)
 
@@ -243,11 +245,25 @@ async def process_action(action: Action):
 @socket.on("action_call")
 async def call_action(sid, action):
     """Handle an action call from the UI."""
-    init_ws_context(sid)
+    context = init_ws_context(sid)
 
     action = Action(**action)
 
-    await process_action(action)
+    try:
+        res = await process_action(action)
+        await context.emitter.send_action_response(
+            id=action.id, status=True, response=res if isinstance(res, str) else None
+        )
+
+    except InterruptedError:
+        await context.emitter.send_action_response(
+            id=action.id, status=False, response="Action interrupted by the user"
+        )
+    except Exception as e:
+        logger.exception(e)
+        await context.emitter.send_action_response(
+            id=action.id, status=False, response="An error occured"
+        )
 
 
 @socket.on("chat_settings_change")

chainlit/step.py

@@ -0,0 +1,393 @@
+import asyncio
+import inspect
+import json
+import uuid
+from datetime import datetime
+from functools import wraps
+from typing import Callable, Dict, List, Optional, TypedDict, Union
+
+from chainlit.config import config
+from chainlit.context import context
+from chainlit.data import get_data_layer
+from chainlit.element import Element
+from chainlit.logger import logger
+from chainlit.telemetry import trace_event
+from chainlit.types import FeedbackDict
+from chainlit_client import BaseGeneration
+from chainlit_client.step import StepType, TrueStepType
+
+
+class StepDict(TypedDict, total=False):
+    name: str
+    type: StepType
+    id: str
+    threadId: str
+    parentId: Optional[str]
+    disableFeedback: bool
+    streaming: bool
+    waitForAnswer: Optional[bool]
+    isError: Optional[bool]
+    metadata: Dict
+    input: str
+    output: str
+    createdAt: Optional[str]
+    start: Optional[str]
+    end: Optional[str]
+    generation: Optional[Dict]
+    showInput: Optional[Union[bool, str]]
+    language: Optional[str]
+    indent: Optional[int]
+    feedback: Optional[FeedbackDict]
+
+
+def step(
+    original_function: Optional[Callable] = None,
+    *,
+    name: Optional[str] = "",
+    type: TrueStepType = "undefined",
+    id: Optional[str] = None,
+    disable_feedback: bool = True,
+    root: bool = False,
+    language: Optional[str] = None,
+    show_input: Union[bool, str] = False,
+):
+    """Step decorator for async and sync functions."""
+
+    def wrapper(func: Callable):
+        nonlocal name
+        if not name:
+            name = func.__name__
+
+        # Handle async decorator
+
+        if inspect.iscoroutinefunction(func):
+
+            @wraps(func)
+            async def async_wrapper(*args, **kwargs):
+                async with Step(
+                    type=type,
+                    name=name,
+                    id=id,
+                    disable_feedback=disable_feedback,
+                    root=root,
+                    language=language,
+                    show_input=show_input,
+                ) as step:
+                    try:
+                        step.input = {"args": args, "kwargs": kwargs}
+                    except:
+                        pass
+                    result = await func(*args, **kwargs)
+                    try:
+                        if result and not step.output:
+                            step.output = result
+                    except:
+                        pass
+                    return result
+
+            return async_wrapper
+        else:
+            # Handle sync decorator
+            @wraps(func)
+            def sync_wrapper(*args, **kwargs):
+                with Step(
+                    type=type,
+                    name=name,
+                    id=id,
+                    disable_feedback=disable_feedback,
+                    root=root,
+                    language=language,
+                    show_input=show_input,
+                ) as step:
+                    try:
+                        step.input = {"args": args, "kwargs": kwargs}
+                    except:
+                        pass
+                    result = func(*args, **kwargs)
+                    try:
+                        if result and not step.output:
+                            step.output = result
+                    except:
+                        pass
+                    return result
+
+            return sync_wrapper
+
+    func = original_function
+    if not func:
+        return wrapper
+    else:
+        return wrapper(func)
+
+
+class Step:
+    # Constructor
+    name: str
+    type: TrueStepType
+    id: str
+    parent_id: Optional[str]
+    disable_feedback: bool
+
+    streaming: bool
+    persisted: bool
+
+    root: bool
+    show_input: Union[bool, str]
+
+    is_error: Optional[bool]
+    metadata: Dict
+    thread_id: str
+    created_at: Union[str, None]
+    start: Union[str, None]
+    end: Union[str, None]
+    generation: Optional[BaseGeneration]
+    language: Optional[str]
+    elements: Optional[List[Element]]
+    fail_on_persist_error: bool
+
+    def __init__(
+        self,
+        name: Optional[str] = config.ui.name,
+        type: TrueStepType = "undefined",
+        id: Optional[str] = None,
+        parent_id: Optional[str] = None,
+        elements: Optional[List[Element]] = None,
+        disable_feedback: bool = True,
+        root: bool = False,
+        language: Optional[str] = None,
+        show_input: Union[bool, str] = False,
+    ):
+        trace_event(f"init {self.__class__.__name__} {type}")
+        self._input = ""
+        self._output = ""
+        self.thread_id = context.session.thread_id
+        self.name = name or ""
+        self.type = type
+        self.id = id or str(uuid.uuid4())
+        self.disable_feedback = disable_feedback
+        self.metadata = {}
+        self.is_error = False
+        self.show_input = show_input
+        self.parent_id = parent_id
+        self.root = root
+
+        self.language = language
+        self.generation = None
+        self.elements = elements or []
+
+        self.created_at = datetime.utcnow().isoformat()
+        self.start = None
+        self.end = None
+
+        self.streaming = False
+        self.persisted = False
+        self.fail_on_persist_error = False
+
+    def _process_content(self, content, set_language=False):
+        if content is None:
+            return ""
+        if isinstance(content, dict):
+            try:
+                processed_content = json.dumps(content, indent=4, ensure_ascii=False)
+                if set_language:
+                    self.language = "json"
+            except TypeError:
+                processed_content = str(content)
+                if set_language:
+                    self.language = "text"
+        elif isinstance(content, str):
+            processed_content = content
+        else:
+            processed_content = str(content)
+            if set_language:
+                self.language = "text"
+        return processed_content
+
+    @property
+    def input(self):
+        return self._input
+
+    @input.setter
+    def input(self, content: Union[Dict, str]):
+        self._input = self._process_content(content, set_language=False)
+
+    @property
+    def output(self):
+        return self._output
+
+    @output.setter
+    def output(self, content: Union[Dict, str]):
+        self._output = self._process_content(content, set_language=True)
+
+    def to_dict(self) -> StepDict:
+        _dict: StepDict = {
+            "name": self.name,
+            "type": self.type,
+            "id": self.id,
+            "threadId": self.thread_id,
+            "parentId": self.parent_id,
+            "disableFeedback": self.disable_feedback,
+            "streaming": self.streaming,
+            "metadata": self.metadata,
+            "input": self.input,
+            "isError": self.is_error,
+            "output": self.output,
+            "createdAt": self.created_at,
+            "start": self.start,
+            "end": self.end,
+            "language": self.language,
+            "showInput": self.show_input,
+            "generation": self.generation.to_dict() if self.generation else None,
+        }
+        return _dict
+
+    async def update(self):
+        """
+        Update a step already sent to the UI.
+        """
+        trace_event("update_step")
+
+        if self.streaming:
+            self.streaming = False
+
+        step_dict = self.to_dict()
+        data_layer = get_data_layer()
+
+        if data_layer:
+            try:
+                asyncio.create_task(data_layer.update_step(step_dict))
+            except Exception as e:
+                if self.fail_on_persist_error:
+                    raise e
+                logger.error(f"Failed to persist step update: {str(e)}")
+
+        tasks = [el.send(for_id=self.id) for el in self.elements]
+        await asyncio.gather(*tasks)
+
+        if not config.features.prompt_playground and "generation" in step_dict:
+            step_dict.pop("generation", None)
+
+        await context.emitter.update_step(step_dict)
+
+        return True
+
+    async def remove(self):
+        """
+        Remove a step already sent to the UI.
+        """
+        trace_event("remove_step")
+
+        step_dict = self.to_dict()
+        data_layer = get_data_layer()
+
+        if data_layer:
+            try:
+                asyncio.create_task(data_layer.delete_step(self.id))
+            except Exception as e:
+                if self.fail_on_persist_error:
+                    raise e
+                logger.error(f"Failed to persist step deletion: {str(e)}")
+
+        await context.emitter.delete_step(step_dict)
+
+        return True
+
+    async def send(self):
+        if self.persisted:
+            return
+
+        if config.code.author_rename:
+            self.name = await config.code.author_rename(self.name)
+
+        if self.streaming:
+            self.streaming = False
+
+        step_dict = self.to_dict()
+
+        data_layer = get_data_layer()
+
+        if data_layer:
+            try:
+                asyncio.create_task(data_layer.create_step(step_dict))
+                self.persisted = True
+            except Exception as e:
+                if self.fail_on_persist_error:
+                    raise e
+                logger.error(f"Failed to persist step creation: {str(e)}")
+
+        tasks = [el.send(for_id=self.id) for el in self.elements]
+        await asyncio.gather(*tasks)
+
+        if not config.features.prompt_playground and "generation" in step_dict:
+            step_dict.pop("generation", None)
+
+        await context.emitter.send_step(step_dict)
+
+        return self.id
+
+    async def stream_token(self, token: str, is_sequence=False):
+        """
+        Sends a token to the UI.
+        Once all tokens have been streamed, call .send() to end the stream and persist the step if persistence is enabled.
+        """
+
+        if not self.streaming:
+            self.streaming = True
+            step_dict = self.to_dict()
+            await context.emitter.stream_start(step_dict)
+
+        if is_sequence:
+            self.output = token
+        else:
+            self.output += token
+
+        assert self.id
+        await context.emitter.send_token(
+            id=self.id, token=token, is_sequence=is_sequence
+        )
+
+    # Handle parameter less decorator
+    def __call__(self, func):
+        return step(
+            original_function=func,
+            type=self.type,
+            name=self.name,
+            id=self.id,
+            parent_id=self.parent_id,
+            thread_id=self.thread_id,
+            disable_feedback=self.disable_feedback,
+        )
+
+    # Handle Context Manager Protocol
+    async def __aenter__(self):
+        self.start = datetime.utcnow().isoformat()
+        if not self.parent_id and not self.root:
+            if current_step := context.current_step:
+                self.parent_id = current_step.id
+            elif context.session.root_message:
+                self.parent_id = context.session.root_message.id
+        context.session.active_steps.append(self)
+        await self.send()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        self.end = datetime.utcnow().isoformat()
+        context.session.active_steps.pop()
+        await self.update()
+
+    def __enter__(self):
+        self.start = datetime.utcnow().isoformat()
+        if not self.parent_id and not self.root:
+            if current_step := context.current_step:
+                self.parent_id = current_step.id
+            elif context.session.root_message:
+                self.parent_id = context.session.root_message.id
+        context.session.active_steps.append(self)
+
+        asyncio.create_task(self.send())
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.end = datetime.utcnow().isoformat()
+        context.session.active_steps.pop()
+        asyncio.create_task(self.update())