PyPI - chainlit - Versions diffs - 1.0.401__py3-none-any.whl → 2.0.4__py3-none-any.whl - Mend

chainlit 1.0.401py3-none-any.whl → 2.0.4py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.

This version of chainlit might be problematic. Click here for more details.

Files changed (113) hide show

chainlit/__init__.py +98 -279
chainlit/_utils.py +8 -0
chainlit/action.py +12 -10
chainlit/{auth.py → auth/__init__.py} +28 -36
chainlit/auth/cookie.py +123 -0
chainlit/auth/jwt.py +39 -0
chainlit/cache.py +4 -6
chainlit/callbacks.py +362 -0
chainlit/chat_context.py +64 -0
chainlit/chat_settings.py +3 -1
chainlit/cli/__init__.py +77 -8
chainlit/config.py +191 -102
chainlit/context.py +42 -13
chainlit/copilot/dist/index.js +8750 -903
chainlit/data/__init__.py +101 -416
chainlit/data/acl.py +6 -2
chainlit/data/base.py +107 -0
chainlit/data/chainlit_data_layer.py +614 -0
chainlit/data/dynamodb.py +590 -0
chainlit/data/literalai.py +500 -0
chainlit/data/sql_alchemy.py +721 -0
chainlit/data/storage_clients/__init__.py +0 -0
chainlit/data/storage_clients/azure.py +81 -0
chainlit/data/storage_clients/azure_blob.py +89 -0
chainlit/data/storage_clients/base.py +26 -0
chainlit/data/storage_clients/gcs.py +88 -0
chainlit/data/storage_clients/s3.py +75 -0
chainlit/data/utils.py +29 -0
chainlit/discord/__init__.py +6 -0
chainlit/discord/app.py +354 -0
chainlit/element.py +91 -33
chainlit/emitter.py +81 -29
chainlit/frontend/dist/assets/DailyMotion-Ce9dQoqZ.js +1 -0
chainlit/frontend/dist/assets/Dataframe-C1XonMcV.js +22 -0
chainlit/frontend/dist/assets/Facebook-DVVt6lrr.js +1 -0
chainlit/frontend/dist/assets/FilePlayer-c7stW4vz.js +1 -0
chainlit/frontend/dist/assets/Kaltura-BmMmgorA.js +1 -0
chainlit/frontend/dist/assets/Mixcloud-Cw8hDmiO.js +1 -0
chainlit/frontend/dist/assets/Mux-DiRZfeUf.js +1 -0
chainlit/frontend/dist/assets/Preview-6Jt2mRHx.js +1 -0
chainlit/frontend/dist/assets/SoundCloud-DKwcT58_.js +1 -0
chainlit/frontend/dist/assets/Streamable-BVdxrEeX.js +1 -0
chainlit/frontend/dist/assets/Twitch-DFqZR7Gu.js +1 -0
chainlit/frontend/dist/assets/Vidyard-0BQAAtVk.js +1 -0
chainlit/frontend/dist/assets/Vimeo-CRFSH0Vu.js +1 -0
chainlit/frontend/dist/assets/Wistia-CKrmdQaG.js +1 -0
chainlit/frontend/dist/assets/YouTube-CQpL-rvU.js +1 -0
chainlit/frontend/dist/assets/index-DQmLRKyv.css +1 -0
chainlit/frontend/dist/assets/index-QdmxtIMQ.js +8665 -0
chainlit/frontend/dist/assets/react-plotly-B9hvVpUG.js +3484 -0
chainlit/frontend/dist/index.html +2 -4
chainlit/haystack/callbacks.py +4 -7
chainlit/input_widget.py +8 -4
chainlit/langchain/callbacks.py +103 -68
chainlit/langflow/__init__.py +1 -0
chainlit/llama_index/callbacks.py +65 -40
chainlit/markdown.py +22 -6
chainlit/message.py +54 -56
chainlit/mistralai/__init__.py +50 -0
chainlit/oauth_providers.py +266 -8
chainlit/openai/__init__.py +10 -18
chainlit/secret.py +1 -1
chainlit/server.py +789 -228
chainlit/session.py +108 -90
chainlit/slack/__init__.py +6 -0
chainlit/slack/app.py +397 -0
chainlit/socket.py +199 -116
chainlit/step.py +141 -89
chainlit/sync.py +2 -1
chainlit/teams/__init__.py +6 -0
chainlit/teams/app.py +338 -0
chainlit/translations/bn.json +244 -0
chainlit/translations/en-US.json +122 -8
chainlit/translations/gu.json +244 -0
chainlit/translations/he-IL.json +244 -0
chainlit/translations/hi.json +244 -0
chainlit/translations/ja.json +242 -0
chainlit/translations/kn.json +244 -0
chainlit/translations/ml.json +244 -0
chainlit/translations/mr.json +244 -0
chainlit/translations/nl-NL.json +242 -0
chainlit/translations/ta.json +244 -0
chainlit/translations/te.json +244 -0
chainlit/translations/zh-CN.json +243 -0
chainlit/translations.py +60 -0
chainlit/types.py +133 -28
chainlit/user.py +14 -3
chainlit/user_session.py +6 -3
chainlit/utils.py +52 -5
chainlit/version.py +3 -2
{chainlit-1.0.401.dist-info → chainlit-2.0.4.dist-info}/METADATA +48 -50
chainlit-2.0.4.dist-info/RECORD +107 -0
chainlit/cli/utils.py +0 -24
chainlit/frontend/dist/assets/index-9711593e.js +0 -723
chainlit/frontend/dist/assets/index-d088547c.css +0 -1
chainlit/frontend/dist/assets/react-plotly-d8762cc2.js +0 -3602
chainlit/playground/__init__.py +0 -2
chainlit/playground/config.py +0 -40
chainlit/playground/provider.py +0 -108
chainlit/playground/providers/__init__.py +0 -13
chainlit/playground/providers/anthropic.py +0 -118
chainlit/playground/providers/huggingface.py +0 -75
chainlit/playground/providers/langchain.py +0 -89
chainlit/playground/providers/openai.py +0 -408
chainlit/playground/providers/vertexai.py +0 -171
chainlit/translations/pt-BR.json +0 -155
chainlit-1.0.401.dist-info/RECORD +0 -66
/chainlit/copilot/dist/assets/{logo_dark-2a3cf740.svg → logo_dark-IkGJ_IwC.svg} +0 -0
/chainlit/copilot/dist/assets/{logo_light-b078e7bc.svg → logo_light-Bb_IPh6r.svg} +0 -0
/chainlit/frontend/dist/assets/{logo_dark-2a3cf740.svg → logo_dark-IkGJ_IwC.svg} +0 -0
/chainlit/frontend/dist/assets/{logo_light-b078e7bc.svg → logo_light-Bb_IPh6r.svg} +0 -0
{chainlit-1.0.401.dist-info → chainlit-2.0.4.dist-info}/WHEEL +0 -0
{chainlit-1.0.401.dist-info → chainlit-2.0.4.dist-info}/entry_points.txt +0 -0

chainlit/frontend/dist/index.html CHANGED Viewed

@@ -4,7 +4,6 @@
     <meta charset="UTF-8" />
     <meta name="viewport" content="width=device-width, initial-scale=1.0" />
     <!-- TAG INJECTION PLACEHOLDER -->
-    <link rel="icon" href="/favicon" />
     <link rel="preconnect" href="https://fonts.googleapis.com" />
     <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin />
     <!-- FONT START -->
@@ -22,11 +21,10 @@
     <script>
       const global = globalThis;
     </script>
-    <script type="module" crossorigin src="/assets/index-9711593e.js"></script>
-    <link rel="stylesheet" href="/assets/index-d088547c.css">
+    <script type="module" crossorigin src="/assets/index-QdmxtIMQ.js"></script>
+    <link rel="stylesheet" crossorigin href="/assets/index-DQmLRKyv.css">
   </head>
   <body>
     <div id="root"></div>
   </body>
 </html>

chainlit/haystack/callbacks.py CHANGED Viewed

@@ -1,14 +1,13 @@
 import re
 from typing import Any, Generic, List, Optional, TypeVar
-from chainlit.context import context
-from chainlit.step import Step
-from chainlit.sync import run_sync
 from haystack.agents import Agent, Tool
 from haystack.agents.agent_step import AgentStep
 from literalai.helper import utc_now
 from chainlit import Message
+from chainlit.step import Step
+from chainlit.sync import run_sync
 T = TypeVar("T")
@@ -68,9 +67,7 @@ class HaystackAgentCallbackHandler:
             self.last_tokens: List[str] = []
             self.answer_reached = False
-        root_message = context.session.root_message
-        parent_id = root_message.id if root_message else None
-        run_step = Step(name=self.agent_name, type="run", parent_id=parent_id)
+        run_step = Step(name=self.agent_name, type="run")
         run_step.start = utc_now()
         run_step.input = kwargs
@@ -135,7 +132,7 @@ class HaystackAgentCallbackHandler:
         tool_result: str,
         tool_name: Optional[str] = None,
         tool_input: Optional[str] = None,
-        **kwargs: Any
+        **kwargs: Any,
     ) -> None:
         # Tool finished, send step with tool_result
         tool_step = self.stack.pop()

chainlit/input_widget.py CHANGED Viewed

@@ -2,8 +2,10 @@ from abc import abstractmethod
 from collections import defaultdict
 from typing import Any, Dict, List, Optional
+from pydantic import Field
+from pydantic.dataclasses import dataclass
 from chainlit.types import InputWidgetType
-from pydantic.dataclasses import Field, dataclass
 @dataclass
@@ -75,7 +77,7 @@ class Select(InputWidget):
     initial: Optional[str] = None
     initial_index: Optional[int] = None
     initial_value: Optional[str] = None
-    values: List[str] = Field(default_factory=lambda: [])
+    values: List[str] = Field(default_factory=list)
     items: Dict[str, str] = Field(default_factory=lambda: defaultdict(dict))
     def __post_init__(
@@ -127,6 +129,7 @@ class TextInput(InputWidget):
     type: InputWidgetType = "textinput"
     initial: Optional[str] = None
     placeholder: Optional[str] = None
+    multiline: bool = False
     def to_dict(self) -> Dict[str, Any]:
         return {
@@ -137,6 +140,7 @@ class TextInput(InputWidget):
             "placeholder": self.placeholder,
             "tooltip": self.tooltip,
             "description": self.description,
+            "multiline": self.multiline,
         }
@@ -165,8 +169,8 @@ class Tags(InputWidget):
     """Useful to create an input for an array of strings."""
     type: InputWidgetType = "tags"
-    initial: List[str] = Field(default_factory=lambda: [])
-    values: List[str] = Field(default_factory=lambda: [])
+    initial: List[str] = Field(default_factory=list)
+    values: List[str] = Field(default_factory=list)
     def to_dict(self) -> Dict[str, Any]:
         return {

chainlit/langchain/callbacks.py CHANGED Viewed

@@ -1,19 +1,20 @@
 import json
 import time
-from typing import Any, Dict, List, Optional, TypedDict, Union
+from typing import Any, Dict, List, Optional, Sequence, Tuple, TypedDict, Union
 from uuid import UUID
-from chainlit.context import context_var
-from chainlit.message import Message
-from chainlit.step import Step
-from langchain.callbacks.tracers.base import BaseTracer
+import pydantic
 from langchain.callbacks.tracers.schemas import Run
 from langchain.schema import BaseMessage
-from langchain.schema.output import ChatGenerationChunk, GenerationChunk
 from langchain_core.outputs import ChatGenerationChunk, GenerationChunk
+from langchain_core.tracers.base import AsyncBaseTracer
 from literalai import ChatGeneration, CompletionGeneration, GenerationMessage
 from literalai.helper import utc_now
-from literalai.step import TrueStepType
+from literalai.observability.step import TrueStepType
+from chainlit.context import context_var
+from chainlit.message import Message
+from chainlit.step import Step
 DEFAULT_ANSWER_PREFIX_TOKENS = ["Final", "Answer", ":"]
@@ -123,6 +124,14 @@ class GenerationHelper:
                 key: self.ensure_values_serializable(value)
                 for key, value in data.items()
             }
+        elif isinstance(data, pydantic.BaseModel):
+            # Fallback to support pydantic v1
+            # https://docs.pydantic.dev/latest/migration/#changes-to-pydanticbasemodel
+            if pydantic.VERSION.startswith("1"):
+                return data.dict()
+            # pydantic v2
+            return data.model_dump()  # pyright: ignore reportAttributeAccessIssue
         elif isinstance(data, list):
             return [self.ensure_values_serializable(item) for item in data]
         elif isinstance(data, (str, int, float, bool, type(None))):
@@ -229,11 +238,28 @@ class GenerationHelper:
         return provider, model, tools, settings
-DEFAULT_TO_IGNORE = ["Runnable", "<lambda>"]
+def process_content(content: Any) -> Tuple[Dict, Optional[str]]:
+    if content is None:
+        return {}, None
+    if isinstance(content, dict):
+        return content, "json"
+    elif isinstance(content, str):
+        return {"content": content}, "text"
+    else:
+        return {"content": str(content)}, "text"
+DEFAULT_TO_IGNORE = [
+    "RunnableSequence",
+    "RunnableParallel",
+    "RunnableAssign",
+    "RunnableLambda",
+    "<lambda>",
+]
 DEFAULT_TO_KEEP = ["retriever", "llm", "agent", "chain", "tool"]
-class LangchainTracer(BaseTracer, GenerationHelper, FinalStreamHelper):
+class LangchainTracer(AsyncBaseTracer, GenerationHelper, FinalStreamHelper):
     steps: Dict[str, Step]
     parent_id_map: Dict[str, str]
     ignored_runs: set
@@ -252,7 +278,7 @@ class LangchainTracer(BaseTracer, GenerationHelper, FinalStreamHelper):
         to_keep: Optional[List[str]] = None,
         **kwargs: Any,
     ) -> None:
-        BaseTracer.__init__(self, **kwargs)
+        AsyncBaseTracer.__init__(self, **kwargs)
         GenerationHelper.__init__(self)
         FinalStreamHelper.__init__(
             self,
@@ -267,8 +293,6 @@ class LangchainTracer(BaseTracer, GenerationHelper, FinalStreamHelper):
         if self.context.current_step:
             self.root_parent_id = self.context.current_step.id
-        elif self.context.session.root_message:
-            self.root_parent_id = self.context.session.root_message.id
         else:
             self.root_parent_id = None
@@ -282,7 +306,7 @@ class LangchainTracer(BaseTracer, GenerationHelper, FinalStreamHelper):
         else:
             self.to_keep = to_keep
-    def on_chat_model_start(
+    async def on_chat_model_start(
         self,
         serialized: Dict[str, Any],
         messages: List[List[BaseMessage]],
@@ -291,8 +315,9 @@ class LangchainTracer(BaseTracer, GenerationHelper, FinalStreamHelper):
         parent_run_id: Optional["UUID"] = None,
         tags: Optional[List[str]] = None,
         metadata: Optional[Dict[str, Any]] = None,
+        name: Optional[str] = None,
         **kwargs: Any,
-    ) -> Any:
+    ) -> Run:
         lc_messages = messages[0]
         self.chat_generations[str(run_id)] = {
             "input_messages": lc_messages,
@@ -301,46 +326,48 @@ class LangchainTracer(BaseTracer, GenerationHelper, FinalStreamHelper):
             "tt_first_token": None,
         }
-        return super().on_chat_model_start(
+        return await super().on_chat_model_start(
             serialized,
             messages,
             run_id=run_id,
             parent_run_id=parent_run_id,
             tags=tags,
             metadata=metadata,
+            name=name,
             **kwargs,
         )
-    def on_llm_start(
+    async def on_llm_start(
         self,
         serialized: Dict[str, Any],
         prompts: List[str],
         *,
         run_id: "UUID",
+        parent_run_id: Optional[UUID] = None,
         tags: Optional[List[str]] = None,
-        parent_run_id: Optional["UUID"] = None,
         metadata: Optional[Dict[str, Any]] = None,
-        name: Optional[str] = None,
         **kwargs: Any,
-    ) -> Run:
-        self.completion_generations[str(run_id)] = {
-            "prompt": prompts[0],
-            "start": time.time(),
-            "token_count": 0,
-            "tt_first_token": None,
-        }
-        return super().on_llm_start(
+    ) -> None:
+        await super().on_llm_start(
             serialized,
             prompts,
             run_id=run_id,
             parent_run_id=parent_run_id,
             tags=tags,
             metadata=metadata,
-            name=name,
             **kwargs,
         )
-    def on_llm_new_token(
+        self.completion_generations[str(run_id)] = {
+            "prompt": prompts[0],
+            "start": time.time(),
+            "token_count": 0,
+            "tt_first_token": None,
+        }
+        return None
+    async def on_llm_new_token(
         self,
         token: str,
         *,
@@ -348,7 +375,14 @@ class LangchainTracer(BaseTracer, GenerationHelper, FinalStreamHelper):
         run_id: "UUID",
         parent_run_id: Optional["UUID"] = None,
         **kwargs: Any,
-    ) -> Run:
+    ) -> None:
+        await super().on_llm_new_token(
+            token=token,
+            chunk=chunk,
+            run_id=run_id,
+            parent_run_id=parent_run_id,
+            **kwargs,
+        )
         if isinstance(chunk, ChatGenerationChunk):
             start = self.chat_generations[str(run_id)]
         else:
@@ -363,24 +397,13 @@ class LangchainTracer(BaseTracer, GenerationHelper, FinalStreamHelper):
             if self.answer_reached:
                 if not self.final_stream:
                     self.final_stream = Message(content="")
-                    self._run_sync(self.final_stream.send())
-                self._run_sync(self.final_stream.stream_token(token))
+                    await self.final_stream.send()
+                await self.final_stream.stream_token(token)
                 self.has_streamed_final_answer = True
             else:
                 self.answer_reached = self._check_if_answer_reached()
-        return super().on_llm_new_token(
-            token,
-            chunk=chunk,
-            run_id=run_id,
-            parent_run_id=parent_run_id,
-        )
-    def _run_sync(self, co):  # TODO: WHAT TO DO WITH THIS?
-        context_var.set(self.context)
-        self.context.loop.create_task(co)
-    def _persist_run(self, run: Run) -> None:
+    async def _persist_run(self, run: Run) -> None:
         pass
     def _get_run_parent_id(self, run: Run):
@@ -431,11 +454,8 @@ class LangchainTracer(BaseTracer, GenerationHelper, FinalStreamHelper):
                 self.ignored_runs.add(str(run.id))
             return ignore, parent_id
-    def _is_annotable(self, run: Run):
-        return run.run_type in ["retriever", "llm"]
-    def _start_trace(self, run: Run) -> None:
-        super()._start_trace(run)
+    async def _start_trace(self, run: Run) -> None:
+        await super()._start_trace(run)
         context_var.set(self.context)
         ignore, parent_id = self._should_ignore_run(run)
@@ -448,40 +468,39 @@ class LangchainTracer(BaseTracer, GenerationHelper, FinalStreamHelper):
         if ignore:
             return
-        step_type: "TrueStepType" = "undefined"
+        step_type: TrueStepType = "undefined"
         if run.run_type == "agent":
             step_type = "run"
         elif run.run_type == "chain":
-            pass
+            if not self.steps:
+                step_type = "run"
         elif run.run_type == "llm":
             step_type = "llm"
         elif run.run_type == "retriever":
-            step_type = "retrieval"
+            step_type = "tool"
         elif run.run_type == "tool":
             step_type = "tool"
         elif run.run_type == "embedding":
             step_type = "embedding"
-        if not self.steps:
-            step_type = "run"
-        disable_feedback = not self._is_annotable(run)
         step = Step(
             id=str(run.id),
             name=run.name,
             type=step_type,
             parent_id=parent_id,
-            disable_feedback=disable_feedback,
         )
         step.start = utc_now()
-        step.input = run.inputs
+        step.input, language = process_content(run.inputs)
+        if language is not None:
+            if step.metadata is None:
+                step.metadata = {}
+            step.metadata["language"] = language
         self.steps[str(run.id)] = step
-        self._run_sync(step.send())
+        await step.send()
-    def _on_run_update(self, run: Run) -> None:
+    async def _on_run_update(self, run: Run) -> None:
         """Process a run upon update."""
         context_var.set(self.context)
@@ -499,6 +518,9 @@ class LangchainTracer(BaseTracer, GenerationHelper, FinalStreamHelper):
             generations = (run.outputs or {}).get("generations", [])
             generation = generations[0][0]
             variables = self.generation_inputs.get(str(run.parent_run_id), {})
+            variables = {
+                k: process_content(v) for k, v in variables.items() if v is not None
+            }
             if message := generation.get("message"):
                 chat_start = self.chat_generations[str(run.id)]
                 duration = time.time() - chat_start["start"]
@@ -529,11 +551,17 @@ class LangchainTracer(BaseTracer, GenerationHelper, FinalStreamHelper):
                             "prompt_id"
                         ]
                         if custom_variables := m.additional_kwargs.get("variables"):
-                            current_step.generation.variables = custom_variables
+                            current_step.generation.variables = {
+                                k: process_content(v)
+                                for k, v in custom_variables.items()
+                                if v is not None
+                            }
                     break
                 current_step.language = "json"
-                current_step.output = json.dumps(message_completion)
+                current_step.output = json.dumps(
+                    message_completion, indent=4, ensure_ascii=False
+                )
             else:
                 completion_start = self.completion_generations[str(run.id)]
                 completion = generation.get("text", "")
@@ -557,32 +585,39 @@ class LangchainTracer(BaseTracer, GenerationHelper, FinalStreamHelper):
             if current_step:
                 current_step.end = utc_now()
-                self._run_sync(current_step.update())
+                await current_step.update()
             if self.final_stream and self.has_streamed_final_answer:
-                self._run_sync(self.final_stream.update())
+                await self.final_stream.update()
             return
         outputs = run.outputs or {}
         output_keys = list(outputs.keys())
         output = outputs
         if output_keys:
             output = outputs.get(output_keys[0], outputs)
         if current_step:
-            current_step.output = output
+            current_step.output = (
+                output[0]
+                if isinstance(output, Sequence)
+                and not isinstance(output, str)
+                and len(output)
+                else output
+            )
             current_step.end = utc_now()
-            self._run_sync(current_step.update())
+            await current_step.update()
-    def _on_error(self, error: BaseException, *, run_id: UUID, **kwargs: Any):
+    async def _on_error(self, error: BaseException, *, run_id: UUID, **kwargs: Any):
         context_var.set(self.context)
         if current_step := self.steps.get(str(run_id), None):
             current_step.is_error = True
             current_step.output = str(error)
             current_step.end = utc_now()
-            self._run_sync(current_step.update())
+            await current_step.update()
     on_llm_error = _on_error
     on_chain_error = _on_error

chainlit/langflow/__init__.py CHANGED Viewed

@@ -8,6 +8,7 @@ if not check_module_version("langflow", "0.1.4"):
 from typing import Dict, Optional, Union
 import httpx
 from chainlit.telemetry import trace_event

chainlit/llama_index/callbacks.py CHANGED Viewed

@@ -1,20 +1,21 @@
 from typing import Any, Dict, List, Optional
-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.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
 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 1.0.401__py3-none-any.whl → 2.0.4__py3-none-any.whl

Potentially problematic release.

chainlit 1.0.401py3-none-any.whl → 2.0.4py3-none-any.whl