"""Aurelius — FastAPI app: source-agnostic search, relate, discover. v2 rewrite: the server no longer knows anything about Wikipedia. It picks a GraphSource by name from the adapter registry and hands it to the source-agnostic GraphNavigator / relate / discover. Adding a data source is an adapter file, not a server change. Endpoints GET /api/sources registered sources + ingest status GET /api/random?source= a demo pair for a source GET /api/relate?source=&a=&b= connection strength + intermediaries GET /api/discover?source=&a= hidden-connection candidates GET /api/health readiness (implies model loaded) WS /ws streamed pathfinding; first msg carries source """ import asyncio import io import json import os import re import time from collections import defaultdict, deque from fastapi import (FastAPI, Request, WebSocket, WebSocketDisconnect, Query, UploadFile, File) from fastapi.middleware.cors import CORSMiddleware from starlette.middleware.base import BaseHTTPMiddleware from config import ( ALLOWED_ORIGINS, MAX_QUERY_LEN, MAX_CONCURRENT_SEARCHES, MAX_SEARCH_SECONDS, RATE_LIMIT_WINDOW_S, RATE_LIMIT_MAX_REQUESTS, NEWS_REFRESH_MINUTES, ) from core import embedding from core import llm from core.navigator import GraphNavigator from core.discovery import relate as relate_query, discover as discover_query from core.source import get_source, list_sources import adapters # noqa: F401 — importing registers every adapter from news_intel import service as news_service app = FastAPI(docs_url=None, redoc_url=None, openapi_url=None) app.add_middleware( CORSMiddleware, allow_origins=ALLOWED_ORIGINS, allow_methods=["GET", "POST"], # POST for the PDF citation upload allow_headers=["*"], ) class SecurityHeadersMiddleware(BaseHTTPMiddleware): async def dispatch(self, request, call_next): resp = await call_next(request) resp.headers["X-Content-Type-Options"] = "nosniff" resp.headers["X-Frame-Options"] = "DENY" resp.headers["Referrer-Policy"] = "no-referrer" return resp app.add_middleware(SecurityHeadersMiddleware) # Concurrency cap — single-threaded event loop, so a plain counter is # race-free. Bounds CPU-bound embedding + upstream API fan-out. _active_searches = 0 # Per-IP sliding-window rate limit — bounds total volume (the concurrency # cap alone can't stop rapid connect/abandon loops). Process-local; move to # Redis if this ever runs multi-replica. _ip_hits: dict[str, deque] = defaultdict(deque) def _rate_ok(ip: str) -> bool: now = time.time() # Sweep stale IPs so the dict can't grow for the life of the process # (one abusive scanner cycling IPs would otherwise leak deques forever). if len(_ip_hits) > 1000: for stale in [k for k, dq in _ip_hits.items() if not dq or now - dq[-1] > RATE_LIMIT_WINDOW_S]: del _ip_hits[stale] dq = _ip_hits[ip] while dq and now - dq[0] > RATE_LIMIT_WINDOW_S: dq.popleft() if len(dq) >= RATE_LIMIT_MAX_REQUESTS: return False dq.append(now) return True # Ingested sources to self-populate at boot when their store is empty — # lets an ephemeral-storage host (Hugging Face Spaces free tier) come up # with finance/news ready without a manual CLI step. e.g. "finance,news". AUTOINGEST = [s.strip() for s in os.getenv("AURELIUS_AUTOINGEST", "").split(",") if s.strip()] @app.on_event("startup") async def _startup_embed_model(): await embedding.load_model() if AUTOINGEST: asyncio.create_task(_autoingest()) if NEWS_REFRESH_MINUTES > 0: asyncio.create_task(_news_refresh_loop()) async def _autoingest(): for name in AUTOINGEST: try: src = _resolve_source(name) if src is None or getattr(src, "ingested", None) is None: continue if src.ingested(): continue print(f"[ingest] auto-ingesting '{name}' (AURELIUS_AUTOINGEST)...") if name == "finance": from ingest.finance import ingest as fin_ingest from ingest.pipeline import embed_and_fuse await fin_ingest() await embed_and_fuse("finance") elif name == "news": from news_intel.pipeline import run_pipeline await run_pipeline() except Exception as e: print(f"[ingest] auto-ingest '{name}' failed: {e}") async def _news_refresh_loop(): """Optional scheduled News Intelligence refresh (env-gated, off by default — free-tier hosts shouldn't burn quota unattended).""" while True: try: stats = await news_service.refresh() print(f"[news] scheduled refresh: {stats}") except Exception as e: print(f"[news] scheduled refresh failed: {e}") await asyncio.sleep(NEWS_REFRESH_MINUTES * 60) def _resolve_source(name: str): try: return get_source(name or "wikipedia") except KeyError: return None # ══════════════════════════════════════════════════════════════ # REST # ══════════════════════════════════════════════════════════════ @app.get("/") async def api_root(): return {"status": "Aurelius backend is running.", "frontend": "Open index.html directly in your browser."} @app.get("/api/health") async def api_health(): return {"status": "ready"} @app.get("/api/sources") async def api_sources(): """Every registered source + whether it's queryable right now (live sources always are; ingested ones need an ingest run).""" out = [] for s in list_sources(): if getattr(s, "hidden", False): continue # e.g. news — used inside Finance, not a standalone graph ingested = getattr(s, "ingested", None) ready = True if ingested is None else bool(ingested()) out.append({ "name": s.name, "description": s.description, "edge_types": list(s.edge_types), "supports_backlinks": s.supports_backlinks, "mode": "ingested" if ingested is not None else "live", "ready": ready, }) return {"sources": out} @app.get("/api/random") async def api_random(source: str = Query("wikipedia")): src = _resolve_source(source) if src is None: return {"error": f"Unknown source '{source}'"} pair = await src.sample_pair() if pair: return {"start": pair[0], "end": pair[1], "source": src.name} # Fallback: two random nodes for an ingested source. if hasattr(src, "store") and getattr(src, "ingested", lambda: False)(): picks = src.store.random_nodes(src.name, 2) if len(picks) == 2: return {"start": picks[0]["title"], "end": picks[1]["title"], "source": src.name} return {"start": "Google", "end": "Mohali", "source": src.name} @app.get("/api/neighbors") async def api_neighbors(request_source: str = Query("wikipedia", alias="source"), q: str = Query(...), limit: int = Query(15)): """Outbound edges of one node — powers click-to-expand in the explorer. Each edge carries type/weight plus a human-readable display so the UI can show *why* the two nodes connect.""" src = _resolve_source(request_source) if src is None: return {"error": f"Unknown source '{request_source}'"} if len(q) > MAX_QUERY_LEN: return {"error": "Query too long."} ref = await src.resolve(q) if not ref: return {"error": f"Cannot find: '{q}'"} try: edges = await src.neighbors(ref) except Exception: return {"error": f"Could not fetch neighbors of '{ref.title}'."} limit = max(1, min(limit, 50)) # Store-backed sources have cheap edge evidence; live sources would pay # one upstream fetch per edge for it, so they return type only. cheap_display = getattr(src, "ingested", None) is not None edges = sorted(edges, key=lambda e: -e.weight) def _kind(node_id: str): # node kind (company/etf/sector/…) rides in stored features — # the UI colors expanded nodes by it (visual clustering). if not cheap_display: return None node = src.store.get_node(src.name, node_id) return (node or {}).get("features", {}).get("kind") out, seen = [], set() for e in edges: if e.dst.id in seen or e.dst.id == ref.id: continue seen.add(e.dst.id) display = (await src.edge_display(ref, e.dst)) if cheap_display else None out.append({"id": e.dst.id, "title": e.dst.title, "type": e.type, "weight": round(e.weight, 3), "kind": _kind(e.dst.id), "display": display or e.type.replace("_", " ")}) if len(out) >= limit: break return {"source": src.name, "node": {"id": ref.id, "title": ref.title, "kind": _kind(ref.id)}, "neighbors": out} @app.get("/api/suggest") async def api_suggest(request_source: str = Query("wikipedia", alias="source"), q: str = Query(...), limit: int = Query(8)): """Domain-aware type-ahead. Each source suggests from its OWN vocabulary (companies/tickers for finance, papers for research, diseases/genes for biology, articles for Wikipedia) instead of the old hardcoded Wikipedia opensearch. Never errors the UI — returns an empty list on any failure.""" src = _resolve_source(request_source) if src is None: return {"error": f"Unknown source '{request_source}'"} if len(q) > MAX_QUERY_LEN: return {"error": "Query too long."} try: items = await src.suggest(q.strip(), limit=max(1, min(limit, 15))) except Exception as e: print(f"[suggest] {src.name}: {e}") items = [] return {"source": src.name, "suggestions": items} @app.get("/api/node") async def api_node(request_source: str = Query("wikipedia", alias="source"), q: str = Query(...)): """One resolved node with its stored features (price series, sector, kind, …) — powers the Compare panel's charts and header facts.""" src = _resolve_source(request_source) if src is None: return {"error": f"Unknown source '{request_source}'"} if len(q) > MAX_QUERY_LEN: return {"error": "Query too long."} ref = await src.resolve(q) if not ref: return {"error": f"Cannot find: '{q}'"} info = await src.node_info(ref, rich=True) return {"source": src.name, "id": ref.id, "title": ref.title, "summary": info.summary, "features": info.features or {}} def _compute_exposure(src, ref, k: int) -> list[dict]: """Weighted 1-/2-hop event propagation over an ingested source's typed edges. Shared by /api/exposure and the LLM compare explainer.""" store = src.store fmt = getattr(src, "format_edge", lambda t, w: t.replace("_", " ")) SKIP = {"based_in", "headquarters_of"} # country hub adds noise, not signal INDIRECT_DECAY = 0.7 agg: dict[str, float] = {} # id → summed exposure best: dict[str, tuple[float, list, str]] = {} # id → (score, via, chain) l1 = [(d, t, w) for d, t, w in store.neighbors(src.name, ref.id) if t not in SKIP and d != ref.id][:40] for d1, t1, w1 in l1: agg[d1] = agg.get(d1, 0.0) + w1 if d1 not in best or w1 > best[d1][0]: best[d1] = (w1, [], fmt(t1, w1)) for d2, t2, w2 in store.neighbors(src.name, d1)[:40]: if d2 in (ref.id, d1) or t2 in SKIP: continue sc = w1 * w2 * INDIRECT_DECAY agg[d2] = agg.get(d2, 0.0) + sc if d2 not in best or sc > best[d2][0]: best[d2] = (sc, [d1], f"{fmt(t1, w1)} → {fmt(t2, w2)}") if not agg: return [] k = max(1, min(k, 25)) ranked = sorted(agg.items(), key=lambda kv: -kv[1])[:k] need_titles = [i for i, _ in ranked] for i, _sc in ranked: need_titles.extend(best[i][1]) titles = store.titles_for(src.name, need_titles) max_sc = ranked[0][1] exposed = [] for node_id, score in ranked: _sc, via, chain = best[node_id] node = store.get_node(src.name, node_id) or {} exposed.append({ "id": node_id, "title": titles.get(node_id, node_id), "kind": (node.get("features") or {}).get("kind"), "score": round(100.0 * score / max_sc, 1), "via": [{"id": v, "title": titles.get(v, v)} for v in via], "chain": chain, }) return exposed def _fin_card(store, node_id: str) -> dict: """Compact facts for a related finance node — enough to render a row (name, ticker, kind, sector, last price, window change).""" node = store.get_node("finance", node_id) or {} f = node.get("features") or {} return { "id": node_id, "title": node.get("title", node_id), "kind": f.get("kind"), "sector": f.get("sector"), "last_price": f.get("last_price"), "change_pct": f.get("window_change_pct"), } @app.get("/api/company") async def api_company(q: str = Query(...), request_source: str = Query("finance", alias="source")): """One company's full research profile, aggregated from the finance graph in a single call: price series + key facts, and every typed relationship grouped into the sections a researcher actually wants — peers, supply chain, correlations, ownership, leadership. The graph becomes a secondary explorer; THIS is the primary finance surface.""" src = _resolve_source(request_source) if src is None or getattr(src, "ingested", None) is None: return {"error": "Company profiles need an ingested finance source."} if not src.ingested(): return {"error": "Finance data has not been ingested yet."} if len(q) > MAX_QUERY_LEN: return {"error": "Query too long."} ref = await src.resolve(q) if not ref: return {"error": f"Cannot find: '{q}'"} store = src.store node = store.get_node("finance", ref.id) or {} feats = node.get("features") or {} # group the node's typed edges into researcher-facing buckets buckets: dict[str, list[tuple[str, float]]] = { "competes_with": [], "supplied_by": [], "supplies": [], "co_moves": [], "macro_correlates": [], "held_by": [], "stake_held_by": [], } ceo_id = country_id = sector_id = None for dst, typ, w in store.neighbors("finance", ref.id): if typ in buckets: buckets[typ].append((dst, w)) elif typ == "led_by": ceo_id = dst elif typ == "based_in": country_id = dst elif typ == "sector_member": sector_id = dst def cards(items, with_corr=False, limit=12): items = sorted(items, key=lambda kv: -kv[1])[:limit] out = [] for did, w in items: c = _fin_card(store, did) if with_corr: c["corr"] = round(w, 2) out.append(c) return out # peers: direct rivals first, then same-sector members peers = [d for d, _ in buckets["competes_with"]] seen = set(peers) | {ref.id} if sector_id: for d, typ, _w in store.neighbors("finance", sector_id): if typ == "has_member" and d not in seen: peers.append(d); seen.add(d) peer_cards = [_fin_card(store, p) for p in peers[:10]] ceo = None if ceo_id: cnode = store.get_node("finance", ceo_id) or {} ceo = {"id": ceo_id, "title": cnode.get("title", ceo_id), "summary": cnode.get("summary", "")} country = None if country_id: country = (store.get_node("finance", country_id) or {}).get( "title", country_id).replace(" (Sector)", "") return { "id": ref.id, "title": node.get("title", ref.id), "kind": feats.get("kind"), "sector": feats.get("sector"), "summary": node.get("summary", ""), "last_price": feats.get("last_price"), "change_pct": feats.get("window_change_pct"), "series": feats.get("series"), "series_days": feats.get("series_days"), "ceo": ceo, "country": country, "peers": peer_cards, "suppliers": cards(buckets["supplied_by"]), "customers": cards(buckets["supplies"]), "correlated": cards(buckets["co_moves"], with_corr=True), "macro": cards(buckets["macro_correlates"], with_corr=True), "etfs": cards(buckets["held_by"]), "investors": cards(buckets["stake_held_by"]), } @app.get("/api/exposure") async def api_exposure(request_source: str = Query("finance", alias="source"), q: str = Query(...), k: int = Query(10)): """Dependency / event-propagation analysis: if this node moves, who feels it? Weighted 1- and 2-hop walk over the stored typed edges; every result carries the strongest chain as evidence.""" src = _resolve_source(request_source) if src is None: return {"error": f"Unknown source '{request_source}'"} if getattr(src, "ingested", None) is None or not src.ingested(): return {"error": "Exposure analysis needs an ingested source."} if len(q) > MAX_QUERY_LEN: return {"error": "Query too long."} ref = await src.resolve(q) if not ref: return {"error": f"Cannot find: '{q}'"} return {"source": src.name, "node": {"id": ref.id, "title": ref.title}, "exposed": _compute_exposure(src, ref, k)} @app.get("/api/paper") async def api_paper(q: str = Query(...), request_source: str = Query("openalex", alias="source"), refs: int = Query(25), cites: int = Query(25)): """A paper's citation dossier: its metadata plus the works it CITES (references) and the works that CITE it (citations), each with authors, year, venue and citation count — the seed for the citation explorer. Replaces the old connect-two-papers pathfinding as the primary research surface.""" src = _resolve_source(request_source) if src is None: return {"error": f"Unknown source '{request_source}'"} if len(q) > MAX_QUERY_LEN: return {"error": "Query too long."} ref = await src.resolve(q) if not ref: return {"error": f"Couldn't find a paper matching '{q}'."} return await _paper_payload(src, ref, refs, cites) async def _paper_payload(src, ref, refs: int, cites: int) -> dict: info = await src.node_info(ref, rich=True) def _card(nref, feats): f = feats or {} return {"id": nref.id, "title": nref.title, "year": f.get("year") or None, "authors": f.get("authors") or [], "venue": f.get("venue") or "", "cited_by_count": f.get("cited_by_count") or 0} try: out_edges = await src.neighbors(ref) except Exception: out_edges = [] in_edges = [] if src.supports_backlinks: try: in_edges = await src.back_neighbors(ref, limit=cites * 2) except Exception: in_edges = [] ref_nodes = [e.dst for e in out_edges][:max(1, min(refs, 60))] cite_nodes = [e.src for e in in_edges][:max(1, min(cites, 60))] ref_infos = await src.node_infos(ref_nodes) if ref_nodes else [] cite_infos = await src.node_infos(cite_nodes) if cite_nodes else [] references = [_card(n, i.features) for n, i in zip(ref_nodes, ref_infos)] citations = [_card(n, i.features) for n, i in zip(cite_nodes, cite_infos)] citations.sort(key=lambda c: -c["cited_by_count"]) # influential first f = info.features or {} return { "source": src.name, "id": ref.id, "title": ref.title, "year": f.get("year") or None, "authors": f.get("authors") or [], "venue": f.get("venue") or "", "cited_by_count": f.get("cited_by_count") or 0, "abstract": info.summary or "", "references": references, "citations": citations, "n_references": len(references), "n_citations": len(citations), } _PDF_DOI_RE = re.compile(r"\b(10\.\d{4,9}/[-._;()/:a-z0-9]+)", re.IGNORECASE) _PDF_ARXIV_RE = re.compile(r"arXiv:\s*(\d{4}\.\d{4,5})", re.IGNORECASE) def _pdf_locator(data: bytes) -> str | None: """Pull a resolvable identifier out of an uploaded PDF: a DOI or arXiv id from the first few pages (where they almost always sit), else the title guessed from the largest first-page line. OpenAlex resolves any of these.""" try: from pypdf import PdfReader reader = PdfReader(io.BytesIO(data)) text = "" for page in reader.pages[:3]: text += "\n" + (page.extract_text() or "") except Exception as e: print(f"[paper] pdf parse failed: {e}") return None m = _PDF_DOI_RE.search(text) if m: return m.group(1).rstrip(".,;)") m = _PDF_ARXIV_RE.search(text) if m: return f"arXiv:{m.group(1)}" # Fallback: the first substantial line is usually the title. for line in text.splitlines(): s = line.strip() if len(s) >= 20 and any(ch.isalpha() for ch in s): return s[:MAX_QUERY_LEN] return None @app.post("/api/paper/upload") async def api_paper_upload(request: Request, file: UploadFile = File(...), request_source: str = Query("openalex", alias="source")): """Resolve a paper from an uploaded PDF: extract its DOI / arXiv id / title, then return the same citation dossier as /api/paper. Rate-limited and size-capped like the other write-ish endpoints.""" ip = request.client.host if request.client else "unknown" if not _rate_ok(ip): return {"error": "Rate limit reached — slow down a moment."} src = _resolve_source(request_source) if src is None: return {"error": f"Unknown source '{request_source}'"} data = await file.read() if not data: return {"error": "Empty file."} if len(data) > 15 * 1024 * 1024: return {"error": "PDF too large (15 MB max)."} locator = _pdf_locator(data) if not locator: return {"error": "Couldn't find a DOI, arXiv id or title in that PDF."} ref = await src.resolve(locator) if not ref: return {"error": f"Found “{locator[:60]}” in the PDF but couldn't match a paper."} payload = await _paper_payload(src, ref, 25, 25) payload["matched_via"] = locator[:80] return payload @app.get("/api/relate") async def api_relate(request_source: str = Query("wikipedia", alias="source"), a: str = Query(...), b: str = Query(...)): src = _resolve_source(request_source) if src is None: return {"error": f"Unknown source '{request_source}'"} if len(a) > MAX_QUERY_LEN or len(b) > MAX_QUERY_LEN: return {"error": "Query too long."} return await relate_query(src, a, b) @app.get("/api/discover") async def api_discover(request_source: str = Query("wikipedia", alias="source"), a: str = Query(...), k: int = Query(12)): src = _resolve_source(request_source) if src is None: return {"error": f"Unknown source '{request_source}'"} if len(a) > MAX_QUERY_LEN: return {"error": "Query too long."} return await discover_query(src, a, k=max(1, min(k, 30))) # ══════════════════════════════════════════════════════════════ # Optional LLM layer — lazy, additive, fallback-first # # Every endpoint here returns 200 with {available, reason} and never 500s # for an LLM problem. The frontend calls these AFTER the fast non-LLM # result is already on screen, so nothing the user relies on can be slowed # or broken by Gemini being down/keyless/throttled. `_ai_guard` folds the # per-IP client rate limit and the Gemini budget into one early return. # ══════════════════════════════════════════════════════════════ def _ai_guard(request: Request) -> dict | None: """None = clear to call the LLM; otherwise the fallback payload.""" ip = request.client.host if request.client else "unknown" if not _rate_ok(ip): return {"available": False, "reason": "rate_limited"} if not llm.available(): return {"available": False, "reason": llm.llm_status()["reason"]} return None @app.get("/api/llm/status") async def api_llm_status(): """Whether AI features can run right now — booleans only, no key leak.""" return llm.llm_status() @app.get("/api/explain/path") async def api_explain_path(request: Request, request_source: str = Query("wikipedia", alias="source"), nodes: str = Query(...), edges: str = Query("")): guard = _ai_guard(request) if guard: return guard node_list = [n for n in nodes.split("|") if n][:12] edge_list = edges.split("|") if edges else [] if len(node_list) < 2: return {"available": False, "reason": "error"} res = await llm.generate(llm.explain_path(request_source, node_list, edge_list), max_output_tokens=300) return {"available": res["ok"], "reason": res["reason"], "text": res["text"]} @app.get("/api/explain/compare") async def api_explain_compare(request: Request, request_source: str = Query("wikipedia", alias="source"), a: str = Query(...), b: str = Query(...)): guard = _ai_guard(request) if guard: return guard if len(a) > MAX_QUERY_LEN or len(b) > MAX_QUERY_LEN: return {"available": False, "reason": "error"} src = _resolve_source(request_source) if src is None: return {"available": False, "reason": "error"} data = await relate_query(src, a, b) if data.get("error"): return {"available": False, "reason": "error"} exposure = None if getattr(src, "ingested", None) is not None and src.ingested(): ref = await src.resolve(a) if ref: exposure = _compute_exposure(src, ref, 6) res = await llm.generate( llm.analyze_relation(request_source, data["a"]["title"], data["b"]["title"], data, exposure), max_output_tokens=340) return {"available": res["ok"], "reason": res["reason"], "text": res["text"]} @app.get("/api/explain/discover") async def api_explain_discover(request: Request, request_source: str = Query("wikipedia", alias="source"), a: str = Query(...)): guard = _ai_guard(request) if guard: return guard if len(a) > MAX_QUERY_LEN: return {"available": False, "reason": "error"} src = _resolve_source(request_source) if src is None: return {"available": False, "reason": "error"} data = await discover_query(src, a, k=8) cands = data.get("candidates", []) if not cands: return {"available": False, "reason": "error"} res = await llm.generate( llm.explain_discovery(request_source, data["a"]["title"], cands), max_output_tokens=300) return {"available": res["ok"], "reason": res["reason"], "text": res["text"]} @app.get("/api/explain/entity") async def api_explain_entity(request: Request, request_source: str = Query("wikipedia", alias="source"), q: str = Query(...)): guard = _ai_guard(request) if guard: return guard if len(q) > MAX_QUERY_LEN: return {"available": False, "reason": "error"} src = _resolve_source(request_source) if src is None: return {"available": False, "reason": "error"} ref = await src.resolve(q) if not ref: return {"available": False, "reason": "error"} info = await src.node_info(ref, rich=True) res = await llm.generate( llm.summarize_entity(request_source, ref.title, {"summary": info.summary, "features": info.features}), max_output_tokens=160) return {"available": res["ok"], "reason": res["reason"], "text": res["text"]} # ══════════════════════════════════════════════════════════════ # News Intelligence — shared service, domain-agnostic # ══════════════════════════════════════════════════════════════ @app.get("/api/news/status") async def api_news_status(): return news_service.status() @app.get("/api/news/search") async def api_news_search(q: str = Query(...), k: int = Query(12)): if len(q) > MAX_QUERY_LEN: return {"error": "Query too long."} return await news_service.search(q, k=max(1, min(k, 30))) @app.get("/api/news/stories") async def api_news_stories(limit: int = Query(12)): return news_service.stories(limit=max(1, min(limit, 50))) @app.get("/api/news/entity") async def api_news_entity(name: str = Query(...), k: int = Query(10)): if len(name) > MAX_QUERY_LEN: return {"error": "Query too long."} return news_service.entity_news(name, k=max(1, min(k, 30))) @app.get("/api/news/summary") async def api_news_summary(request: Request, entity: str = Query(...)): """LLM coverage summary + refined tone for an entity. Additive: the per-article lexicon sentiment is unaffected; this narrates on top and degrades to {available:false} when the LLM is off/throttled.""" guard = _ai_guard(request) if guard: return guard if len(entity) > MAX_QUERY_LEN: return {"available": False, "reason": "error"} return await news_service.coverage_summary(entity) @app.get("/api/news/refresh") async def api_news_refresh(request: Request): # A refresh hits upstream feeds + runs embedding — rate-limit it like # the WS endpoint so a public deployment can't be farmed. ip = request.client.host if request.client else "unknown" if not _rate_ok(ip): return {"error": "Rate limit reached — slow down a moment."} return await news_service.refresh() # ══════════════════════════════════════════════════════════════ # WebSocket — streamed pathfinding # ══════════════════════════════════════════════════════════════ @app.websocket("/ws") async def websocket_endpoint(ws: WebSocket): global _active_searches origin = ws.headers.get("origin") if origin is not None and origin not in ALLOWED_ORIGINS: await ws.close(code=1008) return client_ip = ws.client.host if ws.client else "unknown" if not _rate_ok(client_ip): await ws.accept() await ws.send_text(json.dumps( {"event": "error", "message": "Rate limit reached — slow down a moment."})) await ws.close(code=1008) return await ws.accept() navigator: GraphNavigator | None = None counted = False try: raw = await ws.receive_text() data = json.loads(raw) if "control" in data: return start = data.get("start", "").strip() end = data.get("end", "").strip() source_name = (data.get("source") or "wikipedia").strip() print(f"\n{'='*60}\n[WS] [{source_name}] '{start}' -> '{end}'\n{'='*60}") if not start or not end: await ws.send_text(json.dumps({"event": "error", "message": "Need start and end."})) return if len(start) > MAX_QUERY_LEN or len(end) > MAX_QUERY_LEN: await ws.send_text(json.dumps({"event": "error", "message": "Article name is too long."})) return src = _resolve_source(source_name) if src is None: await ws.send_text(json.dumps({"event": "error", "message": f"Unknown source '{source_name}'."})) return if getattr(src, "ingested", None) is not None and not src.ingested(): await ws.send_text(json.dumps({"event": "error", "message": f"Source '{source_name}' has not been ingested yet."})) return if _active_searches >= MAX_CONCURRENT_SEARCHES: await ws.send_text(json.dumps( {"event": "error", "message": "Server is busy right now — please try again in a moment."})) return _active_searches += 1 counted = True async def emit(event: str, payload: dict): try: await ws.send_text(json.dumps({"event": event, **payload})) except Exception: pass navigator = GraphNavigator(src, start, end, emit) async def listen_controls(): while True: try: msg = await asyncio.wait_for(ws.receive_text(), timeout=0.5) ctrl = json.loads(msg) if navigator: if ctrl.get("control") == "pause": navigator.paused = True; print("[WS] Paused") elif ctrl.get("control") == "resume": navigator.paused = False; print("[WS] Resumed") except asyncio.TimeoutError: pass except Exception: break ctrl_task = asyncio.create_task(listen_controls()) search_task = asyncio.create_task(navigator.run()) done, pending = await asyncio.wait( [ctrl_task, search_task], timeout=MAX_SEARCH_SECONDS, return_when=asyncio.FIRST_COMPLETED, ) for t in pending: t.cancel() # Tell the client when the wall-clock cap fired — otherwise the UI # sits on the last status line with no explanation. if search_task in pending: await emit("not_found", { "message": f"Search stopped at the {MAX_SEARCH_SECONDS}s time limit.", "visited": [], "stats": navigator.stats.to_dict(), }) except WebSocketDisconnect: print("[WS] disconnected") except Exception as e: print(f"[WS] {e}") try: await ws.send_text(json.dumps( {"event": "error", "message": "Something went wrong on the server."})) except Exception: pass finally: if counted: _active_searches -= 1