Spaces:
Sleeping
Sleeping
File size: 7,444 Bytes
d309047 fb272ee d309047 8e93416 d309047 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 | from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Request, Depends, HTTPException
import uuid
import json
import base64
import asyncio
from auth import get_current_user, decode_access_token
from services.session_store import (
create_session, get_session, close_session,
append_transcript, append_action_item, append_summary
)
from ml.stt import transcribe_audio
from ml.embeddings import embed_single
from ml.classifier import classify_sentence
from ml.drift_detector import predict_drift
from ml.rag import add_to_store
from services.ollama_client import summarize_segment
router = APIRouter()
active_connections: dict[str, WebSocket] = {}
session_buffers: dict[str, dict] = {}
SUMMARY_WORD_THRESHOLD = 50 # Generate summary every ~50 words (~30 seconds of speech)
MAX_TRANSCRIPT_CHARS = 500_000
@router.post("/start")
async def start_session():
session_id = str(uuid.uuid4())
create_session(session_id)
session_buffers[session_id] = {
"sentences": [], "timestamps": [], "speakers": [],
"word_count": 0, "last_summary_at": 0, "total_chars": 0,
}
return {"session_id": session_id, "status": "created"}
@router.post("/{session_id}/end")
async def end_session(session_id: str):
close_session(session_id)
session_buffers.pop(session_id, None)
ws = active_connections.pop(session_id, None)
if ws:
try:
await ws.close()
except Exception:
pass
return {"status": "ended"}
@router.get("/{session_id}")
async def get_session_data(session_id: str):
session = get_session(session_id)
if not session:
return {"error": "Session not found"}
return session
@router.websocket("/{session_id}/ws")
async def session_websocket(websocket: WebSocket, session_id: str):
await websocket.accept()
active_connections[session_id] = websocket
app = websocket.app
embedder = app.state.model_cache.get("embedder")
classifier = app.state.model_cache.get("classifier")
try:
while True:
raw = await asyncio.wait_for(websocket.receive_text(), timeout=120)
message = json.loads(raw)
msg_type = message.get("type")
if msg_type == "ping":
await websocket.send_json({"type": "pong"})
elif msg_type == "audio_chunk":
audio_b64 = message.get("data", "")
if len(audio_b64) > 5_000_000:
await websocket.send_json({"type": "error", "message": "Chunk too large"})
continue
audio_bytes = base64.b64decode(audio_b64)
try:
result = await transcribe_audio(audio_bytes)
except Exception as e:
await websocket.send_json({"type": "error", "message": str(e)[:100]})
continue
buf = session_buffers.get(session_id)
if not buf:
continue
for seg in result.get("segments", []):
text = seg.get("text", "").strip()
if not text or len(text) < 3:
continue
buf["total_chars"] += len(text)
if buf["total_chars"] > MAX_TRANSCRIPT_CHARS:
await websocket.send_json({"type": "warning", "message": "Session very long. Consider ending."})
break
ts = float(seg.get("start", 0))
speaker = str(seg.get("speaker", "Speaker"))[:50]
await websocket.send_json({"type": "transcript", "text": text, "timestamp": ts, "speaker": speaker})
try:
embedding = embed_single(text, embedder)
add_to_store(session_id, text, embedding, ts)
except Exception as e:
print(f"RAG Embed Error: {e}")
append_transcript(session_id, {"text": text, "timestamp": ts, "speaker": speaker})
buf["sentences"].append(text)
buf["timestamps"].append(ts)
buf["speakers"].append(speaker)
buf["word_count"] += len(text.split())
try:
clf = classify_sentence(text, classifier)
label = clf.get("label", "general")
conf = clf.get("confidence", 0.0)
if conf >= 0.65 and label != "general":
entry = {"type": label, "text": text, "confidence": conf, "timestamp": ts, "speaker": speaker}
await websocket.send_json(entry)
if label == "action_item":
append_action_item(session_id, entry)
except Exception:
pass
if buf["word_count"] - buf["last_summary_at"] >= SUMMARY_WORD_THRESHOLD:
try:
recent = " ".join(buf["sentences"][-25:])
summary_text = summarize_segment(recent)
await websocket.send_json({"type": "summary", "text": summary_text, "timestamp": ts})
append_summary(session_id, {"text": summary_text, "timestamp": ts})
buf["last_summary_at"] = buf["word_count"]
except Exception as e:
print(f"Summary generation error: {e}")
elif msg_type == "behavior":
data = message.get("data", {})
safe = {
"keyboard_idle_seconds": min(float(data.get("keyboard_idle_seconds", 0)), 600),
"mouse_movement_delta": min(float(data.get("mouse_movement_delta", 0)), 10),
"topic_shift_score": min(float(data.get("topic_shift_score", 0)), 1),
"audio_energy_variance": min(float(data.get("audio_energy_variance", 0.3)), 1),
"time_since_last_ui_action": min(float(data.get("time_since_last_ui_action", 0)), 600),
"words_per_minute_drop": min(float(data.get("words_per_minute_drop", 0)), 1),
"scroll_activity": min(float(data.get("scroll_activity", 0)), 50),
}
try:
drift_result = predict_drift(safe)
if drift_result.get("trigger_nudge"):
session = get_session(session_id)
ctx = ""
if session and session.get("summaries"):
ctx = session["summaries"][-1].get("text", "")
await websocket.send_json({
"type": "drift_nudge",
"drift_probability": drift_result["drift_probability"],
"context": ctx[:300] or "Meeting is in progress.",
})
except Exception:
pass
except asyncio.TimeoutError:
try:
await websocket.close(code=1001)
except Exception:
pass
except WebSocketDisconnect:
pass
except Exception as e:
print(f"WS error [{session_id}]: {e}")
finally:
active_connections.pop(session_id, None)
|