core/pipeline_2/logic.py

"""Pipeline 2: Hybrid RAG — vector search + reranker + Groq LLM.

Architectural improvements over baseline:
  - Query decomposition for multi-hop questions (LLM-driven)
  - Multi-paper full-text fetch (up to 3 papers)
  - Fallback: if top papers lack full text, find the highest-scoring paper that does
  - Structure-aware sentence-level passage extraction with contextual prepend
  - Improved synthesis prompt for completeness and specificity
"""
import json
import logging
import os
import re
import sys
import time
from dataclasses import dataclass
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed

from dotenv import load_dotenv
from groq import Groq, RateLimitError as GroqRateLimitError

# Project root on path so services/ imports resolve
sys.path.append(str(Path(__file__).parents[2]))
# pipeline-2 dir on path so indexer/retriever imports resolve
sys.path.append(str(Path(__file__).parent))

from indexer import build_index
from retriever import HybridRetriever
from services.metrics_service import MetricsService
from services.paper_fetcher import PaperFetcher

load_dotenv()

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

_STORAGE_PATH = str(Path(__file__).parent / "qdrant_storage")
_DEFAULT_MODEL = "llama-3.3-70b-versatile"
_DECOMPOSE_MODEL = "llama-3.1-8b-instant"  # cheap classifier — no need for 70B


@dataclass
class _GroqUsage:
    """Maps Groq usage fields to the shape MetricsService expects."""
    prompt_token_count: int
    candidates_token_count: int


class PipelineRAG:
    """
    Pipeline 2: Hybrid RAG.
    Flow: query decomposition → hybrid search (RRF) → cross-encoder rerank
          → multi-paper full-text fetch → structure-aware chunking → Groq synthesis.

    Run the indexer once before starting the server:
        .venv/bin/python core/pipeline-2/indexer.py
    """

    def __init__(
        self,
        retrieval_top_k: int = 50,
        rerank_top_n: int = 10,
        max_full_text: int = 3,
    ):
        self.retrieval_top_k = retrieval_top_k
        self.rerank_top_n = rerank_top_n
        self.max_full_text = max_full_text

        self.model_name = os.environ.get("GROQ_MODEL", _DEFAULT_MODEL)
        _keys = [os.environ.get(f"GROQ_API_KEY_{i}") for i in range(1, 4)]
        _keys = [k for k in _keys if k]
        if not _keys:
            _keys = [os.environ["GROQ_API_KEY"]]
        self._groq_clients = [Groq(api_key=k) for k in _keys]
        self._groq_key_idx = 0
        logger.info(f"Loaded {len(self._groq_clients)} Groq API key(s).")
        self.fetcher = PaperFetcher()
        self.metrics = MetricsService(self.model_name)

        logger.info("Ensuring Qdrant index exists...")
        build_index(storage_path=_STORAGE_PATH)

        logger.info("Initializing retriever (loading models)...")
        self.retriever = HybridRetriever(storage_path=_STORAGE_PATH)

    # ------------------------------------------------------------------
    # Groq API call with key rotation on rate limit
    # ------------------------------------------------------------------
    def _groq_call(self, **kwargs):
        for _ in range(len(self._groq_clients)):
            try:
                return self._groq_clients[self._groq_key_idx].chat.completions.create(**kwargs)
            except GroqRateLimitError:
                next_idx = (self._groq_key_idx + 1) % len(self._groq_clients)
                if next_idx == self._groq_key_idx:
                    raise
                logger.warning(
                    f"Rate limit on Groq key {self._groq_key_idx + 1}, "
                    f"rotating to key {next_idx + 1}..."
                )
                self._groq_key_idx = next_idx
        raise GroqRateLimitError("All Groq API keys exhausted.")

    # ------------------------------------------------------------------
    # Query Decomposition
    # ------------------------------------------------------------------
    def _decompose_query(self, query: str) -> dict:
        """Use a lightweight LLM to detect true multi-paper questions and decompose them.

        MULTI-HOP = the question explicitly references 2+ DISTINCT papers/methods
        that must each be independently retrieved (e.g. "Both CGCNN and DPMD...").
        Multi-part questions about ONE paper are NOT multi-hop.

        Returns {"is_multi_hop": bool, "sub_queries": [str, ...]}.
        For single-hop questions, sub_queries contains only the original query.
        """
        try:
            response = self._groq_call(
                model=_DECOMPOSE_MODEL,
                messages=[
                    {
                        "role": "system",
                        "content": (
                            "You classify scientific search queries as single-hop or multi-hop.\n\n"
                            "MULTI-HOP = the question EXPLICITLY references 2+ DISTINCT papers or methods "
                            "that must each be independently retrieved. "
                            "Look for: 'Both X and Y', 'X paper and Y paper', 'X review and Y review', "
                            "or a comparison between two named systems from different papers.\n\n"
                            "SINGLE-HOP = anything else, INCLUDING multi-part questions about ONE topic.\n\n"
                            "SINGLE-HOP examples (is_multi_hop=false):\n"
                            "  'What three components does X have and what does each do?' → one paper\n"
                            "  'What does ADMET stand for and why does ADMETlab 2.0 address it?' → one paper\n"
                            "  'What makes DPMD first-principles based and what symmetries does it preserve?' → one paper\n"
                            "  'What two factors drove NLP progress and what is the library goal?' → one paper\n\n"
                            "MULTI-HOP examples (is_multi_hop=true, max 2 sub-queries):\n"
                            "  'Both CGCNN and DPMD apply neural networks... what does each predict?' → two papers\n"
                            "  'The HuggingFace paper and the molecular representations paper both...' → two papers\n"
                            "  'The ML for fluid mechanics review and the ML in materials science review...' → two papers\n\n"
                            "Return ONLY valid JSON, no markdown fences:\n"
                            '{"is_multi_hop": false, "sub_queries": ["original query"]}\n'
                            "or\n"
                            '{"is_multi_hop": true, "sub_queries": ["focused query about paper A", "focused query about paper B"]}'
                        ),
                    },
                    {"role": "user", "content": query},
                ],
                temperature=0.0,
                max_tokens=256,
            )
            raw = response.choices[0].message.content.strip()
            # Strip markdown code fences if the LLM wraps the JSON
            raw = re.sub(r"^```(?:json)?\s*", "", raw)
            raw = re.sub(r"\s*```$", "", raw)
            result = json.loads(raw)
            if not isinstance(result.get("sub_queries"), list) or not result["sub_queries"]:
                return {"is_multi_hop": False, "sub_queries": [query]}

            # Post-LLM sanity checks: override false multi-hop classifications
            if result.get("is_multi_hop") and len(result["sub_queries"]) >= 2:
                subs = result["sub_queries"]
                override = False

                # Check 1: both sub-queries share a capitalized word-pair → same paper
                def _cap_pairs(text):
                    tokens = text.split()
                    return {
                        (tokens[i].lower(), tokens[i + 1].lower())
                        for i in range(len(tokens) - 1)
                        if tokens[i][0].isupper() and tokens[i + 1][0].isupper()
                        and len(tokens[i]) > 1 and len(tokens[i + 1]) > 1
                    }
                if _cap_pairs(subs[0]) & _cap_pairs(subs[1]):
                    override = True

                # Check 2: second sub-query uses "it" as a pronoun and neither sub-query
                # mentions a specific named system (ALL-CAPS or CamelCase) → same topic
                if not override and re.search(r'\bit\b', subs[1], re.IGNORECASE):
                    has_named_system = re.search(
                        r'\b([A-Z]{2,}|[A-Z][a-z]+[A-Z][a-zA-Z]*)\b',
                        subs[0] + " " + subs[1]
                    )
                    if not has_named_system:
                        override = True

                if override:
                    logger.info("Multi-hop override: sub-queries target the same paper — single-hop.")
                    result = {"is_multi_hop": False, "sub_queries": [query]}

            logger.info(f"Query decomposition: multi_hop={result['is_multi_hop']}, "
                        f"sub_queries={result['sub_queries']}")
            return result
        except Exception as e:
            logger.warning(f"Query decomposition failed ({e}), using original query.")
            return {"is_multi_hop": False, "sub_queries": [query]}

    # ------------------------------------------------------------------
    # Structure-aware passage extraction
    # ------------------------------------------------------------------
    def _extract_relevant_passages(
        self, full_text: str, query: str, paper_title: str = "",
        top_k: int = 8, target_tokens: int = 600,
    ) -> str:
        """Sentence-aware passage extraction with contextual prepend.

        1. Split full text on sentence boundaries (preserving scientific structure).
        2. Group sentences into coherent passages (~3-5 sentences each).
        3. Embed passages and rank by cosine similarity to the query.
        4. Prepend paper title to each selected passage for context.
        """
        # Split on sentence boundaries — handles abbreviations reasonably
        sentences = re.split(r'(?<=[.!?])\s+(?=[A-Z])', full_text)
        if not sentences:
            return full_text

        # Group sentences into passages of roughly target_tokens characters
        passages: list[str] = []
        current: list[str] = []
        current_len = 0
        for sent in sentences:
            current.append(sent)
            current_len += len(sent)
            if current_len >= target_tokens:
                passages.append(" ".join(current))
                # Carry the last sentence over for overlap context
                current = [current[-1]] if len(current) > 1 else []
                current_len = len(current[0]) if current else 0
        if current:
            passages.append(" ".join(current))

        if len(passages) <= top_k:
            return full_text

        # Stage 1: bi-encoder to narrow to top-30 candidates (fast)
        texts = [query] + passages
        embeddings = self.retriever.dense.encode(
            texts, normalize_embeddings=True, show_progress_bar=False, batch_size=64
        )
        dense_scores = embeddings[1:] @ embeddings[0]
        stage1_k = min(30, len(passages))
        stage1_indices = dense_scores.argsort()[-stage1_k:][::-1].tolist()

        # Always include passage[0] (intro/abstract) — key summary statements live there
        mandatory = {0}
        stage1_pool = sorted(set(stage1_indices) | mandatory)

        # Stage 2: cross-encoder rerank for accurate passage selection
        pairs = [[query, passages[i]] for i in stage1_pool]
        rerank_scores = self.retriever.reranker.predict(pairs, batch_size=32)
        scored = sorted(zip(rerank_scores, stage1_pool), key=lambda x: x[0], reverse=True)

        # Take top_k-1 by reranker score + always keep passage[0]
        top_by_reranker = [idx for _, idx in scored if idx not in mandatory][: top_k - 1]
        top_indices = sorted(mandatory | set(top_by_reranker))

        # Contextual prepend: attach paper title to each passage
        prefix = f"[Source: {paper_title}]" if paper_title else ""
        selected = []
        for idx in top_indices:
            passage_text = passages[idx].strip()
            if prefix:
                selected.append(f"{prefix}\n{passage_text}")
            else:
                selected.append(passage_text)

        return "\n\n[...]\n\n".join(selected)

    # ------------------------------------------------------------------
    # Retrieval with dedup support for multi-hop
    # ------------------------------------------------------------------
    def _deduplicate_papers(self, papers: list[dict], top_n: int) -> list[dict]:
        """Deduplicate papers by paper_id, keeping the first (highest-ranked) occurrence."""
        seen = set()
        deduped = []
        for p in papers:
            pid = p.get("paper_id")
            if pid and pid not in seen:
                seen.add(pid)
                deduped.append(p)
            if len(deduped) >= top_n:
                break
        return deduped

    # ------------------------------------------------------------------
    # Multi-hop coverage enforcement
    # ------------------------------------------------------------------
    def _ensure_sub_query_coverage(
        self, top_papers: list[dict], sub_queries: list[str], top_n: int
    ) -> list[dict]:
        """Guarantee that each sub-query is represented by at least one paper.

        For each sub-query, score every paper in top_papers with the cross-encoder.
        If no paper scores above a minimal threshold, run a targeted search for that
        sub-query and inject the best hit at the end of the list.
        """
        # bge-reranker-base returns sigmoid-activated [0,1] scores; 0.1 means "weak relevance"
        COVERAGE_THRESHOLD = 0.1

        for sub_q in sub_queries:
            pairs = [[sub_q, f"{p.get('title', '')}\n{p.get('abstract', '')}"]
                     for p in top_papers]
            scores = self.retriever.reranker.predict(pairs)
            if float(max(scores)) > COVERAGE_THRESHOLD:
                continue  # sub-query already covered

            # No paper covers this sub-query — fetch a targeted result and append
            logger.info(f"Coverage gap detected for sub-query: {sub_q[:60]}... Fetching targeted paper.")
            extra = self.retriever.search(sub_q, top_k=5)
            reranked_extra = self.retriever.rerank(sub_q, extra, top_n=1)
            if reranked_extra:
                pid = reranked_extra[0].get("paper_id")
                if pid not in {p.get("paper_id") for p in top_papers}:
                    top_papers.append(reranked_extra[0])

        return self._deduplicate_papers(top_papers, top_n)

    # ------------------------------------------------------------------
    # Context assembly with multi-paper full-text + fallback
    # ------------------------------------------------------------------
    def _build_context(self, top_papers: list[dict], query: str, max_full_text: int = None) -> tuple[str, list[str], list[str]]:
        """Build the context string for LLM synthesis.

        Fetches full text for the top-N papers IN PARALLEL to avoid sequential
        30s-per-paper timeouts stacking up. Falls back to abstract when fetch fails.

        Returns (full_context, abstracts_list, full_text_paper_titles).
        """
        if max_full_text is None:
            max_full_text = self.max_full_text
        papers_to_fetch = top_papers[:max_full_text]
        remaining = top_papers[max_full_text:]

        def _fetch(args: tuple[int, dict]) -> tuple[int, dict, str | None]:
            i, paper = args
            title = paper.get("title", "")
            logger.info(f"Attempting full text fetch for: {title[:60]}")
            text = self.fetcher.fetch_full_text(paper["paper_id"])
            return i, paper, text

        # Parallel fetch — wall time is max(individual fetches) instead of sum
        fetch_results: dict[int, tuple[dict, str | None]] = {}
        with ThreadPoolExecutor(max_workers=self.max_full_text) as pool:
            for i, paper, text in pool.map(_fetch, enumerate(papers_to_fetch)):
                fetch_results[i] = (paper, text)

        context_parts: list[str] = []
        abstracts: list[str] = []
        full_text_titles: list[str] = []

        # Adaptive top_k: distribute passage budget evenly across however many
        # papers actually returned full text, capped at 8 for a single paper.
        full_text_count = sum(
            1 for i in range(len(papers_to_fetch))
            if fetch_results[i][1] and len(fetch_results[i][1]) > 1000
        )
        adaptive_top_k = max(6, 12 // max(1, full_text_count))

        for i in range(len(papers_to_fetch)):
            paper, full_text = fetch_results[i]
            abstract = paper.get("abstract", "")
            title = paper.get("title", "")
            abstracts.append(abstract)

            if full_text and len(full_text) > 1000:
                content = self._extract_relevant_passages(
                    full_text, query, paper_title=title, top_k=adaptive_top_k
                )
                context_parts.append(
                    f"--- PAPER {i} [{title}] (FULL TEXT, RELEVANT PASSAGES) ---\n{content}\n"
                )
                full_text_titles.append(title)
            else:
                context_parts.append(f"--- PAPER {i} [{title}] (ABSTRACT) ---\n{abstract}\n")

        # Remaining papers (beyond max_full_text) always use abstract
        for j, paper in enumerate(remaining, start=len(papers_to_fetch)):
            abstract = paper.get("abstract", "")
            title = paper.get("title", "")
            abstracts.append(abstract)
            context_parts.append(f"--- PAPER {j} [{title}] (ABSTRACT) ---\n{abstract}\n")

        # If no full text was fetched at all, scan remaining for any accessible paper
        if not full_text_titles:
            logger.info("No full text found in top papers — scanning remaining for fallback...")
            for paper in remaining:
                title = paper.get("title", "")
                full_text = self.fetcher.fetch_full_text(paper["paper_id"])
                if full_text and len(full_text) > 1000:
                    content = self._extract_relevant_passages(
                        full_text, query, paper_title=title, top_k=8
                    )
                    context_parts.append(
                        f"--- FALLBACK PAPER [{title}] (FULL TEXT, RELEVANT PASSAGES) ---\n{content}\n"
                    )
                    full_text_titles.append(title)
                    break

        full_context = "\n".join(context_parts)
        return full_context, abstracts, full_text_titles

    # ------------------------------------------------------------------
    # Main entry point
    # ------------------------------------------------------------------
    def run_stream(self, query: str, ground_truth: str = None):
        """Streaming version of Pipeline 2"""
        start = time.time()
        
        yield {"type": "status", "data": "ANALYZING QUERY STRUCTURE..."}
        decomposition = self._decompose_query(query)
        is_multi_hop = decomposition["is_multi_hop"]
        
        yield {"type": "status", "data": f"PLAN: multi_hop={is_multi_hop}"}
        
        yield {"type": "status", "data": "EXECUTING HYBRID VECTOR RETRIEVAL..."}
        if is_multi_hop:
            def _search_sub_query(sub_q: str) -> list[dict]:
                return self.retriever.search(sub_q, top_k=20)
            all_candidates = []
            with ThreadPoolExecutor(max_workers=len(decomposition["sub_queries"])) as pool:
                futures = [pool.submit(_search_sub_query, sq) for sq in decomposition["sub_queries"]]
                for fut in as_completed(futures): all_candidates.extend(fut.result())
            unique_candidates = self._deduplicate_papers(all_candidates, top_n=self.retrieval_top_k)
            
            yield {"type": "status", "data": f"RETRIEVED {len(unique_candidates)} CANDIDATES. RERANKING..."}
            top_papers = self.retriever.rerank(query, unique_candidates, top_n=self.rerank_top_n)
            top_papers = self._ensure_sub_query_coverage(top_papers, decomposition["sub_queries"], top_n=self.rerank_top_n)
        else:
            candidates = self.retriever.search(query, top_k=self.retrieval_top_k)
            yield {"type": "status", "data": f"RETRIEVED {len(candidates)} CANDIDATES. RERANKING..."}
            top_papers = self.retriever.rerank(query, candidates, top_n=self.rerank_top_n)

        sources = [p.get("title", "") for p in top_papers]
        yield {"type": "sources", "data": sources}
        
        yield {"type": "status", "data": "FETCHING FULL-TEXT & BUILDING CONTEXT..."}
        max_ft = self.max_full_text if is_multi_hop else min(2, self.max_full_text)
        full_context, abstracts, _ = self._build_context(top_papers, query, max_full_text=max_ft)

        yield {"type": "status", "data": "SYNTHESIZING RESPONSE VIA GROQ..."}

        system_msg = (
            "You are an expert AI research assistant specializing in scientific literature synthesis. "
            "Answer the query using ONLY the provided paper context. Speak with absolute certainty.\n\n"
            "Rules:\n"
            "1. Answer directly and immediately. Do NOT use any preamble. Start the response with the subject of the question.\n"
            "2. Quote technical terms EXACTLY as they appear.\n"
            "3. Be complete.\n"
            "4. No disclaimers or hedges.\n"
            "5. No square brackets or citations.\n"
            "6. Precise, highly technical language."
        )
        
        full_answer = ""
        try:
            stream = self._groq_call(
                model=self.model_name,
                messages=[{"role": "system", "content": system_msg}, {"role": "user", "content": f"QUERY: {query}\n\nCONTEXT:\n{full_context}"}],
                temperature=0.1,
                max_tokens=1024,
                stream=True
            )
            
            usage = None
            for chunk in stream:
                if chunk.choices and chunk.choices[0].delta.content:
                    token = chunk.choices[0].delta.content
                    full_answer += token
                    yield {"type": "token", "data": token}
                if chunk.x_groq and chunk.x_groq.usage:
                    usage = _GroqUsage(
                        prompt_token_count=chunk.x_groq.usage.prompt_tokens,
                        candidates_token_count=chunk.x_groq.usage.completion_tokens
                    )

            yield {"type": "status", "data": "GENERATION COMPLETE. CALCULATING BENCHMARK METRICS..."}
            stats = self.metrics.process_metrics(
                client=None, query=query, answer=full_answer, context=full_context,
                usage_metadata=usage, start_time=start, abstracts_list=abstracts, ground_truth=ground_truth,
                model_name=self.model_name
            )
            yield {"type": "metrics", "data": stats}
            yield {"type": "status", "data": "PIPELINE EXECUTION SUCCESSFUL."}
            
        except Exception as e:
            logger.error(f"Pipeline 2 stream error: {e}")
            yield {"type": "error", "data": str(e)}

    def run(self, query: str, ground_truth: str = None) -> dict:
        start = time.time()
        logger.info(f"Pipeline 2 query: {query}")

        # Step 1: Query decomposition
        decomposition = self._decompose_query(query)

        # Step 2: Retrieval + Reranking
        if decomposition["is_multi_hop"]:
            # Parallel sub-query retrieval: each sub-query targets one required paper
            def _search_sub_query(sub_q: str) -> list[dict]:
                return self.retriever.search(sub_q, top_k=20)

            all_candidates: list[dict] = []
            with ThreadPoolExecutor(max_workers=len(decomposition["sub_queries"])) as pool:
                futures = [pool.submit(_search_sub_query, sq)
                           for sq in decomposition["sub_queries"]]
                for fut in as_completed(futures):
                    all_candidates.extend(fut.result())

            # Deduplicate then do ONE final rerank against the ORIGINAL question (Fix 3)
            unique_candidates = self._deduplicate_papers(all_candidates, top_n=self.retrieval_top_k)
            top_papers = self.retriever.rerank(query, unique_candidates, top_n=self.rerank_top_n)

            # Fix 4: Mandatory coverage check — ensure each sub-query has a representative paper
            top_papers = self._ensure_sub_query_coverage(
                top_papers, decomposition["sub_queries"], top_n=self.rerank_top_n
            )
        else:
            candidates = self.retriever.search(query, top_k=self.retrieval_top_k)
            top_papers = self.retriever.rerank(query, candidates, top_n=self.rerank_top_n)

        # Step 3: Context assembly — single-hop gets 2 full-text papers (less noise),
        # multi-hop gets 3 (needs evidence from multiple sources)
        max_ft = self.max_full_text if decomposition["is_multi_hop"] else min(2, self.max_full_text)
        full_context, abstracts, full_text_titles = self._build_context(top_papers, query, max_full_text=max_ft)

        # Step 4: LLM Synthesis
        system_msg = (
            "You are an expert AI research assistant specializing in scientific literature synthesis. "
            "Answer the query using ONLY the provided paper context. Speak with absolute certainty.\n\n"
            "Rules:\n"
            "1. Answer directly and immediately. Do NOT use any preamble like 'Based on the provided papers', 'According to the context', or 'It can be inferred'. Start the response with the subject of the question.\n"
            "2. Quote technical terms EXACTLY as they appear in the papers. "
            "For any list of methods, data sources, properties, or named items: copy the exact words. "
            "Do NOT substitute synonyms or rephrase — "
            "e.g. 'large-scale simulations' must not become 'multi-scale simulations'; "
            "'field measurements' must not become 'real data from the field'.\n"
            "3. Be complete — address EVERY part of the question. For list questions, include "
            "every item. For comparison questions, cover EACH paper explicitly.\n"
            "4. Never add disclaimers, caveats, hedges, or uncertainty statements. "
            "Do not state that information is missing or unclear. If a detail is missing, omit it silently.\n"
            "5. Do not use square brackets or cite PAPER numbers in the final text. "
            "6. Prefer concise, precise, and highly technical language."
        )
        user_msg = f"QUERY: {query}\n\nCONTEXT:\n{full_context}"

        try:
            response = self._groq_call(
                model=self.model_name,
                messages=[
                    {"role": "system", "content": system_msg},
                    {"role": "user", "content": user_msg},
                ],
                temperature=0.1,
                max_tokens=1024,
            )

            answer = response.choices[0].message.content
            usage = _GroqUsage(
                prompt_token_count=response.usage.prompt_tokens,
                candidates_token_count=response.usage.completion_tokens,
            )

            stats = self.metrics.process_metrics(
                client=None,
                query=query,
                answer=answer,
                context=full_context,
                usage_metadata=usage,
                start_time=start,
                abstracts_list=abstracts,
                ground_truth=ground_truth,
                model_name=self.model_name
            )

            return {
                "answer": answer,
                "sources": [p.get("title", "") for p in top_papers],
                "metrics": stats,
                "selected_papers": full_text_titles,
                "query_decomposition": decomposition,
            }

        except Exception as e:
            logger.error(f"Pipeline 2 error: {e}")
            return {"error": str(e), "answer": "Error generating response."}