Skip to content

Commit eab9533

Browse files
authored
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.
1 parent dde96bc commit eab9533

File tree

1 file changed

+29
-17
lines changed

1 file changed

+29
-17
lines changed

backend/chainlit/data/__init__.py

Lines changed: 29 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -7,26 +7,38 @@
77
)
88

99
_data_layer: Optional[BaseDataLayer] = None
10+
_data_layer_initialized = False
1011

1112

1213
def get_data_layer():
13-
global _data_layer
14-
print("Getting data layer", _data_layer)
15-
16-
if not _data_layer:
17-
from chainlit.config import config
18-
19-
if config.code.data_layer:
20-
# When @data_layer is configured, call it to get data layer.
21-
_data_layer = config.code.data_layer()
22-
elif api_key := os.environ.get("LITERAL_API_KEY"):
23-
# When LITERAL_API_KEY is defined, use LiteralAI data layer
24-
from .literalai import LiteralDataLayer
25-
26-
# support legacy LITERAL_SERVER variable as fallback
27-
server = os.environ.get("LITERAL_API_URL") or os.environ.get(
28-
"LITERAL_SERVER"
14+
global _data_layer, _data_layer_initialized
15+
16+
if not _data_layer_initialized:
17+
if _data_layer:
18+
# Data layer manually set, warn user that this is deprecated.
19+
import warnings
20+
21+
warnings.warn(
22+
"Setting data layer manually is deprecated. Use @data_layer instead.",
23+
DeprecationWarning,
2924
)
30-
_data_layer = LiteralDataLayer(api_key=api_key, server=server)
25+
26+
else:
27+
from chainlit.config import config
28+
29+
if config.code.data_layer:
30+
# When @data_layer is configured, call it to get data layer.
31+
_data_layer = config.code.data_layer()
32+
elif api_key := os.environ.get("LITERAL_API_KEY"):
33+
# When LITERAL_API_KEY is defined, use LiteralAI data layer
34+
from .literalai import LiteralDataLayer
35+
36+
# support legacy LITERAL_SERVER variable as fallback
37+
server = os.environ.get("LITERAL_API_URL") or os.environ.get(
38+
"LITERAL_SERVER"
39+
)
40+
_data_layer = LiteralDataLayer(api_key=api_key, server=server)
41+
42+
_data_layer_initialized = True
3143

3244
return _data_layer

0 commit comments

Comments
 (0)