TanWenyan/leaudit-platform-backend
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: harden rag streaming and queue status reporting

Browse Source
This commit is contained in:
wren
2026-05-11 20:39:24 +08:00
parent e749abd9ed
commit 2ca44f6312
3 changed files with 71 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastapi_modules/fastapi_leaudit/controllers/documentController.py
+35 -37
View File
@@ -294,64 +294,62 @@ class DocumentController(BaseController):
async def GetQueueStatus():
"""获取文档处理队列状态。"""
from datetime import datetime
from fastapi_admin.config import LEAUDIT_WORKER_CONCURRENCY
async with GetAsyncSession() as Session:
statusRows = (
runStats = (
await Session.execute(
text(
"""
SELECT processing_status, COUNT(*) AS cnt
FROM leaudit_documents
WHERE deleted_at IS NULL AND is_latest_version = true
GROUP BY processing_status
SELECT
COUNT(*) FILTER (WHERE r.status IN ('queued', 'pending', 'retrying'))::int AS pending_tasks,
COUNT(*) FILTER (WHERE r.status = 'running')::int AS processing_tasks
FROM leaudit_audit_runs r
JOIN leaudit_documents d
ON d.current_run_id = r.id
WHERE d.deleted_at IS NULL
AND d.is_latest_version = true
"""
)
)
).mappings().all()
).mappings().first()
waiting = 0
processing = 0
for row in statusRows:
s = str(row["processing_status"] or "")
c = int(row["cnt"] or 0)
if s == "waiting":
waiting = c
elif s in ("processing", "running"):
processing += c
pending = int((runStats or {}).get("pending_tasks") or 0)
processing = int((runStats or {}).get("processing_tasks") or 0)
processingIdsRows: list[int] = []
if processing > 0:
async with GetAsyncSession() as Session:
idRows = (
await Session.execute(
text(
"""
SELECT id FROM leaudit_documents
WHERE deleted_at IS NULL
AND is_latest_version = true
AND processing_status IN ('processing', 'running')
ORDER BY updated_at DESC
LIMIT 50
"""
)
processingIdsRows = (
await Session.execute(
text(
"""
SELECT d.id
FROM leaudit_documents d
JOIN leaudit_audit_runs r
ON r.id = d.current_run_id
WHERE d.deleted_at IS NULL
AND d.is_latest_version = true
AND r.status = 'running'
ORDER BY COALESCE(r.started_at, r.created_at) DESC, d.id DESC
LIMIT 50
"""
)
).fetchall()
processingIdsRows = [int(r[0]) for r in idRows]
)
).fetchall()
processingIds = [int(r[0]) for r in processingIdsRows]
return Result.success(
data=QueueStatusVO(
success=True,
timestamp=datetime.now().isoformat(),
queue={
"pending_tasks": waiting,
"pending_tasks": pending,
"processing_tasks": processing,
"available_slots": max(0, 4 - processing),
"max_concurrent": 4,
"available_slots": max(0, LEAUDIT_WORKER_CONCURRENCY - processing),
"max_concurrent": LEAUDIT_WORKER_CONCURRENCY,
},
documents={
"waiting": waiting,
"waiting": pending,
"processing": processing,
"processing_ids": processingIdsRows,
"processing_ids": processingIds,
},
)
)
fastapi_modules/fastapi_leaudit/rag_engine/generator.py
+4 -2
View File
@@ -77,7 +77,9 @@ async def generate_stream(
if payload == "[DONE]":
break
chunk = json.loads(payload)
delta = chunk.get("choices", [{}])[0].get("delta", {})
choices = chunk.get("choices") or []
first_choice = choices[0] if choices and isinstance(choices[0], dict) else {}
delta = first_choice.get("delta", {})
text = delta.get("content", "")
if text:
yield _sse_line(
@@ -141,4 +143,4 @@ async def generate_stream(
def _sse_line(data: dict) -> str:
return f"data: {json.dumps(data, ensure_ascii=False)}\\n\\n"
return f"data: {json.dumps(data, ensure_ascii=False)}\n\n"