agent.py

"""
agent.py
--------
Agentic pipeline for Journal Topic Modelling.
Uses Mistral API (mistralai v2 SDK) via mistralai.client.Mistral.

Phases:
  1 – Parse & ingest CSV
  2 – Extract topics (batched LLM calls) + label + frequency table
  3 – Title vs Abstract comparison → comparison.csv
  4 – PAJAIS taxonomy mapping → taxonomy_map.json
  5 – Section 7 narrative generation (500+ words) → narrative.txt
"""

import os
import json
import re
import time
from typing import Callable, Any

# ── FIX: correct import for mistralai v2 SDK ─────────────────────────────────
from mistralai import Mistral

from tools import (
    PAJAIS_THEMES_LIST,
    parse_csv,
    extract_text_corpus,
    chunk_corpus_for_topic_extraction,
    merge_and_deduplicate_topics,
    assign_topic_labels,
    build_topic_frequency_table,
    split_corpus_by_field,
    compare_topic_distributions,
    save_comparison_csv,
    map_topics_to_pajais,
    save_taxonomy_json,
    build_narrative_context,
    validate_narrative,
    summarise_run,
)

# ── Client factory ────────────────────────────────────────────────────────────

def get_client() -> Mistral:
    api_key = os.environ.get("MISTRAL_API_KEY", "").strip()
    if not api_key:
        raise ValueError(
            "MISTRAL_API_KEY environment variable is not set. "
            "Add it as a Space Secret in Hugging Face Settings."
        )
    return Mistral(api_key=api_key)


# ── LLM helper ────────────────────────────────────────────────────────────────

def llm_call(
    client: Mistral,
    messages: list,
    temperature: float = 0.2,
    max_tokens: int = 4096,
) -> str:
    """
    Single synchronous LLM call. Returns assistant text content as a string.
    Uses mistralai v2 SDK: client.chat.complete(...)
    """
    response = client.chat.complete(
        model="mistral-large-latest",
        messages=messages,
        temperature=temperature,
        max_tokens=max_tokens,
    )
    # response.choices[0].message.content may be a string or list of content blocks
    content = response.choices[0].message.content
    if isinstance(content, list):
        # Extract text from content blocks
        return " ".join(
            block.text if hasattr(block, "text") else str(block)
            for block in content
        ).strip()
    return str(content).strip()


# ── JSON parser ───────────────────────────────────────────────────────────────

def parse_json_from_response(text: str) -> Any:
    """
    Robustly extract and parse the first JSON object or array from LLM output.
    Handles markdown code fences, leading prose, and trailing text.
    """
    # 1. Direct parse
    try:
        return json.loads(text)
    except Exception:
        pass

    # 2. Strip markdown code fences
    fenced = re.search(r"```(?:json)?\s*([\s\S]+?)```", text)
    if fenced:
        try:
            return json.loads(fenced.group(1).strip())
        except Exception:
            pass

    # 3. Find first [ or { and last matching ] or }
    def try_extract(open_ch, close_ch):
        start = text.find(open_ch)
        if start == -1:
            return None
        end = text.rfind(close_ch)
        if end == -1 or end <= start:
            return None
        try:
            return json.loads(text[start : end + 1])
        except Exception:
            return None

    arr = try_extract("[", "]")
    if arr is not None:
        return arr
    obj = try_extract("{", "}")
    if obj is not None:
        return obj

    raise ValueError(f"No valid JSON found in response: {text[:300]!r}")


# ── Phase 1 ───────────────────────────────────────────────────────────────────

def phase1_ingest(csv_text: str, log: Callable) -> tuple:
    """Parse CSV. Returns (records: list, columns: list)."""
    log("📂 Phase 1 – Parsing CSV…")
    result = parse_csv(csv_text)
    if result["status"] != "ok":
        raise RuntimeError(f"CSV parse failed: {result.get('message','unknown error')}")
    log(f"   ✅ {result['count']} records found | Columns: {result['columns']}")
    return result["records"], result["columns"]


# ── Phase 2 ───────────────────────────────────────────────────────────────────

def phase2_extract_topics(client: Mistral, records: list, log: Callable) -> tuple:
    """
    Extract topics from abstracts+titles, label them, build frequency table.
    Returns (labelled_topics: list, frequency_table: list).
    """
    log("📖 Phase 2 – Building text corpus…")
    corpus_result = extract_text_corpus(records, text_fields=["abstract", "title"])
    corpus = corpus_result["corpus"]
    log(f"   ✅ {corpus_result['doc_count']} documents | {corpus_result['total_words']} total words")

    chunks_result = chunk_corpus_for_topic_extraction(corpus, chunk_size=15)
    chunks = chunks_result["chunks"]
    log(f"   ✂️  {chunks_result['chunk_count']} chunk(s) for topic extraction")

    raw_topic_lists = []
    for i, chunk in enumerate(chunks):
        log(f"   🔍 Extracting topics — chunk {i + 1}/{len(chunks)}…")
        excerpts = "\n".join(
            f"[{doc['id']}] {doc['text'][:400]}" for doc in chunk
        )
        prompt = (
            "You are a research analyst specialising in academic topic modelling.\n"
            f"Analyse the following {len(chunk)} paper excerpts and extract KEY RESEARCH TOPICS.\n\n"
            "Rules:\n"
            "- Return ONLY a JSON array of topic strings (2–5 words each).\n"
            "- Extract 5–15 distinct topics capturing the core themes.\n"
            "- Be specific: prefer 'deep learning image classification' over 'AI'.\n"
            "- Do NOT include author names, journal names, or years.\n"
            "- Output ONLY valid JSON — no preamble, no markdown fences.\n\n"
            f"Excerpts:\n{excerpts}\n\n"
            'Output example: ["topic one", "topic two", "topic three"]'
        )
        try:
            resp = llm_call(client, [{"role": "user", "content": prompt}], temperature=0.1)
            topics = parse_json_from_response(resp)
            if isinstance(topics, list):
                raw_topic_lists.append([str(t) for t in topics if t])
                log(f"      → {len(topics)} topics extracted")
            else:
                log(f"      ⚠️  Unexpected response type for chunk {i + 1}, skipping")
        except Exception as e:
            log(f"      ⚠️  Chunk {i + 1} failed: {e}")
        time.sleep(0.5)  # Respect rate limits

    log("   🔄 Merging and deduplicating topics…")
    merge_result = merge_and_deduplicate_topics(raw_topic_lists)
    unique_topics = merge_result["unique_topics"]
    log(
        f"   ✅ {merge_result['raw_count']} raw → "
        f"{merge_result['unique_count']} unique topics"
    )

    # Label topics via LLM
    log("   🏷️  Labelling topics…")
    label_prompt = (
        "You are a research librarian. Assign a short, clear, title-case label to each topic below.\n\n"
        f"Topics:\n{json.dumps(unique_topics, indent=2)}\n\n"
        "Return a JSON object mapping each exact topic string to its label.\n"
        'Example: {"deep learning image classification": "Deep Learning for Vision"}\n'
        "Output ONLY valid JSON — no preamble, no markdown fences."
    )
    labels_map = {}
    try:
        label_resp = llm_call(
            client,
            [{"role": "user", "content": label_prompt}],
            temperature=0.1,
            max_tokens=3000,
        )
        parsed = parse_json_from_response(label_resp)
        if isinstance(parsed, dict):
            labels_map = parsed
        else:
            log("   ⚠️  Labels response was not a dict; using topics as labels")
    except Exception as e:
        log(f"   ⚠️  Labelling failed ({e}); using topics as their own labels")

    labelled_result = assign_topic_labels(unique_topics, labels_map)
    labelled_topics = labelled_result["labelled_topics"]
    log(f"   ✅ {len(labelled_topics)} labelled topics")

    log("   📊 Building frequency table…")
    freq_result = build_topic_frequency_table(corpus, labelled_topics)
    log(
        f"   ✅ Frequency table: {len(labelled_topics)} topics × "
        f"{freq_result['doc_count']} documents"
    )
    return labelled_topics, freq_result["frequency_table"]


# ── Phase 3 ───────────────────────────────────────────────────────────────────

def phase3_comparison(
    records: list,
    labelled_topics: list,
    log: Callable,
) -> tuple:
    """
    Compare title vs abstract topic distributions.
    Returns (comparison: list, csv_text: str).
    """
    log("⚖️  Phase 3 – Title vs Abstract comparison…")
    split = split_corpus_by_field(records)
    log(
        f"   ✅ {split['title_count']} title docs | "
        f"{split['abstract_count']} abstract docs"
    )

    title_freq = build_topic_frequency_table(
        split["title_corpus"], labelled_topics
    )["frequency_table"]
    abstract_freq = build_topic_frequency_table(
        split["abstract_corpus"], labelled_topics
    )["frequency_table"]

    compare_result = compare_topic_distributions(title_freq, abstract_freq)
    comparison = compare_result["comparison"]
    log(
        f"   ✅ {len(comparison)} topics compared | "
        f"{len(compare_result['title_only'])} title-only | "
        f"{len(compare_result['abstract_only'])} abstract-only | "
        f"{len(compare_result['shared'])} shared"
    )

    csv_result = save_comparison_csv(comparison)
    return comparison, csv_result.get("csv_text", "")


# ── Phase 4 ───────────────────────────────────────────────────────────────────

def phase4_pajais_mapping(
    client: Mistral,
    labelled_topics: list,
    log: Callable,
) -> tuple:
    """
    Map topics to PAJAIS taxonomy, identify NOVEL themes.
    Returns (taxonomy_result: dict, json_text: str).
    """
    log("🗺️  Phase 4 – PAJAIS taxonomy mapping…")
    topics_list = [item["topic"] for item in labelled_topics]

    map_prompt = (
        "You are a PAJAIS (Pacific Asia Journal of the Association for Information Systems) "
        "taxonomy expert.\n\n"
        "Map each topic below to the single most appropriate PAJAIS theme. "
        'If a topic does NOT clearly fit any theme, assign "NOVEL".\n\n'
        f"PAJAIS Themes:\n{json.dumps(PAJAIS_THEMES_LIST, indent=2)}\n\n"
        f"Topics to map:\n{json.dumps(topics_list, indent=2)}\n\n"
        "Return ONLY a JSON object: "
        '{"topic string": "PAJAIS Theme or NOVEL", ...}\n'
        "Output ONLY valid JSON — no preamble, no markdown fences."
    )
    mapping = {}
    try:
        map_resp = llm_call(
            client,
            [{"role": "user", "content": map_prompt}],
            temperature=0.1,
            max_tokens=4096,
        )
        parsed = parse_json_from_response(map_resp)
        if isinstance(parsed, dict):
            mapping = parsed
        else:
            log("   ⚠️  Mapping response was not a dict; all topics marked NOVEL")
    except Exception as e:
        log(f"   ⚠️  PAJAIS mapping failed ({e}); all topics marked NOVEL")

    taxonomy_result = map_topics_to_pajais(labelled_topics, mapping)
    log(
        f"   ✅ {taxonomy_result['mapped_count']} MAPPED | "
        f"{taxonomy_result['novel_count']} NOVEL | "
        f"{taxonomy_result['coverage_pct']}% PAJAIS coverage"
    )

    json_result = save_taxonomy_json(taxonomy_result["taxonomy_map"])
    return taxonomy_result, json_result.get("json_text", "{}")


# ── Phase 5 ───────────────────────────────────────────────────────────────────

def phase5_narrative(
    client: Mistral,
    records: list,
    frequency_table: list,
    comparison: list,
    taxonomy_result: dict,
    log: Callable,
) -> str:
    """
    Generate a 500+ word Section 7 narrative draft.
    Returns the narrative string.
    """
    log("✍️  Phase 5 – Generating Section 7 narrative…")

    ctx_result = build_narrative_context(
        frequency_table, comparison, taxonomy_result, len(records)
    )
    ctx = ctx_result["context"]

    narrative_prompt = (
        "You are an academic researcher writing Section 7 (Discussion and Conclusions) "
        "of a journal article on topic modelling results from a corpus of academic papers.\n\n"
        f"Analysis context:\n{json.dumps(ctx, indent=2)}\n\n"
        "Write a scholarly, coherent 500–600 word Section 7 covering:\n"
        "1. Overview of dominant themes discovered\n"
        "2. Title themes vs abstract themes — what does this divergence reveal?\n"
        "3. NOVEL themes not in PAJAIS and their significance for the field\n"
        "4. Research gaps and implications for future work\n"
        "5. Limitations of this study\n\n"
        "Style: academic but readable; hedged language ('suggests', 'appears to').\n"
        "Format: flowing paragraphs — NO bullet points, NO headers, NO title line.\n"
        "Start directly with the prose. Write AT LEAST 500 words."
    )

    narrative = llm_call(
        client,
        [{"role": "user", "content": narrative_prompt}],
        temperature=0.5,
        max_tokens=1800,
    )

    val = validate_narrative(narrative)
    log(f"   ✅ Narrative: {val['word_count']} words | {val['message']}")

    if not val["valid"]:
        log("   🔁 Narrative too short — extending…")
        extend_prompt = (
            "The narrative below needs extending to reach at least 500 words. "
            "Add 1–2 paragraphs discussing implications and limitations "
            "while maintaining academic tone. Return the FULL extended narrative only.\n\n"
            f"Current narrative:\n{narrative}"
        )
        narrative = llm_call(
            client,
            [{"role": "user", "content": extend_prompt}],
            temperature=0.5,
            max_tokens=2000,
        )
        val2 = validate_narrative(narrative)
        log(f"   ✅ Extended narrative: {val2['word_count']} words")

    return narrative


# ── Main entry point ──────────────────────────────────────────────────────────

def run_topic_modelling_agent(csv_text: str, log: Callable = print) -> dict:
    """
    Full 5-phase agentic pipeline.

    Args:
        csv_text: Raw CSV string (must have 'title' and/or 'abstract' columns).
        log: Callable that accepts a single string (used for live progress updates).

    Returns dict with keys:
        labelled_topics, frequency_table, comparison, comparison_csv,
        taxonomy_result, taxonomy_json, narrative, summary
    """
    client = get_client()
    log("🚀 Journal Topic Modelling Agent — starting")
    log("=" * 55)

    # Phase 1
    records, columns = phase1_ingest(csv_text, log)
    log("")

    # Phase 2
    labelled_topics, frequency_table = phase2_extract_topics(client, records, log)
    log("")

    # Phase 3
    comparison, comparison_csv = phase3_comparison(records, labelled_topics, log)
    log("")

    # Phase 4
    taxonomy_result, taxonomy_json = phase4_pajais_mapping(client, labelled_topics, log)
    log("")

    # Phase 5
    narrative = phase5_narrative(
        client, records, frequency_table, comparison, taxonomy_result, log
    )
    log("")

    # Summary
    summary_result = summarise_run(
        record_count=len(records),
        topic_count=len(labelled_topics),
        mapped_count=taxonomy_result["mapped_count"],
        novel_count=taxonomy_result["novel_count"],
        coverage_pct=taxonomy_result["coverage_pct"],
    )

    log("🎉 All phases complete!")
    log("=" * 55)

    return {
        "labelled_topics": labelled_topics,
        "frequency_table": frequency_table,
        "comparison": comparison,
        "comparison_csv": comparison_csv,
        "taxonomy_result": taxonomy_result,
        "taxonomy_json": taxonomy_json,
        "narrative": narrative,
        "summary": summary_result["summary"],
    }