Spaces:
Sleeping
Sleeping
| 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 | |
| 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"} | |
| 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"} | |
| async def get_session_data(session_id: str): | |
| session = get_session(session_id) | |
| if not session: | |
| return {"error": "Session not found"} | |
| return session | |
| 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) | |