from __future__ import annotations """ WebSocket Handler 클라이언트 메시지: JOIN | BEHAVIOR 서버 Push: SESSION_READY | EVENT_ACK | INTENT_UPDATE | ERROR """ import json import logging from datetime import datetime from fastapi import APIRouter, WebSocket, WebSocketDisconnect from config import settings from core.extractor import get_extractor from core.inference import ( infer_with_behavior, to_customer_context_json, to_probability_dict, to_topn_with_others, ) from data.executor import get_executor from ws.manager import manager router = APIRouter() logger = logging.getLogger(__name__) @router.websocket("/ws") async def websocket_endpoint(websocket: WebSocket) -> None: """WebSocket 엔드포인트. 첫 메시지로 JOIN을 받아 세션에 결합하고 SESSION_READY를 보낸 뒤, 이후 BEHAVIOR 메시지마다 재추론하여 INTENT_UPDATE를 Push한다. 연결 종료 시 매니저에서 연결을 해제한다. Args: websocket: 클라이언트 WebSocket 연결. """ session_id: str | None = None try: await websocket.accept() raw = await websocket.receive_json() if raw.get("type") != "JOIN": await websocket.send_json({"type": "ERROR", "code": "MISSING_JOIN"}) return session_id = raw["session_id"] ex = get_executor() row = ex.fetchone("SELECT id, scenario_id, stage FROM sessions WHERE id = ?", [session_id]) if row is None: await websocket.send_json({"type": "ERROR", "code": "SESSION_NOT_FOUND"}) return scenario_id = row[1] manager._connections[session_id] = websocket await websocket.send_json({ "type": "SESSION_READY", "session_id": session_id, "scenario_id": scenario_id, "stage": row[2], }) # 설문 답변 캐시 (재추론 시 사용) survey_answers = _load_session_answers(session_id) while True: msg = await websocket.receive_json() mtype = msg.get("type") if mtype == "BEHAVIOR": await _handle_behavior(websocket, session_id, scenario_id, survey_answers, msg) else: await websocket.send_json({"type": "ERROR", "code": "UNKNOWN_TYPE"}) except WebSocketDisconnect: pass except Exception as e: logger.exception(f"WS error: {e}") try: await websocket.send_json({"type": "ERROR", "message": str(e)}) except Exception: pass finally: if session_id: manager.disconnect(session_id) async def _handle_behavior( websocket: WebSocket, session_id: str, scenario_id: str, survey_answers: dict[str, str], msg: dict, ) -> None: """BEHAVIOR 메시지를 처리한다. 이벤트를 event_log와 extractor에 적재하고 EVENT_ACK를 보낸다. app_exit이면 세션을 종료(stage='exited')하고 INTENT_UPDATE 없이 반환한다. 그 외에는 누적 행동 기반으로 재추론하여 intent_scores·customer_contexts에 저장하고 INTENT_UPDATE를 Push한다. Args: websocket: 응답을 보낼 WebSocket 연결. session_id: 대상 세션 ID. scenario_id: 세션의 시나리오 ID. survey_answers: 재추론에 사용할 설문 답변 ({question_id: answer_code}). msg: BEHAVIOR 메시지 (behavior_id·event_type·entity 포함). """ behavior_id = msg.get("behavior_id") event_type = msg.get("event_type") entity = msg.get("entity") if not (behavior_id and event_type and entity): await websocket.send_json({"type": "ERROR", "code": "MISSING_FIELDS"}) return ex = get_executor() # step: behavior_id 접두사로 판정 (1-* = step1, 그 외 step2). 시나리오 무관. step = 1 if str(behavior_id).startswith("1-") else 2 # 이벤트 로그 적재 ex.execute( "INSERT INTO event_log (session_id, scenario_id, step, behavior_id, event_type, entity) " "VALUES (?, ?, ?, ?, ?, ?)", [session_id, scenario_id, step, behavior_id, event_type, entity], ) ex.execute( "UPDATE sessions SET last_active_at = CURRENT_TIMESTAMP WHERE id = ?", [session_id], ) # Extractor에 이벤트 추가 (app_exit/navigate_back은 Pattern 집계 영향 약함) get_extractor().add_event(session_id, event_type, entity) await websocket.send_json({ "type": "EVENT_ACK", "behavior_id": behavior_id, "occurred_at": datetime.utcnow().isoformat(), }) # app_exit: 세션 종료. INTENT_UPDATE 미전송. stage='exited' 갱신만. if event_type == "app_exit": ex.execute("UPDATE sessions SET stage = ? WHERE id = ?", ["exited", session_id]) return # Intent 재추론 (Batch + 누적 행동) batch_features, intent_scores = infer_with_behavior(survey_answers, session_id, scenario_id) new_stage = f"step_{step}" if step > 0 else "initial" # Intent Score 적재 score_rows = [ [session_id, new_stage, s.intent_id, s.baseline_score, s.final_score, s.delta_score, s.baseline_rank, s.rank, s.rank_change, s.inference_type] for s in intent_scores ] ex.executemany( "INSERT INTO intent_scores " "(session_id, stage, intent_id, baseline_score, final_score, delta_score, " " baseline_rank, rank, rank_change, inference_type) " "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", score_rows, ) # Customer Context 적재 context = to_customer_context_json( session_id=session_id, stage=new_stage, scenario_id=scenario_id, scores=intent_scores, batch_features=batch_features, ) ex.execute( "INSERT INTO customer_contexts (session_id, scenario_id, stage, context_json) " "VALUES (?, ?, ?, ?)", [session_id, scenario_id, new_stage, json.dumps(context, ensure_ascii=False)], ) ex.execute("UPDATE sessions SET stage = ? WHERE id = ?", [new_stage, session_id]) top_n_cnt = settings.TOP_N_INTENT top_items, others = to_topn_with_others(intent_scores, top_n=top_n_cnt, scenario_id=scenario_id) # 추론 이유(reasoning) 첨부 — 결합 feature(batch+누적 pattern+최신 event) from core.engines import get_engine from core import explain as _explain _eng = get_engine(scenario_id) _combined = {**batch_features, **_eng.pattern_features(session_id), **_eng.event_features(session_id)} _explain.attach_reasoning(_eng, _combined, top_items) all_probabilities = to_probability_dict(intent_scores, scenario_id=scenario_id) await websocket.send_json({ "type": "INTENT_UPDATE", "session_id": session_id, "behavior_id": behavior_id, "stage": new_stage, "computed_at": datetime.utcnow().isoformat(), "top_n": top_items, "others": others, "all_probabilities": all_probabilities, }) def _load_session_answers(session_id: str) -> dict[str, str]: """세션의 설문 답변을 DB에서 로드한다. Args: session_id: 답변을 조회할 세션 ID. Returns: {question_id: answer_code} 형태의 설문 답변 dict. """ ex = get_executor() df = ex.to_pandas( "SELECT question_id, answer_code FROM survey_answers WHERE session_id = ?", [session_id], ) return {r["question_id"]: r["answer_code"] for _, r in df.iterrows()}