chainlit 0.4.1__py3-none-any.whl → 0.4.3__py3-none-any.whl

chainlit/client/local.py

@@ -0,0 +1,257 @@
+from typing import Optional, Dict
+import uuid
+import json
+import os
+
+import asyncio
+import aiofiles
+
+from chainlit.client.base import PaginatedResponse, PageInfo
+
+from .base import BaseClient
+
+from chainlit.logger import logger
+from chainlit.config import config
+from chainlit.element import mime_to_ext
+
+
+class LocalClient(BaseClient):
+    conversation_id: Optional[str] = None
+    lock: asyncio.Lock
+
+    def __init__(self):
+        self.lock = asyncio.Lock()
+
+    def before_write(self, variables: Dict):
+        if "llmSettings" in variables:
+            # Sqlite doesn't support json fields, so we need to serialize it.
+            variables["llmSettings"] = json.dumps(variables["llmSettings"])
+
+        if "forIds" in variables:
+            # Sqlite doesn't support list of primitives, so we need to serialize it.
+            variables["forIds"] = json.dumps(variables["forIds"])
+
+        if "tempId" in variables:
+            del variables["tempId"]
+
+    def after_read(self, variables: Dict):
+        if "llmSettings" in variables:
+            # Sqlite doesn't support json fields, so we need to parse it.
+            variables["llmSettings"] = json.loads(variables["llmSettings"])
+
+    async def is_project_member(self):
+        return True
+
+    async def get_member_role(self):
+        return "OWNER"
+
+    async def get_project_members(self):
+        return []
+
+    async def get_conversation_id(self):
+        self.conversation_id = await self.create_conversation()
+
+        return self.conversation_id
+
+    async def create_conversation(self):
+        from prisma.models import Conversation
+
+        # If we run multiple send concurrently, we need to make sure we don't create multiple conversations.
+        async with self.lock:
+            if self.conversation_id:
+                return self.conversation_id
+
+            res = await Conversation.prisma().create(data={})
+
+            return res.id
+
+    async def delete_conversation(self, conversation_id):
+        from prisma.models import Conversation
+
+        await Conversation.prisma().delete(where={"id": conversation_id})
+
+        return True
+
+    async def get_conversation(self, conversation_id: int):
+        from prisma.models import Conversation
+
+        c = await Conversation.prisma().find_unique_or_raise(
+            where={"id": conversation_id}, include={"messages": True, "elements": True}
+        )
+
+        for m in c.messages:
+            if m.llmSettings:
+                m.llmSettings = json.loads(m.llmSettings)
+
+        for e in c.elements:
+            if e.forIds:
+                e.forIds = json.loads(e.forIds)
+
+        return json.loads(c.json())
+
+    async def get_conversations(self, pagination, filter):
+        from prisma.models import Conversation
+
+        some_messages = {}
+
+        if filter.feedback is not None:
+            some_messages["humanFeedback"] = filter.feedback
+
+        if filter.search is not None:
+            some_messages["content"] = {"contains": filter.search or None}
+
+        if pagination.cursor:
+            cursor = {"id": pagination.cursor}
+        else:
+            cursor = None
+
+        conversations = await Conversation.prisma().find_many(
+            take=pagination.first,
+            skip=1 if pagination.cursor else None,
+            cursor=cursor,
+            include={
+                "messages": {
+                    "take": 1,
+                    "where": {
+                        "authorIsUser": True,
+                    },
+                    "orderBy": [
+                        {
+                            "createdAt": "asc",
+                        }
+                    ],
+                }
+            },
+            where={"messages": {"some": some_messages}},
+            order={
+                "createdAt": "desc",
+            },
+        )
+
+        has_more = len(conversations) == pagination.first
+
+        if has_more:
+            end_cursor = conversations[-1].id
+        else:
+            end_cursor = None
+
+        conversations = [json.loads(c.json()) for c in conversations]
+
+        return PaginatedResponse(
+            pageInfo=PageInfo(hasNextPage=has_more, endCursor=end_cursor),
+            data=conversations,
+        )
+
+    async def create_message(self, variables):
+        from prisma.models import Message
+
+        c_id = await self.get_conversation_id()
+
+        if not c_id:
+            logger.warning("Missing conversation ID, could not persist the message.")
+            return None
+
+        variables = variables.copy()
+
+        variables["conversationId"] = c_id
+
+        self.before_write(variables)
+
+        res = await Message.prisma().create(data=variables)
+        return res.id
+
+    async def get_message(self, message_id):
+        from prisma.models import Message
+
+        res = await Message.prisma().find_first(where={"id": message_id})
+        res = res.dict()
+        self.after_read(res)
+        return res
+
+    async def update_message(self, message_id, variables):
+        from prisma.models import Message
+
+        variables = variables.copy()
+
+        self.before_write(variables)
+
+        await Message.prisma().update(data=variables, where={"id": message_id})
+
+        return True
+
+    async def delete_message(self, message_id):
+        from prisma.models import Message
+
+        await Message.prisma().delete(where={"id": message_id})
+
+        return True
+
+    async def upsert_element(
+        self,
+        variables,
+    ):
+        from prisma.models import Element
+
+        c_id = await self.get_conversation_id()
+
+        if not c_id:
+            logger.warning("Missing conversation ID, could not persist the element.")
+            return None
+
+        variables["conversationId"] = c_id
+
+        self.before_write(variables)
+
+        if "id" in variables:
+            res = await Element.prisma().update(
+                data=variables, where={"id": variables.get("id")}
+            )
+        else:
+            res = await Element.prisma().create(data=variables)
+
+        return res.dict()
+
+    async def get_element(
+        self,
+        conversation_id,
+        element_id,
+    ):
+        from prisma.models import Element
+
+        res = await Element.prisma().find_unique_or_raise(where={"id": element_id})
+        return json.loads(res.json())
+
+    async def upload_element(self, content: bytes, mime: str):
+        c_id = await self.get_conversation_id()
+
+        if not c_id:
+            logger.warning("Missing conversation ID, could not persist the message.")
+            return None
+
+        file_ext = mime_to_ext.get(mime, "bin")
+        file_name = f"{uuid.uuid4()}.{file_ext}"
+
+        sub_path = os.path.join(str(c_id), file_name)
+        full_path = os.path.join(config.project.local_fs_path, sub_path)
+
+        if not os.path.exists(os.path.dirname(full_path)):
+            os.makedirs(os.path.dirname(full_path))
+
+        async with aiofiles.open(full_path, "wb") as out:
+            await out.write(content)
+            await out.flush()
+
+            url = f"/files/{sub_path}"
+            return url
+
+    async def set_human_feedback(self, message_id, feedback):
+        from prisma.models import Message
+
+        await Message.prisma().update(
+            where={"id": message_id},
+            data={
+                "humanFeedback": feedback,
+            },
+        )
+
+        return True

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,4 +1,4 @@
-from typing import Optional, Any, Callable, List, Dict, TYPE_CHECKING
+from typing import Optional, Any, Callable, Union, Literal, List, Dict, TYPE_CHECKING
 import os
 import sys
 import tomli
@@ -10,6 +10,7 @@ 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,15 +108,19 @@ class CodeSettings:
     lc_postprocess: Optional[Callable[[Any], str]] = None
     lc_factory: Optional[Callable[[], Any]] = None
     lc_rename: Optional[Callable[[str], str]] = None
+    llama_index_factory: Optional[Callable[[], Any]] = None
+    langflow_schema: Union[Dict, str] = None
+    client_factory: Optional[Callable[[str], "BaseClient"]] = None
 
     def validate(self):
         requires_one_of = [
             "lc_factory",
+            "llama_index_factory",
             "on_message",
             "on_chat_start",
         ]
 
-        mutually_exclusive = ["lc_factory"]
+        mutually_exclusive = ["lc_factory", "llama_index_factory"]
 
         # Check if at least one of the required attributes is set
         if not any(getattr(self, attr) for attr in requires_one_of):
@@ -136,12 +146,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()
@@ -203,10 +219,20 @@ def load_settings():
             )
 
         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.id:

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