Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 53 additions & 0 deletions backend/agents/kb_builder/_ws_stub.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,53 @@
"""Temporary WebSocket broadcast stub — replace with ``ws_manager`` when M4.1 lands.

Issue #22 (M3.6) wires `ui_render` events into the KB Builder pipeline so the
frontend Activity Feed and Onboarding wizard can render progress in real time.
M4.1 (#23) introduces ``core.ws_manager.WSManager`` which owns the actual
WebSocket fan-out; until that ships, every call site fires through this shim
so the swap is one diff::

# before (M3.6)
from agents.kb_builder._ws_stub import broadcast_stub
await broadcast_stub("ui_render", {...})

# after (M4.1)
from core.ws_manager import ws_manager
await ws_manager.broadcast("ui_render", {...})

The payload shape this stub accepts is the contract documented in issue #23
and reproduced in #22 §1 — keep them aligned so callers do not need to be
edited again at swap time. The stub is intentionally awaitable so callers
already use ``await`` and the eventual swap is purely textual.

This module emits no WebSocket frames and has no side-effects beyond a single
``logging`` call. It must not import ``WebSocket`` or anything from
``core.ws_manager`` to avoid a circular dependency once that module exists.
"""

from __future__ import annotations

import json
import logging
from typing import Any

_log = logging.getLogger("aria.kb_builder.ws_stub")


async def broadcast_stub(event_type: str, payload: dict[str, Any]) -> None:
"""Log a structured representation of a future WebSocket broadcast.

The first positional arg matches ``WSManager.broadcast``'s signature
(``event_type``, ``payload``) so the M4.1 swap is mechanical.

Args:
event_type: One of the event types listed in issue #23 (here, almost
always ``"ui_render"``).
payload: The JSON-serialisable dict that will become the WebSocket
frame's body. Logged as a single line so it can be grepped from
container logs during demo prep.
"""
try:
rendered = json.dumps(payload, default=str, separators=(",", ":"))
except (TypeError, ValueError):
rendered = repr(payload)
_log.info("ws_stub event=%s payload=%s", event_type, rendered)
48 changes: 48 additions & 0 deletions backend/agents/kb_builder/onboarding/service.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
import uuid
from typing import Any

from agents.kb_builder._ws_stub import broadcast_stub
from agents.kb_builder.onboarding import session_store
from agents.kb_builder.onboarding.extraction import extract_patch
from agents.kb_builder.onboarding.questions import QUESTIONS
Expand Down Expand Up @@ -196,6 +197,35 @@ async def submit_onboarding_message(session_id: str, answer: str) -> dict[str, A
)
session.question_index += 1

# M3.6 — emit per-question progress AFTER the MCP write completes
# (issue #22 acceptance #5: "All events are emitted AFTER the
# corresponding MCP write completes (not before)"). Will become
# ``ws_manager.broadcast`` once M4.1 (#23) lands.
await broadcast_stub(
"ui_render",
{
"agent": "kb_builder",
"component": "kb_progress",
"props": {
"cell_id": session.cell_id,
"steps": [
{
"label": f"Question {q['index'] + 1}/{len(QUESTIONS)}",
"status": (
"done"
if q["index"] < session.question_index
else (
"in_progress" if q["index"] == session.question_index else "pending"
)
),
}
for q in QUESTIONS
],
},
"turn_id": None, # set by orchestrator ContextVar after M4.1 (#23)
},
)

if is_final:
# Re-read the row through the repository so the response shape matches
# ``EquipmentKbOut`` (the router serialises it for the client).
Expand All @@ -206,6 +236,24 @@ async def submit_onboarding_message(session_id: str, answer: str) -> dict[str, A
if rec is None:
# Should be impossible — update_equipment_kb just returned success.
raise NotFoundError(f"No equipment_kb row for cell {session.cell_id} after onboarding")
# M3.6 — final card AFTER the DB re-read so the frontend's re-fetch
# via ``GET /api/v1/kb/equipment/{cell_id}`` (M8.2) sees the
# post-onboarding KB. Will become ``ws_manager.broadcast`` after M4.1.
await broadcast_stub(
"ui_render",
{
"agent": "kb_builder",
"component": "equipment_kb_card",
"props": {
"cell_id": session.cell_id,
"highlight_fields": [
"thresholds.vibration_mm_s",
"failure_patterns",
],
},
"turn_id": None, # set by orchestrator ContextVar after M4.1 (#23)
},
)
return {
"session_id": session_id,
"complete": True,
Expand Down
56 changes: 49 additions & 7 deletions backend/modules/kb/router.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
start_onboarding,
submit_onboarding_message,
)
from agents.kb_builder._ws_stub import broadcast_stub
from aria_mcp.client import mcp_client
from core.api_response import created, ok
from core.database import get_db
Expand Down Expand Up @@ -51,6 +52,45 @@ def _ser_failure(r):
return FailureHistoryOut.model_validate(decode_record(r, JSON_FIELDS)).model_dump(mode="json")


# 5 phase labels — kept in sync with issue #22 §2 acceptance criterion
# ("PDF upload emits exactly 5 ui_render events with component kb_progress").
_UPLOAD_PHASES: tuple[str, ...] = (
"Validating PDF",
"Reading pages with Opus vision",
"Extracting thresholds",
"Validating schema",
"Saving knowledge base",
)


def _phase_status(idx: int, active_idx: int) -> str:
if idx < active_idx:
return "done"
if idx == active_idx:
return "in_progress"
return "pending"


def _upload_steps(active_idx: int) -> list[dict[str, str]]:
return [
{"label": label, "status": _phase_status(i, active_idx)}
for i, label in enumerate(_UPLOAD_PHASES)
]


async def _emit_upload_phase(cell_id: int, active_idx: int) -> None:
"""Stub WS broadcast for one PDF-upload phase. See M3.6 (#22) / M4.1 (#23)."""
await broadcast_stub(
"ui_render",
{
"agent": "kb_builder",
"component": "kb_progress",
"props": {"cell_id": cell_id, "steps": _upload_steps(active_idx)},
"turn_id": None, # set by orchestrator ContextVar after M4.1 (#23)
},
)


@router.get("/equipment")
async def list_kb(conn: asyncpg.Connection = Depends(get_db)):
rows = await KbRepository(conn).list()
Expand Down Expand Up @@ -95,22 +135,24 @@ async def upload_pdf(
recomputes completeness.
5. Re-read and serialise via ``EquipmentKbOut``.

Phase log lines stand in for live progress events until M4.1 (#23) lands
the websocket manager — see issue #18 §7.
Phase events are emitted via ``broadcast_stub`` (M3.6 / issue #22). Each
call will become ``ws_manager.broadcast("ui_render", ...)`` once M4.1
(#23) lands; the payload shape is already final, only the transport is
stubbed.
"""
if file.content_type not in ("application/pdf", "application/octet-stream"):
raise HTTPException(400, "File must be a PDF")

lock = _upload_locks.setdefault(cell_id, asyncio.Lock())
async with lock:
log.info("kb_upload[cell=%d] phase=Validating PDF", cell_id)
await _emit_upload_phase(cell_id, 0)
pdf_bytes = await file.read()
if not pdf_bytes:
raise HTTPException(400, "Uploaded file is empty")

log.info("kb_upload[cell=%d] phase=Reading pages with Opus vision", cell_id)
await _emit_upload_phase(cell_id, 1)
try:
log.info("kb_upload[cell=%d] phase=Extracting thresholds", cell_id)
await _emit_upload_phase(cell_id, 2)
kb, raw_markdown = await extract_from_pdf(pdf_bytes, cell_id)
except ValueError as e:
# ValidationError is a subclass of ValueError, so this single
Expand All @@ -122,11 +164,11 @@ async def upload_pdf(
raise HTTPException(413, msg) from e
raise HTTPException(422, f"Extraction failed after retry: {msg}") from e

log.info("kb_upload[cell=%d] phase=Validating schema", cell_id)
await _emit_upload_phase(cell_id, 3)
kb_dict = kb.model_dump(exclude={"kb_meta"})
kb_dict = await bootstrap_thresholds(cell_id, kb_dict)

log.info("kb_upload[cell=%d] phase=Saving knowledge base", cell_id)
await _emit_upload_phase(cell_id, 4)
result = await mcp_client.call_tool(
"update_equipment_kb",
{
Expand Down
Loading
Loading