FIRST = [
    "Bob",
    "Arno",
    "Ronald",
    "Richard",
    "Sara",
    "Mina",
    "Leila",
    "Omid",
    "Nora",
    "Ali",
    "Hana",
    "Reza",
    "Tara",
    "Kian",
    "Dara",
    "Lena",
    "Yara",
    "Ivan",
    "Maya",
    "Noah",
    "Zane",
    "Iris",
    "Elif",
    "Omar",
    "Ravi",
]
LAST = [
    "Joe",
    "Swarzfinger",
    "Coleyman",
    "Preor",
    "Karimi",
    "Novak",
    "Haddad",
    "Silva",
    "Moreau",
    "Tanaka",
    "Okafor",
    "Larsen",
    "Rossi",
    "Weber",
    "Costa",
    "Ahmadi",
]
DEPTS = ["research", "sales", "finance", "engineering", "legal"]
TOPICS = {
    "research": ["experiments", "papers", "models", "hypotheses"],
    "sales": ["clients", "deals", "pipeline", "quotas"],
    "finance": ["budgets", "audits", "forecasts", "invoices"],
    "engineering": ["servers", "code", "deployments", "APIs"],
    "legal": ["contracts", "compliance", "disputes", "policies"],
}

def make_company(n=150, seed=SEED):
    r = np.random.default_rng(seed)
    names = []
    while len(names) < n:
        nm = f"{r.choice(FIRST)} {r.choice(LAST)}"
        if nm not in names:
            names.append(nm)
    dept = r.integers(0, 5, n)
    g = nx.MultiDiGraph()
    for i, nm in enumerate(names):
        # bios mention the department's topics, but 30% of people mention a *random* department's topics (noisy text)
        topic_dept = dept[i] if r.random() > 0.3 else r.integers(0, 5)
        words = r.choice(TOPICS[DEPTS[topic_dept]], 2, replace=False)
        g.add_node(
            i,
            name=nm,
            dept=int(dept[i]),
            bio=f"{nm} enjoys working on {words[0]} and {words[1]}.",
        )
    leads = {d: int(np.flatnonzero(dept == d)[0]) for d in range(5)}
    for i in range(n):
        if i != leads[dept[i]]:
            g.add_edge(i, leads[dept[i]], rel="reports_to")
        for _ in range(r.integers(1, 4)):
            same = r.random() < 0.8
            pool = (
                np.flatnonzero(dept == dept[i]) if same else np.arange(n)
            )
            j = int(r.choice(pool))
            if j != i and not g.has_edge(i, j) and not g.has_edge(j, i):
                g.add_edge(i, j, rel="works_with")
        if r.random() < 0.3:
            j = int(r.integers(n))
            if j != i:
                g.add_edge(i, j, rel="mentors")
    return g

company = make_company()
works_with = nx.Graph(
    [
        (u, v)
        for u, v, d in company.edges(data=True)
        if d["rel"] == "works_with"
    ]
)
facts = [
    f"{company.nodes[u]['name']} {d['rel'].replace('_', ' ')} {company.nodes[v]['name']}."
    for u, v, d in company.edges(data=True)
]
mo.vstack(
    [
        mo.md(
            f"**{company.number_of_nodes()} people, {company.number_of_edges()} typed edges** "
            f"({pd.Series([d['rel'] for *_, d in company.edges(data=True)]).value_counts().to_dict()}). Some facts as text:"
        ),
        mo.md(
            "\n".join(f"- {f}" for f in facts[:6])
            + "\n\nSome bios:\n\n"
            + "\n".join(f"- {company.nodes[i]['bio']}" for i in range(3))
        ),
    ]
)

_edge_list = [(u, v, d["rel"]) for u, v, d in company.edges(data=True)]
fact_vectorizer = TfidfVectorizer().fit(facts + list(questions.question))
fact_matrix = fact_vectorizer.transform(facts)
name_vectorizer = TfidfVectorizer(
    analyzer="char_wb", ngram_range=(2, 4)
).fit([company.nodes[i]["name"] for i in company])
name_matrix = name_vectorizer.transform(
    [company.nodes[i]["name"] for i in company]
)

def supporting(q):
    """Indices of the two facts that answer the question (works_with in either direction)."""
    need = {frozenset((q.x, q.answer)), frozenset((q.answer, q.z))}
    return {
        i
        for i, (u, v, rel) in enumerate(_edge_list)
        if rel == "works_with" and frozenset((u, v)) in need
    }

def vector_rag(q, k):
    sims = (
        (fact_matrix @ fact_vectorizer.transform([q.question]).T)
        .toarray()
        .ravel()
    )
    return set(np.argsort(-sims)[:k])

def seed_entities(question, n_seeds=2):
    """k-NN between the question's name-like spans and node names (char n-grams)."""
    text = question.replace("Who works with both ", "").rstrip("?")
    spans = [s.strip() for s in text.split(" and ")]
    sims = (name_matrix @ name_vectorizer.transform(spans).T).toarray()
    return [int(np.argmax(sims[:, j])) for j in range(len(spans))][
        :n_seeds
    ]

def graph_rag(q, mode):
    seeds = seed_entities(q.question)
    if mode == "2-hop expansion":
        nodes = set(seeds)
        for s in seeds:
            nodes |= set(
                nx.single_source_shortest_path_length(
                    company.to_undirected(as_view=True), s, cutoff=2
                )
            )
        keep = {
            i
            for i, (u, v, _) in enumerate(_edge_list)
            if u in nodes and v in nodes
        }
    else:  # Steiner tree linking the seeds
        tree = (
            steiner_tree(works_with, seeds)
            if all(s in works_with for s in seeds)
            and nx.has_path(works_with, *seeds)
            else nx.Graph()
        )
        keep = {
            i
            for i, (u, v, rel) in enumerate(_edge_list)
            if rel == "works_with" and tree.has_edge(u, v)
        }
    return keep, seeds

rag_k = mo.ui.slider(
    5, 60, step=5, value=10, label="VectorRAG top-k facts", show_value=True
)
rag_k

_undirected = company.to_undirected(as_view=False)
_all_nodes = sorted(company.nodes)

def subgraph_example(q):
    """2-hop subgraph around the seeds as a PyG-style tuple (x, edge_index, nodes, answer position)."""
    seeds = seed_entities(q.question)
    nodes = set(seeds)
    for s in seeds:
        nodes |= set(
            nx.single_source_shortest_path_length(_undirected, s, cutoff=2)
        )
    nodes = sorted(nodes)
    pos = {v: i for i, v in enumerate(nodes)}
    rel_ids = {"works_with": 0, "reports_to": 1, "mentors": 2}
    src, dst, et = [], [], []
    for u, v, d in company.edges(data=True):
        if u in pos and v in pos:
            src += [pos[u], pos[v]]
            dst += [pos[v], pos[u]]
            et += [rel_ids[d["rel"]]] * 2
    x = torch.zeros(len(nodes), 2)
    for s in seeds:
        x[pos[s], 0] = 1.0  # labelling trick: "I am a seed"
    x[:, 1] = 1.0
    return (
        x,
        torch.tensor([src, dst]),
        torch.tensor(et),
        nodes,
        pos.get(q.answer, -1),
    )

class GNNReader(nn.Module):
    """Relation-aware 2-layer GNN (one weight per relation type) + node-scoring head."""

    def __init__(self, d=32, n_rel=3):
        super().__init__()
        self.inp = nn.Linear(2, d)
        self.W = nn.ModuleList(
            nn.ModuleList(
                nn.Linear(d, d, bias=False) for _ in range(n_rel)
            )
            for _ in range(2)
        )
        self.self_w = nn.ModuleList(nn.Linear(d, d) for _ in range(2))
        self.score = nn.Linear(d, 1)

    def forward(self, x, edge_index, edge_type):
        h = F.relu(self.inp(x))
        for W, S in zip(self.W, self.self_w):
            agg = torch.zeros_like(h)
            for r, lin in enumerate(W):
                m = edge_type == r
                agg = agg.index_add(
                    0, edge_index[1, m], lin(h[edge_index[0, m]])
                )
            h = F.relu(S(h) + agg)
        return self.score(h).squeeze(1)

reader_data = [subgraph_example(q) for q in questions.itertuples()]
torch.manual_seed(SEED)
reader = GNNReader()
_opt = torch.optim.Adam(reader.parameters(), lr=0.01)
_train, _test = reader_data[:100], reader_data[100:]
for _ in range(60):
    for _x, _ei, _et, _nodes, _ans in _train:
        _opt.zero_grad()
        F.cross_entropy(
            reader(_x, _ei, _et)[None], torch.tensor([_ans])
        ).backward()
        _opt.step()
with torch.no_grad():
    reader_acc = np.mean(
        [
            int(reader(_x, _ei, _et).argmax()) == _ans
            for _x, _ei, _et, _nodes, _ans in _test
        ]
    )
_sizes = [len(n) for *_, n, _ in reader_data]
mo.md(
    f"GNN reader on 2-hop subgraphs ({np.mean(_sizes):.0f} nodes on average): test accuracy **{reader_acc:.3f}** "
    f"on {len(_test)} unseen questions (random guess ≈ {np.mean([1 / s for s in _sizes]):.3f})."
)

import os

# Point GL_OLLAMA_URL at your Ollama server (e.g. the team reverse proxy); defaults to a local Ollama.
OLLAMA_URL = os.environ.get("GL_OLLAMA_URL", "http://localhost:11434")
LLM_CACHE = CACHE_DIR / "llm_cache.json"
_cache_lock = threading.Lock()

def _load_cache():
    try:
        return json.loads(LLM_CACHE.read_text())
    except (FileNotFoundError, json.JSONDecodeError):
        return {}

def ask_llm(prompt, model, timeout=180):
    """One chat call to Ollama, cached on disk by (model, prompt); thread-safe. Returns the answer (or an error string)."""
    key = hashlib.sha1(f"{model}\n{prompt}".encode()).hexdigest()
    with _cache_lock:
        cached = _load_cache().get(key)
    if cached is not None:
        return cached
    body = json.dumps(
        {
            "model": model,
            "messages": [{"role": "user", "content": prompt}],
            "stream": False,
            "think": False,
            "options": {"temperature": 0},
        }
    ).encode()
    try:
        req = urllib.request.Request(
            f"{OLLAMA_URL}/api/chat",
            data=body,
            headers={"Content-Type": "application/json"},
        )
        answer = json.loads(
            urllib.request.urlopen(req, timeout=timeout).read()
        )["message"]["content"].strip()
    except Exception as e:  # unreachable server etc.: don't cache
        return f"[LLM error: {e}]"
    with _cache_lock:
        cache = _load_cache()
        cache[key] = answer
        tmp = LLM_CACHE.with_suffix(".tmp")
        tmp.write_text(json.dumps(cache))
        tmp.replace(LLM_CACHE)
    return answer

def server_status():
    try:
        tags = json.loads(
            urllib.request.urlopen(
                f"{OLLAMA_URL}/api/tags", timeout=5
            ).read()
        )
        return [m["name"] for m in tags["models"]]
    except Exception:
        return []

_models = server_status()
llm_model = mo.ui.dropdown(
    _models or ["gemma4:latest"],
    value="gemma4:latest"
    if "gemma4:latest" in _models or not _models
    else _models[0],
    label="LLM",
)
llm_n_questions = mo.ui.slider(
    10, 42, step=4, value=30, label="test questions", show_value=True
)
llm_run = mo.ui.run_button(
    label="Ask the LLM (≈ 1–3 min uncached, instant when cached)"
)
mo.vstack(
    [
        mo.md(
            f"Server `{OLLAMA_URL}`: "
            + (
                f"reachable, {len(_models)} models."
                if _models
                else "**not reachable**; cached answers (if any) still work."
            )
        ),
        mo.hstack([llm_model, llm_n_questions, llm_run], justify="start"),
    ]
)

def context_text(q, kind):
    if kind == "no context":
        return ""
    if kind == "all facts (raw text)":
        return "\n".join(facts)
    if kind == "VectorRAG top-10":
        return "\n".join(facts[i] for i in sorted(vector_rag(q, 10)))
    if kind == "GraphRAG Steiner tree":
        return "\n".join(
            facts[i] for i in sorted(graph_rag(q, "steiner")[0])
        )

LLM_CONTEXTS = [
    "no context",
    "all facts (raw text)",
    "VectorRAG top-10",
    "GraphRAG Steiner tree",
]

def llm_prompt(q, kind):
    ctx = context_text(q, kind)
    return (
        (f"Facts about a company:\n{ctx}\n\n" if ctx else "")
        + f"Question: {q.question}\nAnswer with the person's full name only. If you cannot know, answer 'unknown'."
    )

mo.stop(not llm_run.value, mo.md("Press **Ask the LLM**."))

_qs = list(questions.iloc[100 : 100 + llm_n_questions.value].itertuples())
_jobs = [(q, k) for q in _qs for k in LLM_CONTEXTS]
_t = time.time()
with ThreadPoolExecutor(8) as _pool:
    _answers = list(
        _pool.map(
            lambda job: ask_llm(llm_prompt(*job), llm_model.value), _jobs
        )
    )
llm_answers = pd.DataFrame(
    [
        {
            "question": q.question,
            "context": k,
            "gold": company.nodes[q.answer]["name"],
            "llm": a,
            "correct": company.nodes[q.answer]["name"].lower()
            in a.lower(),
        }
        for (q, k), a in zip(_jobs, _answers)
    ]
)
_summary = (
    llm_answers.groupby("context", sort=False)
    .correct.mean()
    .rename("accuracy")
    .to_frame()
)
_summary["context size (facts)"] = [0, len(facts), 10, 2]
mo.vstack(
    [
        mo.md(
            f"{llm_model.value}: {len(_jobs)} calls in {time.time() - _t:.0f} s"
        ),
        _summary.round(3),
        mo.accordion({"all answers": llm_answers}),
    ]
)

_rev = reveal(
    r"""
Measured on 30 held-out questions (answers cached, temperature 0):

| context | facts given | gemma4 (8B) | gemma4:26b |
|---|---|---|---|
| no context | 0 | 0.00 | 0.00 |
| all facts, raw text | 450 | 0.10 | 0.50 |
| VectorRAG top-10 | 10 | 0.03 | 0.37 |
| **GraphRAG Steiner tree** | **2** | **0.97** | **1.00** |

* Without context both models (correctly) say "unknown": the company is fictional.
* Given **all** the facts, the LLM must perform the join itself. The 8B model mostly gives up; the 26B model gets half right.
* VectorRAG is worse than raw text here: its top-10 often misses the bridge fact (§2), and even when both facts are present
  they sit among 8 distractors.
* The **2-fact Steiner-tree context** makes the question trivial for both: 0.97 / 1.00 with a 225× shorter prompt. Structure-aware
  retrieval does the reasoning, so a small, cheap LLM suffices. This is the lecture's STaRK-Prime point: an 8B model plus a GNN beat
  agentic GraphRAG with much larger models.

""",
    predict="Rank the four contexts (none, all 450 facts, VectorRAG top-10, a 2-fact Steiner tree) for the small and the large LLM.",
)
_post = mo.md(r"""
**Try it:** pick another model in the dropdown and press the button again (new answers are cached too).

## 5 · LLM feeds GNN: text-attributed graphs

The reverse direction: an LLM (or any text encoder) turns each node's **text** into features, and a GNN propagates them. Task:
predict each employee's **department** from their bio (which is misleading for 30% of people, by construction) and the graph.
Two text encoders: TF-IDF of the bio, and an **LLM zero-shot guess** ("which department does this bio suggest?", one-hot,
cached). Each is used alone (logistic regression) and as input features to a 2-layer GCN over works_with + reports_to edges.
""")
mo.vstack([_rev, _post])

mo.stop(
    not glem_run.value,
    mo.md(
        "Press the button (150 short LLM calls the first time, ≈ 30 s)."
    ),
)
with ThreadPoolExecutor(8) as _pool:
    _guess = list(_pool.map(llm_department_guess, employee_bios))
text_llm = torch.zeros(len(employee_bios), 5)
for _i, _g in enumerate(_guess):
    if _g >= 0:
        text_llm[_i, _g] = 1.0
_g_und = nx.Graph(
    [
        (u, v)
        for u, v, d in company.edges(data=True)
        if d["rel"] in ("works_with", "reports_to")
    ]
)
_ei = torch.tensor(
    [[u, v] for u, v in _g_und.edges] + [[v, u] for u, v in _g_und.edges]
).t()
_perm = torch.randperm(
    len(employee_bios), generator=torch.Generator().manual_seed(SEED)
)
_tr = torch.zeros(len(employee_bios), dtype=torch.bool)
_tr[_perm[:75]] = True

def _fit_gcn(x):
    torch.manual_seed(SEED)
    c1, c2 = GCNConv(x.size(1), 32), GCNConv(32, 5)
    opt = torch.optim.Adam(
        list(c1.parameters()) + list(c2.parameters()),
        lr=0.01,
        weight_decay=5e-4,
    )
    for _ in range(200):
        opt.zero_grad()
        out = c2(F.relu(c1(x, _ei)), _ei)
        F.cross_entropy(out[_tr], dept_y[_tr]).backward()
        opt.step()
    with torch.no_grad():
        return (
            (c2(F.relu(c1(x, _ei)), _ei).argmax(1)[~_tr] == dept_y[~_tr])
            .float()
            .mean()
            .item()
        )

def _fit_text(x):
    clf = LogisticRegression(max_iter=2000).fit(
        x[_tr].numpy(), dept_y[_tr].numpy()
    )
    return clf.score(x[~_tr].numpy(), dept_y[~_tr].numpy())

glem_results = pd.DataFrame(
    [
        {
            "text encoder": "TF-IDF of bio",
            "text only (logistic)": _fit_text(text_tfidf),
            "text → GCN": _fit_gcn(text_tfidf),
        },
        {
            "text encoder": "LLM zero-shot department guess",
            "text only (logistic)": _fit_text(text_llm),
            "text → GCN": _fit_gcn(text_llm),
        },
    ]
)
mo.vstack(
    [
        mo.md(
            f"LLM guess = true department for **{np.mean([g == int(y) for g, y in zip(_guess, dept_y)]):.2f}** of people "
            "(bios are misleading for ≈ 30% by construction). Test accuracy on 75 held-out people:"
        ),
        glem_results.round(3),
    ]
)

_rev = reveal(
    r"""
Measured (75 held-out employees): from text alone, TF-IDF reaches **0.81** and the LLM's zero-shot guesses **0.83** (the LLM is
right for 82% of bios; the misleading ones fool both). Feed either text representation into a **GCN over the org graph** and accuracy
rises to **1.00 / 0.99**: colleagues and managers correct the misleading bios. That is the LLM → GNN direction: language models
read the text, the GNN reads the structure. GLEM goes further and **co-trains** them, with the GNN's predictions as pseudo-labels
to fine-tune the LM and vice versa.

""",
    predict="Text alone vs. text + a GCN over the org graph: how much can colleagues and managers correct misleading bios?",
)
_post = mo.md(r"""
---

## 6 · Exercises

**Exercise 1 (context budget).** How many times shorter is the Steiner-tree context than the "all facts" context, in facts? → `ex1_ratio`.

**Exercise 2 (better VectorRAG).** Write `split_query_rag(q, k)`: retrieve the top-$k$ facts for **each** name in the question
separately (two TF-IDF queries, e.g. just the name string) and return the union of fact indices. With $k=8$ per name (≤ 16 facts) it should make
≥ 90% of the questions answerable, versus 46% for plain top-10. (Use `fact_vectorizer`, `fact_matrix`, `seed_entities` or plain string splitting.)

**Exercise 3 (think).** Why can't an LLM given the raw edge list be "fixed" just by giving it a longer context window?
Write your answer down, then open the solution.
""")
mo.vstack([_rev, _post])

ex1_ratio = None

def split_query_rag(q, k):
    return None

def _verdict(ok, answered):
    if not answered:
        return mo.callout(mo.md("Not answered yet."), kind="neutral")
    return mo.callout(
        mo.md("Correct!" if ok else "Not yet — try again or open a hint."),
        kind="success" if ok else "warn",
    )

def _check_split():
    try:
        q0 = next(questions.itertuples())
        if split_query_rag(q0, 8) is None:
            return _verdict(False, False)
        rate = np.mean(
            [
                supporting(q) <= set(split_query_rag(q, 8))
                for q in questions.itertuples()
            ]
        )
        return mo.callout(
            mo.md(
                f"Answerable: **{rate:.2f}** of questions with ≤ 16 facts. "
                + ("Correct!" if rate >= 0.9 else "Aim for ≥ 0.90.")
            ),
            kind="success" if rate >= 0.9 else "warn",
        )
    except Exception as e:
        return mo.callout(mo.md(f"Error: `{e}`"), kind="danger")

mo.vstack(
    [
        mo.md("**Exercise 1**"),
        _verdict(
            ex1_ratio is not None
            and abs(ex1_ratio - len(facts) / 2) < 1e-6,
            ex1_ratio is not None,
        ),
        mo.md("**Exercise 2**"),
        _check_split(),
        mo.accordion(
            {
                "Solution 1": mo.md(
                    f"{len(facts)} / 2 = **{len(facts) / 2:.0f}×**."
                ),
                "Solution 2": mo.md("""```python
def split_query_rag(q, k):
    names = q.question.replace("Who works with both ", "").rstrip("?").split(" and ")
    out = set()
    for name in names:
        sims = (fact_matrix @ fact_vectorizer.transform([name]).T).toarray().ravel()
        out |= set(np.argsort(-sims)[:k])
    return out
```
Each entity gets its own budget, so the bridge facts are no longer crowded out by the other name. This is a small step towards
entity-centric (graph) retrieval."""),
                "Solution 3": mo.md(
                    "The facts *are* all there with a long context (0.10 / 0.50 accuracy above). The bottleneck is doing the **join** "
                    "reliably inside the model: finding X's partners, Z's partners and intersecting them among hundreds of lines. "
                    "Structure-aware retrieval (or a GNN) performs the join explicitly and hands the LLM only the result."
                ),
            }
        ),
    ]
)