From eab953392b47f8b292cce9465940b1573cea51e1 Mon Sep 17 00:00:00 2001 From: Mathijs de Bruin Date: Thu, 21 Nov 2024 12:47:14 +0000 Subject: [PATCH] Various improvements to data_layer acquisition logic. (#1534) * Initialise data layer only once. * Give DeprecationWarning when users set _data_layer directly. * Remove debug print statement. --- backend/chainlit/data/__init__.py | 46 +++++++++++++++++++------------ 1 file changed, 29 insertions(+), 17 deletions(-) diff --git a/backend/chainlit/data/__init__.py b/backend/chainlit/data/__init__.py index dfc359c24b..513bc5975d 100644 --- a/backend/chainlit/data/__init__.py +++ b/backend/chainlit/data/__init__.py @@ -7,26 +7,38 @@ ) _data_layer: Optional[BaseDataLayer] = None +_data_layer_initialized = False def get_data_layer(): - global _data_layer - print("Getting data layer", _data_layer) - - if not _data_layer: - 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 api_key := os.environ.get("LITERAL_API_KEY"): - # When LITERAL_API_KEY is defined, use LiteralAI 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" + global _data_layer, _data_layer_initialized + + if not _data_layer_initialized: + if _data_layer: + # Data layer manually set, warn user that this is deprecated. + import warnings + + warnings.warn( + "Setting data layer manually is deprecated. Use @data_layer instead.", + DeprecationWarning, ) - _data_layer = LiteralDataLayer(api_key=api_key, server=server) + + else: + 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 api_key := os.environ.get("LITERAL_API_KEY"): + # When LITERAL_API_KEY is defined, use LiteralAI 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 return _data_layer