mirror of
https://github.com/open-webui/open-webui.git
synced 2026-09-01 19:50:41 +02:00
refac
This commit is contained in:
Timothy Jaeryang Baek
@@ -11,6 +11,7 @@ The scheduler_worker_loop handles all time-based background work:
|
||||
|
||||
Environment:
|
||||
SCHEDULER_POLL_INTERVAL – seconds between polls (default: 10)
|
||||
TIMER_POLL_INTERVAL – seconds between timer polls (default: 1)
|
||||
CALENDAR_ALERT_LOOKAHEAD_MINUTES – default alert window (default: 5)
|
||||
"""
|
||||
|
||||
@@ -43,6 +44,7 @@ from starlette.datastructures import Headers
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
SCHEDULER_POLL_INTERVAL = int(os.getenv('SCHEDULER_POLL_INTERVAL', os.getenv('AUTOMATION_POLL_INTERVAL', '10')))
|
||||
TIMER_POLL_INTERVAL = int(os.getenv('TIMER_POLL_INTERVAL', '1'))
|
||||
CALENDAR_ALERT_LOOKAHEAD_MINUTES = int(os.getenv('CALENDAR_ALERT_LOOKAHEAD_MINUTES', '10'))
|
||||
|
||||
|
||||
@@ -190,10 +192,15 @@ async def scheduler_worker_loop(app) -> None:
|
||||
Runs on every instance. Poll interval is configurable via
|
||||
SCHEDULER_POLL_INTERVAL env var (default: 10 seconds).
|
||||
"""
|
||||
log.info(f'Scheduler worker started (poll interval: {SCHEDULER_POLL_INTERVAL}s)')
|
||||
log.info(
|
||||
f'Scheduler worker started (timer poll interval: {TIMER_POLL_INTERVAL}s, '
|
||||
f'scheduler poll interval: {SCHEDULER_POLL_INTERVAL}s)'
|
||||
)
|
||||
next_scheduler_poll = 0.0
|
||||
|
||||
while True:
|
||||
try:
|
||||
now = time.monotonic()
|
||||
# ── Timers ──
|
||||
try:
|
||||
from open_webui.utils.timers import claim_due_timers, execute_due_timer
|
||||
@@ -203,6 +210,12 @@ async def scheduler_worker_loop(app) -> None:
|
||||
except Exception:
|
||||
log.exception('Scheduler: timer error')
|
||||
|
||||
if now < next_scheduler_poll:
|
||||
await asyncio.sleep(max(1, TIMER_POLL_INTERVAL))
|
||||
continue
|
||||
# Jitter to spread automation/calendar load across instances; timers keep a tight poll.
|
||||
next_scheduler_poll = now + SCHEDULER_POLL_INTERVAL + random.uniform(0, 2)
|
||||
|
||||
# ── Automations ──
|
||||
if await Config.get('automations.enable'):
|
||||
try:
|
||||
@@ -225,8 +238,7 @@ async def scheduler_worker_loop(app) -> None:
|
||||
except Exception:
|
||||
log.exception('Scheduler worker error')
|
||||
|
||||
# Jitter to spread load across instances
|
||||
await asyncio.sleep(SCHEDULER_POLL_INTERVAL + random.uniform(0, 2))
|
||||
await asyncio.sleep(max(1, TIMER_POLL_INTERVAL))
|
||||
|
||||
|
||||
##########################
|
||||
|
||||
@@ -9,13 +9,15 @@ from uuid import uuid4
|
||||
|
||||
from fastapi import Request
|
||||
from fastapi.security import HTTPAuthorizationCredentials
|
||||
from open_webui.internal.db import get_async_db
|
||||
from open_webui.models.chat_messages import ChatMessages
|
||||
from open_webui.models.chats import ChatForm, Chats
|
||||
from open_webui.models.chats import Chat, ChatForm, Chats
|
||||
from open_webui.models.config import Config
|
||||
from open_webui.models.users import UserModel, Users
|
||||
from open_webui.tasks import create_task, has_active_tasks
|
||||
from open_webui.utils.auth import create_token
|
||||
from open_webui.utils.misc import get_message_list
|
||||
from sqlalchemy import select
|
||||
from starlette.datastructures import Headers
|
||||
|
||||
DEFAULT_SUBAGENT_SYSTEM_PROMPT = """You are a sub-agent working on a specific task assigned by the lead agent.
|
||||
@@ -81,145 +83,143 @@ async def process_pending_internal_messages(
|
||||
if await has_active_tasks(source_request.app.state.redis, parent_chat_id):
|
||||
return
|
||||
|
||||
chat = await Chats.get_chat_by_id_and_user_id(parent_chat_id, user_id)
|
||||
user = await Users.get_user_by_id(user_id)
|
||||
if not chat or not user:
|
||||
if not user:
|
||||
return
|
||||
|
||||
history = copy.deepcopy(chat.chat.get('history') or {})
|
||||
messages = history.get('messages') or {}
|
||||
pending = [
|
||||
message
|
||||
for message in messages.values()
|
||||
if message.get('role') == 'user'
|
||||
and not message.get('childrenIds')
|
||||
and (
|
||||
(message.get('meta') or {}).get('async_subagent_result') is True
|
||||
or (
|
||||
(message.get('meta') or {}).get('internal') is True
|
||||
and (message.get('meta') or {}).get('type') == 'timer'
|
||||
async with get_async_db() as db:
|
||||
stmt = select(Chat).where(Chat.id == parent_chat_id, Chat.user_id == user_id)
|
||||
if db.bind.dialect.name == 'postgresql':
|
||||
stmt = stmt.with_for_update()
|
||||
result = await db.execute(stmt)
|
||||
chat = result.scalar_one_or_none()
|
||||
if not chat:
|
||||
return
|
||||
|
||||
history = copy.deepcopy((chat.chat or {}).get('history') or {})
|
||||
messages = history.get('messages') or {}
|
||||
pending = [
|
||||
message
|
||||
for message in messages.values()
|
||||
if message.get('role') == 'user'
|
||||
and not message.get('childrenIds')
|
||||
and (
|
||||
(message.get('meta') or {}).get('async_subagent_result') is True
|
||||
or (
|
||||
(message.get('meta') or {}).get('internal') is True
|
||||
and (message.get('meta') or {}).get('type') == 'timer'
|
||||
)
|
||||
)
|
||||
)
|
||||
]
|
||||
if not pending:
|
||||
return
|
||||
]
|
||||
if not pending:
|
||||
return
|
||||
|
||||
first = pending[0]
|
||||
first_meta = first.get('meta') or {}
|
||||
kind = 'subagent' if first_meta.get('async_subagent_result') is True else 'timer'
|
||||
parent_id = first.get('parentId')
|
||||
if kind == 'timer' and first_meta.get('timer_id'):
|
||||
timer = await Chats.get_chat_by_id(first_meta['timer_id'])
|
||||
run = {**run, **(((timer.meta or {}).get('run') if timer else None) or {})}
|
||||
model_id = first.get('model') or run['model_id']
|
||||
if kind == 'timer':
|
||||
batch = [
|
||||
message
|
||||
for message in pending
|
||||
if message.get('parentId') == parent_id
|
||||
and (message.get('model') or model_id) == model_id
|
||||
and (message.get('meta') or {}).get('internal') is True
|
||||
and (message.get('meta') or {}).get('type') == 'timer'
|
||||
]
|
||||
else:
|
||||
batch = [
|
||||
message
|
||||
for message in pending
|
||||
if message.get('parentId') == parent_id
|
||||
and (message.get('model') or model_id) == model_id
|
||||
and (message.get('meta') or {}).get('async_subagent_result') is True
|
||||
]
|
||||
combined_content = '\n\n'.join(message.get('content', '') for message in batch if message.get('content'))
|
||||
if kind == 'timer':
|
||||
timer_ids = [message['meta']['timer_id'] for message in batch if (message.get('meta') or {}).get('timer_id')]
|
||||
combined_meta = {'internal': True, 'type': 'timer'}
|
||||
if len(timer_ids) == 1:
|
||||
combined_meta['timer_id'] = timer_ids[0]
|
||||
elif timer_ids:
|
||||
combined_meta['timer_ids'] = timer_ids
|
||||
else:
|
||||
delegation_ids = [
|
||||
message['meta']['delegation_id']
|
||||
for message in batch
|
||||
if (message.get('meta') or {}).get('delegation_id')
|
||||
]
|
||||
subagent_chat_ids = [
|
||||
message['meta']['subagent_chat_id']
|
||||
for message in batch
|
||||
if (message.get('meta') or {}).get('subagent_chat_id')
|
||||
]
|
||||
combined_meta = {'async_subagent_result': True}
|
||||
if len(delegation_ids) == 1:
|
||||
combined_meta['delegation_id'] = delegation_ids[0]
|
||||
elif delegation_ids:
|
||||
combined_meta['delegation_ids'] = delegation_ids
|
||||
if len(subagent_chat_ids) == 1:
|
||||
combined_meta['subagent_chat_id'] = subagent_chat_ids[0]
|
||||
elif subagent_chat_ids:
|
||||
combined_meta['subagent_chat_ids'] = subagent_chat_ids
|
||||
|
||||
pending_flag = 'timer_pending' if kind == 'timer' else 'async_subagent_pending'
|
||||
reuse_message = len(batch) == 1 and not (first.get('meta') or {}).get(pending_flag)
|
||||
user_message_id = first['id'] if reuse_message else str(uuid4())
|
||||
if not reuse_message:
|
||||
removed_ids = {message['id'] for message in batch}
|
||||
for message_id in removed_ids:
|
||||
messages.pop(message_id, None)
|
||||
if parent_id and parent_id in messages:
|
||||
messages[parent_id]['childrenIds'] = [
|
||||
child_id for child_id in messages[parent_id].get('childrenIds', []) if child_id not in removed_ids
|
||||
first = pending[0]
|
||||
first_meta = first.get('meta') or {}
|
||||
kind = 'subagent' if first_meta.get('async_subagent_result') is True else 'timer'
|
||||
parent_id = first.get('parentId')
|
||||
if kind == 'timer' and first_meta.get('timer_id'):
|
||||
timer = await Chats.get_chat_by_id(first_meta['timer_id'])
|
||||
run = {**run, **(((timer.meta or {}).get('run') if timer else None) or {})}
|
||||
model_id = first.get('model') or run['model_id']
|
||||
if kind == 'timer':
|
||||
batch = [first]
|
||||
else:
|
||||
batch = [
|
||||
message
|
||||
for message in pending
|
||||
if message.get('parentId') == parent_id
|
||||
and (message.get('model') or model_id) == model_id
|
||||
and (message.get('meta') or {}).get('async_subagent_result') is True
|
||||
]
|
||||
combined_content = '\n\n'.join(message.get('content', '') for message in batch if message.get('content'))
|
||||
if kind == 'timer':
|
||||
timer_ids = [
|
||||
message['meta']['timer_id'] for message in batch if (message.get('meta') or {}).get('timer_id')
|
||||
]
|
||||
combined_meta = {'internal': True, 'type': 'timer'}
|
||||
if len(timer_ids) == 1:
|
||||
combined_meta['timer_id'] = timer_ids[0]
|
||||
elif timer_ids:
|
||||
combined_meta['timer_ids'] = timer_ids
|
||||
else:
|
||||
delegation_ids = [
|
||||
message['meta']['delegation_id']
|
||||
for message in batch
|
||||
if (message.get('meta') or {}).get('delegation_id')
|
||||
]
|
||||
subagent_chat_ids = [
|
||||
message['meta']['subagent_chat_id']
|
||||
for message in batch
|
||||
if (message.get('meta') or {}).get('subagent_chat_id')
|
||||
]
|
||||
combined_meta = {'async_subagent_result': True}
|
||||
if len(delegation_ids) == 1:
|
||||
combined_meta['delegation_id'] = delegation_ids[0]
|
||||
elif delegation_ids:
|
||||
combined_meta['delegation_ids'] = delegation_ids
|
||||
if len(subagent_chat_ids) == 1:
|
||||
combined_meta['subagent_chat_id'] = subagent_chat_ids[0]
|
||||
elif subagent_chat_ids:
|
||||
combined_meta['subagent_chat_ids'] = subagent_chat_ids
|
||||
|
||||
pending_flag = 'timer_pending' if kind == 'timer' else 'async_subagent_pending'
|
||||
reuse_message = len(batch) == 1 and not (first.get('meta') or {}).get(pending_flag)
|
||||
user_message_id = first['id'] if reuse_message else str(uuid4())
|
||||
removed_ids = set()
|
||||
if not reuse_message:
|
||||
removed_ids = {message['id'] for message in batch}
|
||||
for message_id in removed_ids:
|
||||
messages.pop(message_id, None)
|
||||
if parent_id and parent_id in messages:
|
||||
messages[parent_id]['childrenIds'] = [
|
||||
child_id
|
||||
for child_id in messages[parent_id].get('childrenIds', [])
|
||||
if child_id not in removed_ids
|
||||
]
|
||||
|
||||
assistant_message_id = str(uuid4())
|
||||
message_list = get_message_list(messages, parent_id)
|
||||
system_prompt = run.get('system_prompt')
|
||||
user_message = {
|
||||
'id': user_message_id,
|
||||
'parentId': parent_id,
|
||||
'childrenIds': [assistant_message_id],
|
||||
'role': 'user',
|
||||
'content': combined_content,
|
||||
'model': model_id,
|
||||
'meta': combined_meta,
|
||||
'timestamp': int(time.time()),
|
||||
}
|
||||
assistant_message = {
|
||||
'id': assistant_message_id,
|
||||
'parentId': user_message_id,
|
||||
'childrenIds': [],
|
||||
'role': 'assistant',
|
||||
'content': '',
|
||||
'done': False,
|
||||
'model': model_id,
|
||||
'timestamp': int(time.time()),
|
||||
}
|
||||
|
||||
if parent_id and parent_id in messages:
|
||||
parent_children = [
|
||||
child_id for child_id in messages[parent_id].get('childrenIds', []) if child_id != user_message_id
|
||||
]
|
||||
parent_children.append(user_message_id)
|
||||
messages[parent_id]['childrenIds'] = parent_children
|
||||
messages[user_message_id] = {**messages.get(user_message_id, {}), **user_message}
|
||||
messages[assistant_message_id] = assistant_message
|
||||
history['messages'] = messages
|
||||
history['currentId'] = parent_id
|
||||
updated_chat = copy.deepcopy(chat.chat)
|
||||
updated_chat['history'] = history
|
||||
await Chats.update_chat_by_id(parent_chat_id, updated_chat)
|
||||
history['currentId'] = assistant_message_id
|
||||
chat.chat = {**(chat.chat or {}), 'history': history}
|
||||
chat.updated_at = int(time.time())
|
||||
await db.commit()
|
||||
|
||||
if removed_ids:
|
||||
await ChatMessages.delete_message_ids_by_chat_id(parent_chat_id, removed_ids)
|
||||
|
||||
assistant_message_id = str(uuid4())
|
||||
message_list = get_message_list(messages, parent_id)
|
||||
system_prompt = run.get('system_prompt')
|
||||
user_message = {
|
||||
'id': user_message_id,
|
||||
'parentId': parent_id,
|
||||
'childrenIds': [assistant_message_id],
|
||||
'role': 'user',
|
||||
'content': combined_content,
|
||||
'model': model_id,
|
||||
'meta': combined_meta,
|
||||
'timestamp': int(time.time()),
|
||||
}
|
||||
assistant_message = {
|
||||
'id': assistant_message_id,
|
||||
'parentId': user_message_id,
|
||||
'childrenIds': [],
|
||||
'role': 'assistant',
|
||||
'content': '',
|
||||
'done': False,
|
||||
'model': model_id,
|
||||
'timestamp': int(time.time()),
|
||||
}
|
||||
|
||||
if parent_id and parent_id in messages:
|
||||
parent_children = [
|
||||
child_id for child_id in messages[parent_id].get('childrenIds', []) if child_id != user_message_id
|
||||
]
|
||||
parent_children.append(user_message_id)
|
||||
await Chats.upsert_message_to_chat_by_id_and_message_id(
|
||||
parent_chat_id,
|
||||
parent_id,
|
||||
{'childrenIds': parent_children},
|
||||
)
|
||||
await Chats.upsert_message_to_chat_by_id_and_message_id(
|
||||
parent_chat_id,
|
||||
user_message_id,
|
||||
user_message,
|
||||
)
|
||||
await Chats.upsert_message_to_chat_by_id_and_message_id(
|
||||
parent_chat_id,
|
||||
assistant_message_id,
|
||||
assistant_message,
|
||||
)
|
||||
await ChatMessages.upsert_message(user_message_id, parent_chat_id, user_id, user_message)
|
||||
await ChatMessages.upsert_message(assistant_message_id, parent_chat_id, user_id, assistant_message)
|
||||
|
||||
from open_webui.socket.main import sio
|
||||
|
||||
|
||||
@@ -275,45 +275,61 @@ async def execute_due_timer(app, timer_id: str, claim_id: str | None = None) ->
|
||||
user_message_id = str(uuid4())
|
||||
parent_lock = _parent_locks.setdefault(parent_chat_id, asyncio.Lock())
|
||||
async with parent_lock:
|
||||
parent = await Chats.get_chat_by_id_and_user_id(parent_chat_id, timer.user_id)
|
||||
if not parent:
|
||||
await _set_timer_status(timer_id, 'error', timer_error='parent chat no longer exists')
|
||||
return
|
||||
parent_chat = copy.deepcopy(parent.chat or {})
|
||||
history = parent_chat.setdefault('history', {})
|
||||
messages = history.setdefault('messages', {})
|
||||
done_assistants = [
|
||||
message
|
||||
for message in messages.values()
|
||||
if message.get('role') == 'assistant' and message.get('done') is not False
|
||||
]
|
||||
parent_id = (
|
||||
max(done_assistants, key=lambda message: message.get('timestamp', 0)).get('id')
|
||||
if done_assistants
|
||||
else meta.get('parent_message_id')
|
||||
)
|
||||
pending_meta = {'internal': True, 'type': 'timer', 'timer_id': timer_id}
|
||||
if await has_active_tasks(app.state.redis, parent_chat_id):
|
||||
pending_meta['timer_pending'] = True
|
||||
async with get_async_db() as db:
|
||||
stmt = select(Chat).where(Chat.id == parent_chat_id, Chat.user_id == timer.user_id)
|
||||
if db.bind.dialect.name == 'postgresql':
|
||||
stmt = stmt.with_for_update()
|
||||
result = await db.execute(stmt)
|
||||
parent = result.scalar_one_or_none()
|
||||
if not parent:
|
||||
await _set_timer_status(timer_id, 'error', timer_error='parent chat no longer exists')
|
||||
return
|
||||
|
||||
user_message = {
|
||||
'id': user_message_id,
|
||||
'parentId': parent_id,
|
||||
'childrenIds': [],
|
||||
'role': 'user',
|
||||
'content': prompt,
|
||||
'model': model_id,
|
||||
'meta': pending_meta,
|
||||
'timestamp': int(time.time()),
|
||||
}
|
||||
messages[user_message_id] = user_message
|
||||
if parent_id and parent_id in messages:
|
||||
children = messages[parent_id].setdefault('childrenIds', [])
|
||||
if user_message_id not in children:
|
||||
children.append(user_message_id)
|
||||
await Chats.update_chat_by_id(parent_chat_id, parent_chat)
|
||||
parent_chat = copy.deepcopy(parent.chat or {})
|
||||
history = parent_chat.setdefault('history', {})
|
||||
messages = history.setdefault('messages', {})
|
||||
done_assistants = [
|
||||
message
|
||||
for message in messages.values()
|
||||
if message.get('role') == 'assistant' and message.get('done') is not False
|
||||
]
|
||||
parent_id = (
|
||||
max(done_assistants, key=lambda message: message.get('timestamp', 0)).get('id')
|
||||
if done_assistants
|
||||
else meta.get('parent_message_id')
|
||||
)
|
||||
pending_meta = {'internal': True, 'type': 'timer', 'timer_id': timer_id}
|
||||
if await has_active_tasks(app.state.redis, parent_chat_id):
|
||||
pending_meta['timer_pending'] = True
|
||||
|
||||
user_message = {
|
||||
'id': user_message_id,
|
||||
'parentId': parent_id,
|
||||
'childrenIds': [],
|
||||
'role': 'user',
|
||||
'content': prompt,
|
||||
'model': model_id,
|
||||
'meta': pending_meta,
|
||||
'timestamp': int(time.time()),
|
||||
}
|
||||
messages[user_message_id] = user_message
|
||||
if parent_id and parent_id in messages:
|
||||
children = messages[parent_id].setdefault('childrenIds', [])
|
||||
if user_message_id not in children:
|
||||
children.append(user_message_id)
|
||||
|
||||
parent.chat = parent_chat
|
||||
parent.updated_at = int(time.time())
|
||||
timer_row = await db.get(Chat, timer_id)
|
||||
if timer_row:
|
||||
timer_row.meta = {
|
||||
**(timer_row.meta or {}),
|
||||
'timer_status': 'dispatched',
|
||||
'timer_dispatched_at': int(time.time_ns()),
|
||||
}
|
||||
timer_row.updated_at = int(time.time())
|
||||
await db.commit()
|
||||
await ChatMessages.upsert_message(user_message_id, parent_chat_id, timer.user_id, user_message)
|
||||
await _set_timer_status(timer_id, 'dispatched', timer_dispatched_at=int(time.time_ns()))
|
||||
|
||||
if user_message['meta'].get('timer_pending') is True:
|
||||
await sio.emit(
|
||||
|
||||
Reference in New Issue
Block a user