chainlit 1.0.400__py3-none-any.whl → 2.0.3__py3-none-any.whl

chainlit/llama_index/callbacks.py

@@ -1,20 +1,21 @@
 from typing import Any, Dict, List, Optional
 
+from literalai import ChatGeneration, CompletionGeneration, GenerationMessage
+from literalai.helper import utc_now
+from llama_index.core.callbacks import TokenCountingHandler
+from llama_index.core.callbacks.schema import CBEventType, EventPayload
+from llama_index.core.llms import ChatMessage, ChatResponse, CompletionResponse
+from llama_index.core.tools.types import ToolMetadata
+
 from chainlit.context import context_var
 from chainlit.element import Text
 from chainlit.step import Step, StepType
-from literalai import ChatGeneration, CompletionGeneration, GenerationMessage
-from literalai.helper import utc_now
-from llama_index.callbacks import TokenCountingHandler
-from llama_index.callbacks.schema import CBEventType, EventPayload
-from llama_index.llms.base import ChatMessage, ChatResponse, CompletionResponse
 
 DEFAULT_IGNORE = [
     CBEventType.CHUNKING,
     CBEventType.SYNTHESIZE,
     CBEventType.EMBEDDING,
     CBEventType.NODE_PARSING,
-    CBEventType.QUERY,
     CBEventType.TREE,
 ]
 
@@ -34,33 +35,17 @@ class LlamaIndexCallbackHandler(TokenCountingHandler):
             event_starts_to_ignore=event_starts_to_ignore,
             event_ends_to_ignore=event_ends_to_ignore,
         )
-        self.context = context_var.get()
 
         self.steps = {}
 
     def _get_parent_id(self, event_parent_id: Optional[str] = None) -> Optional[str]:
         if event_parent_id and event_parent_id in self.steps:
             return event_parent_id
-        elif self.context.current_step:
-            return self.context.current_step.id
-        elif self.context.session.root_message:
-            return self.context.session.root_message.id
+        elif context_var.get().current_step:
+            return context_var.get().current_step.id
         else:
             return None
 
-    def _restore_context(self) -> None:
-        """Restore Chainlit context in the current thread
-
-        Chainlit context is local to the main thread, and LlamaIndex
-        runs the callbacks in its own threads, so they don't have a
-        Chainlit context by default.
-
-        This method restores the context in which the callback handler
-        has been created (it's always created in the main thread), so
-        that we can actually send messages.
-        """
-        context_var.set(self.context)
-
     def on_event_start(
         self,
         event_type: CBEventType,
@@ -70,26 +55,36 @@ class LlamaIndexCallbackHandler(TokenCountingHandler):
         **kwargs: Any,
     ) -> str:
         """Run when an event starts and return id of event."""
-        self._restore_context()
         step_type: StepType = "undefined"
-        if event_type == CBEventType.RETRIEVE:
-            step_type = "retrieval"
+        step_name: str = event_type.value
+        step_input: Optional[Dict[str, Any]] = payload
+        if event_type == CBEventType.FUNCTION_CALL:
+            step_type = "tool"
+            if payload:
+                metadata: Optional[ToolMetadata] = payload.get(EventPayload.TOOL)
+                if metadata:
+                    step_name = getattr(metadata, "name", step_name)
+                step_input = payload.get(EventPayload.FUNCTION_CALL)
+        elif event_type == CBEventType.RETRIEVE:
+            step_type = "tool"
+        elif event_type == CBEventType.QUERY:
+            step_type = "tool"
         elif event_type == CBEventType.LLM:
             step_type = "llm"
         else:
             return event_id
 
         step = Step(
-            name=event_type.value,
+            name=step_name,
             type=step_type,
             parent_id=self._get_parent_id(parent_id),
             id=event_id,
-            disable_feedback=False,
         )
+
         self.steps[event_id] = step
         step.start = utc_now()
-        step.input = payload or {}
-        self.context.loop.create_task(step.send())
+        step.input = step_input or {}
+        context_var.get().loop.create_task(step.send())
         return event_id
 
     def on_event_end(
@@ -105,37 +100,59 @@ class LlamaIndexCallbackHandler(TokenCountingHandler):
         if payload is None or step is None:
             return
 
-        self._restore_context()
-
         step.end = utc_now()
 
-        if event_type == CBEventType.RETRIEVE:
+        if event_type == CBEventType.FUNCTION_CALL:
+            response = payload.get(EventPayload.FUNCTION_OUTPUT)
+            if response:
+                step.output = f"{response}"
+                context_var.get().loop.create_task(step.update())
+
+        elif event_type == CBEventType.QUERY:
+            response = payload.get(EventPayload.RESPONSE)
+            source_nodes = getattr(response, "source_nodes", None)
+            if source_nodes:
+                source_refs = ", ".join(
+                    [f"Source {idx}" for idx, _ in enumerate(source_nodes)]
+                )
+                step.elements = [
+                    Text(
+                        name=f"Source {idx}",
+                        content=source.text or "Empty node",
+                        display="side",
+                    )
+                    for idx, source in enumerate(source_nodes)
+                ]
+                step.output = f"Retrieved the following sources: {source_refs}"
+                context_var.get().loop.create_task(step.update())
+
+        elif event_type == CBEventType.RETRIEVE:
             sources = payload.get(EventPayload.NODES)
             if sources:
-                source_refs = "\, ".join(
+                source_refs = ", ".join(
                     [f"Source {idx}" for idx, _ in enumerate(sources)]
                 )
                 step.elements = [
                     Text(
                         name=f"Source {idx}",
+                        display="side",
                         content=source.node.get_text() or "Empty node",
                     )
                     for idx, source in enumerate(sources)
                 ]
                 step.output = f"Retrieved the following sources: {source_refs}"
-            self.context.loop.create_task(step.update())
+            context_var.get().loop.create_task(step.update())
 
-        if event_type == CBEventType.LLM:
-            formatted_messages = payload.get(
-                EventPayload.MESSAGES
-            )  # type: Optional[List[ChatMessage]]
+        elif event_type == CBEventType.LLM:
+            formatted_messages = payload.get(EventPayload.MESSAGES)  # type: Optional[List[ChatMessage]]
             formatted_prompt = payload.get(EventPayload.PROMPT)
             response = payload.get(EventPayload.RESPONSE)
 
             if formatted_messages:
                 messages = [
                     GenerationMessage(
-                        role=m.role.value, content=m.content or ""  # type: ignore
+                        role=m.role.value,  # type: ignore
+                        content=m.content or "",
                     )
                     for m in formatted_messages
                 ]
@@ -152,10 +169,13 @@ class LlamaIndexCallbackHandler(TokenCountingHandler):
             step.output = content
 
             token_count = self.total_llm_token_count or None
+            raw_response = response.raw if response else None
+            model = getattr(raw_response, "model", None)
 
             if messages and isinstance(response, ChatResponse):
                 msg: ChatMessage = response.message
                 step.generation = ChatGeneration(
+                    model=model,
                     messages=messages,
                     message_completion=GenerationMessage(
                         role=msg.role.value,  # type: ignore
@@ -165,12 +185,17 @@ class LlamaIndexCallbackHandler(TokenCountingHandler):
                 )
             elif formatted_prompt:
                 step.generation = CompletionGeneration(
+                    model=model,
                     prompt=formatted_prompt,
                     completion=content,
                     token_count=token_count,
                 )
 
-            self.context.loop.create_task(step.update())
+            context_var.get().loop.create_task(step.update())
+
+        else:
+            step.output = payload
+            context_var.get().loop.create_task(step.update())
 
         self.steps.pop(event_id, None)
 

chainlit/markdown.py

@@ -1,7 +1,11 @@
 import os
+from pathlib import Path
+from typing import Optional
 
 from chainlit.logger import logger
 
+from ._utils import is_path_inside
+
 # Default chainlit.md file created if none exists
 DEFAULT_MARKDOWN_STR = """# Welcome to Chainlit! 🚀🤖
 
@@ -30,12 +34,24 @@ def init_markdown(root: str):
             logger.info(f"Created default chainlit markdown file at {chainlit_md_file}")
 
 
-def get_markdown_str(root: str):
+def get_markdown_str(root: str, language: str) -> Optional[str]:
     """Get the chainlit.md file as a string."""
-    chainlit_md_path = os.path.join(root, "chainlit.md")
-    if os.path.exists(chainlit_md_path):
-        with open(chainlit_md_path, "r", encoding="utf-8") as f:
-            chainlit_md = f.read()
-            return chainlit_md
+    root_path = Path(root)
+    translated_chainlit_md_path = root_path / f"chainlit_{language}.md"
+    default_chainlit_md_path = root_path / "chainlit.md"
+
+    if (
+        is_path_inside(translated_chainlit_md_path, root_path)
+        and translated_chainlit_md_path.is_file()
+    ):
+        chainlit_md_path = translated_chainlit_md_path
+    else:
+        chainlit_md_path = default_chainlit_md_path
+        logger.warning(
+            f"Translated markdown file for {language} not found. Defaulting to chainlit.md."
+        )
+
+    if chainlit_md_path.is_file():
+        return chainlit_md_path.read_text(encoding="utf-8")
     else:
         return None

chainlit/message.py

@@ -5,9 +5,13 @@ import uuid
 from abc import ABC
 from typing import Dict, List, Optional, Union, cast
 
+from literalai.helper import utc_now
+from literalai.observability.step import MessageStepType
+
 from chainlit.action import Action
+from chainlit.chat_context import chat_context
 from chainlit.config import config
-from chainlit.context import context
+from chainlit.context import context, local_steps
 from chainlit.data import get_data_layer
 from chainlit.element import ElementBased
 from chainlit.logger import logger
@@ -21,9 +25,6 @@ from chainlit.types import (
     AskSpec,
     FileDict,
 )
-from literalai import BaseGeneration
-from literalai.helper import utc_now
-from literalai.step import MessageStepType
 
 
 class MessageBase(ABC):
@@ -32,57 +33,60 @@ class MessageBase(ABC):
     author: str
     content: str = ""
     type: MessageStepType = "assistant_message"
-    disable_feedback = False
     streaming = False
     created_at: Union[str, None] = None
     fail_on_persist_error: bool = False
     persisted = False
     is_error = False
+    parent_id: Optional[str] = None
     language: Optional[str] = None
+    metadata: Optional[Dict] = None
+    tags: Optional[List[str]] = None
     wait_for_answer = False
-    indent: Optional[int] = None
-    generation: Optional[BaseGeneration] = None
 
     def __post_init__(self) -> None:
         trace_event(f"init {self.__class__.__name__}")
         self.thread_id = context.session.thread_id
 
+        previous_steps = local_steps.get() or []
+        parent_step = previous_steps[-1] if previous_steps else None
+        if parent_step:
+            self.parent_id = parent_step.id
+
         if not getattr(self, "id", None):
             self.id = str(uuid.uuid4())
 
     @classmethod
     def from_dict(self, _dict: StepDict):
         type = _dict.get("type", "assistant_message")
-        message = Message(
+        return Message(
             id=_dict["id"],
+            parent_id=_dict.get("parentId"),
             created_at=_dict["createdAt"],
             content=_dict["output"],
             author=_dict.get("name", config.ui.name),
             type=type,  # type: ignore
-            disable_feedback=_dict.get("disableFeedback", False),
             language=_dict.get("language"),
+            metadata=_dict.get("metadata", {}),
         )
 
-        return message
-
     def to_dict(self) -> StepDict:
         _dict: StepDict = {
             "id": self.id,
             "threadId": self.thread_id,
+            "parentId": self.parent_id,
             "createdAt": self.created_at,
             "start": self.created_at,
             "end": self.created_at,
             "output": self.content,
             "name": self.author,
             "type": self.type,
-            "createdAt": self.created_at,
             "language": self.language,
             "streaming": self.streaming,
-            "disableFeedback": self.disable_feedback,
             "isError": self.is_error,
             "waitForAnswer": self.wait_for_answer,
-            "indent": self.indent,
-            "generation": self.generation.to_dict() if self.generation else None,
+            "metadata": self.metadata or {},
+            "tags": self.tags,
         }
 
         return _dict
@@ -99,6 +103,7 @@ class MessageBase(ABC):
             self.streaming = False
 
         step_dict = self.to_dict()
+        chat_context.add(self)
 
         data_layer = get_data_layer()
         if data_layer:
@@ -107,7 +112,7 @@ class MessageBase(ABC):
             except Exception as e:
                 if self.fail_on_persist_error:
                     raise e
-                logger.error(f"Failed to persist message update: {str(e)}")
+                logger.error(f"Failed to persist message update: {e!s}")
 
         await context.emitter.update_step(step_dict)
 
@@ -118,7 +123,7 @@ class MessageBase(ABC):
         Remove a message already sent to the UI.
         """
         trace_event("remove_message")
-
+        chat_context.remove(self)
         step_dict = self.to_dict()
         data_layer = get_data_layer()
         if data_layer:
@@ -127,7 +132,7 @@ class MessageBase(ABC):
             except Exception as e:
                 if self.fail_on_persist_error:
                     raise e
-                logger.error(f"Failed to persist message deletion: {str(e)}")
+                logger.error(f"Failed to persist message deletion: {e!s}")
 
         await context.emitter.delete_step(step_dict)
 
@@ -143,7 +148,7 @@ class MessageBase(ABC):
             except Exception as e:
                 if self.fail_on_persist_error:
                     raise e
-                logger.error(f"Failed to persist message creation: {str(e)}")
+                logger.error(f"Failed to persist message creation: {e!s}")
 
         return step_dict
 
@@ -160,30 +165,31 @@ class MessageBase(ABC):
             self.streaming = False
 
         step_dict = await self._create()
+        chat_context.add(self)
         await context.emitter.send_step(step_dict)
 
-        return self.id
+        return self
 
     async def stream_token(self, token: str, is_sequence=False):
         """
         Sends a token to the UI. This is useful for streaming messages.
         Once all tokens have been streamed, call .send() to end the stream and persist the message 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.content = token
         else:
             self.content += token
 
         assert self.id
-        await context.emitter.send_token(
-            id=self.id, token=token, is_sequence=is_sequence
-        )
+
+        if not self.streaming:
+            self.streaming = True
+            step_dict = self.to_dict()
+            await context.emitter.stream_start(step_dict)
+        else:
+            await context.emitter.send_token(
+                id=self.id, token=token, is_sequence=is_sequence
+            )
 
 
 class Message(MessageBase):
@@ -192,29 +198,28 @@ class Message(MessageBase):
 
     Args:
         content (Union[str, Dict]): The content of the message.
-        author (str, optional): The author of the message, this will be used in the UI. Defaults to the chatbot name (see config).
+        author (str, optional): The author of the message, this will be used in the UI. Defaults to the assistant name (see config).
         language (str, optional): Language of the code is the content is code. See https://react-code-blocks-rajinwonderland.vercel.app/?path=/story/codeblock--supported-languages for a list of supported languages.
         actions (List[Action], optional): A list of actions to send with the message.
         elements (List[ElementBased], optional): A list of elements to send with the message.
-        disable_feedback (bool, optional): Hide the feedback buttons for this specific message
     """
 
     def __init__(
         self,
         content: Union[str, Dict],
-        author: str = config.ui.name,
+        author: Optional[str] = None,
         language: Optional[str] = None,
         actions: Optional[List[Action]] = None,
         elements: Optional[List[ElementBased]] = None,
-        disable_feedback: bool = False,
         type: MessageStepType = "assistant_message",
-        generation: Optional[BaseGeneration] = None,
+        metadata: Optional[Dict] = None,
+        tags: Optional[List[str]] = None,
         id: Optional[str] = None,
+        parent_id: Optional[str] = None,
         created_at: Union[str, None] = None,
     ):
         time.sleep(0.001)
         self.language = language
-        self.generation = generation
         if isinstance(content, dict):
             try:
                 self.content = json.dumps(content, indent=4, ensure_ascii=False)
@@ -231,18 +236,23 @@ class Message(MessageBase):
         if id:
             self.id = str(id)
 
+        if parent_id:
+            self.parent_id = str(parent_id)
+
         if created_at:
             self.created_at = created_at
 
-        self.author = author
+        self.metadata = metadata
+        self.tags = tags
+
+        self.author = author or config.ui.name
         self.type = type
         self.actions = actions if actions is not None else []
         self.elements = elements if elements is not None else []
-        self.disable_feedback = disable_feedback
 
         super().__post_init__()
 
-    async def send(self) -> str:
+    async def send(self):
         """
         Send the message to the UI and persist it in the cloud if a project ID is configured.
         Return the ID of the message.
@@ -250,8 +260,6 @@ class Message(MessageBase):
         trace_event("send_message")
         await super().send()
 
-        context.session.root_message = self
-
         # Create tasks for all actions and elements
         tasks = [action.send(for_id=self.id) for action in self.actions]
         tasks.extend(element.send(for_id=self.id) for element in self.elements)
@@ -259,7 +267,7 @@ class Message(MessageBase):
         # Run all tasks concurrently
         await asyncio.gather(*tasks)
 
-        return self.id
+        return self
 
     async def update(self):
         """
@@ -294,9 +302,7 @@ class ErrorMessage(MessageBase):
 
     Args:
         content (str): Text displayed above the upload button.
-        author (str, optional): The author of the message, this will be used in the UI. Defaults to the chatbot name (see config).
-        parent_id (str, optional): If provided, the message will be nested inside the parent in the UI.
-        indent (int, optional): If positive, the message will be nested in the UI.
+        author (str, optional): The author of the message, this will be used in the UI. Defaults to the assistant name (see config).
     """
 
     def __init__(
@@ -307,7 +313,7 @@ class ErrorMessage(MessageBase):
     ):
         self.content = content
         self.author = author
-        self.type = "system_message"
+        self.type = "assistant_message"
         self.is_error = True
         self.fail_on_persist_error = fail_on_persist_error
 
@@ -337,8 +343,7 @@ class AskUserMessage(AskMessageBase):
 
     Args:
         content (str): The content of the prompt.
-        author (str, optional): The author of the message, this will be used in the UI. Defaults to the chatbot name (see config).
-        disable_feedback (bool, optional): Hide the feedback buttons for this specific message
+        author (str, optional): The author of the message, this will be used in the UI. Defaults to the assistant name (see config).
         timeout (int, optional): The number of seconds to wait for an answer before raising a TimeoutError.
         raise_on_timeout (bool, optional): Whether to raise a socketio TimeoutError if the user does not answer in time.
     """
@@ -348,7 +353,6 @@ class AskUserMessage(AskMessageBase):
         content: str,
         author: str = config.ui.name,
         type: MessageStepType = "assistant_message",
-        disable_feedback: bool = False,
         timeout: int = 60,
         raise_on_timeout: bool = False,
     ):
@@ -356,7 +360,6 @@ class AskUserMessage(AskMessageBase):
         self.author = author
         self.timeout = timeout
         self.type = type
-        self.disable_feedback = disable_feedback
         self.raise_on_timeout = raise_on_timeout
 
         super().__post_init__()
@@ -402,8 +405,7 @@ class AskFileMessage(AskMessageBase):
         accept (Union[List[str], Dict[str, List[str]]]): List of mime type to accept like ["text/csv", "application/pdf"] or a dict like {"text/plain": [".txt", ".py"]}.
         max_size_mb (int, optional): Maximum size per file in MB. Maximum value is 100.
         max_files (int, optional): Maximum number of files to upload. Maximum value is 10.
-        author (str, optional): The author of the message, this will be used in the UI. Defaults to the chatbot name (see config).
-        disable_feedback (bool, optional): Hide the feedback buttons for this specific message
+        author (str, optional): The author of the message, this will be used in the UI. Defaults to the assistant name (see config).
         timeout (int, optional): The number of seconds to wait for an answer before raising a TimeoutError.
         raise_on_timeout (bool, optional): Whether to raise a socketio TimeoutError if the user does not answer in time.
     """
@@ -416,7 +418,6 @@ class AskFileMessage(AskMessageBase):
         max_files=1,
         author=config.ui.name,
         type: MessageStepType = "assistant_message",
-        disable_feedback: bool = False,
         timeout=90,
         raise_on_timeout=False,
     ):
@@ -428,7 +429,6 @@ class AskFileMessage(AskMessageBase):
         self.author = author
         self.timeout = timeout
         self.raise_on_timeout = raise_on_timeout
-        self.disable_feedback = disable_feedback
 
         super().__post_init__()
 
@@ -492,14 +492,12 @@ class AskActionMessage(AskMessageBase):
         content: str,
         actions: List[Action],
         author=config.ui.name,
-        disable_feedback=False,
         timeout=90,
         raise_on_timeout=False,
     ):
         self.content = content
         self.actions = actions
         self.author = author
-        self.disable_feedback = disable_feedback
         self.timeout = timeout
         self.raise_on_timeout = raise_on_timeout
 
@@ -542,7 +540,7 @@ class AskActionMessage(AskMessageBase):
         if res is None:
             self.content = "Timed out: no action was taken"
         else:
-            self.content = f'**Selected action:** {res["label"]}'
+            self.content = f"**Selected:** {res['label']}"
 
         self.wait_for_answer = False
 

chainlit/mistralai/__init__.py

@@ -0,0 +1,50 @@
+import asyncio
+from typing import Union
+
+from literalai import ChatGeneration, CompletionGeneration
+from literalai.helper import timestamp_utc
+
+from chainlit.context import get_context
+from chainlit.step import Step
+
+
+def instrument_mistralai():
+    from literalai.instrumentation.mistralai import instrument_mistralai
+
+    def on_new_generation(
+        generation: Union["ChatGeneration", "CompletionGeneration"], timing
+    ):
+        context = get_context()
+
+        parent_id = None
+        if context.current_step:
+            parent_id = context.current_step.id
+
+        step = Step(
+            name=generation.model if generation.model else generation.provider,
+            type="llm",
+            parent_id=parent_id,
+        )
+        step.generation = generation
+        # Convert start/end time from seconds to milliseconds
+        step.start = (
+            timestamp_utc(timing.get("start"))
+            if timing.get("start", None) is not None
+            else None
+        )
+        step.end = (
+            timestamp_utc(timing.get("end"))
+            if timing.get("end", None) is not None
+            else None
+        )
+
+        if isinstance(generation, ChatGeneration):
+            step.input = generation.messages
+            step.output = generation.message_completion  # type: ignore
+        else:
+            step.input = generation.prompt
+            step.output = generation.completion
+
+        asyncio.create_task(step.send())
+
+    instrument_mistralai(None, on_new_generation)