426 lines
14 KiB
Python
426 lines
14 KiB
Python
"""
|
||
Automation utilities.
|
||
|
||
RRULE helpers, worker loop, and execution logic.
|
||
Follows the utils/<feature>.py pattern (cf. utils/channels.py, utils/task.py).
|
||
|
||
Environment:
|
||
AUTOMATION_POLL_INTERVAL – seconds between polls (default: 10)
|
||
"""
|
||
|
||
import asyncio
|
||
import logging
|
||
import os
|
||
import random
|
||
import time
|
||
from datetime import datetime
|
||
from typing import Optional
|
||
from uuid import uuid4
|
||
from zoneinfo import ZoneInfo
|
||
|
||
from dateutil.rrule import rrulestr
|
||
from fastapi import Request
|
||
from starlette.datastructures import Headers
|
||
|
||
from open_webui.models.automations import Automations, AutomationRuns, AutomationModel
|
||
from open_webui.models.chats import ChatForm, Chats
|
||
from open_webui.models.users import Users
|
||
from open_webui.utils.task import prompt_template
|
||
from open_webui.internal.db import get_db
|
||
|
||
log = logging.getLogger(__name__)
|
||
|
||
AUTOMATION_POLL_INTERVAL = int(os.getenv('AUTOMATION_POLL_INTERVAL', '10'))
|
||
|
||
|
||
####################
|
||
# RRULE Helpers
|
||
####################
|
||
|
||
|
||
def _parse_rule(s: str):
|
||
"""Parse RRULE with clock-aligned DTSTART for sub-daily frequencies.
|
||
|
||
MINUTELY/HOURLY rules use a fixed epoch DTSTART (2000-01-01 00:00)
|
||
so intervals snap to clock boundaries (e.g. every 5min = :00, :05, :10).
|
||
"""
|
||
raw = s.replace('RRULE:', '')
|
||
parts = dict(p.split('=', 1) for p in raw.split(';') if '=' in p)
|
||
freq = parts.get('FREQ', '')
|
||
|
||
if freq in ('MINUTELY', 'HOURLY'):
|
||
epoch = datetime(2000, 1, 1, 0, 0, 0)
|
||
return rrulestr(s, dtstart=epoch, ignoretz=True)
|
||
return rrulestr(s, ignoretz=True)
|
||
|
||
|
||
def validate_rrule(s: str) -> None:
|
||
"""Raise ValueError if the RRULE is malformed or exhausted."""
|
||
try:
|
||
rule = _parse_rule(s)
|
||
except Exception as e:
|
||
raise ValueError(f'Invalid RRULE: {e}')
|
||
if rule.after(datetime.now()) is None:
|
||
raise ValueError('RRULE has no future occurrences')
|
||
|
||
|
||
def next_run_ns(s: str, tz: str = None) -> Optional[int]:
|
||
"""Next occurrence as epoch nanoseconds, respecting user timezone."""
|
||
now = datetime.now(ZoneInfo(tz)) if tz else datetime.now()
|
||
dt = _parse_rule(s).after(now.replace(tzinfo=None))
|
||
if dt is None:
|
||
return None
|
||
if tz:
|
||
dt = dt.replace(tzinfo=ZoneInfo(tz))
|
||
return int(dt.timestamp() * 1_000_000_000)
|
||
|
||
|
||
def next_n_runs_ns(s: str, n: int = 5, tz: str = None) -> list[int]:
|
||
"""Compute next N occurrences for UI preview."""
|
||
rule = _parse_rule(s)
|
||
result = []
|
||
dt = datetime.now()
|
||
for _ in range(n):
|
||
dt = rule.after(dt)
|
||
if not dt:
|
||
break
|
||
if tz:
|
||
dt_tz = dt.replace(tzinfo=ZoneInfo(tz))
|
||
result.append(int(dt_tz.timestamp() * 1_000_000_000))
|
||
else:
|
||
result.append(int(dt.timestamp() * 1_000_000_000))
|
||
return result
|
||
|
||
|
||
############################
|
||
# Worker Loop
|
||
############################
|
||
|
||
|
||
async def automation_worker_loop(app) -> None:
|
||
"""Poll for due automations, claim, fire-and-forget execute.
|
||
|
||
Runs on every instance. Poll interval is configurable via
|
||
AUTOMATION_POLL_INTERVAL env var (default: 10 seconds).
|
||
"""
|
||
log.info(
|
||
f'Automation worker started (poll interval: {AUTOMATION_POLL_INTERVAL}s)'
|
||
)
|
||
while True:
|
||
try:
|
||
with get_db() as db:
|
||
batch = Automations.claim_due(
|
||
int(time.time_ns()), limit=10, db=db
|
||
)
|
||
if batch:
|
||
log.info(f'Claimed {len(batch)} due automation(s)')
|
||
for automation in batch:
|
||
asyncio.create_task(execute_automation(app, automation))
|
||
except Exception:
|
||
log.exception('Automation worker error')
|
||
|
||
# Jitter to spread load across instances
|
||
await asyncio.sleep(
|
||
AUTOMATION_POLL_INTERVAL + random.uniform(0, 2)
|
||
)
|
||
|
||
|
||
##########################
|
||
# Execute
|
||
####################
|
||
|
||
|
||
def _build_request(app) -> Request:
|
||
"""Build a minimal ASGI Request for chat_completion.
|
||
|
||
Mirrors the mock-request pattern used in main.py lifespan
|
||
(model pre-fetch, tool server init) for consistency.
|
||
"""
|
||
scope = {
|
||
"type": "http",
|
||
"asgi": {"version": "3.0", "spec_version": "2.0"},
|
||
"method": "POST",
|
||
"path": "/api/v1/automations/internal",
|
||
"query_string": b"",
|
||
"headers": Headers({}).raw,
|
||
"client": ("127.0.0.1", 0),
|
||
"server": ("127.0.0.1", 80),
|
||
"scheme": "http",
|
||
"app": app,
|
||
}
|
||
request = Request(scope)
|
||
# Ensure request.state is initialized with required attributes
|
||
request.state.token = None
|
||
request.state.enable_api_keys = False
|
||
return request
|
||
|
||
|
||
def _resolve_model_tool_ids(app, model_id: str) -> list[str]:
|
||
"""Read model-attached tool_ids from model config.
|
||
|
||
The frontend does this in Chat.svelte (model.info.meta.toolIds).
|
||
The backend never auto-resolves them, so we must do it explicitly.
|
||
"""
|
||
models = getattr(app.state, "MODELS", {})
|
||
model = models.get(model_id, {})
|
||
tool_ids = model.get("info", {}).get("meta", {}).get("toolIds", [])
|
||
return list(tool_ids) if tool_ids else []
|
||
|
||
|
||
def _resolve_model_features(app, model_id: str) -> dict:
|
||
"""Read model default features from model config.
|
||
|
||
The frontend does this in Chat.svelte (model.info.meta.defaultFeatureIds
|
||
+ model.info.meta.capabilities). Enables features like web_search,
|
||
code_interpreter, image_generation when the model has them as defaults
|
||
AND the capability is enabled AND the admin has enabled the feature.
|
||
"""
|
||
models = getattr(app.state, "MODELS", {})
|
||
model = models.get(model_id, {})
|
||
meta = model.get("info", {}).get("meta", {})
|
||
|
||
default_feature_ids = meta.get("defaultFeatureIds", [])
|
||
if not default_feature_ids:
|
||
return {}
|
||
|
||
capabilities = meta.get("capabilities", {})
|
||
config = app.state.config
|
||
features = {}
|
||
|
||
# code_interpreter is excluded: it requires the frontend event emitter
|
||
# and does not work in headless backend execution.
|
||
feature_checks = {
|
||
"web_search": getattr(config, "ENABLE_WEB_SEARCH", False),
|
||
"image_generation": getattr(config, "ENABLE_IMAGE_GENERATION", False),
|
||
}
|
||
|
||
for feature_id in default_feature_ids:
|
||
if feature_id in feature_checks:
|
||
# Feature must be: in defaultFeatureIds + capability enabled + admin enabled
|
||
if capabilities.get(feature_id) and feature_checks[feature_id]:
|
||
features[feature_id] = True
|
||
|
||
return features
|
||
|
||
|
||
def _resolve_model_filter_ids(app, model_id: str) -> list[str]:
|
||
"""Read model default filter_ids from model config."""
|
||
models = getattr(app.state, "MODELS", {})
|
||
model = models.get(model_id, {})
|
||
filter_ids = model.get("info", {}).get("meta", {}).get("defaultFilterIds", [])
|
||
return list(filter_ids) if filter_ids else []
|
||
|
||
|
||
async def _set_terminal_cwd(
|
||
app, server_id: str, user, cwd: str, chat_id: str
|
||
) -> None:
|
||
"""Set the working directory on a terminal server via the proxy.
|
||
|
||
Routes through the open-webui terminal proxy endpoint so that
|
||
auth headers, orchestrator policy routing, and X-User-Id are
|
||
handled correctly — same path the frontend uses.
|
||
"""
|
||
import aiohttp
|
||
|
||
connections = getattr(
|
||
getattr(app, 'state', None), 'config', None
|
||
)
|
||
if connections is None:
|
||
return
|
||
connections = getattr(connections, 'TERMINAL_SERVER_CONNECTIONS', None) or []
|
||
connection = next((c for c in connections if c.get('id') == server_id), None)
|
||
if connection is None:
|
||
log.warning(f'Terminal server {server_id} not found for CWD set')
|
||
return
|
||
|
||
base_url = (connection.get('url') or '').rstrip('/')
|
||
if not base_url:
|
||
return
|
||
|
||
# Build target URL — route through orchestrator policy if configured
|
||
policy_id = connection.get('policy_id')
|
||
if connection.get('server_type') == 'orchestrator' and policy_id:
|
||
target_url = f'{base_url}/p/{policy_id}/files/cwd'
|
||
else:
|
||
target_url = f'{base_url}/files/cwd'
|
||
|
||
headers = {'Content-Type': 'application/json', 'X-User-Id': user.id}
|
||
if chat_id:
|
||
headers['X-Session-Id'] = chat_id
|
||
|
||
auth_type = connection.get('auth_type', 'bearer')
|
||
if auth_type == 'bearer':
|
||
headers['Authorization'] = f'Bearer {connection.get("key", "")}'
|
||
|
||
try:
|
||
async with aiohttp.ClientSession(
|
||
timeout=aiohttp.ClientTimeout(total=10)
|
||
) as session:
|
||
async with session.post(
|
||
target_url,
|
||
json={'path': cwd},
|
||
headers=headers,
|
||
) as resp:
|
||
if resp.status != 200:
|
||
body = await resp.text()
|
||
log.warning(
|
||
f'Failed to set terminal CWD to {cwd}: '
|
||
f'HTTP {resp.status} — {body[:200]}'
|
||
)
|
||
except Exception as e:
|
||
log.warning(f'Failed to set terminal CWD: {e}')
|
||
|
||
|
||
async def execute_automation(app, automation: AutomationModel) -> None:
|
||
"""Execute an automation through the full chat completion pipeline.
|
||
|
||
Creates a real chat, then calls chat_completion exactly like the frontend:
|
||
session_id + chat_id + message_id → async task → pipeline handles everything
|
||
(filters, model params, knowledge/RAG, tools, DB saves, webhooks).
|
||
"""
|
||
try:
|
||
user = Users.get_user_by_id(automation.user_id)
|
||
if not user:
|
||
_record_run(automation.id, "error", error="User not found")
|
||
return
|
||
|
||
prompt = prompt_template(automation.data["prompt"], user)
|
||
model_id = automation.data["model_id"]
|
||
terminal_config = automation.data.get("terminal")
|
||
|
||
# Generate proper UUIDs for messages (same as frontend)
|
||
user_msg_id = str(uuid4())
|
||
assistant_msg_id = str(uuid4())
|
||
|
||
# Create the chat with user message (same structure as frontend)
|
||
chat = Chats.insert_new_chat(
|
||
automation.user_id,
|
||
ChatForm(
|
||
chat={
|
||
"title": automation.name,
|
||
"models": [model_id],
|
||
"history": {
|
||
"currentId": assistant_msg_id,
|
||
"messages": {
|
||
user_msg_id: {
|
||
"id": user_msg_id,
|
||
"parentId": None,
|
||
"role": "user",
|
||
"content": prompt,
|
||
"childrenIds": [assistant_msg_id],
|
||
"timestamp": int(time.time()),
|
||
"models": [model_id],
|
||
},
|
||
assistant_msg_id: {
|
||
"id": assistant_msg_id,
|
||
"parentId": user_msg_id,
|
||
"role": "assistant",
|
||
"content": "",
|
||
"done": False,
|
||
"model": model_id,
|
||
"childrenIds": [],
|
||
"timestamp": int(time.time()),
|
||
},
|
||
},
|
||
},
|
||
"messages": [
|
||
{"role": "user", "content": prompt},
|
||
],
|
||
"meta": {"automation_id": automation.id},
|
||
}
|
||
),
|
||
)
|
||
|
||
if not chat:
|
||
_record_run(automation.id, "error", error="Failed to create chat")
|
||
return
|
||
|
||
# Notify frontend to refresh chat list
|
||
from open_webui.socket.main import sio
|
||
|
||
await sio.emit(
|
||
"events",
|
||
{
|
||
"chat_id": chat.id,
|
||
"message_id": user_msg_id,
|
||
"data": {"type": "chat:list"},
|
||
},
|
||
room=f"user:{automation.user_id}",
|
||
)
|
||
|
||
# Resolve model defaults (frontend does this, backend doesn't)
|
||
tool_ids = _resolve_model_tool_ids(app, model_id)
|
||
features = _resolve_model_features(app, model_id)
|
||
filter_ids = _resolve_model_filter_ids(app, model_id)
|
||
|
||
# If a terminal is linked, set the CWD before building the payload
|
||
terminal_id = None
|
||
if terminal_config and terminal_config.get("server_id"):
|
||
terminal_id = terminal_config["server_id"]
|
||
cwd = terminal_config.get("cwd")
|
||
if cwd:
|
||
await _set_terminal_cwd(app, terminal_id, user, cwd, chat.id)
|
||
|
||
# Build the same payload the frontend sends to /api/chat/completions
|
||
form_data = {
|
||
"model": model_id,
|
||
"messages": [{"role": "user", "content": prompt}],
|
||
"stream": True,
|
||
"chat_id": chat.id,
|
||
"id": assistant_msg_id,
|
||
"parent_id": user_msg_id,
|
||
"session_id": f"automation:{automation.id}",
|
||
"background_tasks": {},
|
||
}
|
||
if tool_ids:
|
||
form_data["tool_ids"] = tool_ids
|
||
if features:
|
||
form_data["features"] = features
|
||
if filter_ids:
|
||
form_data["filter_ids"] = filter_ids
|
||
if terminal_id:
|
||
form_data["terminal_id"] = terminal_id
|
||
|
||
# Call the full chat completion pipeline (same as POST /api/chat/completions).
|
||
# The handler reference is stored on app.state to avoid circular imports.
|
||
request = _build_request(app)
|
||
await app.state.CHAT_COMPLETION_HANDLER(request, form_data, user=user)
|
||
|
||
# Notify user
|
||
from open_webui.socket.main import sio
|
||
|
||
await sio.emit(
|
||
"automation:result",
|
||
{
|
||
"automation_id": automation.id,
|
||
"name": automation.name,
|
||
"chat_id": chat.id,
|
||
"status": "success",
|
||
},
|
||
room=f"user:{automation.user_id}",
|
||
)
|
||
|
||
_record_run(automation.id, "success", chat_id=chat.id)
|
||
|
||
except Exception as e:
|
||
log.exception(f"Automation {automation.id} failed")
|
||
_record_run(automation.id, "error", error=str(e)[:4000])
|
||
|
||
|
||
####################
|
||
# Internals
|
||
####################
|
||
|
||
|
||
def _record_run(
|
||
automation_id: str,
|
||
status: str,
|
||
chat_id: str = None,
|
||
error: str = None,
|
||
):
|
||
"""Insert a run record into automation_run."""
|
||
with get_db() as db:
|
||
AutomationRuns.insert(
|
||
automation_id, status, chat_id=chat_id, error=error, db=db
|
||
)
|