ai-station/.venv/lib/python3.12/site-packages/chainlit/data/utils.py

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