Make callbacks available at the main module level again, but lazily

cypress/e2e/author_rename/main.py

@@ -1,7 +1,6 @@
 from langchain import LLMMathChain, OpenAI
 
 import chainlit as cl
-from chainlit.langchain.callbacks import AsyncLangchainCallbackHandler
 
 
 @cl.author_rename
@@ -14,6 +13,6 @@ def rename(orig_author: str):
 async def main(message: str):
     llm = OpenAI(temperature=0)
     llm_math = LLMMathChain.from_llm(llm=llm)
-    res = await llm_math.acall(message, callbacks=[AsyncLangchainCallbackHandler()])
+    res = await llm_math.acall(message, callbacks=[cl.AsyncLangchainCallbackHandler()])
 
     await cl.Message(content="Hello").send()

cypress/e2e/haystack_cb/main.py

@@ -3,7 +3,6 @@
 from haystack.nodes import PromptNode
 
 import chainlit as cl
-from chainlit.haystack.callbacks import HaystackAgentCallbackHandler
 
 
 @cl.on_chat_start
@@ -13,7 +12,7 @@ async def start():
     fake_prompt_node = PromptNode(model_name_or_path="gpt-3.5-turbo", api_key="fakekey")
 
     agent = Agent(fake_prompt_node)
-    cb = HaystackAgentCallbackHandler(agent)
+    cb = cl.HaystackAgentCallbackHandler(agent)
 
     cb.on_agent_start(name="agent")
 

cypress/e2e/langchain_cb/main_async.py

@@ -1,14 +1,13 @@
 from langchain.schema import Generation, LLMResult, SystemMessage
 
 import chainlit as cl
-from chainlit.langchain.callbacks import AsyncLangchainCallbackHandler
 
 
 @cl.on_chat_start
 async def main():
     await cl.Message(content="AsyncLangchainCb").send()
 
-    acb = AsyncLangchainCallbackHandler()
+    acb = cl.AsyncLangchainCallbackHandler()
 
     await acb.on_chain_start(serialized={"id": ["TestChain1"]}, inputs={})
 

cypress/e2e/langchain_cb/main_sync.py

@@ -1,14 +1,13 @@
 from langchain.schema import Generation, LLMResult, SystemMessage
 
 import chainlit as cl
-from chainlit.langchain.callbacks import LangchainCallbackHandler
 
 
 @cl.on_chat_start
 async def main():
     await cl.Message(content="AsyncLangchainCb").send()
 
-    cb = LangchainCallbackHandler()
+    cb = cl.LangchainCallbackHandler()
 
     cb.on_chain_start(serialized={"id": ["TestChain1"]}, inputs={})
 

cypress/e2e/llama_index_cb/main.py

@@ -2,14 +2,13 @@
 from llama_index.schema import NodeWithScore, TextNode
 
 import chainlit as cl
-from chainlit.llama_index.callbacks import LlamaIndexCallbackHandler
 
 
 @cl.on_chat_start
 async def start():
     await cl.Message(content="LlamaIndexCb").send()
 
-    cb = LlamaIndexCallbackHandler()
+    cb = cl.LlamaIndexCallbackHandler()
 
     cb.start_trace()
 

src/chainlit/__init__.py

@@ -32,7 +32,7 @@
 from chainlit.telemetry import trace
 from chainlit.types import LLMSettings
 from chainlit.user_session import user_session
-from chainlit.utils import wrap_user_function
+from chainlit.utils import make_module_getattr, wrap_user_function
 from chainlit.version import __version__
 
 env_found = load_dotenv(dotenv_path=os.path.join(os.getcwd(), ".env"))
@@ -181,6 +181,15 @@ def sleep(duration: int):
     return asyncio.sleep(duration)
 
 
+__getattr__ = make_module_getattr(
+    {
+        "LangchainCallbackHandler": "chainlit.langchain.callbacks",
+        "AsyncLangchainCallbackHandler": "chainlit.langchain.callbacks",
+        "LlamaIndexCallbackHandler": "chainlit.llama_index.callbacks",
+        "HaystackAgentCallbackHandler": "chainlit.haystack.callbacks",
+    }
+)
+
 __all__ = [
     "user_session",
     "LLMSettings",
@@ -213,4 +222,8 @@ def sleep(duration: int):
     "run_sync",
     "make_async",
     "cache",
+    "LangchainCallbackHandler",
+    "AsyncLangchainCallbackHandler",
+    "LlamaIndexCallbackHandler",
+    "HaystackAgentCallbackHandler",
 ]

src/chainlit/utils.py

@@ -1,3 +1,4 @@
+import importlib
 import inspect
 from typing import Callable
 
@@ -49,3 +50,18 @@ async def wrapper(*args):
                 await emitter.task_end()
 
     return wrapper
+
+
+def make_module_getattr(registry):
+    """Leverage PEP 562 to make imports lazy in an __init__.py
+
+    The registry must be a dictionary with the items to import as keys and the
+    modules they belong to as a value.
+    """
+
+    def __getattr__(name):
+        module_path = registry[name]
+        module = importlib.import_module(module_path, __package__)
+        return getattr(module, name)
+
+    return __getattr__