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

chainlit/data/storage_clients/azure.py

@@ -0,0 +1,81 @@
+from typing import TYPE_CHECKING, Any, Dict, Optional, Union
+
+from azure.storage.filedatalake import (
+    ContentSettings,
+    DataLakeFileClient,
+    DataLakeServiceClient,
+    FileSystemClient,
+)
+
+from chainlit.data.storage_clients.base import BaseStorageClient
+from chainlit.logger import logger
+
+if TYPE_CHECKING:
+    from azure.core.credentials import (
+        AzureNamedKeyCredential,
+        AzureSasCredential,
+        TokenCredential,
+    )
+
+
+class AzureStorageClient(BaseStorageClient):
+    """
+    Class to enable Azure Data Lake Storage (ADLS) Gen2
+
+    parms:
+        account_url: "https://<your_account>.dfs.core.windows.net"
+        credential: Access credential (AzureKeyCredential)
+        sas_token: Optionally include SAS token to append to urls
+    """
+
+    def __init__(
+        self,
+        account_url: str,
+        container: str,
+        credential: Optional[
+            Union[
+                str,
+                Dict[str, str],
+                "AzureNamedKeyCredential",
+                "AzureSasCredential",
+                "TokenCredential",
+            ]
+        ],
+        sas_token: Optional[str] = None,
+    ):
+        try:
+            self.data_lake_client = DataLakeServiceClient(
+                account_url=account_url, credential=credential
+            )
+            self.container_client: FileSystemClient = (
+                self.data_lake_client.get_file_system_client(file_system=container)
+            )
+            self.sas_token = sas_token
+            logger.info("AzureStorageClient initialized")
+        except Exception as e:
+            logger.warn(f"AzureStorageClient initialization error: {e}")
+
+    async def upload_file(
+        self,
+        object_key: str,
+        data: Union[bytes, str],
+        mime: str = "application/octet-stream",
+        overwrite: bool = True,
+    ) -> Dict[str, Any]:
+        try:
+            file_client: DataLakeFileClient = self.container_client.get_file_client(
+                object_key
+            )
+            content_settings = ContentSettings(content_type=mime)
+            file_client.upload_data(
+                data, overwrite=overwrite, content_settings=content_settings
+            )
+            url = (
+                f"{file_client.url}{self.sas_token}"
+                if self.sas_token
+                else file_client.url
+            )
+            return {"object_key": object_key, "url": url}
+        except Exception as e:
+            logger.warn(f"AzureStorageClient, upload_file error: {e}")
+            return {}

chainlit/data/storage_clients/azure_blob.py

@@ -0,0 +1,89 @@
+from datetime import datetime, timedelta
+from typing import Any, Dict, Union
+
+from azure.storage.blob import BlobSasPermissions, ContentSettings, generate_blob_sas
+from azure.storage.blob.aio import BlobServiceClient as AsyncBlobServiceClient
+
+from chainlit.data.storage_clients.base import EXPIRY_TIME, BaseStorageClient
+from chainlit.logger import logger
+
+
+class AzureBlobStorageClient(BaseStorageClient):
+    def __init__(self, container_name: str, storage_account: str, storage_key: str):
+        self.container_name = container_name
+        self.storage_account = storage_account
+        self.storage_key = storage_key
+        connection_string = (
+            f"DefaultEndpointsProtocol=https;"
+            f"AccountName={storage_account};"
+            f"AccountKey={storage_key};"
+            f"EndpointSuffix=core.windows.net"
+        )
+        self.service_client = AsyncBlobServiceClient.from_connection_string(
+            connection_string
+        )
+        self.container_client = self.service_client.get_container_client(
+            self.container_name
+        )
+        logger.info("AzureBlobStorageClient initialized")
+
+    async def get_read_url(self, object_key: str) -> str:
+        if not self.storage_key:
+            raise Exception("Not using Azure Storage")
+
+        sas_permissions = BlobSasPermissions(read=True)
+        start_time = datetime.now()
+        expiry_time = start_time + timedelta(seconds=EXPIRY_TIME)
+
+        sas_token = generate_blob_sas(
+            account_name=self.storage_account,
+            container_name=self.container_name,
+            blob_name=object_key,
+            account_key=self.storage_key,
+            permission=sas_permissions,
+            start=start_time,
+            expiry=expiry_time,
+        )
+
+        return f"https://{self.storage_account}.blob.core.windows.net/{self.container_name}/{object_key}?{sas_token}"
+
+    async def upload_file(
+        self,
+        object_key: str,
+        data: Union[bytes, str],
+        mime: str = "application/octet-stream",
+        overwrite: bool = True,
+    ) -> Dict[str, Any]:
+        try:
+            blob_client = self.container_client.get_blob_client(object_key)
+
+            if isinstance(data, str):
+                data = data.encode("utf-8")
+
+            content_settings = ContentSettings(content_type=mime)
+
+            await blob_client.upload_blob(
+                data, overwrite=overwrite, content_settings=content_settings
+            )
+
+            properties = await blob_client.get_blob_properties()
+
+            return {
+                "path": object_key,
+                "size": properties.size,
+                "last_modified": properties.last_modified,
+                "etag": properties.etag,
+                "content_type": properties.content_settings.content_type,
+            }
+
+        except Exception as e:
+            raise Exception(f"Failed to upload file to Azure Blob Storage: {e!s}")
+
+    async def delete_file(self, object_key: str) -> bool:
+        try:
+            blob_client = self.container_client.get_blob_client(blob=object_key)
+            await blob_client.delete_blob()
+            return True
+        except Exception as e:
+            logger.warn(f"AzureBlobStorageClient, delete_file error: {e}")
+            return False

chainlit/data/storage_clients/base.py

@@ -0,0 +1,26 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Union
+
+EXPIRY_TIME = 3600
+
+
+class BaseStorageClient(ABC):
+    """Base class for non-text data persistence like Azure Data Lake, S3, Google Storage, etc."""
+
+    @abstractmethod
+    async def upload_file(
+        self,
+        object_key: str,
+        data: Union[bytes, str],
+        mime: str = "application/octet-stream",
+        overwrite: bool = True,
+    ) -> Dict[str, Any]:
+        pass
+
+    @abstractmethod
+    async def delete_file(self, object_key: str) -> bool:
+        pass
+
+    @abstractmethod
+    async def get_read_url(self, object_key: str) -> str:
+        pass

chainlit/data/storage_clients/gcs.py

@@ -0,0 +1,88 @@
+import base64
+from typing import Any, Dict, Union
+
+from google.cloud import storage  # type: ignore
+from google.oauth2 import service_account
+
+from chainlit import make_async
+from chainlit.data.storage_clients.base import EXPIRY_TIME, BaseStorageClient
+from chainlit.logger import logger
+
+
+class GCSStorageClient(BaseStorageClient):
+    def __init__(
+        self, project_id: str, client_email: str, private_key: str, bucket_name: str
+    ):
+        private_key = base64.b64decode(private_key).decode("utf-8")
+
+        credentials = service_account.Credentials.from_service_account_info(
+            {
+                "type": "service_account",
+                "project_id": project_id,
+                "private_key": private_key,
+                "client_email": client_email,
+                "token_uri": "https://oauth2.googleapis.com/token",
+            }
+        )
+
+        self.client = storage.Client(project=project_id, credentials=credentials)
+        self.bucket = self.client.bucket(bucket_name)
+        logger.info("GCSStorageClient initialized")
+
+    def sync_get_read_url(self, object_key: str) -> str:
+        return self.bucket.blob(object_key).generate_signed_url(
+            version="v4", expiration=EXPIRY_TIME, method="GET"
+        )
+
+    async def get_read_url(self, object_key: str) -> str:
+        return await make_async(self.sync_get_read_url)(object_key)
+
+    def sync_upload_file(
+        self,
+        object_key: str,
+        data: Union[bytes, str],
+        mime: str = "application/octet-stream",
+        overwrite: bool = True,
+    ) -> Dict[str, Any]:
+        try:
+            blob = self.bucket.blob(object_key)
+
+            if not overwrite and blob.exists():
+                raise Exception(
+                    f"File {object_key} already exists and overwrite is False"
+                )
+
+            if isinstance(data, str):
+                data = data.encode("utf-8")
+
+            blob.upload_from_string(data, content_type=mime)
+
+            return {
+                "object_key": object_key,
+                "url": f"gs://{self.bucket.name}/{object_key}",
+            }
+
+        except Exception as e:
+            raise Exception(f"Failed to upload file to GCS: {e!s}")
+
+    async def upload_file(
+        self,
+        object_key: str,
+        data: Union[bytes, str],
+        mime: str = "application/octet-stream",
+        overwrite: bool = True,
+    ) -> Dict[str, Any]:
+        return await make_async(self.sync_upload_file)(
+            object_key, data, mime, overwrite
+        )
+
+    def sync_delete_file(self, object_key: str) -> bool:
+        try:
+            self.bucket.blob(object_key).delete()
+            return True
+        except Exception as e:
+            logger.warn(f"GCSStorageClient, delete_file error: {e}")
+            return False
+
+    async def delete_file(self, object_key: str) -> bool:
+        return await make_async(self.sync_delete_file)(object_key)

chainlit/data/storage_clients/s3.py

@@ -0,0 +1,75 @@
+from typing import Any, Dict, Union
+
+import boto3  # type: ignore
+
+from chainlit import make_async
+from chainlit.data.storage_clients.base import EXPIRY_TIME, BaseStorageClient
+from chainlit.logger import logger
+
+
+class S3StorageClient(BaseStorageClient):
+    """
+    Class to enable Amazon S3 storage provider
+    """
+
+    def __init__(self, bucket: str, **kwargs: Any):
+        try:
+            self.bucket = bucket
+            self.client = boto3.client("s3", **kwargs)
+            logger.info("S3StorageClient initialized")
+        except Exception as e:
+            logger.warn(f"S3StorageClient initialization error: {e}")
+
+    def sync_get_read_url(self, object_key: str) -> str:
+        try:
+            url = self.client.generate_presigned_url(
+                "get_object",
+                Params={"Bucket": self.bucket, "Key": object_key},
+                ExpiresIn=EXPIRY_TIME,
+            )
+            return url
+        except Exception as e:
+            logger.warn(f"S3StorageClient, get_read_url error: {e}")
+            return object_key
+
+    async def get_read_url(self, object_key: str) -> str:
+        return await make_async(self.sync_get_read_url)(object_key)
+
+    def sync_upload_file(
+        self,
+        object_key: str,
+        data: Union[bytes, str],
+        mime: str = "application/octet-stream",
+        overwrite: bool = True,
+    ) -> Dict[str, Any]:
+        try:
+            self.client.put_object(
+                Bucket=self.bucket, Key=object_key, Body=data, ContentType=mime
+            )
+            url = f"https://{self.bucket}.s3.amazonaws.com/{object_key}"
+            return {"object_key": object_key, "url": url}
+        except Exception as e:
+            logger.warn(f"S3StorageClient, upload_file error: {e}")
+            return {}
+
+    async def upload_file(
+        self,
+        object_key: str,
+        data: Union[bytes, str],
+        mime: str = "application/octet-stream",
+        overwrite: bool = True,
+    ) -> Dict[str, Any]:
+        return await make_async(self.sync_upload_file)(
+            object_key, data, mime, overwrite
+        )
+
+    def sync_delete_file(self, object_key: str) -> bool:
+        try:
+            self.client.delete_object(Bucket=self.bucket, Key=object_key)
+            return True
+        except Exception as e:
+            logger.warn(f"S3StorageClient, delete_file error: {e}")
+            return False
+
+    async def delete_file(self, object_key: str) -> bool:
+        return await make_async(self.sync_delete_file)(object_key)

chainlit/data/utils.py

@@ -0,0 +1,29 @@
+import functools
+from collections import deque
+
+from chainlit.context import context
+from chainlit.session import WebsocketSession
+
+
+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

chainlit/discord/__init__.py

@@ -0,0 +1,6 @@
+import importlib.util
+
+if importlib.util.find_spec("discord") is None:
+    raise ValueError(
+        "The discord package is required to integrate Chainlit with a Discord app. Run `pip install discord --upgrade`"
+    )