mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-04 16:13:23 +00:00
arc53-machine
5d0992eef8
run_agent_headless records each unattended run under its endpoint; the scheduler passes its run id and the webhook worker its task id so Logs rows can find their trace, while the LLM's own request id stays untouched for quota counts. /api/search and MCP search_docs record their retrieval, and a graph build records every extraction call under one step.
459 lines
17 KiB
Python
459 lines
17 KiB
Python
"""Body of ``execute_scheduled_run`` — runs a single agent execution.
|
|
|
|
Not a DURABLE_TASK: agent runs have side effects (messages, CRM writes)
|
|
and blind auto-retry would double them. Failures after agent.gen starts
|
|
are terminal and recorded; only the pre-start load is retry-safe.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from datetime import datetime, timezone
|
|
from typing import Any, Dict, Optional
|
|
|
|
from celery.exceptions import SoftTimeLimitExceeded
|
|
from sqlalchemy import text as sql_text
|
|
|
|
from docsgpt.agents.headless_runner import run_agent_headless
|
|
from docsgpt.core.settings import settings
|
|
from docsgpt.events.publisher import publish_user_event
|
|
from docsgpt.quotas.service import QuotaExceededError
|
|
from docsgpt.storage.db.base_repository import row_to_dict
|
|
from docsgpt.storage.db.engine import get_engine
|
|
from docsgpt.storage.db.repositories.conversations import (
|
|
ConversationsRepository,
|
|
)
|
|
from docsgpt.storage.db.repositories.schedule_runs import (
|
|
ScheduleRunsRepository,
|
|
)
|
|
from docsgpt.storage.db.repositories.schedules import SchedulesRepository
|
|
from docsgpt.storage.db.repositories.token_usage import TokenUsageRepository
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# Cap output verbatim in the run log; beyond the cap we keep the head and stamp output_truncated.
|
|
_OUTPUT_CAP_CHARS = 24_000
|
|
|
|
|
|
def _agent_config_for_schedule(schedule: Dict[str, Any]) -> Optional[Dict[str, Any]]:
|
|
"""Resolve the agent row (agent-bound) or build an ephemeral classic config.
|
|
|
|
For agentless schedules (``agent_id IS NULL``), the worker constructs an
|
|
in-memory agent shape carrying just enough fields for ``run_agent_headless``:
|
|
classic agent type, system-default retriever/chunks/prompt, no source, and
|
|
the optional ``model_id`` override. The runtime toolset is rebuilt by
|
|
``ToolExecutor`` at fire time (current ``user_tools`` + non-disabled,
|
|
non-headless-excluded defaults), so a snapshot here would be dead code.
|
|
"""
|
|
if schedule.get("agent_id"):
|
|
engine = get_engine()
|
|
with engine.connect() as conn:
|
|
row = conn.execute(
|
|
sql_text("SELECT * FROM agents WHERE id = CAST(:id AS uuid)"),
|
|
{"id": str(schedule["agent_id"])},
|
|
).fetchone()
|
|
return row_to_dict(row) if row is not None else None
|
|
return _ephemeral_agent_for_agentless(schedule)
|
|
|
|
|
|
def _ephemeral_agent_for_agentless(
|
|
schedule: Dict[str, Any],
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""Build an agent-shaped config for a schedule with no parent agent."""
|
|
# ``agent_config["tools"]`` is intentionally omitted: ``run_agent_headless``
|
|
# never reads it. The runtime toolset is rebuilt by
|
|
# ``ToolExecutor._get_user_tools(owner)`` at fire time — same dereference
|
|
# the agent-bound path uses, so a tool added/disabled after creation is
|
|
# reflected. Headless mode there filters chat-only tools (``scheduler``).
|
|
user_id = schedule.get("user_id")
|
|
if not user_id:
|
|
return None
|
|
return {
|
|
"id": None,
|
|
"user_id": user_id,
|
|
"agent_type": "classic",
|
|
"retriever": "classic",
|
|
"chunks": 2,
|
|
"prompt_id": "default",
|
|
"source_id": None,
|
|
"default_model_id": schedule.get("model_id") or "",
|
|
}
|
|
|
|
|
|
def _load_chat_history(schedule: Dict[str, Any]) -> list:
|
|
"""Originating conversation history (one-time only; recurring has none)."""
|
|
origin = schedule.get("origin_conversation_id")
|
|
if not origin or schedule.get("trigger_type") != "once":
|
|
return []
|
|
user_id = schedule.get("user_id")
|
|
if not user_id:
|
|
return []
|
|
try:
|
|
engine = get_engine()
|
|
with engine.connect() as conn:
|
|
conv = ConversationsRepository(conn).get_any(str(origin), user_id)
|
|
if conv is None:
|
|
return []
|
|
messages = ConversationsRepository(conn).get_messages(str(conv["id"]))
|
|
except Exception:
|
|
logger.exception("scheduler: failed loading chat history")
|
|
return []
|
|
history: list = []
|
|
for msg in messages:
|
|
if msg.get("prompt") and msg.get("response"):
|
|
history.append({
|
|
"prompt": msg["prompt"],
|
|
"response": msg["response"],
|
|
})
|
|
return history
|
|
|
|
|
|
def _publish_run_event(
|
|
event_type: str, run: Dict[str, Any], schedule: Dict[str, Any], **extra: Any
|
|
) -> None:
|
|
"""Best-effort SSE publish for a scheduler run state transition."""
|
|
user_id = run.get("user_id") or schedule.get("user_id")
|
|
if not user_id:
|
|
return
|
|
agent_id_raw = schedule.get("agent_id")
|
|
payload = {
|
|
"run_id": str(run["id"]),
|
|
"schedule_id": str(schedule["id"]),
|
|
"agent_id": str(agent_id_raw) if agent_id_raw else None,
|
|
"trigger_type": schedule.get("trigger_type"),
|
|
"status": run.get("status"),
|
|
**extra,
|
|
}
|
|
try:
|
|
publish_user_event(
|
|
user_id,
|
|
event_type,
|
|
payload,
|
|
scope={"kind": "schedule", "id": str(schedule["id"])},
|
|
)
|
|
except Exception:
|
|
logger.exception(
|
|
"scheduler: SSE publish failed event=%s run=%s",
|
|
event_type, run.get("id"),
|
|
)
|
|
|
|
|
|
def _publish_message_appended(
|
|
user_id: str,
|
|
conversation_id: str,
|
|
message: Dict[str, Any],
|
|
schedule_id: str,
|
|
run_id: str,
|
|
) -> None:
|
|
"""SSE message-appended event for a one-time run's chat turn."""
|
|
try:
|
|
publish_user_event(
|
|
user_id,
|
|
"schedule.message.appended",
|
|
{
|
|
"conversation_id": str(conversation_id),
|
|
"message_id": str(message["id"]),
|
|
"schedule_id": str(schedule_id),
|
|
"run_id": str(run_id),
|
|
"position": int(message.get("position", 0)),
|
|
},
|
|
scope={"kind": "conversation", "id": str(conversation_id)},
|
|
)
|
|
except Exception:
|
|
logger.exception(
|
|
"scheduler: message.appended publish failed run=%s", run_id,
|
|
)
|
|
|
|
|
|
def _append_one_time_turn(
|
|
schedule: Dict[str, Any],
|
|
run: Dict[str, Any],
|
|
outcome: Dict[str, Any],
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""Insert an assistant turn in the originating conversation (once only)."""
|
|
origin = schedule.get("origin_conversation_id")
|
|
if not origin:
|
|
return None
|
|
engine = get_engine()
|
|
user_id = schedule.get("user_id")
|
|
metadata = {
|
|
"scheduled": True,
|
|
"schedule_id": str(schedule["id"]),
|
|
"run_id": str(run["id"]),
|
|
"scheduled_run_at": (
|
|
run.get("scheduled_for")
|
|
if isinstance(run.get("scheduled_for"), str)
|
|
else None
|
|
),
|
|
}
|
|
with engine.begin() as conn:
|
|
conv = ConversationsRepository(conn).get_any(str(origin), user_id)
|
|
if conv is None:
|
|
return None
|
|
message = ConversationsRepository(conn).append_message(
|
|
str(conv["id"]),
|
|
{
|
|
"prompt": schedule.get("instruction") or "",
|
|
"response": outcome.get("answer") or "",
|
|
"thought": outcome.get("thought") or "",
|
|
"sources": outcome.get("sources") or [],
|
|
"tool_calls": outcome.get("tool_calls") or [],
|
|
"model_id": outcome.get("model_id"),
|
|
"metadata": metadata,
|
|
},
|
|
)
|
|
return message
|
|
|
|
|
|
def execute_scheduled_run_body(run_id: str, celery_task_id: Optional[str]) -> Dict[str, Any]:
|
|
"""Execute one scheduled run by id; returns a result dict for tracing."""
|
|
if not settings.POSTGRES_URI:
|
|
return {"status": "skipped", "reason": "POSTGRES_URI not set"}
|
|
|
|
engine = get_engine()
|
|
|
|
with engine.connect() as conn:
|
|
run = ScheduleRunsRepository(conn).get_internal(run_id)
|
|
if run is None:
|
|
return {"status": "skipped", "reason": "run not found"}
|
|
schedule = SchedulesRepository(conn).get_internal(str(run["schedule_id"]))
|
|
if schedule is None:
|
|
return {"status": "skipped", "reason": "schedule not found"}
|
|
|
|
# Refuse non-runnable terminal states; manual run-now bypasses.
|
|
if run.get("status") != "pending":
|
|
return {"status": "skipped", "reason": f"run status={run.get('status')}"}
|
|
if schedule.get("status") in {"cancelled", "completed"} and run.get(
|
|
"trigger_source"
|
|
) != "manual":
|
|
with engine.begin() as conn:
|
|
ScheduleRunsRepository(conn).update(
|
|
run_id,
|
|
{
|
|
"status": "skipped",
|
|
"finished_at": datetime.now(timezone.utc),
|
|
"error_type": "internal",
|
|
"error": "schedule no longer active",
|
|
},
|
|
)
|
|
return {"status": "skipped", "reason": "schedule terminal"}
|
|
|
|
agent_config = _agent_config_for_schedule(schedule)
|
|
if agent_config is None:
|
|
with engine.begin() as conn:
|
|
updated = ScheduleRunsRepository(conn).update(
|
|
run_id,
|
|
{
|
|
"status": "failed",
|
|
"finished_at": datetime.now(timezone.utc),
|
|
"error_type": "internal",
|
|
"error": "agent missing",
|
|
},
|
|
)
|
|
SchedulesRepository(conn).bump_failure_count(str(schedule["id"]))
|
|
_publish_run_event("schedule.run.failed", updated or run, schedule,
|
|
error="agent missing")
|
|
return {"status": "failed", "reason": "agent missing"}
|
|
|
|
with engine.begin() as conn:
|
|
if not ScheduleRunsRepository(conn).mark_running(run_id, celery_task_id):
|
|
return {"status": "skipped", "reason": "lost race to mark_running"}
|
|
|
|
started = datetime.now(timezone.utc)
|
|
instruction = schedule.get("instruction") or ""
|
|
allowlist = schedule.get("tool_allowlist") or []
|
|
chat_history = _load_chat_history(schedule)
|
|
outcome: Dict[str, Any]
|
|
error_type: Optional[str] = None
|
|
error_text: Optional[str] = None
|
|
timed_out = False
|
|
try:
|
|
outcome = run_agent_headless(
|
|
agent_config,
|
|
instruction,
|
|
tool_allowlist=allowlist,
|
|
model_id_override=schedule.get("model_id"),
|
|
endpoint="schedule",
|
|
conversation_id=schedule.get("origin_conversation_id"),
|
|
chat_history=chat_history,
|
|
# Links the run's execution trace to its Logs row.
|
|
request_id=str(run_id),
|
|
)
|
|
except SoftTimeLimitExceeded:
|
|
timed_out = True
|
|
outcome = {"answer": "", "tool_calls": [], "sources": [], "thought": ""}
|
|
error_type = "timeout"
|
|
error_text = "run exceeded soft time limit"
|
|
except QuotaExceededError as exc:
|
|
# The owner's usage quota is spent; the run never started.
|
|
outcome = {"answer": "", "tool_calls": [], "sources": [], "thought": ""}
|
|
error_type = "budget_exceeded"
|
|
error_text = str(exc)
|
|
except Exception as exc:
|
|
outcome = {"answer": "", "tool_calls": [], "sources": [], "thought": ""}
|
|
error_type = "agent_error"
|
|
error_text = str(exc)
|
|
logger.exception("scheduler: agent run failed run=%s", run_id)
|
|
|
|
finished = datetime.now(timezone.utc)
|
|
|
|
# Honour the outcome the runner reported: it classifies headless denials
|
|
# (``tool_not_allowed``) and mid-stream failures (``stream_error``) alike,
|
|
# and owns the message. Re-deriving the denial case here only duplicated
|
|
# the runner's rule while leaving a failed stream — which returns normally,
|
|
# with an empty answer — to fall through to "success".
|
|
if error_type is None and outcome.get("error_type"):
|
|
error_type = str(outcome["error_type"])
|
|
error_text = str(outcome.get("error") or "") or error_type
|
|
|
|
prompt_tokens = int(outcome.get("prompt_tokens", 0) or 0)
|
|
generated_tokens = int(outcome.get("generated_tokens", 0) or 0)
|
|
used_tokens = prompt_tokens + generated_tokens
|
|
if (
|
|
schedule.get("token_budget") is not None
|
|
and int(schedule["token_budget"]) > 0
|
|
and used_tokens > int(schedule["token_budget"])
|
|
):
|
|
error_type = "budget_exceeded"
|
|
error_text = (
|
|
f"used {used_tokens} tokens exceeds budget "
|
|
f"{schedule['token_budget']}"
|
|
)
|
|
|
|
# Backstop for silent failures that raise nothing and emit no error event:
|
|
# a run that produced no answer, did no work and burned no completion
|
|
# tokens did not do its job, whatever the absence of an exception suggests.
|
|
# "Work" is deliberately broad — tool calls for a chat agent, completed
|
|
# nodes for a workflow — so a schedule whose whole purpose is a side effect
|
|
# (post a message, file a ticket) is never flagged for staying quiet.
|
|
if (
|
|
error_type is None
|
|
and not (outcome.get("answer") or "").strip()
|
|
and not (outcome.get("tool_calls") or [])
|
|
and not int(outcome.get("steps_completed") or 0)
|
|
and generated_tokens == 0
|
|
):
|
|
error_type = "empty_output"
|
|
error_text = "run produced no answer, tool calls or completion tokens"
|
|
|
|
answer = outcome.get("answer") or ""
|
|
truncated = False
|
|
if len(answer) > _OUTPUT_CAP_CHARS:
|
|
answer = answer[:_OUTPUT_CAP_CHARS]
|
|
truncated = True
|
|
|
|
new_status = (
|
|
"timeout" if timed_out else ("failed" if error_type else "success")
|
|
)
|
|
|
|
with engine.begin() as conn:
|
|
update_fields: Dict[str, Any] = {
|
|
"status": new_status,
|
|
"started_at": started,
|
|
"finished_at": finished,
|
|
"output": answer or None,
|
|
"output_truncated": truncated,
|
|
"prompt_tokens": prompt_tokens,
|
|
"generated_tokens": generated_tokens,
|
|
}
|
|
if error_type:
|
|
update_fields["error_type"] = error_type
|
|
update_fields["error"] = error_text
|
|
updated_run = ScheduleRunsRepository(conn).update(run_id, update_fields)
|
|
if used_tokens > 0:
|
|
agent_id_raw = schedule.get("agent_id")
|
|
try:
|
|
TokenUsageRepository(conn).insert(
|
|
user_id=schedule.get("user_id"),
|
|
api_key=None,
|
|
prompt_tokens=prompt_tokens,
|
|
generated_tokens=generated_tokens,
|
|
timestamp=finished,
|
|
agent_id=str(agent_id_raw) if agent_id_raw else None,
|
|
source="schedule",
|
|
request_id=str(run_id),
|
|
model_id=outcome.get("model_id"),
|
|
)
|
|
except Exception:
|
|
logger.exception(
|
|
"scheduler: token_usage insert failed run=%s", run_id,
|
|
)
|
|
schedules_repo = SchedulesRepository(conn)
|
|
autopaused = False
|
|
if new_status == "success":
|
|
schedules_repo.reset_failure_count(str(schedule["id"]))
|
|
elif new_status in ("failed", "timeout"):
|
|
count = schedules_repo.bump_failure_count(str(schedule["id"]))
|
|
if (
|
|
settings.SCHEDULE_AUTOPAUSE_FAILURES > 0
|
|
and count >= settings.SCHEDULE_AUTOPAUSE_FAILURES
|
|
and schedule.get("trigger_type") == "recurring"
|
|
):
|
|
autopaused = schedules_repo.autopause(str(schedule["id"]))
|
|
# Once: terminal-flip on cron-fired runs only; manual runs on a
|
|
# still-active once-schedule leave the future cadence intact.
|
|
if (
|
|
schedule.get("trigger_type") == "once"
|
|
and run.get("trigger_source") != "manual"
|
|
and schedule.get("status") == "active"
|
|
):
|
|
schedules_repo.update_internal(
|
|
str(schedule["id"]),
|
|
{"status": "completed", "next_run_at": None},
|
|
)
|
|
|
|
appended: Optional[Dict[str, Any]] = None
|
|
if (
|
|
schedule.get("trigger_type") == "once"
|
|
and new_status == "success"
|
|
and schedule.get("origin_conversation_id")
|
|
):
|
|
try:
|
|
appended = _append_one_time_turn(schedule, updated_run or run, outcome)
|
|
except Exception:
|
|
logger.exception(
|
|
"scheduler: append turn failed run=%s", run_id,
|
|
)
|
|
if appended is not None:
|
|
with engine.begin() as conn:
|
|
ScheduleRunsRepository(conn).update(
|
|
run_id,
|
|
{
|
|
"conversation_id": str(appended["conversation_id"]),
|
|
"message_id": str(appended["id"]),
|
|
},
|
|
)
|
|
_publish_message_appended(
|
|
schedule.get("user_id"),
|
|
str(appended["conversation_id"]),
|
|
appended,
|
|
str(schedule["id"]),
|
|
run_id,
|
|
)
|
|
|
|
if new_status == "success":
|
|
_publish_run_event("schedule.run.completed", updated_run or run, schedule)
|
|
else:
|
|
_publish_run_event(
|
|
"schedule.run.failed",
|
|
updated_run or run,
|
|
schedule,
|
|
error_type=error_type,
|
|
error=error_text,
|
|
)
|
|
|
|
if autopaused:
|
|
_publish_run_event(
|
|
"schedule.autopaused",
|
|
updated_run or run,
|
|
schedule,
|
|
consecutive_failure_count=settings.SCHEDULE_AUTOPAUSE_FAILURES,
|
|
)
|
|
|
|
return {
|
|
"status": new_status,
|
|
"run_id": run_id,
|
|
"error_type": error_type,
|
|
}
|