aarush2109/bertopic-modelling
0
1"""2agent.py — LangGraph-based topic analysis agent (§11).33-LLM Council for topic modelling, 4 sheets, triple-agreement tracking.4"""5from __future__ import annotations6import json, logging, os, re, time7from dataclasses import dataclass, field, asdict8from typing import TypedDict, Optional9from collections import Counter10import pandas as pd, numpy as np, requests11from groq import Groq12from langgraph.graph import StateGraph, END13 14logging.basicConfig(level=logging.INFO, format="%(levelname)s | %(message)s")15logger = logging.getLogger(__name__)16 17GROQ_MODEL = "llama-3.1-8b-instant"18MISTRAL_MODEL = "mistral-small-latest"19 20# ---------------------------------------------------------------------------21# LangGraph state22# ---------------------------------------------------------------------------23class PipelineState(TypedDict, total=False):24 filepath: str25 groq_key: str26 mistral_key: str27 gemini_key: str28 n_trials: int29 topic_data: dict30 interpretations: dict31 sheets: dict # {1: [...], 2: [...], 3: [...], 4: [...]}32 agreement_rates: dict33 mismatch_table: list34 json_path: str35 csv_path: str36 error: str37 38# ---------------------------------------------------------------------------39# API helpers40# ---------------------------------------------------------------------------41def _parse(raw: str) -> dict:42 raw = raw.strip().replace("```json","").replace("```","").strip()43 s, e = raw.find("{"), raw.rfind("}")+144 if s != -1 and e > 0: raw = raw[s:e]45 try: return json.loads(raw)46 except: return {}47 48def _groq(client, prompt):49 try:50 r = client.chat.completions.create(model=GROQ_MODEL,51 messages=[{"role":"user","content":prompt}], temperature=0.2, timeout=15)52 return _parse(r.choices[0].message.content)53 except Exception as e: logger.warning("Groq: %s",e); return {}54 55def _mistral(prompt, key):56 if not key: return {}57 try:58 r = requests.post("https://api.mistral.ai/v1/chat/completions",59 headers={"Authorization":f"Bearer {key}","Content-Type":"application/json"},60 json={"model":MISTRAL_MODEL,"messages":[{"role":"user","content":prompt}],61 "temperature":0.2}, timeout=15)62 return _parse(r.json()["choices"][0]["message"]["content"])63 except Exception as e: logger.warning("Mistral: %s",e); return {}64 65def _gemini(prompt, key):66 if not key: return {}67 model = "gemini-2.5-flash"68 MAX_RETRIES = 369 waits = [4, 5, 5]70 for attempt in range(MAX_RETRIES):71 try:72 r = requests.post(73 f"https://generativelanguage.googleapis.com/v1beta/models/"74 f"{model}:generateContent?key={key}",75 headers={"Content-Type":"application/json"},76 json={"contents":[{"parts":[{"text":prompt}]}],77 "generationConfig":{"temperature":0.2}}, timeout=60)78 d = r.json()79 if "candidates" not in d:80 err = d.get("error",{})81 msg = err.get("message","") if isinstance(err,dict) else str(err)82 if "quota" in msg.lower() or "rate" in msg.lower():83 w = waits[attempt]84 logger.warning("Gemini rate-limited. Retry %d/%d after %ds", attempt+1, MAX_RETRIES, w)85 time.sleep(w); continue86 logger.warning("Gemini attempt %d: %s", attempt+1, msg)87 return {"label": "Unavailable"}88 return _parse(d["candidates"][0]["content"]["parts"][0]["text"])89 except Exception as e:90 w = waits[attempt]91 logger.warning("Gemini error. Retry %d/%d after %ds", attempt+1, MAX_RETRIES, w)92 time.sleep(w)93 logger.warning("Gemini unavailable for topic. Continuing with remaining models.")94 return {"label": "Unavailable"}95 96# ---------------------------------------------------------------------------97# Topic labelling prompt98# ---------------------------------------------------------------------------99def _label_prompt(keyphrases, rep_docs):100 kp = ", ".join(k[0] if isinstance(k,tuple) else k for k in keyphrases[:5])101 ab = " | ".join(a[:250] for a in rep_docs[:3])102 return f"""You are a research topic classifier.103A SPECTER-2 + HDBSCAN pipeline produced a topic cluster.104KEYPHRASES: {kp}105REPRESENTATIVE ABSTRACTS: {ab}106Return ONLY valid JSON:107{{108 "label": "<5-8 word topic label>",109 "description": "<one sentence description>",110 "pacis_match": "<closest PAJAIS 2019 category, or NOVEL if none>",111 "confidence": <0.0-1.0>112}}"""113 114# ---------------------------------------------------------------------------115# Defence prompt for disagreements116# ---------------------------------------------------------------------------117def _defence_prompt(keyphrases, rep_docs, votes):118 kp = ", ".join(k[0] if isinstance(k,tuple) else k for k in keyphrases[:5])119 v_str = "\n".join(f" LLM{i+1}: {v.get('label','?')}" for i,v in enumerate(votes))120 return f"""Resolve this labelling disagreement.121KEYPHRASES: {kp}122Votes:\n{v_str}123Pick the best label or synthesise a better one.124Return ONLY JSON: {{"label":"...","description":"...","pacis_match":"...","confidence":0.0}}"""125 126# ---------------------------------------------------------------------------127# Grounding check128# ---------------------------------------------------------------------------129def _grounding(label, keyphrases):130 if not label or not keyphrases: return {"verdict":"FAIL", "score":0, "precision":0.0, "recall":0.0, "matched":[]}131 lt = set(re.findall(r"\b[a-z]{3,}\b", label.lower()))132 kt = set()133 for k in keyphrases:134 kt.update(re.findall(r"\b[a-z]{3,}\b", (k[0] if isinstance(k,tuple) else k).lower()))135 noise = {"the","and","for","with","using","based","from","that","are","this"}136 lt -= noise; kt -= noise137 m = list(lt & kt)138 precision = len(m)/max(len(lt),1)139 recall = len(m)/max(len(kt),1)140 return {"verdict":"PASS" if m else "FAIL", "score":precision, "precision":precision, "recall":recall, "matched":m}141 142def _clean(s):143 s = str(s or "").replace("\n"," ").strip()144 return s[:60].rsplit(" ",1)[0] if len(s)>60 else s145 146# ---------------------------------------------------------------------------147# LangGraph node: run topic modelling148# ---------------------------------------------------------------------------149def embed_and_cluster(state: PipelineState) -> dict:150 from tools import run_topic_modeling151 try:152 td = run_topic_modeling(state["filepath"], state.get("n_trials", 50))153 return {"topic_data": td}154 except Exception as e:155 return {"error": str(e)}156 157# ---------------------------------------------------------------------------158# LangGraph node: LLM Council — 4 sheets for topic modelling159# ---------------------------------------------------------------------------160def llm_council(state: PipelineState) -> dict:161 td = state["topic_data"]162 if not td: return {"error": "No topic data"}163 client = Groq(api_key=state["groq_key"], max_retries=0)164 mk, gk = state["mistral_key"], state["gemini_key"]165 166 sheets = {1:[], 2:[], 3:[], 4:[]} # 1=Groq, 2=Mistral, 3=Gemini, 4=Consolidated167 interps = {}168 169 for cid in sorted(td["keyphrases"].keys()):170 kps = td["keyphrases"][cid]171 rds = td["representative_docs"].get(cid, [])172 sw = td["membership"].get(cid, {"strong":0,"weak":0})173 prompt = _label_prompt(kps, rds)174 175 s1 = _groq(client, prompt); time.sleep(1)176 s2 = _mistral(prompt, mk); time.sleep(1)177 s3 = _gemini(prompt, gk); time.sleep(4) # respect Gemini free-tier rate limit178 votes = [s1, s2, s3]179 180 # Sheets 1-3181 for si, (sheet_n, resp) in enumerate([(1,s1),(2,s2),(3,s3)]):182 sheets[sheet_n].append({"cluster":cid, **{k:resp.get(k,"—")183 for k in ["label","description","pacis_match","confidence"]}})184 185 # Sheet 4: consolidate186 valid = [v for v in votes if v and "label" in v]187 labels_l = [_clean(v.get("label","")).lower() for v in valid]188 counts = Counter(labels_l)189 190 if any(c>=3 for c in counts.values()):191 agreement = "Triple"192 winner = max(counts, key=counts.get)193 best = next(v for v in valid if _clean(v["label"]).lower()==winner)194 elif any(c>=2 for c in counts.values()):195 agreement = "Two"196 winner = max(counts, key=counts.get)197 best = next(v for v in valid if _clean(v["label"]).lower()==winner)198 else:199 agreement = "Single"200 d = _groq(client, _defence_prompt(kps, rds, votes))201 best = d if d and "label" in d else (valid[0] if valid else {})202 203 label = _clean(best.get("label",""))204 gc = _grounding(label, kps)205 if gc["verdict"]=="FAIL" and valid:206 label = _clean(valid[0].get("label",""))207 gc = _grounding(label, kps)208 209 consistency_score = 1.0 if agreement == "Triple" else (0.67 if agreement == "Two" else 0.33)210 211 cp = td.get("cluster_persistence",{}).get(cid, 0.0)212 sheets[4].append({"cluster":cid, "label":label, "agreement":agreement,213 "description":best.get("description",""),214 "pacis_match":best.get("pacis_match",""),215 "strong":sw["strong"], "weak":sw["weak"],216 "persistence":round(cp,4), "grounding":gc["verdict"],217 "consistency": round(consistency_score, 4),218 "precision": round(gc.get("precision",0), 4),219 "recall": round(gc.get("recall",0), 4)})220 221 interps[cid] = {"label":label, "agreement":agreement,222 "strong":sw["strong"], "weak":sw["weak"],223 "persistence":cp, "description":best.get("description",""),224 "pacis_match":best.get("pacis_match",""),225 "keyphrases":[k[0] if isinstance(k,tuple) else k for k in kps[:5]],226 "consistency": round(consistency_score, 4),227 "precision": round(gc.get("precision",0), 4),228 "recall": round(gc.get("recall",0), 4)}229 230 logger.info("Cluster %d → %s [%s]", cid, label, agreement)231 232 # Agreement rate on labels233 total = len(sheets[4]) or 1234 n_triple = sum(1 for r in sheets[4] if r.get("agreement")=="Triple")235 n_two = sum(1 for r in sheets[4] if r.get("agreement")=="Two")236 237 from sklearn.metrics import cohen_kappa_score238 l1 = [_clean(r.get("label","")).lower() for r in sheets[1]]239 l2 = [_clean(r.get("label","")).lower() for r in sheets[2]]240 l3 = [_clean(r.get("label","")).lower() for r in sheets[3]]241 k12 = cohen_kappa_score(l1, l2) if len(l1) == len(l2) and l1 else 0242 k23 = cohen_kappa_score(l2, l3) if len(l2) == len(l3) and l2 else 0243 k13 = cohen_kappa_score(l1, l3) if len(l1) == len(l3) and l1 else 0244 avg_kappa = (k12 + k23 + k13) / 3 if l1 else 0245 246 avg_precision = sum(r.get("precision",0) for r in sheets[4]) / total247 avg_recall = sum(r.get("recall",0) for r in sheets[4]) / total248 avg_consistency = sum(r.get("consistency",0) for r in sheets[4]) / total249 250 rates = {251 "triple": round(n_triple / total * 100),252 "two_or_more": round((n_triple + n_two) / total * 100),253 "single": round((total - n_triple - n_two) / total * 100),254 "cohen_kappa": round(avg_kappa, 4),255 "global_precision": round(avg_precision, 4),256 "global_recall": round(avg_recall, 4),257 "global_consistency": round(avg_consistency, 4)258 }259 260 # Append the global cohen_kappa to each topic so it appears in the JSON output261 for r in sheets[4]:262 r["cohen_kappa"] = rates["cohen_kappa"]263 if r["cluster"] in interps:264 interps[r["cluster"]]["cohen_kappa"] = rates["cohen_kappa"]265 266 # Save outputs — 4 separate sheet files267 sheet_paths = {}268 names = {1:"sheet1_groq",2:"sheet2_mistral",3:"sheet3_gemini",4:"sheet4_consolidated"}269 for sn, name in names.items():270 path = f"{name}.csv"271 pd.DataFrame(sheets[sn]).to_csv(path, index=False)272 sheet_paths[sn] = path273 with open("topics.json","w") as f: json.dump(sheets[4], f, indent=2)274 275 return {"interpretations":interps, "sheets":sheets,276 "agreement_rates":rates, "sheet_paths":sheet_paths,277 "json_path":"topics.json"}278 279# ---------------------------------------------------------------------------280# LangGraph node: build mismatch table281# ---------------------------------------------------------------------------282def build_mismatch(state: PipelineState) -> dict:283 from tools import build_mismatch_table284 td = state["topic_data"]285 interps = state.get("interpretations", {})286 labels_map = {cid: v["label"] for cid, v in interps.items()}287 mt = build_mismatch_table(td["keyphrases"], labels_map)288 return {"mismatch_table": mt}289 290# ---------------------------------------------------------------------------291# Build the LangGraph292# ---------------------------------------------------------------------------293def build_graph() -> StateGraph:294 g = StateGraph(PipelineState)295 g.add_node("embed_and_cluster", embed_and_cluster)296 g.add_node("llm_council", llm_council)297 g.add_node("build_mismatch", build_mismatch)298 g.set_entry_point("embed_and_cluster")299 g.add_edge("embed_and_cluster", "llm_council")300 g.add_edge("llm_council", "build_mismatch")301 g.add_edge("build_mismatch", END)302 return g.compile()303 304# Compiled graph — importable305pipeline_graph = build_graph()306 307def run_pipeline(filepath, groq_key, mistral_key, gemini_key, n_trials=50):308 """Convenience wrapper."""309 result = pipeline_graph.invoke({310 "filepath": filepath,311 "groq_key": groq_key,312 "mistral_key": mistral_key,313 "gemini_key": gemini_key,314 "n_trials": n_trials,315 })316 return result317