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

chainlit/data/__init__.py

@@ -1,428 +1,113 @@
-import functools
-import json
 import os
-from collections import deque
-from typing import TYPE_CHECKING, Dict, List, Optional, Union
+import warnings
+from typing import Optional
 
-import aiofiles
-from chainlit.config import config
-from chainlit.context import context
-from chainlit.logger import logger
-from chainlit.session import WebsocketSession
-from chainlit.types import Feedback, Pagination, ThreadDict, ThreadFilter
-from chainlit.user import PersistedUser, User, UserDict
-from literalai import Attachment
-from literalai import Feedback as ClientFeedback
-from literalai import PageInfo, PaginatedResponse
-from literalai import Step as ClientStep
-from literalai.step import StepDict as ClientStepDict
-from literalai.thread import NumberListFilter, StringFilter, StringListFilter
-from literalai.thread import ThreadFilter as ClientThreadFilter
+from .base import BaseDataLayer
+from .utils import (
+    queue_until_user_message as queue_until_user_message,  # TODO: Consider deprecating re-export.; Redundant alias tells type checkers to STFU.
+)
 
-if TYPE_CHECKING:
-    from chainlit.element import Element, ElementDict
-    from chainlit.step import FeedbackDict, StepDict
+_data_layer: Optional[BaseDataLayer] = None
+_data_layer_initialized = False
 
-_data_layer = None
 
+def get_data_layer():
+    global _data_layer, _data_layer_initialized
 
-def queue_until_user_message():
-    def decorator(method):
-        @functools.wraps(method)
-        async def wrapper(self, *args, **kwargs):
-            if (
-                isinstance(context.session, WebsocketSession)
-                and not context.session.has_first_interaction
-            ):
-                # Queue the method invocation waiting for the first user message
-                queues = context.session.thread_queues
-                method_name = method.__name__
-                if method_name not in queues:
-                    queues[method_name] = deque()
-                queues[method_name].append((method, self, args, kwargs))
-
-            else:
-                # Otherwise, Execute the method immediately
-                return await method(self, *args, **kwargs)
-
-        return wrapper
-
-    return decorator
-
-
-class BaseDataLayer:
-    """Base class for data persistence."""
-
-    async def get_user(self, identifier: str) -> Optional["PersistedUser"]:
-        return None
-
-    async def create_user(self, user: "User") -> Optional["PersistedUser"]:
-        pass
-
-    async def upsert_feedback(
-        self,
-        feedback: Feedback,
-    ) -> str:
-        return ""
-
-    @queue_until_user_message()
-    async def create_element(self, element_dict: "ElementDict"):
-        pass
-
-    async def get_element(
-        self, thread_id: str, element_id: str
-    ) -> Optional["ElementDict"]:
-        pass
-
-    @queue_until_user_message()
-    async def delete_element(self, element_id: str):
-        pass
-
-    @queue_until_user_message()
-    async def create_step(self, step_dict: "StepDict"):
-        pass
-
-    @queue_until_user_message()
-    async def update_step(self, step_dict: "StepDict"):
-        pass
-
-    @queue_until_user_message()
-    async def delete_step(self, step_id: str):
-        pass
-
-    async def get_thread_author(self, thread_id: str) -> str:
-        return ""
-
-    async def delete_thread(self, thread_id: str):
-        pass
-
-    async def list_threads(
-        self, pagination: "Pagination", filters: "ThreadFilter"
-    ) -> "PaginatedResponse[ThreadDict]":
-        return PaginatedResponse(
-            data=[], pageInfo=PageInfo(hasNextPage=False, endCursor=None)
-        )
-
-    async def get_thread(self, thread_id: str) -> "Optional[ThreadDict]":
-        return None
-
-    async def update_thread(
-        self,
-        thread_id: str,
-        name: Optional[str] = None,
-        user_id: Optional[str] = None,
-        metadata: Optional[Dict] = None,
-        tags: Optional[List[str]] = None,
-    ):
-        pass
-
-    async def delete_user_session(self, id: str) -> bool:
-        return True
-
-
-class ChainlitDataLayer:
-    def __init__(self, api_key: str, server: Optional[str]):
-        from literalai import LiteralClient
-
-        self.client = LiteralClient(api_key=api_key, url=server)
-        logger.info("Chainlit data layer initialized")
-
-    def attachment_to_element_dict(self, attachment: Attachment) -> "ElementDict":
-        metadata = attachment.metadata or {}
-        return {
-            "chainlitKey": None,
-            "display": metadata.get("display", "side"),
-            "language": metadata.get("language"),
-            "page": metadata.get("page"),
-            "size": metadata.get("size"),
-            "type": metadata.get("type", "file"),
-            "forId": attachment.step_id,
-            "id": attachment.id or "",
-            "mime": attachment.mime,
-            "name": attachment.name or "",
-            "objectKey": attachment.object_key,
-            "url": attachment.url,
-            "threadId": attachment.thread_id,
-        }
-
-    def feedback_to_feedback_dict(
-        self, feedback: Optional[ClientFeedback]
-    ) -> "Optional[FeedbackDict]":
-        if not feedback:
-            return None
-        return {
-            "id": feedback.id or "",
-            "forId": feedback.step_id or "",
-            "value": feedback.value or 0,  # type: ignore
-            "comment": feedback.comment,
-            "strategy": "BINARY",
-        }
-
-    def step_to_step_dict(self, step: ClientStep) -> "StepDict":
-        metadata = step.metadata or {}
-        input = (step.input or {}).get("content") or (
-            json.dumps(step.input) if step.input and step.input != {} else ""
-        )
-        output = (step.output or {}).get("content") or (
-            json.dumps(step.output) if step.output and step.output != {} else ""
-        )
-        return {
-            "createdAt": step.created_at,
-            "id": step.id or "",
-            "threadId": step.thread_id or "",
-            "parentId": step.parent_id,
-            "feedback": self.feedback_to_feedback_dict(step.feedback),
-            "start": step.start_time,
-            "end": step.end_time,
-            "type": step.type or "undefined",
-            "name": step.name or "",
-            "generation": step.generation.to_dict() if step.generation else None,
-            "input": input,
-            "output": output,
-            "showInput": metadata.get("showInput", False),
-            "disableFeedback": metadata.get("disableFeedback", False),
-            "indent": metadata.get("indent"),
-            "language": metadata.get("language"),
-            "isError": metadata.get("isError", False),
-            "waitForAnswer": metadata.get("waitForAnswer", False),
-            "feedback": self.feedback_to_feedback_dict(step.feedback),
-        }
-
-    async def get_user(self, identifier: str) -> Optional[PersistedUser]:
-        user = await self.client.api.get_user(identifier=identifier)
-        if not user:
-            return None
-        return PersistedUser(
-            id=user.id or "",
-            identifier=user.identifier or "",
-            metadata=user.metadata,
-            createdAt=user.created_at or "",
-        )
+    if not _data_layer_initialized:
+        if _data_layer:
+            # Data layer manually set, warn user that this is deprecated.
 
-    async def create_user(self, user: User) -> Optional[PersistedUser]:
-        _user = await self.client.api.get_user(identifier=user.identifier)
-        if not _user:
-            _user = await self.client.api.create_user(
-                identifier=user.identifier, metadata=user.metadata
+            warnings.warn(
+                "Setting data layer manually is deprecated. Use @data_layer instead.",
+                DeprecationWarning,
             )
-        elif _user.id:
-            await self.client.api.update_user(id=_user.id, metadata=user.metadata)
-        return PersistedUser(
-            id=_user.id or "",
-            identifier=_user.identifier or "",
-            metadata=_user.metadata,
-            createdAt=_user.created_at or "",
-        )
 
-    async def upsert_feedback(
-        self,
-        feedback: Feedback,
-    ):
-        if feedback.id:
-            await self.client.api.update_feedback(
-                id=feedback.id,
-                update_params={
-                    "comment": feedback.comment,
-                    "strategy": feedback.strategy,
-                    "value": feedback.value,
-                },
-            )
-            return feedback.id
         else:
-            created = await self.client.api.create_feedback(
-                step_id=feedback.forId,
-                value=feedback.value,
-                comment=feedback.comment,
-                strategy=feedback.strategy,
-            )
-            return created.id or ""
-
-    @queue_until_user_message()
-    async def create_element(self, element: "Element"):
-        metadata = {
-            "size": element.size,
-            "language": element.language,
-            "display": element.display,
-            "type": element.type,
-            "page": getattr(element, "page", None),
-        }
-
-        if not element.for_id:
-            return
-
-        object_key = None
-
-        if not element.url:
-            if element.path:
-                async with aiofiles.open(element.path, "rb") as f:
-                    content = await f.read()  # type: Union[bytes, str]
-            elif element.content:
-                content = element.content
-            else:
-                raise ValueError("Either path or content must be provided")
-            uploaded = await self.client.api.upload_file(
-                content=content, mime=element.mime, thread_id=element.thread_id
-            )
-            object_key = uploaded["object_key"]
-
-        await self.client.api.send_steps(
-            [
-                {
-                    "id": element.for_id,
-                    "threadId": element.thread_id,
-                    "attachments": [
-                        {
-                            "id": element.id,
-                            "name": element.name,
-                            "metadata": metadata,
-                            "mime": element.mime,
-                            "url": element.url,
-                            "objectKey": object_key,
-                        }
-                    ],
-                }
-            ]
-        )
-
-    async def get_element(
-        self, thread_id: str, element_id: str
-    ) -> Optional["ElementDict"]:
-        attachment = await self.client.api.get_attachment(id=element_id)
-        if not attachment:
-            return None
-        return self.attachment_to_element_dict(attachment)
-
-    @queue_until_user_message()
-    async def delete_element(self, element_id: str):
-        await self.client.api.delete_attachment(id=element_id)
-
-    @queue_until_user_message()
-    async def create_step(self, step_dict: "StepDict"):
-        metadata = {
-            "disableFeedback": step_dict.get("disableFeedback"),
-            "isError": step_dict.get("isError"),
-            "waitForAnswer": step_dict.get("waitForAnswer"),
-            "language": step_dict.get("language"),
-            "showInput": step_dict.get("showInput"),
-        }
-
-        step: ClientStepDict = {
-            "createdAt": step_dict.get("createdAt"),
-            "startTime": step_dict.get("start"),
-            "endTime": step_dict.get("end"),
-            "generation": step_dict.get("generation"),
-            "id": step_dict.get("id"),
-            "parentId": step_dict.get("parentId"),
-            "name": step_dict.get("name"),
-            "threadId": step_dict.get("threadId"),
-            "type": step_dict.get("type"),
-            "metadata": metadata,
-        }
-        if step_dict.get("input"):
-            step["input"] = {"content": step_dict.get("input")}
-        if step_dict.get("output"):
-            step["output"] = {"content": step_dict.get("output")}
-
-        await self.client.api.send_steps([step])
-
-    @queue_until_user_message()
-    async def update_step(self, step_dict: "StepDict"):
-        await self.create_step(step_dict)
-
-    @queue_until_user_message()
-    async def delete_step(self, step_id: str):
-        await self.client.api.delete_step(id=step_id)
+            from chainlit.config import config
+
+            if config.code.data_layer:
+                # When @data_layer is configured, call it to get data layer.
+                _data_layer = config.code.data_layer()
+            elif database_url := os.environ.get("DATABASE_URL"):
+                # Default to Chainlit data layer if DATABASE_URL specified.
+                from .chainlit_data_layer import ChainlitDataLayer
+
+                if os.environ.get("LITERAL_API_KEY"):
+                    warnings.warn(
+                        "Both LITERAL_API_KEY and DATABASE_URL specified. Ignoring Literal AI data layer and relying on data layer pointing to DATABASE_URL."
+                    )
+                bucket_name = os.environ.get("BUCKET_NAME")
+
+                # AWS S3
+                aws_region = os.getenv("APP_AWS_REGION")
+                aws_access_key = os.getenv("APP_AWS_ACCESS_KEY")
+                aws_secret_key = os.getenv("APP_AWS_SECRET_KEY")
+                dev_aws_endpoint = os.getenv("DEV_AWS_ENDPOINT")
+                is_using_s3 = bool(aws_access_key and aws_secret_key and aws_region)
+
+                # Google Cloud Storage
+                gcs_project_id = os.getenv("APP_GCS_PROJECT_ID")
+                gcs_client_email = os.getenv("APP_GCS_CLIENT_EMAIL")
+                gcs_private_key = os.getenv("APP_GCS_PRIVATE_KEY")
+                is_using_gcs = bool(
+                    gcs_project_id and gcs_client_email and gcs_private_key
+                )
+
+                # Azure Storage
+                azure_storage_account = os.getenv("APP_AZURE_STORAGE_ACCOUNT")
+                azure_storage_key = os.getenv("APP_AZURE_STORAGE_ACCESS_KEY")
+                is_using_azure = bool(azure_storage_account and azure_storage_key)
+
+                storage_client = None
+
+                if sum([is_using_s3, is_using_gcs, is_using_azure]) > 1:
+                    warnings.warn(
+                        "Multiple storage configurations detected. Please use only one."
+                    )
+                elif is_using_s3:
+                    from chainlit.data.storage_clients.s3 import S3StorageClient
+
+                    storage_client = S3StorageClient(
+                        bucket=bucket_name,
+                        region_name=aws_region,
+                        aws_access_key_id=aws_access_key,
+                        aws_secret_access_key=aws_secret_key,
+                        endpoint_url=dev_aws_endpoint,
+                    )
+                elif is_using_gcs:
+                    from chainlit.data.storage_clients.gcs import GCSStorageClient
+
+                    storage_client = GCSStorageClient(
+                        project_id=gcs_project_id,
+                        client_email=gcs_client_email,
+                        private_key=gcs_private_key,
+                        bucket_name=bucket_name,
+                    )
+                elif is_using_azure:
+                    from chainlit.data.storage_clients.azure_blob import (
+                        AzureBlobStorageClient,
+                    )
+
+                    storage_client = AzureBlobStorageClient(
+                        container_name=bucket_name,
+                        storage_account=azure_storage_account,
+                        storage_key=azure_storage_key,
+                    )
+
+                _data_layer = ChainlitDataLayer(
+                    database_url=database_url, storage_client=storage_client
+                )
+            elif api_key := os.environ.get("LITERAL_API_KEY"):
+                # When LITERAL_API_KEY is defined, use Literal AI data layer
+                from .literalai import LiteralDataLayer
+
+                # support legacy LITERAL_SERVER variable as fallback
+                server = os.environ.get("LITERAL_API_URL") or os.environ.get(
+                    "LITERAL_SERVER"
+                )
+                _data_layer = LiteralDataLayer(api_key=api_key, server=server)
+
+        _data_layer_initialized = True
 
-    async def get_thread_author(self, thread_id: str) -> str:
-        thread = await self.get_thread(thread_id)
-        if not thread:
-            return ""
-        user = thread.get("user")
-        if not user:
-            return ""
-        return user.get("identifier") or ""
-
-    async def delete_thread(self, thread_id: str):
-        await self.client.api.delete_thread(id=thread_id)
-
-    async def list_threads(
-        self, pagination: "Pagination", filters: "ThreadFilter"
-    ) -> "PaginatedResponse[ThreadDict]":
-        if not filters.userIdentifier:
-            raise ValueError("userIdentifier is required")
-
-        client_filters = ClientThreadFilter(
-            participantsIdentifier=StringListFilter(
-                operator="in", value=[filters.userIdentifier]
-            ),
-        )
-        if filters.search:
-            client_filters.search = StringFilter(operator="ilike", value=filters.search)
-        if filters.feedback:
-            client_filters.feedbacksValue = NumberListFilter(
-                operator="in", value=[filters.feedback]
-            )
-        return await self.client.api.list_threads(
-            first=pagination.first, after=pagination.cursor, filters=client_filters
-        )
-
-    async def get_thread(self, thread_id: str) -> "Optional[ThreadDict]":
-        thread = await self.client.api.get_thread(id=thread_id)
-        if not thread:
-            return None
-        elements = []  # List[ElementDict]
-        steps = []  # List[StepDict]
-        if thread.steps:
-            for step in thread.steps:
-                if config.ui.hide_cot and step.parent_id:
-                    continue
-                for attachment in step.attachments:
-                    elements.append(self.attachment_to_element_dict(attachment))
-                if not config.features.prompt_playground and step.generation:
-                    step.generation = None
-                steps.append(self.step_to_step_dict(step))
-
-        user = None  # type: Optional["UserDict"]
-
-        if thread.user:
-            user = {
-                "id": thread.user.id or "",
-                "identifier": thread.user.identifier or "",
-                "metadata": thread.user.metadata,
-            }
-
-        return {
-            "createdAt": thread.created_at or "",
-            "id": thread.id,
-            "name": thread.name or None,
-            "steps": steps,
-            "elements": elements,
-            "metadata": thread.metadata,
-            "user": user,
-            "tags": thread.tags,
-        }
-
-    async def update_thread(
-        self,
-        thread_id: str,
-        name: Optional[str] = None,
-        user_id: Optional[str] = None,
-        metadata: Optional[Dict] = None,
-        tags: Optional[List[str]] = None,
-    ):
-        await self.client.api.upsert_thread(
-            thread_id=thread_id,
-            name=name,
-            participant_id=user_id,
-            metadata=metadata,
-            tags=tags,
-        )
-
-
-if api_key := os.environ.get("LITERAL_API_KEY"):
-    server = os.environ.get("LITERAL_SERVER")
-    _data_layer = ChainlitDataLayer(api_key=api_key, server=server)
-
-
-def get_data_layer():
     return _data_layer

chainlit/data/acl.py

@@ -1,14 +1,18 @@
-from chainlit.data import get_data_layer
 from fastapi import HTTPException
 
+from chainlit.data import get_data_layer
+
 
 async def is_thread_author(username: str, thread_id: str):
     data_layer = get_data_layer()
     if not data_layer:
-        raise HTTPException(status_code=401, detail="Unauthorized")
+        raise HTTPException(status_code=400, detail="Data layer not initialized")
 
     thread_author = await data_layer.get_thread_author(thread_id)
 
+    if not thread_author:
+        raise HTTPException(status_code=404, detail="Thread not found")
+
     if thread_author != username:
         raise HTTPException(status_code=401, detail="Unauthorized")
     else:

chainlit/data/base.py

@@ -0,0 +1,107 @@
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Dict, List, Optional
+
+from chainlit.types import (
+    Feedback,
+    PaginatedResponse,
+    Pagination,
+    ThreadDict,
+    ThreadFilter,
+)
+
+from .utils import queue_until_user_message
+
+if TYPE_CHECKING:
+    from chainlit.element import Element, ElementDict
+    from chainlit.step import StepDict
+    from chainlit.user import PersistedUser, User
+
+
+class BaseDataLayer(ABC):
+    """Base class for data persistence."""
+
+    @abstractmethod
+    async def get_user(self, identifier: str) -> Optional["PersistedUser"]:
+        pass
+
+    @abstractmethod
+    async def create_user(self, user: "User") -> Optional["PersistedUser"]:
+        pass
+
+    @abstractmethod
+    async def delete_feedback(
+        self,
+        feedback_id: str,
+    ) -> bool:
+        pass
+
+    @abstractmethod
+    async def upsert_feedback(
+        self,
+        feedback: Feedback,
+    ) -> str:
+        pass
+
+    @queue_until_user_message()
+    @abstractmethod
+    async def create_element(self, element: "Element"):
+        pass
+
+    @abstractmethod
+    async def get_element(
+        self, thread_id: str, element_id: str
+    ) -> Optional["ElementDict"]:
+        pass
+
+    @queue_until_user_message()
+    @abstractmethod
+    async def delete_element(self, element_id: str, thread_id: Optional[str] = None):
+        pass
+
+    @queue_until_user_message()
+    @abstractmethod
+    async def create_step(self, step_dict: "StepDict"):
+        pass
+
+    @queue_until_user_message()
+    @abstractmethod
+    async def update_step(self, step_dict: "StepDict"):
+        pass
+
+    @queue_until_user_message()
+    @abstractmethod
+    async def delete_step(self, step_id: str):
+        pass
+
+    @abstractmethod
+    async def get_thread_author(self, thread_id: str) -> str:
+        return ""
+
+    @abstractmethod
+    async def delete_thread(self, thread_id: str):
+        pass
+
+    @abstractmethod
+    async def list_threads(
+        self, pagination: "Pagination", filters: "ThreadFilter"
+    ) -> "PaginatedResponse[ThreadDict]":
+        pass
+
+    @abstractmethod
+    async def get_thread(self, thread_id: str) -> "Optional[ThreadDict]":
+        pass
+
+    @abstractmethod
+    async def update_thread(
+        self,
+        thread_id: str,
+        name: Optional[str] = None,
+        user_id: Optional[str] = None,
+        metadata: Optional[Dict] = None,
+        tags: Optional[List[str]] = None,
+    ):
+        pass
+
+    @abstractmethod
+    async def build_debug_url(self) -> str:
+        pass