30 lines
966 B
Python
30 lines
966 B
Python
|
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
|