Modular RAG MCP Server

import json
import logging
import re
from pathlib import Path
from typing import Any, Dict, List, Optional

from src.libs.llm.base_llm import BaseLLM
from src.libs.reranker.base_reranker import BaseReranker

logger = logging.getLogger(__name__)


class LLMReranker(BaseReranker):
    """
    Reranker implementation that uses an LLM to reorder candidates.
    It reads a prompt template from a file, injects the query and candidates,
    and parses the LLM's response to determine the new order.
    """

    def __init__(self, llm: BaseLLM, prompt_path: str = "config/prompts/rerank.txt"):
        """
        Initialize the LLM Reranker.

        Args:
            llm: An instance of BaseLLM to perform the reranking.
            prompt_path: Path to the prompt template file.
        """
        self.llm = llm
        self.prompt_path = prompt_path
        self._prompt_template = self._load_prompt_template()

    def _load_prompt_template(self) -> str:
        """Load prompt template from file or return default fallback."""
        try:
            return Path(self.prompt_path).read_text(encoding="utf-8")
        except Exception as e:
            logger.warning(
                f"Failed to load rerank prompt from {self.prompt_path}: {e}. Using default."
            )
            return (
                "Rank the following passages based on relevance to the query: {query}.\n"
                "Passages:\n{candidates}\n"
                "Return a JSON list of indices of the top passages, e.g. [0, 2, 1]."
            )

    def rerank(
        self,
        query: str,
        candidates: List[Any],
        top_k: Optional[int] = None,
        trace: Optional[Any] = None,
    ) -> List[Any]:
        """
        Rerank candidates using LLM.

        Args:
            query: The search query.
            candidates: List of candidate objects.
            top_k: Number of results to return.
            trace: Trace context (unused in this implementation, passed for interface compatibility).

        Returns:
            Reordered list of candidates.
        """
        if not candidates:
            return []

        # Format candidates
        # Try to extract text content: 'page_content', 'text', or str()
        candidate_texts = []
        for i, cand in enumerate(candidates):
            # Safe text extraction
            text = ""
            if hasattr(cand, "page_content"):
                text = cand.page_content
            elif hasattr(cand, "text"):
                text = cand.text
            elif isinstance(cand, dict) and "text" in cand:
                text = cand["text"]
            else:
                text = str(cand)

            # Truncate text to avoid excessive token usage (heuristic)
            text = text[:300].replace("\n", " ")
            candidate_texts.append(f"[{i}] {text}")

        candidates_str = "\n".join(candidate_texts)

        # Construct prompt
        try:
            prompt = self._prompt_template.format(
                query=query, candidates=candidates_str
            )
        except KeyError as e:
            logger.error(
                f"Prompt template format error: {e}. Missing keys in template."
            )
            # Fallback to simple concatenation if template is broken
            prompt = f"Query: {query}\nCandidates:\n{candidates_str}\nRank them as JSON list of indices."

        try:
            # Call LLM
            messages = [{"role": "user", "content": prompt}]
            response = self.llm.chat(messages)

            # Parse response
            ranked_indices = self._parse_indices(response)

            if not ranked_indices:
                logger.warning(
                    "LLM returned no valid indices. Returning original order."
                )
                if top_k:
                    return candidates[:top_k]
                return candidates

            # Reorder
            # Filter indices to ensure they are valid
            valid_indices = [i for i in ranked_indices if 0 <= i < len(candidates)]

            # Create result list
            ranked_results = [candidates[i] for i in valid_indices]

            # Append missing candidates at the end (if any) to preserve recall
            seen_indices = set(valid_indices)
            for i in range(len(candidates)):
                if i not in seen_indices:
                    ranked_results.append(candidates[i])

            if top_k:
                ranked_results = ranked_results[:top_k]

            return ranked_results

        except Exception as e:
            logger.error(f"LLM reranking failed: {e}")
            # Fallback: return original order
            if top_k:
                return candidates[:top_k]
            return candidates

    def _parse_indices(self, response: str) -> List[int]:
        """Parse JSON list of indices from LLM response."""
        try:
            # Clean up response (remove markdown code blocks if present)
            clean_response = response.strip()
            if clean_response.startswith("```"):
                clean_response = clean_response.split("```")[1]
                if clean_response.startswith("json"):
                    clean_response = clean_response[4:]

            # Try to find something that looks like a JSON list
            match = re.search(r"\[[\d,\s]+\]", clean_response)
            if match:
                return json.loads(match.group(0))
        except Exception as e:
            logger.debug(
                f"Failed to parse indices from response: {response[:100]}... Error: {e}"
            )
        return []