chainlit 0.4.0__py3-none-any.whl → 0.4.2__py3-none-any.whl

chainlit/client/utils.py

@@ -0,0 +1,23 @@
+from fastapi import HTTPException, Request
+
+from chainlit.config import config
+from chainlit.client.base import BaseClient
+from chainlit.client.local import LocalClient
+from chainlit.client.cloud import CloudClient
+
+
+async def get_client(request: Request) -> BaseClient:
+    auth_header = request.headers.get("Authorization")
+
+    db = config.project.database
+
+    if db == "local":
+        client = LocalClient()
+    elif db == "cloud":
+        client = CloudClient(config.project.id, auth_header)
+    elif db == "custom":
+        client = await config.code.client_factory()
+    else:
+        raise HTTPException(status_code=500, detail="Invalid database type")
+
+    return client

chainlit/config.py

@@ -1,15 +1,16 @@
-from typing import Optional, Any, Callable, List, Dict, TYPE_CHECKING
+from typing import Optional, Any, Callable, Literal, List, Dict, TYPE_CHECKING
 import os
 import sys
 import tomli
 from pydantic.dataclasses import dataclass
 from dataclasses_json import dataclass_json
-from importlib import machinery
+from importlib import util
 from chainlit.logger import logger
 from chainlit.version import __version__
 
 if TYPE_CHECKING:
     from chainlit.action import Action
+    from chainlit.client.base import BaseClient
 
 PACKAGE_ROOT = os.path.dirname(__file__)
 
@@ -26,10 +27,15 @@ DEFAULT_CONFIG_STR = f"""[project]
 public = true
 
 # The project ID (found on https://cloud.chainlit.io).
-# If provided, all the message data will be stored in the cloud.
-# The project ID is required when public is set to false.
+# The project ID is required when public is set to false or when using the cloud database.
 #id = ""
 
+# Uncomment if you want to persist the chats.
+# local will create a database in your .chainlit directory (requires node.js installed).
+# cloud will use the Chainlit cloud database.
+# custom will load use your custom client.
+# database = "local"
+
 # Whether to enable telemetry (default: true). No personal data is collected.
 enable_telemetry = true
 
@@ -102,6 +108,32 @@ class CodeSettings:
     lc_postprocess: Optional[Callable[[Any], str]] = None
     lc_factory: Optional[Callable[[], Any]] = None
     lc_rename: Optional[Callable[[str], str]] = None
+    client_factory: Optional[Callable[[str], "BaseClient"]] = None
+
+    def validate(self):
+        requires_one_of = [
+            "lc_factory",
+            "on_message",
+            "on_chat_start",
+        ]
+
+        mutually_exclusive = ["lc_factory"]
+
+        # Check if at least one of the required attributes is set
+        if not any(getattr(self, attr) for attr in requires_one_of):
+            raise ValueError(
+                f"Module should at least expose one of {', '.join(requires_one_of)} function"
+            )
+
+        # Check if any mutually exclusive attributes are set together
+        for i, attr1 in enumerate(mutually_exclusive):
+            for attr2 in mutually_exclusive[i + 1 :]:
+                if getattr(self, attr1) and getattr(self, attr2):
+                    raise ValueError(
+                        f"Module should not expose both {attr1} and {attr2} functions"
+                    )
+
+        return True
 
 
 @dataclass_json
@@ -111,12 +143,18 @@ class ProjectSettings:
     id: Optional[str] = None
     # Whether the app is available to anonymous users or only to team members.
     public: bool = True
+    # Storage type
+    database: Optional[Literal["local", "cloud", "custom"]] = None
     # Whether to enable telemetry. No personal data is collected.
     enable_telemetry: bool = True
     # List of environment variables to be provided by each user to use the app. If empty, no environment variables will be asked to the user.
     user_env: List[str] = None
     # Path to the local langchain cache database
     lc_cache_path: str = None
+    # Path to the local chat db
+    local_db_path: str = None
+    # Path to the local file system
+    local_fs_path: str = None
 
 
 @dataclass()
@@ -145,35 +183,26 @@ def init_config(log=False):
         logger.info(f"Config file already exists at {config_file}")
 
 
-def reset_module_config():
-    if not config:
-        return
-
-    config.code = CodeSettings(action_callbacks={})
-
-
 def load_module(target: str):
     """Load the specified module."""
 
-    # Reset the config fields that belonged to the previous module
-    reset_module_config()
-
     # Get the target's directory
     target_dir = os.path.dirname(os.path.abspath(target))
 
     # Add the target's directory to the Python path
     sys.path.insert(0, target_dir)
 
-    loader = machinery.SourceFileLoader(target, target)
-    config.code.module = loader.load_module()
+    spec = util.spec_from_file_location(target, target)
+    module = util.module_from_spec(spec)
+    spec.loader.exec_module(module)
 
     # Remove the target's directory from the Python path
     sys.path.pop(0)
 
+    config.code.validate()
 
-def load_config():
-    """Load the configuration from the config file."""
-    init_config()
+
+def load_settings():
     with open(config_file, "rb") as f:
         toml_dict = tomli.load(f)
         # Load project settings
@@ -187,23 +216,57 @@ def load_config():
             )
 
         lc_cache_path = os.path.join(config_dir, ".langchain.db")
+        local_db_path = os.path.join(config_dir, "chat.db")
+        local_fs_path = os.path.join(config_dir, "chat_files")
+
+        os.environ[
+            "LOCAL_DB_PATH"
+        ] = f"file:{local_db_path}?socket_timeout=10&connection_limit=1"
 
         project_settings = ProjectSettings(
-            lc_cache_path=lc_cache_path, **project_config
+            lc_cache_path=lc_cache_path,
+            local_db_path=local_db_path,
+            local_fs_path=local_fs_path,
+            **project_config,
         )
+
         ui_settings = UISettings(**ui_settings)
 
-        if not project_settings.public and not project_settings.project_id:
+        if not project_settings.public and not project_settings.id:
             raise ValueError("Project ID is required when public is set to false.")
 
-        config = ChainlitConfig(
-            chainlit_server=chainlit_server,
-            chainlit_prod_url=chainlit_prod_url,
-            ui=ui_settings,
-            run=RunSettings(),
-            project=project_settings,
-            code=CodeSettings(action_callbacks={}),
-        )
+        return {
+            "ui": ui_settings,
+            "project": project_settings,
+            "code": CodeSettings(action_callbacks={}),
+        }
+
+
+def reload_config():
+    """Reload the configuration from the config file."""
+    global config
+    if config is None:
+        return
+
+    settings = load_settings()
+
+    config.code = settings["code"]
+    config.ui = settings["ui"]
+    config.project = settings["project"]
+
+
+def load_config():
+    """Load the configuration from the config file."""
+    init_config()
+
+    settings = load_settings()
+
+    config = ChainlitConfig(
+        chainlit_server=chainlit_server,
+        chainlit_prod_url=chainlit_prod_url,
+        run=RunSettings(),
+        **settings,
+    )
 
     return config
 

chainlit/context.py

@@ -0,0 +1,29 @@
+import contextvars
+from typing import TYPE_CHECKING
+from asyncio import AbstractEventLoop
+
+if TYPE_CHECKING:
+    from chainlit.emitter import ChainlitEmitter
+
+
+class ChainlitContextException(Exception):
+    def __init__(self, msg="Chainlit context not found", *args, **kwargs):
+        super().__init__(msg, *args, **kwargs)
+
+
+emitter_var = contextvars.ContextVar("emitter")
+loop_var = contextvars.ContextVar("loop")
+
+
+def get_emitter() -> "ChainlitEmitter":
+    try:
+        return emitter_var.get()
+    except LookupError:
+        raise ChainlitContextException()
+
+
+def get_loop() -> AbstractEventLoop:
+    try:
+        return loop_var.get()
+    except LookupError:
+        raise ChainlitContextException()

chainlit/db/__init__.py

@@ -0,0 +1,35 @@
+import os
+from chainlit.logger import logger
+from chainlit.config import config, PACKAGE_ROOT
+
+SCHEMA_PATH = os.path.join(PACKAGE_ROOT, "db/prisma/schema.prisma")
+
+
+def db_push():
+    from prisma.cli.prisma import run
+    import prisma
+    from importlib import reload
+
+    args = ["db", "push", f"--schema={SCHEMA_PATH}"]
+    env = {"LOCAL_DB_PATH": os.environ.get("LOCAL_DB_PATH")}
+    run(args, env=env)
+
+    # Without this the client will fail to initialize the first time.
+    reload(prisma)
+
+
+def init_local_db():
+    use_local_db = config.project.database == "local"
+    if use_local_db:
+        if not os.path.exists(config.project.local_db_path):
+            db_push()
+
+
+def migrate_local_db():
+    use_local_db = config.project.database == "local"
+    if use_local_db:
+        if os.path.exists(config.project.local_db_path):
+            db_push()
+            logger.info(f"Local db migrated")
+    else:
+        logger.info(f"Local db does not exist, skipping migration")

chainlit/db/prisma/schema.prisma

@@ -0,0 +1,48 @@
+datasource db {
+  provider = "sqlite"
+  url      = env("LOCAL_DB_PATH")
+}
+
+generator client {
+  provider             = "prisma-client-py"
+  recursive_type_depth = 5
+}
+
+model Conversation {
+  id        Int       @id @default(autoincrement())
+  createdAt DateTime  @default(now())
+  messages  Message[]
+  elements  Element[]
+}
+
+model Element {
+  id             Int          @id @default(autoincrement())
+  createdAt      DateTime     @default(now())
+  conversationId Int
+  conversation   Conversation @relation(fields: [conversationId], references: [id], onDelete: Cascade)
+  type           String
+  url            String
+  name           String
+  display        String
+  size           String?
+  language       String?
+  forIds         String?
+}
+
+model Message {
+  id             Int          @id @default(autoincrement())
+  createdAt      DateTime     @default(now())
+  conversationId Int
+  conversation   Conversation @relation(fields: [conversationId], references: [id], onDelete: Cascade)
+  authorIsUser   Boolean      @default(false)
+  isError        Boolean      @default(false)
+  waitForAnswer  Boolean      @default(false)
+  indent         Int          @default(0)
+  author         String
+  content        String
+  humanFeedback  Int          @default(0)
+  language       String?
+  prompt         String?
+  // Sqlite does not support JSON
+  llmSettings    String?
+}

chainlit/element.py

@@ -1,22 +1,27 @@
 from pydantic.dataclasses import dataclass
-from dataclasses_json import dataclass_json
-from typing import Dict, Union, Any
+from typing import Dict, List, Union, Any
 import uuid
 import aiofiles
 from io import BytesIO
 
-from chainlit.emitter import get_emitter, BaseClient
+from chainlit.context import get_emitter
+from chainlit.client.base import BaseClient
 from chainlit.telemetry import trace_event
 from chainlit.types import ElementType, ElementDisplay, ElementSize
 
 type_to_mime = {
-    "image": "binary/octet-stream",
+    "image": "image/png",
     "text": "text/plain",
     "pdf": "application/pdf",
 }
 
+mime_to_ext = {
+    "image/png": "png",
+    "text/plain": "txt",
+    "application/pdf": "pdf",
+}
+
 
-@dataclass_json
 @dataclass
 class Element:
     # Name of the element, this will be used to reference the element in the UI.
@@ -34,20 +39,35 @@ class Element:
     # The ID of the element. This is set automatically when the element is sent to the UI if cloud is enabled.
     id: int = None
     # The ID of the element if cloud is disabled.
-    tempId: str = None
+    temp_id: str = None
     # The ID of the message this element is associated with.
-    forId: str = None
+    for_ids: List[str] = None
 
     def __post_init__(self) -> None:
         trace_event(f"init {self.__class__.__name__}")
         self.emitter = get_emitter()
-        if not self.emitter:
-            raise RuntimeError("Element should be instantiated in a Chainlit context")
+        self.for_ids = []
+        self.temp_id = str(uuid.uuid4())
 
         if not self.url and not self.path and not self.content:
             raise ValueError("Must provide url, path or content to instantiate element")
 
-        self.tempId = uuid.uuid4().hex
+    def to_dict(self) -> Dict:
+        _dict = {
+            "tempId": self.temp_id,
+            "type": self.type,
+            "url": self.url,
+            "name": self.name,
+            "display": self.display,
+            "size": getattr(self, "size", None),
+            "language": getattr(self, "language", None),
+            "forIds": getattr(self, "for_ids", None),
+        }
+
+        if self.id:
+            _dict["id"] = self.id
+
+        return _dict
 
     async def preprocess_content(self):
         pass
@@ -56,29 +76,15 @@ class Element:
         if self.path:
             async with aiofiles.open(self.path, "rb") as f:
                 self.content = await f.read()
-                await self.preprocess_content()
-        elif self.content:
-            await self.preprocess_content()
         else:
             raise ValueError("Must provide path or content to load element")
 
-    async def persist(self, client: BaseClient, for_id: str = None):
-        if not self.url and self.content:
+    async def persist(self, client: BaseClient):
+        if not self.url and self.content and not self.id:
             self.url = await client.upload_element(
                 content=self.content, mime=type_to_mime[self.type]
             )
-
-        size = getattr(self, "size", None)
-        language = getattr(self, "language", None)
-        element = await client.create_element(
-            name=self.name,
-            url=self.url,
-            type=self.type,
-            display=self.display,
-            size=size,
-            language=language,
-            for_id=for_id,
-        )
+        element = await client.upsert_element(self.to_dict())
         return element
 
     async def before_emit(self, element: Dict) -> Dict:
@@ -90,22 +96,34 @@ class Element:
         if not self.content and not self.url and self.path:
             await self.load()
 
-        # Cloud is enabled, upload the element to S3
-        if self.emitter.client and not self.id:
-            element = await self.persist(self.emitter.client, for_id)
-            self.id = element["id"]
+        await self.preprocess_content()
+
+        if for_id:
+            self.for_ids.append(for_id)
+
+        # We have a client, persist the element
+        if self.emitter.client:
+            element = await self.persist(self.emitter.client)
+            self.id = element and element.get("id")
 
         elif not self.url and not self.content:
             raise ValueError("Must provide url or content to send element")
 
         element = self.to_dict()
-        if for_id:
-            element["forId"] = for_id
+
+        element["content"] = self.content
 
         if self.emitter.emit and element:
-            trace_event(f"send {self.__class__.__name__}")
-            element = await self.before_emit(element)
-            await self.emitter.emit("element", element)
+            if len(self.for_ids) > 1:
+                trace_event(f"update {self.__class__.__name__}")
+                await self.emitter.emit(
+                    "update_element",
+                    {"id": self.id or self.temp_id, "forIds": self.for_ids},
+                )
+            else:
+                trace_event(f"send {self.__class__.__name__}")
+                element = await self.before_emit(element)
+                await self.emitter.emit("element", element)
 
 
 @dataclass
@@ -187,8 +205,3 @@ class Pyplot(Element):
         self.content = image.getvalue()
 
         super().__post_init__()
-
-    async def before_emit(self, element: Dict) -> Dict:
-        # Prevent the figure from being serialized
-        del element["figure"]
-        return element

chainlit/emitter.py

@@ -1,9 +1,8 @@
 from typing import Union, Dict
 from chainlit.session import Session
 from chainlit.types import AskSpec
-from chainlit.client import BaseClient
+from chainlit.client.base import BaseClient
 from socketio.exceptions import TimeoutError
-import inspect
 
 
 class ChainlitEmitter:
@@ -129,31 +128,3 @@ class ChainlitEmitter:
     def send_token(self, id: Union[str, int], token: str):
         """Send a message token to the UI."""
         return self.emit("stream_token", {"id": id, "token": token})
-
-
-def get_emitter() -> Union[ChainlitEmitter, None]:
-    """
-    Get the Chainlit Emitter instance from the current call stack.
-    This unusual approach is necessary because:
-     - we need to get the right Emitter instance with the right websocket connection
-     - to preserve a lean developer experience, we do not pass the Emitter instance to every function call
-
-    What happens is that we set __chainlit_emitter__ in the local variables when we receive a websocket message.
-    Then we can retrieve it from the call stack when we need it, even if the developer's code has no idea about it.
-    """
-    attr = "__chainlit_emitter__"
-    candidates = [i[0].f_locals.get(attr) for i in inspect.stack()]
-    emitter = None
-    for candidate in candidates:
-        if candidate:
-            emitter = candidate
-            break
-
-    return emitter
-
-
-def get_emit_fn():
-    emitter = get_emitter()
-    if emitter:
-        return emitter.emit
-    return None