app/modules/search/graph.py

from langchain_groq import ChatGroq
from langgraph.graph import StateGraph, START, END
from typing import TypedDict, List, Dict, Literal, Optional
from pydantic import BaseModel, Field
from dotenv import load_dotenv
from app.core.logging_config import get_node_logger
from app.modules.search.providers.arxiv import search_arxiv
from app.modules.search.providers.semantic_scholar import search_semantic_scholar
from app.modules.search.providers.openalex import search_openalex
from app.modules.search.reranking import rerank_by_relevance
import re
import unicodedata
import os


load_dotenv()

second_api_key = os.getenv("SECOND_GROQ_API_KEY")
api_key = os.getenv("GROQ_API_KEY")
# Setting up model
llm = ChatGroq(model = 'openai/gpt-oss-120b',api_key=second_api_key)
llm_2 = ChatGroq(model='openai/gpt-oss-120b',api_key=second_api_key)

# State of the Agent
class ResearchSearchState(TypedDict):
    run_id : str
    ResearchIntent : str
    Open_Alex_paper: List[dict]
    Semantic_Scholar_paper : List[dict]
    arXiv_paper: List[dict]
    aggregator : dict
    reranked_papers : dict
    clustered_papers : List[dict]
    # Per-source status so the UI can surface WHY a source came back empty
    # (e.g. Semantic Scholar rate-limited) instead of it looking silently blank.
    # Distinct keys per source -> no concurrent-write conflict across the
    # parallel fetch nodes.
    arxiv_status : dict
    semantic_scholar_status : dict
    open_alex_status : dict

class ClusterItem(BaseModel):
    cluster_id: int
    label: str = Field(description="Short, specific approach name")
    rationale: str = Field(description="2-3 sentence explanation of what unifies this cluster")
    paper_ids: List[int]

class ClusteringOutput(BaseModel):
    clusters: List[ClusterItem]

# Defining Nodes
def arXiv_papers(state: ResearchSearchState)-> ResearchSearchState:
    logger = get_node_logger(state["run_id"], "search.arxiv")
    prompt = f"""You are converting a Research Intent into a search query for the arXiv API.

arXiv search is PURELY LEXICAL — it does not understand natural language, intent,
or reasoning. It only matches literal keyword overlap using field-prefixed terms.

RULES FOR THE QUERY YOU PRODUCE:
- Use only these field prefixes: ti: (title), abs: (abstract), all: (all fields).
  Default to all: unless a term is clearly a proper technical name that belongs in a title.
- Combine terms using ONLY uppercase AND, OR, ANDNOT. No other operators exist.
- Do NOT nest field prefixes inside parentheses (e.g. all:(ti:(x)) is invalid/ambiguous
  on arXiv and must be avoided).
- Keep the structure FLAT: field:term AND field:term AND (term1 OR term2) is fine;
  deeper nesting is not.
- Extract only concrete technical noun phrases from the Problem and Objective sections.
  Do NOT use Additional Context — arXiv cannot reason about constraints, exclusions,
  or preferences, only literal keyword matches.
- Multi-word technical phrases should be joined with + between words (arXiv API convention),
  e.g. reinforcement+learning.
- Produce 4-6 core technical terms maximum. More terms narrow results without arXiv
  being able to judge relevance the way a ranked search engine would.
- Output ONLY the final query string in arXiv API syntax. No explanation, no preamble.

Research Intent:
{state['ResearchIntent']}"""
    try:
        query = llm.invoke(prompt).content
    except Exception:
        logger.exception("LLM call failed while generating the arXiv query")
        raise
    query = query.strip().strip('"').strip("'")
    logger.info("arXiv query: %s", query)
    papers = search_arxiv(query=query, logger=logger)
    logger.info("arXiv returned %d paper(s)", len(papers))
    return {'arXiv_paper': papers, 'arxiv_status': {"state": "ok", "count": len(papers)}}

def semantic_scholar_papers(state : ResearchSearchState)-> ResearchSearchState:
    logger = get_node_logger(state["run_id"], "search.semantic_scholar")
    prompt = f"""You are converting a Research Intent into a search query for the Semantic Scholar Search API. This endpoint matches terms ONLY against each paper's title and abstract.

CRITICAL BEHAVIOR TO UNDERSTAND FIRST:
This is an EXACT MATCH system, not a semantic/fuzzy search. There is no "close enough."
Bare words or quoted phrases placed next to each other with no operator between them
are implicitly ANDed — ALL of them must appear, verbatim, in the same paper's title or
abstract. This means every additional required term you add narrows the result set
further and increases the risk of zero results. A SINGLE required phrase that doesn't
match the field's actual standard terminology (e.g. requiring "human-operated" when
the literature actually says "human-driven") can zero out the entire query, even if
every other term in it is correct.

SYNTAX AVAILABLE:
- Bare word or phrase with no operator: still REQUIRED (ANDed with everything else),
  not optional. There is no implicit OR between adjacent terms.
- "exact phrase": double-quote a multi-word phrase that must appear exactly as
  written, word-for-word. Only use this for terminology you are highly confident
  is the field's standard phrasing.
- +term: explicitly required (behaves the same as a bare term — the + is for
  clarity/emphasis, not a different matching rule).
- -term: term must NOT appear. Use only for genuine exclusions the researcher
  explicitly stated.
- (a | b | c): matches ANY of the listed alternatives — use this to hedge against
  terminology uncertainty for a single concept.
- Parentheses group | and + clauses; keep grouping to 1 level deep.

RULES FOR CONSTRUCTING THE QUERY:
- Identify at most 2 CORE concepts from Problem + Objective that are essential to
  the research question, and require each with a quoted exact phrase (using + or
  bare — same effect). Do not require more than 2 concepts — since all required
  terms are ANDed, each additional one multiplies the risk of zero results.
- For ANY concept where you are not fully certain of the field's standard exact
  wording (e.g. how a specific comparison, method, or entity type is typically
  phrased in academic writing), do NOT lock it to a single quoted phrase. Instead,
  express it as an OR group of 2-3 plausible standard phrasings, e.g.
  ("human-driven" | "human-operated" | "manually driven"). This is especially
  important for adjective/description phrases (how something is characterized),
  which vary more across papers than core technical nouns do.
- Use (a | b) groups only for genuine synonymous variants of ONE concept — do not
  invent alternative concepts, methods, or approaches not present in the Research
  Intent.
- Do NOT use - exclusions unless Additional Context explicitly states something
  should be excluded.
- Keep the total query under ~15 words worth of terms. Favor 2 required core
  concepts plus 1 hedged OR group over stacking many required exact phrases.
- Output ONLY the final query string. No explanation, no preamble.

Research Intent:
{state['ResearchIntent']}"""
    try:
        query = llm.invoke(prompt).content
    except Exception:
        logger.exception("LLM call failed while generating the Semantic Scholar query")
        raise
    query = query.strip().strip('"').strip("'")
    logger.info("Semantic Scholar query: %s", query)
    status: dict = {}
    papers = search_semantic_scholar(query=query, logger=logger, status_out=status)
    logger.info("Semantic Scholar returned %d paper(s)", len(papers))
    return {'Semantic_Scholar_paper': papers, 'semantic_scholar_status': status}

def Open_alex_papers(state : ResearchSearchState)-> ResearchSearchState:
    logger = get_node_logger(state["run_id"], "search.openalex")
    prompt = f"""You are converting a Research Intent into a search query for the OpenAlex /works
search endpoint, which searches titles, abstracts, and available fulltext.

CRITICAL BEHAVIOR TO ACCOUNT FOR:
Words not separated by explicit boolean operators are treated as AND by default.
This means a long plain-keyword query risks ZERO results if any single term is too
narrow, because ALL terms must match. You must actively defend against this.

SYNTAX AVAILABLE:
- Uppercase AND, OR, NOT only (lowercase will not be recognized as operators).
- "exact phrase" for phrases that must appear together.
- Parentheses for grouping.

RULES FOR CONSTRUCTING THE QUERY:
- Identify 2-4 DISTINCT core concepts from Problem + Objective (not more — each
  AND-ed concept multiplies the risk of zero results).
- For each core concept, group any synonym/near-synonym terms together with OR
  inside parentheses, e.g. ("fuel efficiency" OR "fuel consumption").
- Join the distinct concept groups with AND.
- Use quotes around every multi-word phrase — unquoted multi-word phrases will be
  split into individual AND-ed words, which is almost always too strict.
- Do NOT use NOT unless Additional Context explicitly names something to exclude.
- Do NOT add synonyms or terms not implied by the Research Intent, even to "help"
  recall — only group terms that are genuinely equivalent phrasings of the same idea.
- Output ONLY the final query string. No explanation, no preamble.

Research Intent:
{state['ResearchIntent']}"""
    try:
        query = llm.invoke(prompt).content
    except Exception:
        logger.exception("LLM call failed while generating the OpenAlex query")
        raise
    query = query.strip().strip('"').strip("'")
    logger.info("OpenAlex query: %s", query)
    status: dict = {}
    papers = search_openalex(query=query, logger=logger, status_out=status)
    logger.info("OpenAlex returned %d paper(s)", len(papers))
    return {'Open_Alex_paper': papers, 'open_alex_status': status}

def normalize_title(title: str) -> str:
    """
    Produce a normalized string used ONLY for matching titles across sources.
    Never stored on the paper record itself — purely a dedup key.
    """
    if not title:
        return ""

    # Normalize unicode (curly quotes, accented chars) to a comparable ASCII-ish form
    text = unicodedata.normalize("NFKD", title)
    text = "".join(c for c in text if not unicodedata.combining(c))

    text = text.lower()
    text = re.sub(r"[^\w\s]", "", text)   # strip punctuation (colons, hyphens, periods, etc.)
    text = re.sub(r"\s+", " ", text)      # collapse whitespace
    return text.strip()

def merge_paper_entry(existing: dict, new: dict) -> dict:
    """
    Merge a newly-encountered paper into an existing aggregated entry
    for the same normalized title. No information is dropped —
    source-varying fields (source, citation_count, url, pdf_url)
    are kept per-source in dicts.
    """
    new_source = new.get("source")

    # --- title: keep the longer/more complete one ---
    existing_title = existing.get("title") or ""
    new_title = new.get("title") or ""
    merged_title = new_title if len(new_title) > len(existing_title) else existing_title

    # --- authors: union, order-preserving ---
    merged_authors = list(existing.get("authors") or [])
    for author in (new.get("authors") or []):
        if author not in merged_authors:
            merged_authors.append(author)

    # --- abstract: keep the longer one ---
    existing_abstract = existing.get("abstract")
    new_abstract = new.get("abstract")
    if existing_abstract and new_abstract:
        merged_abstract = existing_abstract if len(existing_abstract) >= len(new_abstract) else new_abstract
    else:
        merged_abstract = existing_abstract or new_abstract

    # --- year: first non-null wins (should agree across sources) ---
    merged_year = existing.get("year") if existing.get("year") is not None else new.get("year")

    # --- source-varying fields: dict keyed by source, nothing overwritten ---
    merged_citation_count = dict(existing.get("citation_count") or {})
    merged_citation_count[new_source] = new.get("citation_count")

    merged_url = dict(existing.get("url") or {})
    merged_url[new_source] = new.get("url")

    merged_pdf_url = dict(existing.get("pdf_url") or {})
    merged_pdf_url[new_source] = new.get("pdf_url")

    merged_sources = list(existing.get("source") or [])
    if new_source not in merged_sources:
        merged_sources.append(new_source)

    return {
        "title": merged_title,
        "authors": merged_authors or None,
        "abstract": merged_abstract,
        "year": merged_year,
        "source": merged_sources,
        "citation_count": merged_citation_count,
        "url": merged_url,
        "pdf_url": merged_pdf_url,
    }

def aggregator_node(state: ResearchSearchState) -> ResearchSearchState:
    logger = get_node_logger(state["run_id"], "search.aggregate")
    arxiv_papers = state.get("arXiv_paper", []) or []
    openalex_papers = state.get("Open_Alex_paper", []) or []
    semantic_scholar_papers_list = state.get("Semantic_Scholar_paper", []) or []
    all_papers = arxiv_papers + openalex_papers + semantic_scholar_papers_list

    aggregated: dict[str, dict] = {}

    for paper in all_papers:
        key = normalize_title(paper.get("title"))
        if not key:
            continue

        if key in aggregated:
            aggregated[key] = merge_paper_entry(aggregated[key], paper)
        else:
            # first time seeing this title — wrap single-source fields
            # into the same dict-keyed-by-source shape used by merges,
            # so every entry has a UNIFORM shape whether it came from
            # 1 source or all 3
            source = paper.get("source")
            aggregated[key] = {
                "title": paper.get("title"),
                "authors": paper.get("authors"),
                "abstract": paper.get("abstract"),
                "year": paper.get("year"),
                "source": [source],
                "citation_count": {source: paper.get("citation_count")},
                "url": {source: paper.get("url")},
                "pdf_url": {source: paper.get("pdf_url")},
            }

    logger.info(
        "Aggregated %d raw records (arXiv=%d, OpenAlex=%d, SemanticScholar=%d) into %d unique papers after dedup",
        len(all_papers), len(arxiv_papers), len(openalex_papers), len(semantic_scholar_papers_list), len(aggregated),
    )
    return {"aggregator": aggregated}

def reranker_node(state: ResearchSearchState) -> ResearchSearchState:
    """
    Reranks the aggregated papers by semantic similarity to the Research Intent,
    keeps the top 15, and returns the FULL records (not just title:abstract)
    for those top 15 — in the same order the reranker ranked them.
    """
    logger = get_node_logger(state["run_id"], "search.rerank")
    aggregator = state["aggregator"]

    # --- Build the {normalized_title: abstract} view the reranker expects ---
    papers_for_rerank: dict[str, str] = {
        normalized_title: record.get("abstract")
        for normalized_title, record in aggregator.items()
    }
    logger.info("Reranking %d aggregated papers by semantic similarity to the intent", len(papers_for_rerank))
    # --- Run the rerank ---
    top_papers = rerank_by_relevance(
        research_intent=state["ResearchIntent"],
        papers=papers_for_rerank,
        top_n=15,
        logger=logger,
    )

    # --- Rebuild full records for only the titles that survived reranking,
    #     preserving the ranked order (top_papers is an ordered dict) ---
    reranked_full: dict[str, dict] = {}
    for normalized_title in top_papers:
        record = aggregator.get(normalized_title)
        if record is not None:
            reranked_full[normalized_title] = record

    dropped_count = len(aggregator) - len(reranked_full)
    logger.info(
        "Kept top %d of %d papers after semantic reranking (%d dropped)",
        len(reranked_full), len(aggregator), dropped_count,
    )

    return {"reranked_papers": reranked_full}

def build_clustered_output(
    clustering_result: ClusteringOutput,
    id_to_title: dict[int, str],
    aggregator: dict,
) -> list[dict]:
    """
    Converts {cluster_id, label, rationale, paper_ids} clusters into:
    [{"label": ..., "rationale": ..., "papers": {normalized_title: full_aggregator_record, ...}}, ...]
    """
    output = []

    for cluster in clustering_result.clusters:
        papers_in_cluster = {}

        for pid in cluster.paper_ids:
            normalized_title = id_to_title.get(pid)
            if normalized_title is None:
                continue  # LLM hallucinated an ID that was never sent — skip, don't crash

            record = aggregator.get(normalized_title)
            if record is None:
                continue  # defensive: shouldn't happen if id_to_title is built correctly

            papers_in_cluster[normalized_title] = record

        output.append({
            "label": cluster.label,
            "rationale": cluster.rationale,
            "papers": papers_in_cluster,
        })

    return output

def prepare_papers_for_clustering(aggregator: dict) -> tuple[list[dict], dict[int, str]]:
    """
    Convert the aggregator dict into the (id, title, abstract) form needed for the
    clustering prompt, plus an id -> normalized_title map to reverse-lookup results.

    Papers with no abstract are excluded from clustering input (title alone is too
    weak a signal for approach-based grouping) but returned separately so they can
    still be surfaced to the user as "not clustered" rather than silently dropped.
    """
    clustering_input = []
    id_to_title: dict[int, str] = {}
    excluded_no_abstract: list[str] = []

    paper_id = 0
    for normalized_title, record in aggregator.items():
        abstract = record.get("abstract")

        if not abstract or not abstract.strip():
            excluded_no_abstract.append(normalized_title)
            continue

        clustering_input.append({
            "id": paper_id,
            "title": record.get("title") or normalized_title,
            "abstract": abstract,
        })
        id_to_title[paper_id] = normalized_title
        paper_id += 1

    return clustering_input, id_to_title

def clustering_node(state: ResearchSearchState) -> ResearchSearchState:
    logger = get_node_logger(state["run_id"], "search.cluster")
    clustering_input , id_to_title = prepare_papers_for_clustering(state['reranked_papers'])
    logger.info("Clustering %d paper(s) with usable abstracts into approach groups", len(clustering_input))

    prompt = f"""You are helping a researcher make sense of a set of retrieved papers by organizing
them into methodologically distinct APPROACH CLUSTERS.

You will receive:
1. A RESEARCH INTENT (Problem / Objective / Additional Context) — this defines what
   the researcher is actually trying to figure out.
2. A numbered list of papers, each with a title and abstract.

YOUR GOAL:
All these papers were retrieved because they already share a common TOPIC with the
Research Intent — that overlap is not useful information on its own. Your job is to
find the more useful signal underneath: the DIFFERENT APPROACHES, METHODS, or
STRATEGIES these papers take toward addressing that shared topic. Think of each
cluster as answering the question: "If the researcher wanted to pursue THIS
direction, which papers would be their starting point?"

USE THE RESEARCH INTENT TO GUIDE WHAT COUNTS AS A MEANINGFUL DISTINCTION:
- The Objective tells you what decision or comparison the researcher actually cares
  about — let it sharpen which methodological differences matter versus which are
  noise. If the Objective is about comparing two paradigms, treat "which paradigm
  does this paper represent" as a primary clustering axis, not an incidental detail.
- The Additional Context may reveal known approaches the researcher is already aware
  of, or constraints (like avoiding full retraining) — clusters that map onto these
  should be called out explicitly, since they connect directly to something the
  researcher already flagged as relevant.

FOR EACH CLUSTER, PRODUCE:
- cluster_id: a sequential integer starting at 0, unique per cluster.
- label: a short, specific name for the actual approach (not a vague topic word) —
  e.g. "Reinforcement-learning-based car-following control policies," not "RL papers."
- rationale: 2-3 sentences explaining WHAT unifies these papers methodologically,
  and WHY this represents a distinct direction the researcher could pursue or
  compare against, relative to the other clusters.
- paper_ids: the list of paper IDs belonging to this cluster.

SIZING THE CLUSTERS:
- Aim for 4-7 clusters when the paper count comfortably supports that many distinct
  methodological groupings. This is a target, not a quota — with a small paper set,
  fewer, larger clusters are correct; do not manufacture extra clusters just to hit
  the range, and do not force a paper into a cluster it doesn't methodologically fit.
- If two clusters would end up nearly identical in meaning, merge them instead of
  keeping both.
- If some papers do not fit coherently into any clear methodological group, place
  them together in one final cluster labeled "Other / Mixed Approaches" rather than
  forcing a false fit elsewhere. One honest catch-all beats several padded clusters.

HARD CONSTRAINTS:
- Every paper ID in the input must appear in exactly one cluster, including the
  catch-all if used. None may be omitted or duplicated across clusters.
- Do not invent methods, results, or claims not supported by the given title/abstract
  text — ground every rationale only in what's actually stated.

RESEARCH INTENT:
{state["ResearchIntent"]}

PAPERS:
{clustering_input}
"""

    structured_llm = llm_2.with_structured_output(ClusteringOutput)

    try:
        clustering_result = structured_llm.invoke(prompt)
    except Exception as e:
        logger.exception("Structured LLM invocation failed during clustering: %s: %s", type(e).__name__, e)
        return {"clustered_papers": []}

    clustered = build_clustered_output(clustering_result, id_to_title, state["aggregator"])
    total_clustered_papers = sum(len(c["papers"]) for c in clustered)
    logger.info("Clustering produced %d cluster(s) covering %d paper(s)", len(clustered), total_clustered_papers)

    return {"clustered_papers": clustered}

# Building The Graph
builder = StateGraph(ResearchSearchState)

builder.add_node('arxiv',arXiv_papers)
builder.add_node('semantic_scholar',semantic_scholar_papers)
builder.add_node('open_alex',Open_alex_papers)
builder.add_node('aggregation',aggregator_node)
builder.add_node('reranker', reranker_node)
builder.add_node('clustering',clustering_node)

builder.add_edge(START,'arxiv')
builder.add_edge(START,'semantic_scholar')
builder.add_edge(START,'open_alex')
builder.add_edge('arxiv','aggregation')
builder.add_edge('semantic_scholar','aggregation')
builder.add_edge('open_alex','aggregation')
builder.add_edge('aggregation','reranker')
builder.add_edge('reranker','clustering')
builder.add_edge('clustering',END)

graph = builder.compile()

if __name__ == "__main__":
    # Manual standalone test — only runs when you execute this file directly
    # (`python -m app.modules.search.graph`), never on import.
    input_research_intent_5 = """Problem:
Large language models frequently produce fluent, confident-sounding answers that contain factual errors, and existing hallucination-detection methods disagree substantially with each other when applied to the same model outputs, making it unclear which detection approach should be trusted for deployment decisions.

Objective:
Identify hallucination-detection methods for large language models that show consistent, reliable performance across different model families and task types, rather than being effective only in the narrow setting they were originally evaluated on.

Additional Context:
- Focus on post-hoc detection methods applied to generated text, not training-time interventions like RLHF.
- Known approaches the researcher is aware of: self-consistency sampling, retrieval-based fact verification, and uncertainty/entropy-based methods.
- Interested primarily in open-domain question answering and summarization tasks, not code generation."""

    result = graph.invoke({"ResearchIntent": input_research_intent_5, "run_id": "standalone-test"})
    print(result["clustered_papers"])