mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-14 16:08:41 +00:00
ff0a6768c2
* feat(subagents): add capacity controls and durable batches * fix(helm): sync subagent config schema version * fix(subagents): preserve batch history without worker * fix(subagents): support explicit factory runtimes * fix: address durable batch review findings
131 lines
5.2 KiB
Python
131 lines
5.2 KiB
Python
"""Owner-scoped progress and control API for durable subagent batches."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from collections.abc import AsyncIterator
|
|
|
|
from fastapi import APIRouter, HTTPException, Query, Request
|
|
from fastapi.responses import StreamingResponse
|
|
|
|
from app.gateway.authz import require_permission
|
|
from app.gateway.deps import (
|
|
get_current_user,
|
|
get_subagent_batch_repo,
|
|
get_subagent_batch_service,
|
|
)
|
|
from deerflow.utils.thread_id import ThreadId
|
|
|
|
router = APIRouter(prefix="/api/threads/{thread_id}/subagent-batches", tags=["subagent-batches"])
|
|
_ITEM_STATUSES = {"pending", "queued", "leased", "running", "succeeded", "failed", "cancelled"}
|
|
|
|
|
|
async def _user_id(request: Request) -> str:
|
|
user_id = await get_current_user(request)
|
|
if user_id is None:
|
|
raise HTTPException(status_code=401, detail="Authentication required")
|
|
return user_id
|
|
|
|
|
|
async def _owned_batch(request: Request, thread_id: str, batch_id: str) -> tuple[object, str, dict]:
|
|
repo = get_subagent_batch_repo(request)
|
|
user_id = await _user_id(request)
|
|
batch = await repo.get_batch(batch_id, user_id=user_id)
|
|
if batch is None or batch["thread_id"] != thread_id:
|
|
raise HTTPException(status_code=404, detail="Subagent batch not found")
|
|
return repo, user_id, batch
|
|
|
|
|
|
@router.get("")
|
|
@require_permission("threads", "read", owner_check=True)
|
|
async def list_batches(thread_id: ThreadId, request: Request, limit: int = Query(20, ge=1, le=100)) -> list[dict]:
|
|
repo = get_subagent_batch_repo(request)
|
|
return await repo.list_by_thread(thread_id, user_id=await _user_id(request), limit=limit)
|
|
|
|
|
|
@router.get("/{batch_id}")
|
|
@require_permission("threads", "read", owner_check=True)
|
|
async def get_batch(thread_id: ThreadId, batch_id: str, request: Request) -> dict:
|
|
_repo, _user_id_value, batch = await _owned_batch(request, thread_id, batch_id)
|
|
return batch
|
|
|
|
|
|
@router.get("/{batch_id}/items")
|
|
@require_permission("threads", "read", owner_check=True)
|
|
async def list_batch_items(
|
|
thread_id: ThreadId,
|
|
batch_id: str,
|
|
request: Request,
|
|
offset: int = Query(0, ge=0),
|
|
limit: int = Query(100, ge=1, le=500),
|
|
status: str | None = Query(None),
|
|
) -> list[dict]:
|
|
if status is not None and status not in _ITEM_STATUSES:
|
|
raise HTTPException(status_code=422, detail="Unknown batch item status")
|
|
repo, user_id, _batch = await _owned_batch(request, thread_id, batch_id)
|
|
return await repo.list_items(batch_id, user_id=user_id, offset=offset, limit=limit, status=status) or []
|
|
|
|
|
|
@router.post("/{batch_id}/pause")
|
|
@require_permission("threads", "write", owner_check=True)
|
|
async def pause_batch(thread_id: ThreadId, batch_id: str, request: Request) -> dict:
|
|
repo, user_id, _batch = await _owned_batch(request, thread_id, batch_id)
|
|
return await repo.pause_batch(batch_id, user_id=user_id)
|
|
|
|
|
|
@router.post("/{batch_id}/resume")
|
|
@require_permission("threads", "write", owner_check=True)
|
|
async def resume_batch(thread_id: ThreadId, batch_id: str, request: Request) -> dict:
|
|
repo, user_id, _batch = await _owned_batch(request, thread_id, batch_id)
|
|
return await repo.resume_batch(batch_id, user_id=user_id)
|
|
|
|
|
|
@router.post("/{batch_id}/cancel")
|
|
@require_permission("threads", "write", owner_check=True)
|
|
async def cancel_batch(thread_id: ThreadId, batch_id: str, request: Request) -> dict:
|
|
if not getattr(request.app.state, "subagent_batches_available", False):
|
|
raise HTTPException(status_code=503, detail="Subagent batch worker is not running")
|
|
_repo, user_id, _batch = await _owned_batch(request, thread_id, batch_id)
|
|
result = await get_subagent_batch_service(request).cancel_batch(batch_id=batch_id, user_id=user_id)
|
|
if result is None:
|
|
raise HTTPException(status_code=404, detail="Subagent batch not found")
|
|
return result
|
|
|
|
|
|
@router.post("/{batch_id}/items/{item_id}/retry")
|
|
@require_permission("threads", "write", owner_check=True)
|
|
async def retry_batch_item(thread_id: ThreadId, batch_id: str, item_id: str, request: Request) -> dict:
|
|
repo, user_id, _batch = await _owned_batch(request, thread_id, batch_id)
|
|
item = await repo.retry_item(batch_id, item_id, user_id=user_id)
|
|
if item is None:
|
|
raise HTTPException(status_code=409, detail="Only failed items can be retried")
|
|
return item
|
|
|
|
|
|
@router.get("/{batch_id}/results.jsonl")
|
|
@require_permission("threads", "read", owner_check=True)
|
|
async def export_batch_results(thread_id: ThreadId, batch_id: str, request: Request) -> StreamingResponse:
|
|
repo, user_id, _batch = await _owned_batch(request, thread_id, batch_id)
|
|
|
|
async def lines() -> AsyncIterator[bytes]:
|
|
offset = 0
|
|
while True:
|
|
page = await repo.list_items(
|
|
batch_id,
|
|
user_id=user_id,
|
|
offset=offset,
|
|
limit=500,
|
|
include_result=True,
|
|
)
|
|
if not page:
|
|
break
|
|
for item in page:
|
|
yield (json.dumps(item, ensure_ascii=False, default=str) + "\n").encode()
|
|
offset += len(page)
|
|
|
|
return StreamingResponse(
|
|
lines(),
|
|
media_type="application/x-ndjson",
|
|
headers={"Content-Disposition": f'attachment; filename="{batch_id}-results.jsonl"'},
|
|
)
|