forked from Chainlit/chainlit
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Support SQLAlchemy for custom data layer (Chainlit#836)
- adds custom, direct database, data layer using SQLAlchemy with support for a wide-range of SQL dialects - configures ADLS or S3 as the blob storage provider - duplicated `PageInfo` and `PaginatedResponse` from literal SDK into backend/chainlit/types.py and updated typing
- Loading branch information
Showing
7 changed files
with
663 additions
and
6 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Large diffs are not rendered by default.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,58 @@ | ||
from chainlit.data import BaseStorageClient | ||
from chainlit.logger import logger | ||
from typing import TYPE_CHECKING, Optional, Dict, Union, Any | ||
from azure.storage.filedatalake import DataLakeServiceClient, FileSystemClient, DataLakeFileClient, ContentSettings | ||
import boto3 # type: ignore | ||
|
||
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 {} | ||
|
||
class S3StorageClient(BaseStorageClient): | ||
""" | ||
Class to enable Amazon S3 storage provider | ||
""" | ||
def __init__(self, bucket: str): | ||
try: | ||
self.bucket = bucket | ||
self.client = boto3.client("s3") | ||
logger.info("S3StorageClient initialized") | ||
except Exception as e: | ||
logger.warn(f"S3StorageClient 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: | ||
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 {} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,31 @@ | ||
from typing import List, Optional | ||
|
||
import chainlit.data as cl_data | ||
from chainlit.data.sql_alchemy import SQLAlchemyDataLayer | ||
from chainlit.data.storage_clients import AzureStorageClient | ||
from literalai.helper import utc_now | ||
|
||
import chainlit as cl | ||
|
||
storage_client = AzureStorageClient(account_url="<your_account_url>", container="<your_container>") | ||
|
||
cl_data._data_layer = SQLAlchemyDataLayer(conninfo="<your conninfo>", storage_provider=storage_client) | ||
|
||
|
||
@cl.on_chat_start | ||
async def main(): | ||
await cl.Message("Hello, send me a message!", disable_feedback=True).send() | ||
|
||
|
||
@cl.on_message | ||
async def handle_message(): | ||
await cl.sleep(2) | ||
await cl.Message("Ok!").send() | ||
|
||
|
||
@cl.password_auth_callback | ||
def auth_callback(username: str, password: str) -> Optional[cl.User]: | ||
if (username, password) == ("admin", "admin"): | ||
return cl.User(identifier="admin") | ||
else: | ||
return None |