File size: 35,569 Bytes
75ce203
 
 
 
 
 
2bd66df
75ce203
 
 
 
 
 
 
2bd66df
 
 
658d200
2bd66df
75ce203
658d200
75ce203
 
2bd66df
658d200
 
2bd66df
d6382b1
2bd66df
d6382b1
 
 
75ce203
 
d6382b1
75ce203
19c6bad
75ce203
 
 
 
 
2bd66df
d6382b1
 
 
 
 
658d200
d6382b1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
75ce203
 
d6382b1
2bd66df
75ce203
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2bd66df
 
 
 
75ce203
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2bd66df
 
 
75ce203
2bd66df
 
 
 
75ce203
 
2bd66df
 
 
 
 
 
 
75ce203
 
 
 
 
 
19c84c2
 
75ce203
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
590ff18
 
 
 
 
 
 
 
 
75ce203
 
 
 
 
 
 
 
590ff18
75ce203
 
 
 
590ff18
75ce203
 
 
658d200
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
19c84c2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
19c6bad
 
 
19c84c2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
19c6bad
19c84c2
 
 
 
 
 
 
 
 
 
19c6bad
19c84c2
 
 
 
 
 
 
 
 
19c6bad
 
 
658d200
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
19c6bad
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
19c84c2
19c6bad
 
19c84c2
 
658d200
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
75ce203
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
19c6bad
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2bd66df
75ce203
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
19c6bad
 
 
 
 
 
 
 
 
 
 
 
 
75ce203
 
 
 
 
 
 
 
 
 
 
 
2bd66df
 
 
 
d6382b1
 
 
 
75ce203
 
 
 
 
 
 
 
 
d6382b1
 
2bd66df
75ce203
d6382b1
2bd66df
75ce203
2bd66df
 
 
 
 
 
75ce203
 
 
2bd66df
75ce203
2bd66df
d6382b1
75ce203
 
 
 
 
 
 
 
 
 
d6382b1
 
 
 
75ce203
d6382b1
 
 
2bd66df
75ce203
 
 
 
 
 
 
2bd66df
 
 
 
75ce203
2bd66df
 
 
75ce203
2bd66df
75ce203
2bd66df
 
 
 
 
 
 
 
d6382b1
 
 
2bd66df
 
 
75ce203
 
 
 
 
 
 
2bd66df
 
 
 
 
 
d6382b1
75ce203
2bd66df
 
d6382b1
 
 
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
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
"""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