SSE status streaming for doc analysis was added.
This commit is contained in:
parent
89e0db7baa
commit
73187742f2
10 changed files with 670 additions and 6 deletions
|
|
@ -85,6 +85,12 @@ API_METRICS_PORT=9100
|
|||
B2B_DEFAULT_RATE_LIMIT_RPS=3 # per API key; mirrors Ollama Pro concurrency
|
||||
CORS_ORIGINS= # comma-separated, future web SPA
|
||||
|
||||
# SSE streaming of analysis status (GET /api/v1/reports/{id}/events):
|
||||
# poll tick for DB status changes and max stream lifetime before the
|
||||
# server sends a `timeout` event and closes (client reconnects/polls).
|
||||
SSE_POLL_INTERVAL_SECONDS=1.0
|
||||
SSE_MAX_STREAM_SECONDS=300.0
|
||||
|
||||
# --- Auth (JWT + Telegram identity verification) ---
|
||||
# Telegram bot token is also used by the API to verify Login Widget / Mini App signatures.
|
||||
TELEGRAM_BOT_TOKEN=
|
||||
|
|
|
|||
|
|
@ -123,7 +123,7 @@ Control plane:
|
|||
| Doc retention | TTL purge of raw docs after N days | — |
|
||||
| Redis | Kept (rate limit/sessions future) | (no longer the queue) |
|
||||
| Auth | Per-adapter `service_tokens`, revocable | single `SERVICE_TOKEN` |
|
||||
| Report delivery | Polling now (fine-grained stage), SSE/webhook later | — |
|
||||
| Report delivery | Polling + SSE (`GET /reports/{id}/events`, web-контракт `analysisResultSchema`); webhook later | — |
|
||||
| Sync `/analyze` | No | — |
|
||||
| Doc status | Fine-grained `queued→extracting→prescreening→ocr→analyzing→done\|failed` | coarse status |
|
||||
| Extra tables | jobs, service_tokens, invoices(stub) | 3-table plan |
|
||||
|
|
@ -233,7 +233,7 @@ DealDocumentScreening/
|
|||
│ │ ├── routes/
|
||||
│ │ │ ├── health.py (/healthz, /readyz)
|
||||
│ │ │ ├── documents.py (POST /api/v1/documents — upload→MinIO→publish, reserve)
|
||||
│ │ │ ├── reports.py (GET /api/v1/reports/{id} — 202+stage or 200+md)
|
||||
│ │ │ ├── reports.py (GET /api/v1/reports/{id} — 202+stage or 200+md; /events — SSE)
|
||||
│ │ │ ├── me.py (GET /api/v1/me — credits balance)
|
||||
│ │ │ ├── metrics.py (/metrics — prometheus)
|
||||
│ │ │ └── b2b.py (X-API-Key: POST /analyze, GET /b2b/reports, /b2b/usage, /b2b/keys CRUD)
|
||||
|
|
@ -1455,6 +1455,7 @@ Health/metrics exempt from auth.
|
|||
| POST | `/api/v1/documents` | user JWT | multipart → reserve credit → MinIO put → row `queued` → publish `DocumentUploaded` → `202 {document_id, correlation_id}`. `402` if no credit. `400` bad mime/size. |
|
||||
| GET | `/api/v1/documents/{id}` | user JWT | status + stage + filename (for polling UI) |
|
||||
| GET | `/api/v1/reports/{document_id}` | user JWT | `202 {status, stage}` while not done; `200 {markdown, findings, ...}` when done |
|
||||
| GET | `/api/v1/reports/{document_id}/events` | user JWT | SSE-стрим статусов: безымянные кадры с payload `analysisResultSchema` (camelCase, `pending/processing/completed/failed`); терминальные `completed/failed` закрывают стрим; служебные события `timeout`/`error`; keep-alive комментарии между кадрами |
|
||||
| GET | `/api/v1/me` | user JWT | `{telegram_id, credits_left}` |
|
||||
|
||||
### B2B endpoints (`X-API-Key`, `api/routes/b2b.py`)
|
||||
|
|
@ -1495,7 +1496,11 @@ To grant admin access: set `users.role = 'admin'` (or the value of
|
|||
|
||||
No synchronous `/analyze` (locked). Adapters poll `/reports/{id}`; the
|
||||
fine-grained `stage` field powers a progress signal in the bot ("Extracting
|
||||
text…", "Analyzing…"). SSE/webhook added later.
|
||||
text…", "Analyzing…"). Web clients can instead subscribe to the SSE stream
|
||||
`GET /api/v1/reports/{id}/events`: every frame carries the normalized
|
||||
`analysisResultSchema` payload (`pending/processing/completed/failed`,
|
||||
issues, derived riskScore/summary) — see `api/schemas/analysis_result.py`.
|
||||
Webhook added later.
|
||||
|
||||
---
|
||||
|
||||
|
|
|
|||
|
|
@ -325,6 +325,9 @@ async def fetch_document_status_for_user(
|
|||
"prompt_tokens": status.prompt_tokens,
|
||||
"eval_tokens": status.eval_tokens,
|
||||
"latency_ms": status.latency_ms,
|
||||
"mime": status.mime,
|
||||
"bytes": status.bytes_,
|
||||
"report_created_at": status.report_created_at,
|
||||
}
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -544,6 +544,69 @@ TTL `BILLING_RETURN_TOKEN_TTL_MINUTES`). `403` — нет/просрочен/ч
|
|||
|
||||
`404` — отчёт не найден (или чужой).
|
||||
|
||||
### `GET /api/v1/reports/{document_id}/events`
|
||||
|
||||
SSE-стрим статусов анализа (`text/event-stream`). Аутентификация и
|
||||
проверка владельца — до открытия стрима (`401`/`404` как обычные HTTP
|
||||
ошибки). Альтернатива опросу `GET /reports/{document_id}` для web-SPA:
|
||||
сервер сам поллит БД (`SSE_POLL_INTERVAL_SECONDS`) и шлёт событие при
|
||||
смене статуса.
|
||||
|
||||
Каждый кадр данных (безымянное событие → `onmessage` в `EventSource`)
|
||||
несёт payload контракта web-клиента `analysisResultSchema` (camelCase;
|
||||
см. [`api/schemas/analysis_result.py`](../schemas/analysis_result.py)):
|
||||
|
||||
```json
|
||||
{
|
||||
"id": "uuid",
|
||||
"fileName": "contract.pdf",
|
||||
"fileType": "application/pdf",
|
||||
"fileSize": 12345,
|
||||
"status": "pending | processing | completed | failed",
|
||||
"issues": [
|
||||
{
|
||||
"id": "penalties-0",
|
||||
"severity": "critical | warning | info",
|
||||
"category": "penalties",
|
||||
"title": "Неустойки / штрафы",
|
||||
"description": "<риск>\n\nРекомендация: <рекомендация>",
|
||||
"fragment": "<цитата> (п. <section_ref>)",
|
||||
"lineNumber": null
|
||||
}
|
||||
],
|
||||
"summary": "Критичных: 1, предупреждений: 1, замечаний: 0",
|
||||
"riskScore": 6,
|
||||
"createdAt": "2026-09-02T17:35:34.018348+00:00",
|
||||
"completedAt": "2026-09-02T17:36:10.104222+00:00"
|
||||
}
|
||||
```
|
||||
|
||||
Маппинг:
|
||||
|
||||
| Поле | Источник |
|
||||
| ---- | -------- |
|
||||
| `status` | `queued → pending`; `extracting/prescreening/ocr/analyzing → processing`; `done/manual_review → completed`; `failed → failed` |
|
||||
| `issues` | `reports.content_json.findings` (severity: `high→critical`, `medium→warning`, `low→info`) |
|
||||
| `riskScore` | вес находок: critical=4, warning=2, info=1, cap 10 (не хранится, вычисляется) |
|
||||
| `summary` | строка с количеством находок (не хранится, вычисляется) |
|
||||
| `completedAt` | `reports.created_at`; `null`, пока нет отчёта |
|
||||
|
||||
События стрима:
|
||||
|
||||
- безымянные кадры — снимок при подключении, затем по одному на смену
|
||||
статуса; кадр с `status: completed | failed` терминальный, стрим
|
||||
закрывается (клиенту нужно вызвать `es.close()`, иначе `EventSource`
|
||||
переподключится);
|
||||
- `timeout` — стрим жил дольше `SSE_MAX_STREAM_SECONDS`; клиент
|
||||
переподключается (`EventSource` делает это сам) или падает на опрос;
|
||||
- `error` — документ исчез посреди стрима.
|
||||
|
||||
Между кадрами без изменений шлются комментарии `: keep-alive`, чтобы
|
||||
прокси не рвали соединение. Заголовок `X-Accel-Buffering: no` отключает
|
||||
буферизацию nginx.
|
||||
|
||||
`404` — отчёт не найден (или чужой).
|
||||
|
||||
---
|
||||
|
||||
## b2b
|
||||
|
|
|
|||
|
|
@ -1,10 +1,16 @@
|
|||
"""Report polling endpoint."""
|
||||
"""Report polling endpoint + SSE status streaming."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from collections.abc import AsyncIterator
|
||||
from uuid import UUID
|
||||
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from fastapi import APIRouter, HTTPException, Request
|
||||
from fastapi.responses import StreamingResponse
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
||||
|
||||
from src.contract_check.api.deps import (
|
||||
AsyncSessionDep,
|
||||
|
|
@ -12,9 +18,29 @@ from src.contract_check.api.deps import (
|
|||
fetch_document_status_for_user,
|
||||
)
|
||||
from src.contract_check.api.schemas import ReportInProgressResponse, ReportResponse
|
||||
from src.contract_check.api.schemas.analysis_result import AnalysisResult, to_analysis_result
|
||||
from src.contract_check.core.config import get_settings
|
||||
from src.contract_check.core.logging import get_logger
|
||||
|
||||
log = get_logger(__name__)
|
||||
|
||||
router = APIRouter(prefix="/api/v1/reports", tags=["reports"])
|
||||
|
||||
_TERMINAL_ANALYSIS_STATUSES = frozenset({"completed", "failed"})
|
||||
|
||||
_SSE_HEADERS = {
|
||||
"Cache-Control": "no-cache",
|
||||
# Disable proxy buffering (nginx) so events are flushed immediately.
|
||||
"X-Accel-Buffering": "no",
|
||||
}
|
||||
|
||||
|
||||
def _sse_event(data: str, event: str | None = None) -> str:
|
||||
"""Format one SSE frame; unnamed events arrive as `onmessage` in EventSource."""
|
||||
if event is None:
|
||||
return f"data: {data}\n\n"
|
||||
return f"event: {event}\ndata: {data}\n\n"
|
||||
|
||||
|
||||
@router.get("/{document_id}", response_model=ReportResponse | ReportInProgressResponse)
|
||||
async def get_report(
|
||||
|
|
@ -45,3 +71,93 @@ async def get_report(
|
|||
eval_tokens=row["eval_tokens"],
|
||||
latency_ms=row["latency_ms"],
|
||||
)
|
||||
|
||||
|
||||
@router.get("/{document_id}/events")
|
||||
async def stream_report_events(
|
||||
request: Request,
|
||||
session: AsyncSessionDep,
|
||||
document_id: UUID,
|
||||
user: CurrentUserDep,
|
||||
) -> StreamingResponse:
|
||||
"""Server-Sent Events stream of the document analysis status.
|
||||
|
||||
Auth and ownership are checked before the stream starts (401/404 as
|
||||
regular HTTP errors). Every data frame carries an `AnalysisResult`
|
||||
payload (the web client's `analysisResultSchema` — parse each event's
|
||||
`data` with it):
|
||||
|
||||
- unnamed frames (`onmessage`) — initial snapshot, then one per
|
||||
status/stage change while processing; `status: completed | failed`
|
||||
frames are terminal and the stream closes after them;
|
||||
- ``timeout`` — the stream hit ``SSE_MAX_STREAM_SECONDS``; the client
|
||||
should reconnect (EventSource does automatically) or fall back to
|
||||
polling ``GET /api/v1/reports/{document_id}``;
|
||||
- ``error`` — the document disappeared mid-stream.
|
||||
|
||||
Keep-alive comment lines (``: keep-alive``) are sent on polls without
|
||||
changes so proxies keep the connection open.
|
||||
"""
|
||||
settings = get_settings()
|
||||
|
||||
row = await fetch_document_status_for_user(session, document_id, user.user_id)
|
||||
if row is None:
|
||||
raise HTTPException(status_code=404, detail="report not found")
|
||||
|
||||
factory: async_sessionmaker[AsyncSession] = request.app.state.db_session_factory
|
||||
# The request-scoped session would hold a pooled connection for the whole
|
||||
# stream; release it — each poll below opens a short-lived session instead.
|
||||
await session.close()
|
||||
|
||||
async def event_stream() -> AsyncIterator[str]:
|
||||
result = to_analysis_result(row)
|
||||
|
||||
def _terminal(current: AnalysisResult) -> bool:
|
||||
return current.status in _TERMINAL_ANALYSIS_STATUSES
|
||||
|
||||
if _terminal(result):
|
||||
yield _sse_event(result.dump_json())
|
||||
return
|
||||
|
||||
yield _sse_event(result.dump_json())
|
||||
|
||||
deadline = time.monotonic() + settings.sse_max_stream_seconds
|
||||
while time.monotonic() < deadline:
|
||||
await asyncio.sleep(settings.sse_poll_interval_seconds)
|
||||
if await request.is_disconnected():
|
||||
return
|
||||
|
||||
try:
|
||||
async with factory() as poll_session:
|
||||
fresh_row = await fetch_document_status_for_user(
|
||||
poll_session, document_id, user.user_id
|
||||
)
|
||||
except Exception:
|
||||
log.exception("sse_status_poll_failed", document_id=str(document_id))
|
||||
yield ": keep-alive\n\n"
|
||||
continue
|
||||
|
||||
if fresh_row is None:
|
||||
yield _sse_event(json.dumps({"detail": "document not found"}), event="error")
|
||||
return
|
||||
|
||||
fresh = to_analysis_result(fresh_row)
|
||||
if fresh.status != result.status:
|
||||
result = fresh
|
||||
if _terminal(result):
|
||||
yield _sse_event(result.dump_json())
|
||||
return
|
||||
yield _sse_event(result.dump_json())
|
||||
else:
|
||||
yield ": keep-alive\n\n"
|
||||
|
||||
yield _sse_event(
|
||||
json.dumps({"detail": "stream timeout, reconnect or poll"}),
|
||||
event="timeout",
|
||||
)
|
||||
|
||||
return StreamingResponse(
|
||||
event_stream(),
|
||||
media_type="text/event-stream",
|
||||
headers=_SSE_HEADERS,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -7,6 +7,12 @@ live in `api/routes/*` except tiny inline query/path params.
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
from src.contract_check.api.schemas.analysis_result import (
|
||||
AnalysisIssue,
|
||||
AnalysisResult,
|
||||
AnalysisStatus,
|
||||
IssueSeverity,
|
||||
)
|
||||
from src.contract_check.api.schemas.auth import (
|
||||
AuthResponse,
|
||||
ForgotPasswordRequest,
|
||||
|
|
@ -148,6 +154,11 @@ __all__ = [
|
|||
"DocumentListResponse",
|
||||
"ReportInProgressResponse",
|
||||
"ReportResponse",
|
||||
# analysis-result (SSE events contract)
|
||||
"AnalysisIssue",
|
||||
"AnalysisResult",
|
||||
"AnalysisStatus",
|
||||
"IssueSeverity",
|
||||
# b2b
|
||||
"CreateApiKeyRequest",
|
||||
"ApiKeyResponse",
|
||||
|
|
|
|||
160
src/contract_check/api/schemas/analysis_result.py
Normal file
160
src/contract_check/api/schemas/analysis_result.py
Normal file
|
|
@ -0,0 +1,160 @@
|
|||
"""Frontend analysis-result contract (camelCase) for the SSE events endpoint.
|
||||
|
||||
Mirrors the web client's Zod schema (`analysisResultSchema` / `analysisIssueSchema`):
|
||||
|
||||
analysisResultSchema = z.object({
|
||||
id, fileName, fileType, fileSize,
|
||||
status: 'pending' | 'processing' | 'completed' | 'failed',
|
||||
issues: analysisIssue[],
|
||||
summary, riskScore (0-10), createdAt, completedAt: string | null,
|
||||
})
|
||||
|
||||
Backend document statuses are normalized to that enum; report findings
|
||||
(`core.analysis.report_schema.Finding`) are mapped to `analysisIssueSchema`.
|
||||
`riskScore` and `summary` are not stored — they are derived from the issue
|
||||
list on completion (severity-weighted score, counts-based summary line).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime as dt
|
||||
from typing import Any, Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from src.contract_check.core.analysis.analyzer import checklist_title
|
||||
from src.contract_check.core.logging import get_logger
|
||||
|
||||
log = get_logger(__name__)
|
||||
|
||||
AnalysisStatus = Literal["pending", "processing", "completed", "failed"]
|
||||
IssueSeverity = Literal["critical", "warning", "info"]
|
||||
|
||||
# Backend documents.status -> analysisResultSchema.status.
|
||||
_STATUS_MAP: dict[str, AnalysisStatus] = {
|
||||
"queued": "pending",
|
||||
"extracting": "processing",
|
||||
"prescreening": "processing",
|
||||
"ocr": "processing",
|
||||
"analyzing": "processing",
|
||||
"done": "completed",
|
||||
# Terminal without a report; surfaced to the client as a completed,
|
||||
# empty-issues result (prescreen routed it away from deep analysis).
|
||||
"manual_review": "completed",
|
||||
"failed": "failed",
|
||||
}
|
||||
|
||||
# Finding.severity -> analysisIssueSchema.severity.
|
||||
_SEVERITY_MAP: dict[str, IssueSeverity] = {
|
||||
"high": "critical",
|
||||
"medium": "warning",
|
||||
"low": "info",
|
||||
}
|
||||
|
||||
# riskScore weights (capped at 10) per issue severity.
|
||||
_RISK_WEIGHTS: dict[IssueSeverity, int] = {"critical": 4, "warning": 2, "info": 1}
|
||||
|
||||
|
||||
class AnalysisIssue(BaseModel):
|
||||
"""One contract risk finding in the frontend contract shape."""
|
||||
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
id: str
|
||||
severity: IssueSeverity
|
||||
category: str
|
||||
title: str
|
||||
description: str
|
||||
fragment: str | None = None
|
||||
line_number: int | None = Field(default=None, serialization_alias="lineNumber")
|
||||
|
||||
|
||||
class AnalysisResult(BaseModel):
|
||||
"""analysisResultSchema mirror — the SSE event payload for /reports/{id}/events."""
|
||||
|
||||
model_config = ConfigDict(populate_by_name=True)
|
||||
|
||||
id: str
|
||||
file_name: str = Field(serialization_alias="fileName")
|
||||
file_type: str = Field(default="", serialization_alias="fileType")
|
||||
file_size: int = Field(default=0, serialization_alias="fileSize")
|
||||
status: AnalysisStatus
|
||||
issues: list[AnalysisIssue] = Field(default_factory=list)
|
||||
summary: str = ""
|
||||
risk_score: int = Field(default=0, ge=0, le=10, serialization_alias="riskScore")
|
||||
created_at: str = Field(serialization_alias="createdAt")
|
||||
completed_at: str | None = Field(default=None, serialization_alias="completedAt")
|
||||
|
||||
def dump_json(self) -> str:
|
||||
return self.model_dump_json(by_alias=True)
|
||||
|
||||
|
||||
def _iso(value: Any) -> str:
|
||||
if isinstance(value, dt.datetime):
|
||||
return value.isoformat()
|
||||
return str(value)
|
||||
|
||||
|
||||
def _map_finding(finding: dict[str, Any], index: int) -> AnalysisIssue:
|
||||
checklist_id = str(finding.get("checklist_id") or "unknown")
|
||||
severity = _SEVERITY_MAP.get(str(finding.get("severity") or ""), "info")
|
||||
risk = str(finding.get("risk") or "").strip()
|
||||
recommendation = str(finding.get("recommendation") or "").strip()
|
||||
description = risk or checklist_title(checklist_id)
|
||||
if recommendation:
|
||||
description = f"{description}\n\nРекомендация: {recommendation}"
|
||||
quote = str(finding.get("quote") or "").strip()
|
||||
section_ref = str(finding.get("section_ref") or "").strip()
|
||||
fragment = f"{quote} (п. {section_ref})" if quote and section_ref else (quote or None)
|
||||
return AnalysisIssue(
|
||||
id=f"{checklist_id}-{index}",
|
||||
severity=severity,
|
||||
category=checklist_id,
|
||||
title=checklist_title(checklist_id),
|
||||
description=description,
|
||||
fragment=fragment,
|
||||
)
|
||||
|
||||
|
||||
def to_analysis_result(row: dict[str, Any]) -> AnalysisResult:
|
||||
"""Build an AnalysisResult from a `fetch_document_status_for_user` row."""
|
||||
doc_status = str(row["status"])
|
||||
status = _STATUS_MAP.get(doc_status)
|
||||
if status is None:
|
||||
log.warning("unknown_document_status_mapped", status=doc_status)
|
||||
status = "processing" if doc_status not in ("done",) else "completed"
|
||||
|
||||
issues: list[AnalysisIssue] = []
|
||||
risk_score = 0
|
||||
summary = ""
|
||||
|
||||
if status == "completed" and doc_status == "done":
|
||||
content = row.get("content_json") or {}
|
||||
findings = content.get("findings", []) if isinstance(content, dict) else []
|
||||
issues = [_map_finding(f, i) for i, f in enumerate(findings)]
|
||||
risk_score = min(10, sum(_RISK_WEIGHTS[i.severity] for i in issues))
|
||||
counts = {"critical": 0, "warning": 0, "info": 0}
|
||||
for issue in issues:
|
||||
counts[issue.severity] += 1
|
||||
summary = (
|
||||
"Риски не найдены"
|
||||
if not issues
|
||||
else (
|
||||
f"Критичных: {counts['critical']}, "
|
||||
f"предупреждений: {counts['warning']}, "
|
||||
f"замечаний: {counts['info']}"
|
||||
)
|
||||
)
|
||||
|
||||
return AnalysisResult(
|
||||
id=str(row["id"]),
|
||||
file_name=str(row.get("filename") or ""),
|
||||
file_type=str(row.get("mime") or ""),
|
||||
file_size=int(row.get("bytes") or 0),
|
||||
status=status,
|
||||
issues=issues,
|
||||
summary=summary,
|
||||
risk_score=risk_score,
|
||||
created_at=_iso(row["created_at"]),
|
||||
completed_at=_iso(row["report_created_at"]) if row.get("report_created_at") else None,
|
||||
)
|
||||
|
|
@ -110,6 +110,10 @@ class Settings(BaseSettings):
|
|||
api_metrics_port: int = 9100
|
||||
b2b_default_rate_limit_rps: int = 3
|
||||
|
||||
# SSE streaming of report status (GET /api/v1/reports/{id}/events).
|
||||
sse_poll_interval_seconds: float = 1.0
|
||||
sse_max_stream_seconds: float = 300.0
|
||||
|
||||
# Comma-separated browser origins allowed to call the API (CORS).
|
||||
# Empty disables CORS entirely (no browser clients).
|
||||
cors_origins: Annotated[list[str], NoDecode] = []
|
||||
|
|
|
|||
|
|
@ -32,6 +32,9 @@ class DocumentStatus:
|
|||
prompt_tokens: int | None
|
||||
eval_tokens: int | None
|
||||
latency_ms: int | None
|
||||
mime: str | None = None
|
||||
bytes_: int | None = None
|
||||
report_created_at: Any = None
|
||||
|
||||
|
||||
class DocumentRepository:
|
||||
|
|
@ -79,7 +82,8 @@ class DocumentRepository:
|
|||
text(
|
||||
"SELECT d.id, d.status, d.stage, d.filename, d.created_at, "
|
||||
" r.markdown, r.content_json, r.model_used, "
|
||||
" r.prompt_tokens, r.eval_tokens, r.latency_ms "
|
||||
" r.prompt_tokens, r.eval_tokens, r.latency_ms, "
|
||||
" d.mime, d.bytes, r.created_at "
|
||||
"FROM documents d "
|
||||
"LEFT JOIN reports r ON r.document_id = d.id "
|
||||
"WHERE d.id = :d AND d.user_id = :u"
|
||||
|
|
@ -101,6 +105,9 @@ class DocumentRepository:
|
|||
prompt_tokens=row[8],
|
||||
eval_tokens=row[9],
|
||||
latency_ms=row[10],
|
||||
mime=row[11],
|
||||
bytes_=row[12],
|
||||
report_created_at=row[13],
|
||||
)
|
||||
|
||||
async def exists(self, document_id: uuid.UUID) -> bool:
|
||||
|
|
|
|||
289
tests/integration/test_reports_sse.py
Normal file
289
tests/integration/test_reports_sse.py
Normal file
|
|
@ -0,0 +1,289 @@
|
|||
"""Integration tests for the SSE report-events endpoint (GET /reports/{id}/events).
|
||||
|
||||
Run against the Docker Compose infrastructure (`docker compose up -d`).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
|
||||
from tests.integration.conftest import user_token
|
||||
|
||||
pytestmark = pytest.mark.integration
|
||||
|
||||
_FINDINGS = [
|
||||
{
|
||||
"checklist_id": "penalties",
|
||||
"severity": "high",
|
||||
"quote": "Неустойка 10% от суммы договора",
|
||||
"section_ref": "5.2",
|
||||
"risk": "Односторонняя неустойка",
|
||||
"recommendation": "Согласовать взаимную ответственность",
|
||||
},
|
||||
{
|
||||
"checklist_id": "jurisdiction",
|
||||
"severity": "medium",
|
||||
"quote": "Споры рассматриваются в суде г. Минск",
|
||||
"section_ref": "9.1",
|
||||
"risk": "Неудобная подсудность",
|
||||
"recommendation": "",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def pdf_bytes(tmp_path: Path) -> bytes:
|
||||
import pymupdf
|
||||
|
||||
doc = pymupdf.open()
|
||||
page = doc.new_page()
|
||||
page.insert_text((72, 72), "Договор. Стороны обязуются.")
|
||||
path = tmp_path / "contract.pdf"
|
||||
doc.save(str(path))
|
||||
doc.close()
|
||||
return path.read_bytes()
|
||||
|
||||
|
||||
async def _create_user_with_credits(db_session, telegram_id: int) -> None: # noqa: ANN001
|
||||
await db_session.execute(
|
||||
text(
|
||||
"INSERT INTO users (telegram_id, credits_left) VALUES (:t, 10) "
|
||||
"ON CONFLICT (telegram_id) DO UPDATE SET credits_left = 10"
|
||||
),
|
||||
{"t": telegram_id},
|
||||
)
|
||||
await db_session.commit()
|
||||
|
||||
|
||||
async def _upload(client: httpx.AsyncClient, token: str, pdf_bytes: bytes) -> str:
|
||||
"""Upload a document under user JWT and return its id."""
|
||||
resp = await client.post(
|
||||
"/api/v1/documents",
|
||||
headers={"Authorization": f"Bearer {token}"},
|
||||
files={"file": ("contract.pdf", pdf_bytes, "application/pdf")},
|
||||
)
|
||||
assert resp.status_code in (200, 202), resp.text
|
||||
return resp.json()["document_id"]
|
||||
|
||||
|
||||
def _parse_events(raw: str) -> list[tuple[str, dict[str, Any]]]:
|
||||
"""Parse an SSE body into (event, data) tuples, ignoring comment lines."""
|
||||
events: list[tuple[str, dict[str, Any]]] = []
|
||||
for chunk in raw.split("\n\n"):
|
||||
chunk = chunk.strip()
|
||||
if not chunk or chunk.startswith(":"):
|
||||
continue
|
||||
event = "message"
|
||||
data = ""
|
||||
for line in chunk.split("\n"):
|
||||
if line.startswith("event:"):
|
||||
event = line[len("event:") :].strip()
|
||||
elif line.startswith("data:"):
|
||||
data = line[len("data:") :].strip()
|
||||
if data:
|
||||
events.append((event, json.loads(data)))
|
||||
return events
|
||||
|
||||
|
||||
async def _first_events(
|
||||
client: httpx.AsyncClient,
|
||||
token: str,
|
||||
document_id: str,
|
||||
count: int,
|
||||
) -> list[tuple[str, dict[str, Any]]]:
|
||||
"""Open the SSE stream and collect the first `count` parsed events."""
|
||||
collected: list[tuple[str, dict[str, Any]]] = []
|
||||
async with client.stream(
|
||||
"GET",
|
||||
f"/api/v1/reports/{document_id}/events",
|
||||
headers={"Authorization": f"Bearer {token}"},
|
||||
) as response:
|
||||
assert response.status_code == 200
|
||||
assert response.headers["content-type"].startswith("text/event-stream")
|
||||
buffer = ""
|
||||
async for chunk in response.aiter_text():
|
||||
buffer += chunk
|
||||
collected = _parse_events(buffer)
|
||||
if len(collected) >= count:
|
||||
break
|
||||
return collected[:count]
|
||||
|
||||
|
||||
async def test_sse_requires_auth(client: httpx.AsyncClient) -> None:
|
||||
response = await client.get(f"/api/v1/reports/{uuid.uuid4()}/events")
|
||||
assert response.status_code == 401
|
||||
|
||||
|
||||
async def test_sse_unknown_document_returns_404(
|
||||
client: httpx.AsyncClient,
|
||||
infra: dict[str, str],
|
||||
) -> None:
|
||||
telegram_id = 999100101
|
||||
token = await user_token(client, infra, telegram_id)
|
||||
response = await client.get(
|
||||
f"/api/v1/reports/{uuid.uuid4()}/events",
|
||||
headers={"Authorization": f"Bearer {token}"},
|
||||
)
|
||||
assert response.status_code == 404
|
||||
|
||||
|
||||
async def test_sse_streams_progress_then_completed(
|
||||
client: httpx.AsyncClient,
|
||||
db_session, # noqa: ANN001
|
||||
infra: dict[str, str],
|
||||
pdf_bytes: bytes,
|
||||
) -> None:
|
||||
telegram_id = 999100102
|
||||
await _create_user_with_credits(db_session, telegram_id)
|
||||
token = await user_token(client, infra, telegram_id)
|
||||
document_id = await _upload(client, token, pdf_bytes)
|
||||
|
||||
# Simulate the pipeline finishing while the client is subscribed: the
|
||||
# stream should observe the row change on its next poll tick.
|
||||
async def _finish() -> None:
|
||||
await asyncio.sleep(1.5)
|
||||
await db_session.execute(
|
||||
text(
|
||||
"INSERT INTO reports (document_id, content_json, markdown, model_used) "
|
||||
"VALUES (:d, :cj, :md, 'test-model')"
|
||||
),
|
||||
{
|
||||
"d": document_id,
|
||||
"cj": json.dumps({"findings": _FINDINGS}),
|
||||
"md": "# Отчёт",
|
||||
},
|
||||
)
|
||||
await db_session.execute(
|
||||
text("UPDATE documents SET status = 'done' WHERE id = :d"),
|
||||
{"d": document_id},
|
||||
)
|
||||
await db_session.commit()
|
||||
|
||||
task = asyncio.create_task(_finish())
|
||||
try:
|
||||
events = await _first_events(client, token, document_id, 2)
|
||||
finally:
|
||||
await task
|
||||
|
||||
(first_name, first) = events[0]
|
||||
(last_name, last) = events[-1]
|
||||
assert first_name == last_name == "message"
|
||||
|
||||
# Progress frame: analysisResultSchema with empty result fields.
|
||||
assert first["id"] == document_id
|
||||
assert first["fileName"] == "contract.pdf"
|
||||
assert first["fileType"] == "application/pdf"
|
||||
assert first["fileSize"] == len(pdf_bytes)
|
||||
assert first["status"] in ("pending", "processing")
|
||||
assert first["issues"] == []
|
||||
assert first["riskScore"] == 0
|
||||
assert first["completedAt"] is None
|
||||
assert first["createdAt"]
|
||||
|
||||
# Terminal frame: full mapped result.
|
||||
assert last["status"] == "completed"
|
||||
assert last["completedAt"] is not None
|
||||
assert last["riskScore"] == 6 # critical(4) + warning(2)
|
||||
assert last["summary"] == "Критичных: 1, предупреждений: 1, замечаний: 0"
|
||||
|
||||
issues = last["issues"]
|
||||
assert len(issues) == 2
|
||||
assert issues[0]["id"] == "penalties-0"
|
||||
assert issues[0]["severity"] == "critical"
|
||||
assert issues[0]["category"] == "penalties"
|
||||
assert issues[0]["title"] == "Неустойки / штрафы"
|
||||
assert "Односторонняя неустойка" in issues[0]["description"]
|
||||
assert "Рекомендация: Согласовать" in issues[0]["description"]
|
||||
assert issues[0]["fragment"] == "Неустойка 10% от суммы договора (п. 5.2)"
|
||||
assert issues[0]["lineNumber"] is None
|
||||
assert issues[1]["severity"] == "warning"
|
||||
assert "Рекомендация" not in issues[1]["description"]
|
||||
|
||||
|
||||
async def test_sse_failed_document_emits_failed_and_closes(
|
||||
client: httpx.AsyncClient,
|
||||
db_session, # noqa: ANN001
|
||||
infra: dict[str, str],
|
||||
pdf_bytes: bytes,
|
||||
) -> None:
|
||||
telegram_id = 999100103
|
||||
await _create_user_with_credits(db_session, telegram_id)
|
||||
token = await user_token(client, infra, telegram_id)
|
||||
document_id = await _upload(client, token, pdf_bytes)
|
||||
|
||||
await db_session.execute(
|
||||
text("UPDATE documents SET status = 'failed', stage = 'analyze' WHERE id = :d"),
|
||||
{"d": document_id},
|
||||
)
|
||||
await db_session.commit()
|
||||
|
||||
events = await _first_events(client, token, document_id, 1)
|
||||
name, payload = events[0]
|
||||
assert name == "message"
|
||||
assert payload["id"] == document_id
|
||||
assert payload["status"] == "failed"
|
||||
assert payload["issues"] == []
|
||||
assert payload["completedAt"] is None
|
||||
|
||||
|
||||
async def test_sse_manual_review_maps_to_completed(
|
||||
client: httpx.AsyncClient,
|
||||
db_session, # noqa: ANN001
|
||||
infra: dict[str, str],
|
||||
pdf_bytes: bytes,
|
||||
) -> None:
|
||||
telegram_id = 999100105
|
||||
await _create_user_with_credits(db_session, telegram_id)
|
||||
token = await user_token(client, infra, telegram_id)
|
||||
document_id = await _upload(client, token, pdf_bytes)
|
||||
|
||||
await db_session.execute(
|
||||
text("UPDATE documents SET status = 'manual_review' WHERE id = :d"),
|
||||
{"d": document_id},
|
||||
)
|
||||
await db_session.commit()
|
||||
|
||||
events = await _first_events(client, token, document_id, 1)
|
||||
_name, payload = events[0]
|
||||
assert payload["status"] == "completed"
|
||||
assert payload["issues"] == []
|
||||
assert payload["riskScore"] == 0
|
||||
assert payload["summary"] == ""
|
||||
|
||||
|
||||
async def test_sse_emits_keepalive_comments_while_unchanged(
|
||||
client: httpx.AsyncClient,
|
||||
db_session, # noqa: ANN001
|
||||
infra: dict[str, str],
|
||||
pdf_bytes: bytes,
|
||||
) -> None:
|
||||
"""While the status does not change, only comment lines arrive."""
|
||||
telegram_id = 999100104
|
||||
await _create_user_with_credits(db_session, telegram_id)
|
||||
token = await user_token(client, infra, telegram_id)
|
||||
document_id = await _upload(client, token, pdf_bytes)
|
||||
|
||||
raw = ""
|
||||
async with client.stream(
|
||||
"GET",
|
||||
f"/api/v1/reports/{document_id}/events",
|
||||
headers={"Authorization": f"Bearer {token}"},
|
||||
) as response:
|
||||
assert response.status_code == 200
|
||||
async for chunk in response.aiter_text():
|
||||
raw += chunk
|
||||
if raw.count(": keep-alive") >= 2:
|
||||
break
|
||||
|
||||
assert "data: " in raw
|
||||
assert raw.count(": keep-alive") >= 2
|
||||
# Only one data frame (the initial snapshot) while nothing changes.
|
||||
assert len(_parse_events(raw)) == 1
|
||||
Loading…
Add table
Reference in a new issue