Blame for agents.py

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.")