"""WebSocket endpoint for real-time grant query streaming.""" import asyncio import json import logging import time from typing import Dict, Any from fastapi import APIRouter, WebSocket, WebSocketDisconnect from uuid import uuid4 from src.logging.logger import get_logger from src.analyzer.data_loader import load_current_grants, load_past_winners from src.analyzer.chat.chat_tools import ChatTools from src.analyzer.chat.query_router import route from src.analyzer.llm_client import LLMClient from src.analyzer.config import load_config logger = get_logger() router = APIRouter(prefix="/ws", tags=["websocket"]) # Global cache for initialized tools _chat_tools: ChatTools | None = None _llm_client: LLMClient | None = None def get_chat_tools() -> ChatTools: """Lazy load chat tools on first use.""" global _chat_tools if _chat_tools is None: try: from pathlib import Path snapshots_dir = Path("data/snapshots") history_xlsx = Path("data/IUK-141025-InnovateUKFundedProjects-FY2015-16topresent.xlsx") current = load_current_grants(snapshots_dir, limit=100) past = load_past_winners(history_xlsx=history_xlsx) _chat_tools = ChatTools(current, past) logging.info(f"Loaded ChatTools with {len(current)} current and {len(past)} past grants") except Exception as e: logging.error(f"Failed to initialize ChatTools: {e}") raise return _chat_tools def get_llm_client() -> LLMClient: """Lazy load LLM client on first use.""" global _llm_client if _llm_client is None: try: config = load_config() _llm_client = LLMClient(config) logging.info("Initialized LLM client for WebSocket streaming") except Exception as e: logging.error(f"Failed to initialize LLM client: {e}") raise return _llm_client @router.websocket("/query") async def websocket_query_endpoint(websocket: WebSocket): """ WebSocket endpoint for real-time grant query streaming. Protocol: Client sends: {"query": "user question", "session_id": "optional-id"} Server streams: - {"type": "metadata", "session_id": "...", "query": "..."} - {"type": "intent", "intent": "search"} - {"type": "token", "token": "word"} - {"type": "citations", "citations": [...]} - {"type": "done", "latency_ms": 1234} - {"type": "error", "error": "message"} """ await websocket.accept() session_id = str(uuid4()) try: # Initialize tools and client tools = get_chat_tools() llm_client = get_llm_client() logging.info(f"WebSocket connection established: {session_id}") while True: # Receive query from client try: data = await websocket.receive_json() except WebSocketDisconnect: logging.info(f"WebSocket disconnected: {session_id}") break query = data.get("query", "") session_id = data.get("session_id", session_id) if not query: await websocket.send_json({ "type": "error", "error": "Query is required" }) continue start_time = time.time() try: # Send metadata await websocket.send_json({ "type": "metadata", "session_id": session_id, "query": query }) # Route the query routed = route(query, use_llm=True) intent = str(routed.get("intent") or "general") args = routed.get("args") or {} logging.info(f"WebSocket {session_id}: Intent={intent}, Query={query}") # Send intent await websocket.send_json({ "type": "intent", "intent": intent }) citations_list = [] # Handle different intents if intent in {"search", "list", "list_grants"}: # Search for grants - non-streaming response keyword = args.get("keyword") or args.get("query") or "" results = tools.list_grants(keyword=keyword, limit=5) if results: answer_text = f"Found {len(results)} grants matching your query:\n" for grant in results: title = grant.get("title", "Unknown") grant_id = grant.get("id") or grant.get("grant_id", "") answer_text += f"\n- **{title}** (ID: {grant_id})" if grant_id: citations_list.append({ "grant_id": grant_id, "title": title, "url": grant.get("url") }) else: answer_text = f"No grants found matching '{keyword}'." # Stream the complete answer token by token words = answer_text.split() for word in words: await websocket.send_json({ "type": "token", "token": word + " " }) await asyncio.sleep(0.01) # Small delay for visual effect elif intent == "summarize" and args.get("grant_id"): # Summarize a specific grant grant_id = args.get("grant_id") result = tools.summarize_grant(grant_id) summary = result.get("summary_md", "No summary available") # Stream summary token by token words = summary.split() for word in words: await websocket.send_json({ "type": "token", "token": word + " " }) await asyncio.sleep(0.01) citations_list.append({ "grant_id": grant_id, "title": result.get("title", grant_id) }) elif intent == "compare" and args.get("grant_id_a") and args.get("grant_id_b"): # Compare two grants result = tools.compare_grants(args["grant_id_a"], args["grant_id_b"]) comparison = result.get("comparison_md", "Comparison unavailable") # Stream comparison token by token words = comparison.split() for word in words: await websocket.send_json({ "type": "token", "token": word + " " }) await asyncio.sleep(0.01) citations_list.extend([ {"grant_id": args["grant_id_a"], "title": f"Grant {args['grant_id_a']}"}, {"grant_id": args["grant_id_b"], "title": f"Grant {args['grant_id_b']}"} ]) else: # Default: Use LLM streaming for general queries results = tools.list_grants(keyword=query, limit=5) context = f"User query: {query}\n\n" if results: context += "Relevant grants:\n" for grant in results: title = grant.get("title", "Unknown") grant_id = grant.get("id") or grant.get("grant_id", "") context += f"- {title} (ID: {grant_id})\n" if grant_id: citations_list.append({ "grant_id": grant_id, "title": title, "url": grant.get("url") }) # Stream LLM response token by token messages = [ {"role": "system", "content": "You are a UK grant analyst. Answer questions about grants concisely and accurately."}, {"role": "user", "content": context} ] # Get streaming response from LLM stream_generator = llm_client.chat( messages, stream=True, max_tokens=1200, temperature=0.3 ) # Stream each token via WebSocket for token in stream_generator: await websocket.send_json({ "type": "token", "token": token }) await asyncio.sleep(0) # Allow other tasks to run # Send citations if citations_list: await websocket.send_json({ "type": "citations", "citations": citations_list }) # Send completion latency_ms = int((time.time() - start_time) * 1000) await websocket.send_json({ "type": "done", "latency_ms": latency_ms }) except Exception as e: logging.error(f"Error processing WebSocket query: {e}", exc_info=True) await websocket.send_json({ "type": "error", "error": str(e) }) except Exception as e: logging.error(f"WebSocket connection error: {e}", exc_info=True) finally: logging.info(f"WebSocket connection closed: {session_id}")