""" agent.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. 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, Optional import pandas as pd from tools import ( load_csv, preprocess_text, run_topic_modeling, generate_labels, compare_themes, create_taxonomy_map, ) # --------------------------------------------------------------------------- # 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) # 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. 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() # ----------------------------------------------------------------- # Public interface # ----------------------------------------------------------------- def run(self) -> AgentResult: """ Execute the full pipeline step by step. Returns ------- AgentResult Structured results including all DataFrames, taxonomy, and file paths. """ logger.info("=" * 60) logger.info("TopicAgent — Starting pipeline") logger.info("=" * 60) try: # Step 1: Load CSV self._step_load_csv() if self.use_specter: self._step_specter_pipeline() else: # Step 2: Preprocess text self._step_preprocess() # Step 3: Topic modeling on Titles self._step_model_titles() # Step 4: Topic modeling on Abstracts self._step_model_abstracts() # Step 5: Generate labels self._step_generate_labels() # Step 6: Build combined topics table self._step_combine_topics() # Step 7: Compare themes self._step_compare_themes() # Step 8: Create taxonomy map self._step_taxonomy_map() # Step 9: Export outputs (TCCM not used if specter, but export is) if not self.use_specter: self._step_tccm() 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 # ----------------------------------------------------------------- # Pipeline steps # ----------------------------------------------------------------- def _step_load_csv(self): """Step 1 — Ingest CSV dataset.""" logger.info("Step 1/9: Loading CSV …") self._df = load_csv(self.csv_path) self._result.steps_completed.append("load_csv") logger.info(" → %d papers loaded.", len(self._df)) def _step_preprocess(self): """Step 2 — Preprocess Title and Abstract text.""" logger.info("Step 2/9: 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") logger.info(" → Titles preprocessed: %d docs", len(self._titles_clean)) logger.info(" → Abstracts preprocessed: %d docs", len(self._abstracts_clean)) def _step_model_titles(self): """Step 3 — Topic modeling on Titles (clusters + sub-topics).""" logger.info("Step 3/9: 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, max_clusters=self.max_clusters, ) self._result.steps_completed.append("topic_modeling_titles") n_clusters = self._title_topics_df["cluster_id"].nunique() if "cluster_id" in self._title_topics_df.columns else len(self._title_topics_df) logger.info(" → %d title topics in %d clusters.", len(self._title_topics_df), n_clusters) def _step_model_abstracts(self): """Step 4 — Topic modeling on Abstracts (clusters + sub-topics).""" logger.info("Step 4/9: 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, max_clusters=self.max_clusters, ) self._result.steps_completed.append("topic_modeling_abstracts") n_clusters = self._abstract_topics_df["cluster_id"].nunique() if "cluster_id" in self._abstract_topics_df.columns else len(self._abstract_topics_df) logger.info(" → %d abstract topics in %d clusters.", len(self._abstract_topics_df), n_clusters) def _step_generate_labels(self): """Step 5 — Generate human-readable labels.""" logger.info("Step 5/9: 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") logger.info(" → Labels generated for all topics.") def _step_combine_topics(self): """Step 6 — Combine title and abstract topics into one table.""" logger.info("Step 6/9: 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") logger.info(" → Combined table: %d topics total.", len(combined)) def _step_compare_themes(self): """Step 7 — Compare title vs abstract themes.""" logger.info("Step 7/9: 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") logger.info(" → Comparison table: %d rows.", len(comparison)) def _step_taxonomy_map(self): """Step 8 — Create taxonomy map (MAPPED / NOVEL).""" logger.info("Step 8/9: Building taxonomy map …") # Use the combined topics for taxonomy taxonomy = create_taxonomy_map(self._result.combined_topics) self._result.taxonomy_map = taxonomy self._result.steps_completed.append("create_taxonomy_map") logger.info( " → MAPPED: %d, NOVEL: %d", taxonomy["metadata"]["mapped_count"], taxonomy["metadata"]["novel_count"], ) def _step_export(self): """Step 9 — Export all outputs to disk.""" logger.info("Step 9/9: 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 logger.info(" → Saved: %s", 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 logger.info(" → Saved: %s", 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 logger.info(" → Saved: %s", 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 logger.info(" → Saved: %s", 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 logger.info(" → Saved: %s", abstract_path) self._result.steps_completed.append("export") logger.info(" → All outputs exported successfully.") def _step_tccm(self): logger.info("Step X: TCCM Clustering phase ...") # tccm implementation self._result.steps_completed.append("tccm_extraction") def _step_specter_pipeline(self): """Run SPECTER2 Pipeline with two-level hierarchy (clusters + sub-topics).""" logger.info("Running SPECTER2 Pipeline...") from run_specter import ( generate_specter2_embeddings, store_in_chroma, cluster_with_hdbscan, c_tf_idf, extract_top_n_words_per_topic, get_council_labels, extract_subtopics_within_clusters, ) import pandas as pd # 1. Combine Abstract + Title papers = self._df[['Title', 'Abstract']].to_dict('records') titles = [p['Title'] or "" for p in papers] abstracts = [p['Abstract'] or "" for p in papers] combined = [t + " " + a for t, a in zip(titles, abstracts)] # 2. Embed embeddings = generate_specter2_embeddings( [{'title': t, 'abstract': a} for t, a in zip(titles, abstracts)] ) # 3. Store store_in_chroma( embeddings, documents=combined, metadatas=[{'id': i} for i in range(len(combined))], ) # 4. Cluster (produces high-level groups) labels, _ = cluster_with_hdbscan(embeddings) self._df['cluster_id'] = labels # 5. Extract cluster-level keywords docs_per_topic = self._df[self._df['cluster_id'] != -1].groupby( ['cluster_id'], as_index=False ).agg({'Title': ' '.join}) docs_per_topic = docs_per_topic.rename(columns={'cluster_id': 'topic_id'}) tf_idf, count = c_tf_idf(docs_per_topic['Title'].values, m=len(combined)) top_n_words = extract_top_n_words_per_topic(tf_idf, count, docs_per_topic, n=10) # Build cluster-level summary cluster_rows = [] for cluster_id, words in top_n_words.items(): if cluster_id != -1: kw = ", ".join([w for w, _ in words]) cluster_rows.append({"topic_id": int(cluster_id), "keywords": kw}) self._clusters_df = pd.DataFrame(cluster_rows) n_clusters = len(self._clusters_df) logger.info(" → %d clusters formed", n_clusters) # 6. Extract sub-topics within clusters (≥100 topics) cluster_labels_arr = self._df['cluster_id'].values self._combined_topics_df = extract_subtopics_within_clusters( self._df, cluster_labels_arr, min_topics=self.min_topics, n_keywords=10, ) # Add cluster labels cluster_label_map = {} for _, row in self._clusters_df.iterrows(): cid = int(row['topic_id']) kws = [kw.strip().title() for kw in row['keywords'].split(',')][:3] cluster_label_map[cid] = ' / '.join(kws) self._combined_topics_df['cluster_label'] = ( self._combined_topics_df['cluster_id'].map(cluster_label_map).fillna('') ) logger.info( " → %d topics extracted from %d clusters", len(self._combined_topics_df), n_clusters, ) # 7. Label topics if self.use_llm_labels and self.groq_api_key: labels_list = get_council_labels(self._combined_topics_df, self.groq_api_key) self._combined_topics_df['label'] = labels_list else: self._combined_topics_df['label'] = [ ", ".join(row['keywords'].split(",")[:3]).title() for _, row in self._combined_topics_df.iterrows() ] self._result.combined_topics = self._combined_topics_df