from tools import merge_similar_clusters, extract_subtopics_within_clusters """ agents.py — TopicAgent orchestrates the end-to-end topic modeling workflow. This module defines the TopicAgent class, which: 1. Loads and validates the CSV dataset. 2. Preprocesses text for Titles and Abstracts separately. 3. Runs topic modeling on each corpus (≥100 topics guaranteed). 4. Generates human-readable labels for every topic. 5. Compares dominant themes across Title and Abstract topics. 6. Produces a taxonomy map (MAPPED / NOVEL classification). 7. Exports structured outputs: topics table, comparison CSV, taxonomy JSON. OR (SPECTER mode): 1. Loads CSV 2. Generates SPECTER2 embeddings + stores in ChromaDB 3. Adaptive HDBSCAN clustering (5-140 clusters) 4. c-TF-IDF keyword extraction + representative sentences 5. AI Council labeling with debate transcripts 6. Taxonomy mapping 7. Export Usage: agent = TopicAgent(csv_path="dataset.csv") results = agent.run() """ import os import json import logging from dataclasses import dataclass, field from typing import Dict, Any, List, Optional import numpy as np import pandas as pd from tools import ( load_csv, preprocess_text, run_topic_modeling, generate_labels, compare_themes, create_taxonomy_map, generate_specter2_embeddings, store_in_chroma, load_embeddings_from_chroma, adaptive_hdbscan_clustering, c_tf_idf, extract_top_n_words_per_topic, extract_representative_sentences, get_council_debate, build_cluster_graph, ) # --------------------------------------------------------------------------- # Logging # --------------------------------------------------------------------------- logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Structured result container # --------------------------------------------------------------------------- @dataclass class AgentResult: """Container for all outputs produced by the TopicAgent.""" # Core dataframes title_topics: pd.DataFrame = field(default_factory=pd.DataFrame) abstract_topics: pd.DataFrame = field(default_factory=pd.DataFrame) combined_topics: pd.DataFrame = field(default_factory=pd.DataFrame) comparison: pd.DataFrame = field(default_factory=pd.DataFrame) # Taxonomy map (dict serialisable to JSON) taxonomy_map: Dict[str, Any] = field(default_factory=dict) # SPECTER/Cluster-specific data cluster_data: Dict[str, Any] = field(default_factory=dict) # Holds: embeddings_2d, cluster_labels, cluster_details, representative_sentences, # cluster_graph_edges, n_clusters, hdbscan_params council_debates: Dict[int, Dict[str, str]] = field(default_factory=dict) # {topic_id: {expert, analyst, anti_hallucination, final_label, ...}} # Pipeline mode: "default" or "specter" pipeline_mode: str = "default" # Execution metadata status: str = "pending" steps_completed: list = field(default_factory=list) errors: list = field(default_factory=list) # File paths of exported artefacts exported_files: Dict[str, str] = field(default_factory=dict) # --------------------------------------------------------------------------- # TopicAgent # --------------------------------------------------------------------------- class TopicAgent: """ Orchestrates the research-paper topic modeling pipeline. Parameters ---------- csv_path : str Path to the input CSV file. output_dir : str Directory to write output files. min_topics : int Minimum number of topics to generate per source (default 100). use_llm_labels : bool Whether to use Groq LLM for label generation. use_specter : bool Whether to use SPECTER2 + HDBSCAN clustering pipeline. groq_api_key : str, optional API key for Groq (used only when use_llm_labels is True). """ def __init__( self, csv_path: str, output_dir: str = "outputs", min_topics: int = 100, max_clusters: int = 25, use_llm_labels: bool = False, use_specter: bool = False, groq_api_key: Optional[str] = None, ): self.csv_path = csv_path self.output_dir = output_dir self.min_topics = min_topics self.max_clusters = max_clusters self.use_llm_labels = use_llm_labels self.use_specter = use_specter self.groq_api_key = groq_api_key # Ensure output directory exists os.makedirs(self.output_dir, exist_ok=True) self._result = AgentResult() if use_specter: self._result.pipeline_mode = "specter" # Internal state self._df = None self._embeddings = None self._cluster_labels = None self._reduced_2d = None # ----------------------------------------------------------------- # Public interface # ----------------------------------------------------------------- def run(self) -> AgentResult: """Execute the full pipeline. Use sub-steps for progress display.""" logger.info("=" * 60) logger.info("TopicAgent — Starting pipeline (mode=%s)", self._result.pipeline_mode) logger.info("=" * 60) try: self._step_load_csv() if self.use_specter: self._step_specter_embed() self._step_specter_cluster() self._step_specter_extract() self._step_merge_similar_clusters() self._step_specter_subtopics() self._step_specter_label() self._step_specter_taxonomy() else: self._step_preprocess() self._step_model_titles() self._step_model_abstracts() self._step_generate_labels() self._step_combine_topics() self._step_compare_themes() self._step_taxonomy_map() self._step_export() self._result.status = "success" logger.info("Pipeline completed successfully.") except Exception as exc: self._result.status = "failed" self._result.errors.append(str(exc)) logger.error("Pipeline failed: %s", exc, exc_info=True) return self._result # ----------------------------------------------------------------- # Shared steps # ----------------------------------------------------------------- def _step_load_csv(self): """Step 1 — Ingest CSV dataset.""" logger.info("Step 1: Loading CSV …") self._df = load_csv(self.csv_path) self._result.steps_completed.append("load_csv") logger.info(" → %d papers loaded.", len(self._df)) # ----------------------------------------------------------------- # DEFAULT pipeline steps # ----------------------------------------------------------------- def _step_preprocess(self): """Step 2 — Preprocess Title and Abstract text.""" logger.info("Step 2: Preprocessing text …") self._titles_clean = preprocess_text(self._df["Title"].tolist()) self._abstracts_clean = preprocess_text(self._df["Abstract"].tolist()) self._result.steps_completed.append("preprocess_text") def _step_model_titles(self): """Step 3 — Topic modeling on Titles.""" logger.info("Step 3: Topic modeling on Titles …") self._title_topics_df, self._title_model = run_topic_modeling( self._titles_clean, source_label="Titles", min_topics=self.min_topics, ) self._result.steps_completed.append("topic_modeling_titles") def _step_model_abstracts(self): """Step 4 — Topic modeling on Abstracts.""" logger.info("Step 4: Topic modeling on Abstracts …") self._abstract_topics_df, self._abstract_model = run_topic_modeling( self._abstracts_clean, source_label="Abstracts", min_topics=self.min_topics, ) self._result.steps_completed.append("topic_modeling_abstracts") def _step_generate_labels(self): """Step 5 — Generate human-readable labels.""" logger.info("Step 5: Generating topic labels …") self._title_topics_df = generate_labels( self._title_topics_df, use_llm=self.use_llm_labels, groq_api_key=self.groq_api_key, ) self._abstract_topics_df = generate_labels( self._abstract_topics_df, use_llm=self.use_llm_labels, groq_api_key=self.groq_api_key, ) self._result.title_topics = self._title_topics_df self._result.abstract_topics = self._abstract_topics_df self._result.steps_completed.append("generate_labels") def _step_combine_topics(self): """Step 6 — Combine title and abstract topics.""" logger.info("Step 6: Building combined topics table …") combined = pd.concat([self._title_topics_df, self._abstract_topics_df], ignore_index=True) combined["global_id"] = range(len(combined)) self._result.combined_topics = combined self._result.steps_completed.append("combine_topics") def _step_compare_themes(self): """Step 7 — Compare title vs abstract themes.""" logger.info("Step 7: Comparing title vs abstract themes …") comparison = compare_themes(self._title_topics_df, self._abstract_topics_df) self._result.comparison = comparison self._result.steps_completed.append("compare_themes") def _step_taxonomy_map(self): """Step 8 — Create taxonomy map.""" logger.info("Step 8: Building taxonomy map …") taxonomy = create_taxonomy_map(self._result.combined_topics) self._result.taxonomy_map = taxonomy self._result.steps_completed.append("create_taxonomy_map") # ----------------------------------------------------------------- # SPECTER pipeline steps (discrete, callable from app.py) # ----------------------------------------------------------------- def _step_specter_embed(self): """SPECTER Step 2 — Generate embeddings and store in ChromaDB.""" logger.info(">>> Clustering pipeline started") logger.info("SPECTER Step 2: Generating SPECTER2 embeddings …") titles = [str(t) if t else "" for t in self._df["Title"].tolist()] abstracts = [str(a) if a else "" for a in self._df["Abstract"].tolist()] combined_docs = [t + " " + a for t, a in zip(titles, abstracts)] # Check if embeddings already cached in ChromaDB with correct count try: cached_emb, cached_docs, _ = load_embeddings_from_chroma() if len(cached_emb) == len(self._df): logger.info("Loaded embeddings from vector DB — %d vectors", len(cached_emb)) self._embeddings = cached_emb self._combined_docs = combined_docs self._result.steps_completed.append("specter_embed_cached") return except (ValueError, Exception) as e: logger.info("No cached embeddings found (%s), generating fresh …", e) # Generate fresh embeddings papers = [{"title": t, "abstract": a} for t, a in zip(titles, abstracts)] self._embeddings = generate_specter2_embeddings(papers) self._combined_docs = combined_docs # Store in ChromaDB metadatas = [{"id": i, "title": titles[i][:200]} for i in range(len(combined_docs))] store_in_chroma(self._embeddings, documents=combined_docs, metadatas=metadatas) logger.info("Loaded embeddings from vector DB — %d vectors", len(self._embeddings)) # Validate — no nulls null_count = int(np.isnan(self._embeddings).any(axis=1).sum()) if null_count > 0: logger.warning("Found %d embeddings with NaN values, replacing with zeros", null_count) self._embeddings = np.nan_to_num(self._embeddings, nan=0.0) self._result.steps_completed.append("specter_embed") def _step_specter_cluster(self): """SPECTER Step 3 — Adaptive HDBSCAN clustering (10–40 clusters).""" logger.info("Running HDBSCAN (targeting %d–%d clusters)...", 10, self.max_clusters) labels, reduced_2d, n_clusters, params = adaptive_hdbscan_clustering( self._embeddings, min_clusters=10, max_clusters=self.max_clusters, ) self._cluster_labels = np.array(labels) self._reduced_2d = reduced_2d self._df["cluster_id"] = labels logger.info("Clusters generated: %d", n_clusters) # Safeguard: if < 5 clusters, raise error if n_clusters < 5: raise RuntimeError( f"Clustering produced only {n_clusters} clusters (minimum 5 required). " "Try adjusting parameters or adding more data." ) self._result.cluster_data["n_clusters"] = n_clusters self._result.cluster_data["hdbscan_params"] = params self._result.cluster_data["embeddings_2d"] = reduced_2d.tolist() self._result.cluster_data["cluster_labels"] = labels.tolist() if hasattr(labels, 'tolist') else list(labels) self._result.steps_completed.append("specter_cluster") def _step_specter_extract(self): """SPECTER Step 4 — Extract cluster-level keywords and representative sentences.""" logger.info("SPECTER Step 4: Extracting cluster keywords and representative sentences …") # c-TF-IDF keyword extraction at CLUSTER level df_with_clusters = self._df.copy() df_with_clusters["cluster_id"] = self._cluster_labels # Group documents by cluster for c-TF-IDF docs_per_topic = df_with_clusters.groupby("cluster_id", as_index=False).agg( {"Title": " ".join, "Abstract": " ".join} ) docs_per_topic = docs_per_topic[docs_per_topic["cluster_id"] != -1] docs_per_topic = docs_per_topic.rename(columns={"cluster_id": "topic_id"}) # Use combined Title + Abstract for better keyword extraction combined_text = (docs_per_topic["Title"] + " " + docs_per_topic["Abstract"]).values tf_idf, count = c_tf_idf(combined_text, m=len(self._df)) top_n_words = extract_top_n_words_per_topic(tf_idf, count, docs_per_topic, n=10) # Build CLUSTER-level DataFrame (these are the high-level groups) cluster_rows = [] for cluster_id, words in top_n_words.items(): if cluster_id == -1: continue kw = ", ".join([w for w, _ in words]) n_docs = int((self._cluster_labels == cluster_id).sum()) cluster_rows.append({ "topic_id": int(cluster_id), "keywords": kw, "source": "SPECTER2+HDBSCAN", "n_documents": n_docs, }) self._clusters_df = pd.DataFrame(cluster_rows) logger.info(" → %d clusters with keywords extracted", len(self._clusters_df)) # Extract representative sentences rep_sentences = extract_representative_sentences( self._df, self._embeddings, self._cluster_labels, top_n=5, ) self._result.cluster_data["representative_sentences"] = rep_sentences # Build cluster graph edges = build_cluster_graph(self._clusters_df, self._embeddings, self._cluster_labels) self._result.cluster_data["graph_edges"] = edges # Store per-cluster details cluster_details = {} for _, row in self._clusters_df.iterrows(): cid = int(row["topic_id"]) cluster_details[cid] = { "cluster_id": cid, "keywords": row["keywords"], "n_documents": int(row["n_documents"]), "representatives": rep_sentences.get(cid, []), } self._result.cluster_data["cluster_details"] = cluster_details self._result.steps_completed.append("specter_extract") def _step_specter_subtopics(self): """SPECTER Step 5 — Extract ≥100 fine-grained sub-topics within each cluster.""" logger.info( "SPECTER Step 5: Extracting sub-topics (target ≥ %d topics from %d clusters) …", self.min_topics, len(self._clusters_df), ) self._topics_df = extract_subtopics_within_clusters( self._df, self._cluster_labels, min_topics=self.min_topics, n_keywords=10, ) # Add cluster_label from the clusters_df cluster_label_map = {} for _, row in self._clusters_df.iterrows(): cid = int(row["topic_id"]) # Heuristic label from top 3 keywords kws = [kw.strip().title() for kw in row["keywords"].split(",")][:3] cluster_label_map[cid] = " / ".join(kws) self._topics_df["cluster_label"] = self._topics_df["cluster_id"].map(cluster_label_map).fillna("") logger.info( " → %d topics extracted from %d clusters", len(self._topics_df), self._topics_df["cluster_id"].nunique(), ) self._result.steps_completed.append("specter_subtopics") def _step_specter_label(self): """SPECTER Step 6 — Label clusters and topics.""" logger.info("SPECTER Step 6: Labeling clusters and topics …") # --- Label CLUSTERS (high-level) --- if self.use_llm_labels and self.groq_api_key: logger.info("Running AI Council debate for %d clusters …", len(self._clusters_df)) cluster_labels_list = [] debates = {} for _, row in self._clusters_df.iterrows(): cid = int(row["topic_id"]) debate = get_council_debate(row["keywords"], self.groq_api_key) debates[cid] = debate cluster_labels_list.append(debate["final_label"]) logger.info(" Cluster %d → '%s'", cid, debate["final_label"]) self._clusters_df["label"] = cluster_labels_list self._result.council_debates = debates # Update cluster_label in topics_df with LLM labels cluster_label_map = dict(zip(self._clusters_df["topic_id"].astype(int), self._clusters_df["label"])) self._topics_df["cluster_label"] = self._topics_df["cluster_id"].map(cluster_label_map).fillna("") else: logger.info("Using heuristic labels (LLM not enabled)") self._clusters_df["label"] = [ ", ".join(row["keywords"].split(",")[:3]).strip().title() for _, row in self._clusters_df.iterrows() ] # --- Label TOPICS (fine-grained) --- if self.use_llm_labels and self.groq_api_key: logger.info("Generating LLM labels for %d topics …", len(self._topics_df)) from tools import _generate_labels_llm topic_labels_df = self._topics_df.copy() topic_labels_df = _generate_labels_llm(topic_labels_df, self.groq_api_key) self._topics_df["label"] = topic_labels_df["label"] else: self._topics_df["label"] = [ ", ".join(row["keywords"].split(",")[:3]).strip().title() for _, row in self._topics_df.iterrows() ] # Populate result — combined_topics now has the TWO-LEVEL hierarchy self._result.combined_topics = self._topics_df self._result.title_topics = self._clusters_df # Cluster-level view self._result.abstract_topics = pd.DataFrame() # Not applicable for SPECTER self._result.steps_completed.append("specter_label") def _step_merge_similar_clusters(self): """SPECTER Step 4b — Merge semantically similar clusters before sub-topic extraction.""" from tools import merge_similar_clusters if not hasattr(self, "_clusters_df") or self._clusters_df.empty: return logger.info("SPECTER Step 4b: Post-merging similar clusters...") new_clusters_df, new_labels = merge_similar_clusters( self._clusters_df, self._cluster_labels, self._df, threshold=0.75, groq_api_key=self.groq_api_key ) self._clusters_df = new_clusters_df self._cluster_labels = new_labels n_clusters = len(self._clusters_df) self._result.cluster_data["cluster_labels"] = new_labels.tolist() if hasattr(new_labels, 'tolist') else list(new_labels) self._result.cluster_data["n_clusters"] = n_clusters logger.info(" → After merging: %d clusters remain", n_clusters) self._result.steps_completed.append("specter_merge") def _step_specter_taxonomy(self): """SPECTER Step 6 — Build taxonomy map for clustered topics.""" logger.info("SPECTER Step 6: Building taxonomy map …") taxonomy = create_taxonomy_map(self._result.combined_topics) self._result.taxonomy_map = taxonomy self._result.steps_completed.append("specter_taxonomy") # ----------------------------------------------------------------- # Export (shared) # ----------------------------------------------------------------- def _step_export(self): """Export all outputs to disk.""" logger.info("Exporting outputs …") # (a) Combined topics table CSV topics_path = os.path.join(self.output_dir, "topics_table.csv") self._result.combined_topics.to_csv(topics_path, index=False) self._result.exported_files["topics_table"] = topics_path # (b) Comparison CSV comparison_path = os.path.join(self.output_dir, "comparison.csv") self._result.comparison.to_csv(comparison_path, index=False) self._result.exported_files["comparison"] = comparison_path # (c) Taxonomy map JSON taxonomy_path = os.path.join(self.output_dir, "taxonomy_map.json") with open(taxonomy_path, "w", encoding="utf-8") as f: json.dump(self._result.taxonomy_map, f, indent=2, ensure_ascii=False) self._result.exported_files["taxonomy_map"] = taxonomy_path # (d) Title topics CSV title_path = os.path.join(self.output_dir, "title_topics.csv") self._result.title_topics.to_csv(title_path, index=False) self._result.exported_files["title_topics"] = title_path # (e) Abstract topics CSV abstract_path = os.path.join(self.output_dir, "abstract_topics.csv") self._result.abstract_topics.to_csv(abstract_path, index=False) self._result.exported_files["abstract_topics"] = abstract_path # (f) Cluster data JSON (SPECTER mode) if self._result.pipeline_mode == "specter" and self._result.cluster_data: cluster_path = os.path.join(self.output_dir, "cluster_data.json") # Don't save embeddings_2d to JSON (too large), save separately cluster_export = { k: v for k, v in self._result.cluster_data.items() if k != "embeddings_2d" } with open(cluster_path, "w", encoding="utf-8") as f: json.dump(cluster_export, f, indent=2, ensure_ascii=False, default=str) self._result.exported_files["cluster_data"] = cluster_path # (g) Council debates JSON if self._result.council_debates: debates_path = os.path.join(self.output_dir, "council_debates.json") with open(debates_path, "w", encoding="utf-8") as f: json.dump(self._result.council_debates, f, indent=2, ensure_ascii=False, default=str) self._result.exported_files["council_debates"] = debates_path self._result.steps_completed.append("export") logger.info(" → All outputs exported successfully.")