feat: automation
This commit is contained in:
Timothy Jaeryang Baek
@@ -97,6 +97,7 @@ from open_webui.routers import (
|
||||
utils,
|
||||
scim,
|
||||
terminals,
|
||||
automations,
|
||||
)
|
||||
|
||||
from open_webui.routers.retrieval import (
|
||||
@@ -648,6 +649,9 @@ async def lifespan(app: FastAPI):
|
||||
asyncio.create_task(periodic_usage_pool_cleanup())
|
||||
asyncio.create_task(periodic_session_pool_cleanup())
|
||||
|
||||
from open_webui.utils.automations import automation_worker_loop
|
||||
asyncio.create_task(automation_worker_loop(app))
|
||||
|
||||
if app.state.config.ENABLE_BASE_MODELS_CACHE:
|
||||
try:
|
||||
await get_all_models(
|
||||
@@ -1522,6 +1526,7 @@ if ENABLE_ADMIN_ANALYTICS:
|
||||
app.include_router(analytics.router, prefix='/api/v1/analytics', tags=['analytics'])
|
||||
app.include_router(utils.router, prefix='/api/v1/utils', tags=['utils'])
|
||||
app.include_router(terminals.router, prefix='/api/v1/terminals', tags=['terminals'])
|
||||
app.include_router(automations.router, prefix='/api/v1/automations', tags=['automations'])
|
||||
|
||||
# SCIM 2.0 API for identity management
|
||||
if ENABLE_SCIM:
|
||||
|
||||
@@ -0,0 +1,60 @@
|
||||
"""add automation tables
|
||||
|
||||
Revision ID: d4e5f6a7b8c9
|
||||
Revises: f1e2d3c4b5a6
|
||||
Create Date: 2026-03-30
|
||||
"""
|
||||
|
||||
from typing import Union
|
||||
|
||||
from alembic import op
|
||||
import sqlalchemy as sa
|
||||
|
||||
revision: str = "d4e5f6a7b8c9"
|
||||
down_revision: Union[str, None] = "a3dd5bedd151"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade():
|
||||
op.create_table(
|
||||
"automation",
|
||||
sa.Column("id", sa.Text(), primary_key=True),
|
||||
sa.Column("user_id", sa.Text(), nullable=False),
|
||||
sa.Column("name", sa.Text(), nullable=False),
|
||||
sa.Column("data", sa.JSON(), nullable=False),
|
||||
sa.Column("meta", sa.JSON(), nullable=True),
|
||||
sa.Column(
|
||||
"is_active",
|
||||
sa.Boolean(),
|
||||
nullable=False,
|
||||
server_default=sa.text("1"),
|
||||
),
|
||||
sa.Column("last_run_at", sa.BigInteger(), nullable=True),
|
||||
sa.Column("next_run_at", sa.BigInteger(), nullable=True),
|
||||
sa.Column("created_at", sa.BigInteger(), nullable=False),
|
||||
sa.Column("updated_at", sa.BigInteger(), nullable=False),
|
||||
)
|
||||
op.create_index("ix_automation_next_run", "automation", ["next_run_at"])
|
||||
|
||||
op.create_table(
|
||||
"automation_run",
|
||||
sa.Column("id", sa.Text(), primary_key=True),
|
||||
sa.Column("automation_id", sa.Text(), nullable=False),
|
||||
sa.Column("chat_id", sa.Text(), nullable=True),
|
||||
sa.Column("status", sa.Text(), nullable=False),
|
||||
sa.Column("error", sa.Text(), nullable=True),
|
||||
sa.Column("created_at", sa.BigInteger(), nullable=False),
|
||||
)
|
||||
op.create_index(
|
||||
"ix_automation_run_automation_id",
|
||||
"automation_run",
|
||||
["automation_id"],
|
||||
)
|
||||
|
||||
|
||||
def downgrade():
|
||||
op.drop_index("ix_automation_run_automation_id")
|
||||
op.drop_table("automation_run")
|
||||
op.drop_index("ix_automation_next_run")
|
||||
op.drop_table("automation")
|
||||
@@ -0,0 +1,323 @@
|
||||
import time
|
||||
import logging
|
||||
from typing import Optional
|
||||
from uuid import uuid4
|
||||
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
from sqlalchemy import Column, Text, JSON, Boolean, BigInteger, Index, select
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from open_webui.internal.db import Base, get_db, get_db_context
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
####################
|
||||
# Automation DB Schema
|
||||
####################
|
||||
|
||||
|
||||
class Automation(Base):
|
||||
__tablename__ = "automation"
|
||||
|
||||
id = Column(Text, primary_key=True)
|
||||
user_id = Column(Text, nullable=False)
|
||||
name = Column(Text, nullable=False)
|
||||
data = Column(JSON, nullable=False) # {prompt, model_id, rrule}
|
||||
meta = Column(JSON, nullable=True)
|
||||
is_active = Column(Boolean, nullable=False, default=True)
|
||||
last_run_at = Column(BigInteger, nullable=True)
|
||||
next_run_at = Column(BigInteger, nullable=True)
|
||||
|
||||
created_at = Column(BigInteger, nullable=False)
|
||||
updated_at = Column(BigInteger, nullable=False)
|
||||
|
||||
__table_args__ = (
|
||||
Index("ix_automation_next_run", "next_run_at"),
|
||||
)
|
||||
|
||||
|
||||
class AutomationRun(Base):
|
||||
__tablename__ = "automation_run"
|
||||
|
||||
id = Column(Text, primary_key=True)
|
||||
automation_id = Column(Text, nullable=False)
|
||||
chat_id = Column(Text, nullable=True)
|
||||
status = Column(Text, nullable=False) # success | error
|
||||
error = Column(Text, nullable=True)
|
||||
created_at = Column(BigInteger, nullable=False)
|
||||
|
||||
__table_args__ = (
|
||||
Index("ix_automation_run_automation_id", "automation_id"),
|
||||
)
|
||||
|
||||
|
||||
####################
|
||||
# Pydantic Models
|
||||
####################
|
||||
|
||||
|
||||
class AutomationData(BaseModel):
|
||||
prompt: str
|
||||
model_id: str
|
||||
rrule: str
|
||||
|
||||
|
||||
class AutomationModel(BaseModel):
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
id: str
|
||||
user_id: str
|
||||
name: str
|
||||
data: dict
|
||||
meta: Optional[dict] = None
|
||||
is_active: bool
|
||||
last_run_at: Optional[int] = None
|
||||
next_run_at: Optional[int] = None
|
||||
|
||||
created_at: int
|
||||
updated_at: int
|
||||
|
||||
|
||||
class AutomationRunModel(BaseModel):
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
id: str
|
||||
automation_id: str
|
||||
chat_id: Optional[str] = None
|
||||
status: str
|
||||
error: Optional[str] = None
|
||||
created_at: int
|
||||
|
||||
|
||||
class AutomationForm(BaseModel):
|
||||
name: str
|
||||
data: AutomationData
|
||||
meta: Optional[dict] = None
|
||||
is_active: Optional[bool] = True
|
||||
|
||||
|
||||
class AutomationResponse(AutomationModel):
|
||||
last_run: Optional[AutomationRunModel] = None
|
||||
next_runs: Optional[list[int]] = None
|
||||
|
||||
|
||||
####################
|
||||
# AutomationTable
|
||||
####################
|
||||
|
||||
|
||||
class AutomationTable:
|
||||
|
||||
def insert(
|
||||
self,
|
||||
user_id: str,
|
||||
form: AutomationForm,
|
||||
next_run_at: int,
|
||||
db: Optional[Session] = None,
|
||||
) -> AutomationModel:
|
||||
with get_db_context(db) as db:
|
||||
now = int(time.time_ns())
|
||||
row = Automation(
|
||||
id=str(uuid4()),
|
||||
user_id=user_id,
|
||||
name=form.name,
|
||||
data=form.data.model_dump(),
|
||||
meta=form.meta,
|
||||
is_active=form.is_active,
|
||||
next_run_at=next_run_at,
|
||||
created_at=now,
|
||||
updated_at=now,
|
||||
)
|
||||
db.add(row)
|
||||
db.commit()
|
||||
db.refresh(row)
|
||||
return AutomationModel.model_validate(row)
|
||||
|
||||
def get_by_id(
|
||||
self, id: str, db: Optional[Session] = None
|
||||
) -> Optional[AutomationModel]:
|
||||
with get_db_context(db) as db:
|
||||
row = db.get(Automation, id)
|
||||
return AutomationModel.model_validate(row) if row else None
|
||||
|
||||
def get_by_user(
|
||||
self, user_id: str, db: Optional[Session] = None
|
||||
) -> list[AutomationModel]:
|
||||
with get_db_context(db) as db:
|
||||
rows = (
|
||||
db.query(Automation)
|
||||
.filter_by(user_id=user_id)
|
||||
.order_by(Automation.created_at.desc())
|
||||
.all()
|
||||
)
|
||||
return [AutomationModel.model_validate(r) for r in rows]
|
||||
|
||||
def get_all(self, db: Optional[Session] = None) -> list[AutomationModel]:
|
||||
with get_db_context(db) as db:
|
||||
rows = (
|
||||
db.query(Automation)
|
||||
.order_by(Automation.created_at.desc())
|
||||
.all()
|
||||
)
|
||||
return [AutomationModel.model_validate(r) for r in rows]
|
||||
|
||||
def update_by_id(
|
||||
self,
|
||||
id: str,
|
||||
form: AutomationForm,
|
||||
next_run_at: int,
|
||||
db: Optional[Session] = None,
|
||||
) -> Optional[AutomationModel]:
|
||||
with get_db_context(db) as db:
|
||||
row = db.get(Automation, id)
|
||||
if not row:
|
||||
return None
|
||||
row.name = form.name
|
||||
row.data = form.data.model_dump()
|
||||
row.meta = form.meta
|
||||
if form.is_active is not None:
|
||||
row.is_active = form.is_active
|
||||
row.next_run_at = next_run_at
|
||||
row.updated_at = int(time.time_ns())
|
||||
db.commit()
|
||||
db.refresh(row)
|
||||
return AutomationModel.model_validate(row)
|
||||
|
||||
def toggle(
|
||||
self,
|
||||
id: str,
|
||||
next_run_at: Optional[int],
|
||||
db: Optional[Session] = None,
|
||||
) -> Optional[AutomationModel]:
|
||||
with get_db_context(db) as db:
|
||||
row = db.get(Automation, id)
|
||||
if not row:
|
||||
return None
|
||||
row.is_active = not row.is_active
|
||||
row.next_run_at = next_run_at if row.is_active else None
|
||||
row.updated_at = int(time.time_ns())
|
||||
db.commit()
|
||||
db.refresh(row)
|
||||
return AutomationModel.model_validate(row)
|
||||
|
||||
def delete(self, id: str, db: Optional[Session] = None) -> bool:
|
||||
with get_db_context(db) as db:
|
||||
row = db.get(Automation, id)
|
||||
if not row:
|
||||
return False
|
||||
db.delete(row)
|
||||
db.commit()
|
||||
return True
|
||||
|
||||
def claim_due(
|
||||
self, now_ns: int, limit: int = 10, db: Optional[Session] = None
|
||||
) -> list[AutomationModel]:
|
||||
"""
|
||||
Atomically claim due automations for execution.
|
||||
|
||||
Advances next_run_at immediately so the row can never be
|
||||
double-claimed. On PostgreSQL, uses FOR UPDATE SKIP LOCKED
|
||||
for zero-contention distributed work claiming.
|
||||
"""
|
||||
with get_db_context(db) as db:
|
||||
stmt = (
|
||||
select(Automation)
|
||||
.where(
|
||||
Automation.is_active == True,
|
||||
Automation.next_run_at <= now_ns,
|
||||
)
|
||||
.order_by(Automation.next_run_at)
|
||||
.limit(limit)
|
||||
)
|
||||
|
||||
if db.bind.dialect.name == "postgresql":
|
||||
stmt = stmt.with_for_update(skip_locked=True)
|
||||
|
||||
rows = db.execute(stmt).scalars().all()
|
||||
|
||||
from open_webui.utils.automations import next_run_ns
|
||||
|
||||
for row in rows:
|
||||
row.last_run_at = now_ns
|
||||
row.next_run_at = next_run_ns(row.data.get("rrule", ""))
|
||||
|
||||
db.commit()
|
||||
|
||||
return [AutomationModel.model_validate(r) for r in rows]
|
||||
|
||||
|
||||
####################
|
||||
# AutomationRunTable
|
||||
####################
|
||||
|
||||
|
||||
class AutomationRunTable:
|
||||
|
||||
def insert(
|
||||
self,
|
||||
automation_id: str,
|
||||
status: str,
|
||||
chat_id: Optional[str] = None,
|
||||
error: Optional[str] = None,
|
||||
db: Optional[Session] = None,
|
||||
) -> AutomationRunModel:
|
||||
with get_db_context(db) as db:
|
||||
row = AutomationRun(
|
||||
id=str(uuid4()),
|
||||
automation_id=automation_id,
|
||||
chat_id=chat_id,
|
||||
status=status,
|
||||
error=error,
|
||||
created_at=int(time.time_ns()),
|
||||
)
|
||||
db.add(row)
|
||||
db.commit()
|
||||
db.refresh(row)
|
||||
return AutomationRunModel.model_validate(row)
|
||||
|
||||
def get_latest(
|
||||
self, automation_id: str, db: Optional[Session] = None
|
||||
) -> Optional[AutomationRunModel]:
|
||||
with get_db_context(db) as db:
|
||||
row = (
|
||||
db.query(AutomationRun)
|
||||
.filter_by(automation_id=automation_id)
|
||||
.order_by(AutomationRun.created_at.desc())
|
||||
.first()
|
||||
)
|
||||
return AutomationRunModel.model_validate(row) if row else None
|
||||
|
||||
def get_by_automation(
|
||||
self,
|
||||
automation_id: str,
|
||||
skip: int = 0,
|
||||
limit: int = 50,
|
||||
db: Optional[Session] = None,
|
||||
) -> list[AutomationRunModel]:
|
||||
with get_db_context(db) as db:
|
||||
rows = (
|
||||
db.query(AutomationRun)
|
||||
.filter_by(automation_id=automation_id)
|
||||
.order_by(AutomationRun.created_at.desc())
|
||||
.offset(skip)
|
||||
.limit(limit)
|
||||
.all()
|
||||
)
|
||||
return [AutomationRunModel.model_validate(r) for r in rows]
|
||||
|
||||
def delete_by_automation(
|
||||
self, automation_id: str, db: Optional[Session] = None
|
||||
) -> int:
|
||||
with get_db_context(db) as db:
|
||||
count = (
|
||||
db.query(AutomationRun)
|
||||
.filter_by(automation_id=automation_id)
|
||||
.delete()
|
||||
)
|
||||
db.commit()
|
||||
return count
|
||||
|
||||
|
||||
Automations = AutomationTable()
|
||||
AutomationRuns = AutomationRunTable()
|
||||
@@ -0,0 +1,235 @@
|
||||
import asyncio
|
||||
import logging
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, status
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from open_webui.models.automations import (
|
||||
Automations,
|
||||
AutomationRuns,
|
||||
AutomationForm,
|
||||
AutomationModel,
|
||||
AutomationResponse,
|
||||
AutomationRunModel,
|
||||
)
|
||||
from open_webui.utils.automations import (
|
||||
validate_rrule,
|
||||
next_run_ns,
|
||||
next_n_runs_ns,
|
||||
execute_automation,
|
||||
)
|
||||
from open_webui.utils.auth import get_verified_user, get_admin_user
|
||||
from open_webui.internal.db import get_session
|
||||
from open_webui.constants import ERROR_MESSAGES
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
############################
|
||||
# Helpers
|
||||
############################
|
||||
|
||||
|
||||
def check_automation_access(automation, user):
|
||||
if not automation:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
if user.role != 'admin' and user.id != automation.user_id:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail=ERROR_MESSAGES.UNAUTHORIZED,
|
||||
)
|
||||
|
||||
|
||||
def enrich_automation(
|
||||
automation: AutomationModel, db: Session, tz: str = None
|
||||
) -> AutomationResponse:
|
||||
last_run = AutomationRuns.get_latest(automation.id, db=db)
|
||||
return AutomationResponse(
|
||||
**automation.model_dump(),
|
||||
last_run=last_run,
|
||||
next_runs=next_n_runs_ns(automation.data['rrule'], tz=tz),
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
# GetAutomations
|
||||
############################
|
||||
|
||||
|
||||
@router.get('/', response_model=list[AutomationResponse])
|
||||
async def get_automations(
|
||||
request: Request,
|
||||
user=Depends(get_verified_user),
|
||||
db: Session = Depends(get_session),
|
||||
):
|
||||
automations = Automations.get_by_user(user.id, db=db)
|
||||
return [enrich_automation(automation, db, tz=user.timezone) for automation in automations]
|
||||
|
||||
|
||||
############################
|
||||
# CreateNewAutomation
|
||||
############################
|
||||
|
||||
|
||||
@router.post('/create', response_model=AutomationResponse)
|
||||
async def create_new_automation(
|
||||
request: Request,
|
||||
form_data: AutomationForm,
|
||||
user=Depends(get_verified_user),
|
||||
db: Session = Depends(get_session),
|
||||
):
|
||||
try:
|
||||
validate_rrule(form_data.data.rrule)
|
||||
except ValueError as e:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=str(e),
|
||||
)
|
||||
|
||||
tz = user.timezone
|
||||
automation = Automations.insert(
|
||||
user.id, form_data, next_run_ns(form_data.data.rrule, tz=tz), db=db
|
||||
)
|
||||
return enrich_automation(automation, db, tz=tz)
|
||||
|
||||
|
||||
############################
|
||||
# GetAutomationById
|
||||
############################
|
||||
|
||||
|
||||
@router.get('/{id}', response_model=AutomationResponse)
|
||||
async def get_automation_by_id(
|
||||
request: Request,
|
||||
id: str,
|
||||
user=Depends(get_verified_user),
|
||||
db: Session = Depends(get_session),
|
||||
):
|
||||
automation = Automations.get_by_id(id, db=db)
|
||||
check_automation_access(automation, user)
|
||||
return enrich_automation(automation, db, tz=user.timezone)
|
||||
|
||||
|
||||
############################
|
||||
# UpdateAutomationById
|
||||
############################
|
||||
|
||||
|
||||
@router.post('/{id}/update', response_model=AutomationResponse)
|
||||
async def update_automation_by_id(
|
||||
request: Request,
|
||||
id: str,
|
||||
form_data: AutomationForm,
|
||||
user=Depends(get_verified_user),
|
||||
db: Session = Depends(get_session),
|
||||
):
|
||||
automation = Automations.get_by_id(id, db=db)
|
||||
check_automation_access(automation, user)
|
||||
|
||||
try:
|
||||
validate_rrule(form_data.data.rrule)
|
||||
except ValueError as e:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=str(e),
|
||||
)
|
||||
|
||||
tz = user.timezone
|
||||
updated = Automations.update_by_id(
|
||||
id, form_data, next_run_ns(form_data.data.rrule, tz=tz), db=db
|
||||
)
|
||||
return enrich_automation(updated, db, tz=tz)
|
||||
|
||||
|
||||
############################
|
||||
# ToggleAutomationById
|
||||
############################
|
||||
|
||||
|
||||
@router.post('/{id}/toggle', response_model=AutomationResponse)
|
||||
async def toggle_automation_by_id(
|
||||
request: Request,
|
||||
id: str,
|
||||
user=Depends(get_verified_user),
|
||||
db: Session = Depends(get_session),
|
||||
):
|
||||
automation = Automations.get_by_id(id, db=db)
|
||||
check_automation_access(automation, user)
|
||||
toggled = Automations.toggle(
|
||||
id, next_run_ns(automation.data['rrule'], tz=user.timezone), db=db
|
||||
)
|
||||
return enrich_automation(toggled, db, tz=user.timezone)
|
||||
|
||||
|
||||
############################
|
||||
# RunAutomationById
|
||||
############################
|
||||
|
||||
|
||||
@router.post('/{id}/run')
|
||||
async def run_automation_by_id(
|
||||
request: Request,
|
||||
id: str,
|
||||
user=Depends(get_verified_user),
|
||||
db: Session = Depends(get_session),
|
||||
):
|
||||
automation = Automations.get_by_id(id, db=db)
|
||||
check_automation_access(automation, user)
|
||||
asyncio.create_task(execute_automation(request.app, automation))
|
||||
return enrich_automation(automation, db, tz=user.timezone)
|
||||
|
||||
|
||||
############################
|
||||
# DeleteAutomationById
|
||||
############################
|
||||
|
||||
|
||||
@router.delete('/{id}/delete')
|
||||
async def delete_automation_by_id(
|
||||
request: Request,
|
||||
id: str,
|
||||
user=Depends(get_verified_user),
|
||||
db: Session = Depends(get_session),
|
||||
):
|
||||
automation = Automations.get_by_id(id, db=db)
|
||||
check_automation_access(automation, user)
|
||||
AutomationRuns.delete_by_automation(id, db=db)
|
||||
return Automations.delete(id, db=db)
|
||||
|
||||
|
||||
############################
|
||||
# GetAutomationRuns
|
||||
############################
|
||||
|
||||
|
||||
@router.get('/{id}/runs', response_model=list[AutomationRunModel])
|
||||
async def get_automation_runs(
|
||||
request: Request,
|
||||
id: str,
|
||||
skip: int = 0,
|
||||
limit: int = 50,
|
||||
user=Depends(get_verified_user),
|
||||
db: Session = Depends(get_session),
|
||||
):
|
||||
automation = Automations.get_by_id(id, db=db)
|
||||
check_automation_access(automation, user)
|
||||
return AutomationRuns.get_by_automation(id, skip=skip, limit=limit, db=db)
|
||||
|
||||
|
||||
############################
|
||||
# Admin Endpoints
|
||||
############################
|
||||
|
||||
|
||||
@router.get('/admin/all', response_model=list[AutomationModel])
|
||||
async def get_all_automations(
|
||||
request: Request,
|
||||
user=Depends(get_admin_user),
|
||||
db: Session = Depends(get_session),
|
||||
):
|
||||
return Automations.get_all(db=db)
|
||||
@@ -0,0 +1,338 @@
|
||||
"""
|
||||
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 validate_rrule(s: str) -> None:
|
||||
"""Raise ValueError if the RRULE is malformed or exhausted."""
|
||||
try:
|
||||
rule = rrulestr(s, ignoretz=True)
|
||||
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 = rrulestr(s, ignoretz=True).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 = rrulestr(s, ignoretz=True)
|
||||
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 = {}
|
||||
|
||||
feature_checks = {
|
||||
"web_search": getattr(config, "ENABLE_WEB_SEARCH", False),
|
||||
"image_generation": getattr(config, "ENABLE_IMAGE_GENERATION", False),
|
||||
"code_interpreter": getattr(config, "ENABLE_CODE_INTERPRETER", 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 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"]
|
||||
|
||||
# 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": f"[Automation] {automation.name}",
|
||||
"models": [model_id],
|
||||
"history": {
|
||||
"currentId": user_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": "",
|
||||
"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)
|
||||
|
||||
# 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
|
||||
|
||||
# Call chat_completion — returns {'status': True, 'task_id': '...'}
|
||||
# The background task handles everything: LLM, tools, DB, webhooks.
|
||||
from open_webui.main import chat_completion as main_chat_completion
|
||||
|
||||
request = _build_request(app)
|
||||
await main_chat_completion(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
|
||||
)
|
||||
@@ -2750,16 +2750,22 @@ async def process_chat_payload(request, form_data, user, metadata, model):
|
||||
def get_event_emitter_and_caller(metadata):
|
||||
event_emitter = None
|
||||
event_caller = None
|
||||
if (
|
||||
'session_id' in metadata
|
||||
and metadata['session_id']
|
||||
and 'chat_id' in metadata
|
||||
and metadata['chat_id']
|
||||
and 'message_id' in metadata
|
||||
and metadata['message_id']
|
||||
):
|
||||
|
||||
# event_emitter only needs user_id + chat_id + message_id.
|
||||
# It broadcasts to user:{user_id} room AND persists to DB,
|
||||
# so it works for backend-initiated calls (automations, API).
|
||||
if metadata.get('chat_id') and metadata.get('message_id'):
|
||||
event_emitter = get_event_emitter(metadata)
|
||||
|
||||
# event_caller needs session_id — it calls back to a specific
|
||||
# websocket session (used by direct tools, pyodide code interpreter).
|
||||
if (
|
||||
metadata.get('session_id')
|
||||
and metadata.get('chat_id')
|
||||
and metadata.get('message_id')
|
||||
):
|
||||
event_caller = get_event_call(metadata)
|
||||
|
||||
return event_emitter, event_caller
|
||||
|
||||
|
||||
@@ -3206,7 +3212,9 @@ async def streaming_chat_response_handler(response, ctx):
|
||||
]
|
||||
|
||||
# Standard streaming response handler
|
||||
if event_emitter and event_caller:
|
||||
# event_caller is optional — only needed for direct (client-side) tools
|
||||
# and pyodide code interpreter. Server-side tools work without it.
|
||||
if event_emitter:
|
||||
task_id = str(uuid4()) # Create a unique task ID.
|
||||
model_id = form_data.get('model', '')
|
||||
|
||||
|
||||
Reference in New Issue
Block a user