backend/rag/multi_generator.py

import json
from groq import Groq
from backend.config import GROQ_API_KEY, GROQ_MODEL, GROQ_TIMEOUT
from backend.rag.generator import GeneratedAnswer, _build_citations, _get_groq_client
from backend.rag.multi_retriever import MultiSourceResult
from backend.rag.retriever import RetrievedChunk
from backend.core.schemas import UnifiedChunkMetadata
from backend.database.connection import get_connection

def _fetch_unified_metadata(chunk_id: str) -> dict:
    """Fetch unified_metadata for a specific chunk from the database."""
    try:
        with get_connection() as conn:
            cursor = conn.cursor(dictionary=True)
            cursor.execute("SELECT unified_metadata FROM chunks WHERE id = %s", (chunk_id,))
            row = cursor.fetchone()
            if row and row['unified_metadata']:
                # If it's already a dict (mysql-connector does this for JSON), return it
                # Otherwise parse it.
                meta = row['unified_metadata']
                if isinstance(meta, str):
                    return json.loads(meta)
                return meta
    except Exception as e:
        print(f"[MultiGenerator] Error fetching metadata for {chunk_id}: {e}")
    return {}

def _format_chunk_for_prompt(rank: int, chunk: RetrievedChunk) -> str:
    # Try to get unified metadata for better citation
    meta_dict = _fetch_unified_metadata(chunk.chunk_id)
    
    # Extract specific fields as requested
    section_id = meta_dict.get("section_id")
    case_name = meta_dict.get("case_name")
    para_range = meta_dict.get("para_range")
    
    ref_detail = ""
    if section_id:
        ref_detail = f"Section {section_id}"
    elif case_name:
        ref_detail = case_name
    elif para_range:
        ref_detail = f"para {para_range}"
    else:
        # Fallback to page number or timestamp if available
        if chunk.page_number:
            ref_detail = f"page {chunk.page_number}"
        elif chunk.timestamp_s is not None:
            ref_detail = f"at {chunk.timestamp_s}s"

    header = f"[{rank}] {chunk.source_title}"
    if ref_detail:
        header += f" — {ref_detail}"
        
    return f"{header}\n{chunk.chunk_text}"

def build_single_source_prompt(question: str, result: MultiSourceResult, history: list[dict] | None, image_context: str | None = None, is_legal: bool = False) -> list[dict]:
    if is_legal:
        system_prompt = """You are a legal information assistant for Indian law. Answer using ONLY the provided context.
        Structure your answer as:
        ANSWER: [clear explanation]
        LEGAL BASIS: [exact quote from source]
        CITATIONS: [numbered list: Document | Section/Para | Court | Date]
        AMENDMENTS: [any amendments to cited sections]
        Never give legal advice. State only what the law says."""
    else:
        system_prompt = """You are an expert research assistant. Answer accurately using ONLY the provided context.
        Structure your answer as:
        ANSWER: [clear, detailed explanation with inline citations like [Source 1]]
        KEY CONCEPTS: [list the main concepts from the source relevant to the question]
        CITATIONS: [numbered list: Document | Page/Section]
        Quote directly from the source when relevant. If the context doesn't contain the answer, say so clearly."""

    context_parts = []
    for i, chunk in enumerate(result.all_chunks, start=1):
        context_parts.append(_format_chunk_for_prompt(i, chunk))
    
    context_block = "\n\n".join(context_parts)
    
    if image_context:
        system_prompt = f"{system_prompt}\n\n{image_context}"

    messages = [
        {"role": "system", "content": f"{system_prompt}\n\nRETRIEVED CONTEXT:\n{context_block}"}
    ]
    
    if history:
        # CRITICAL: Only inject user/assistant roles — system messages break Groq API
        valid_history = [m for m in history if m.get('role') in ('user', 'assistant') and m.get('content')]
        messages.extend(valid_history[-12:])  # Last 6 turns = 12 messages
        
    messages.append({"role": "user", "content": question})
    return messages

def build_comparison_prompt(question: str, result: MultiSourceResult, history: list[dict] | None, image_context: str | None = None, is_legal: bool = False) -> list[dict]:
    if is_legal:
        system_prompt = """You are a legal analyst. Compare the provided sources objectively.
        Structure your answer as:
        QUERY: [restate what is being compared]
        SOURCE A — {first_source}:
        [what source A says, with exact quote]
        SOURCE B — {second_source}:
        [what source B says, with exact quote]
        KEY DIFFERENCES:
        [bullet points of substantive differences]
        KEY SIMILARITIES:
        [bullet points of shared principles]
        CITATIONS: [numbered, one per claim]
        Do not take sides. Report what each source states."""
    else:
        system_prompt = """You are a research analyst. Compare the provided sources objectively and thoroughly.
        Structure your answer EXACTLY as:

        ## Comparison Overview
        [1-2 sentence summary of what is being compared]

        ## {first_source}
        [Key points and explanation from this source, with page references]

        ## {second_source}
        [Key points and explanation from this source, with page references]

        ## Key Differences
        [Bullet points of the most important differences between the sources]

        ## Key Similarities
        [Bullet points of shared concepts or principles]

        ## Citations
        [Numbered list, one per factual claim]

        Be specific. Quote directly from sources. Do not introduce outside knowledge."""

    context_parts = []
    for title, chunks in result.source_groups.items():
        group_context = "\n".join([_format_chunk_for_prompt(i+1, c) for i, c in enumerate(chunks)])
        context_parts.append(f"=== SOURCE: {title} ===\n{group_context}")
    
    context_block = "\n\n".join(context_parts)
    
    # Fill placeholders in system prompt
    titles = list(result.source_groups.keys())
    s_prompt = system_prompt
    if len(titles) >= 2:
        s_prompt = s_prompt.replace("{first_source}", titles[0]).replace("{second_source}", titles[1])
    elif len(titles) == 1:
        s_prompt = s_prompt.replace("{first_source}", titles[0]).replace("{second_source}", "")
    # Also handle old-style placeholders
    if len(titles) >= 2:
        s_prompt = s_prompt.replace("{first source title}", titles[0]).replace("{second source title}", titles[1])
    
    if image_context:
        s_prompt = f"{s_prompt}\n\n{image_context}"

    messages = [
        {"role": "system", "content": f"{s_prompt}\n\nRETRIEVED CONTEXT:\n{context_block}"}
    ]
    
    if history:
        valid_history = [m for m in history if m.get('role') in ('user', 'assistant') and m.get('content')]
        messages.extend(valid_history[-12:])  # Last 6 turns
        
    messages.append({"role": "user", "content": question})
    return messages

def build_synthesis_prompt(question: str, result: MultiSourceResult, history: list[dict] | None, image_context: str | None = None, is_legal: bool = False) -> list[dict]:
    if is_legal:
        system_prompt = """You are a legal research synthesizer for Indian law. Consolidate and summarize information from MULTIPLE sources.
        Structure your answer as:
        CONSOLIDATED LEGAL VIEW: [comprehensive answer weaving sources]
        BY STATUTE/CASE: [for each source, key points and citations]
        RULES: Cite every claim. Quote exactly when mentioning sections. Do not give legal advice."""
    else:
        system_prompt = """You are a research synthesizer. Your job is to consolidate and summarize information from MULTIPLE sources.

Structure your answer EXACTLY as:
## Consolidated Answer
[A comprehensive 2-3 paragraph answer that weaves information from ALL sources]

## By Source
[For each source: source name in bold, then 2-3 key points from that source with page/section references]

## Common Themes
[Bullet list of themes found across sources]

## Key Differences
[Where sources differ or contradict, if any]

## Citations
[Numbered list with source title and reference for each claim]

RULES: Cite EVERY claim. Reference specific pages/sections. Do NOT guess."""

    context_parts = []
    for title, chunks in result.source_groups.items():
        # Use top 4 chunks per source for better coverage
        group_context = "\n".join([_format_chunk_for_prompt(i+1, c) for i, c in enumerate(chunks[:4])])
        context_parts.append(f"=== SOURCE: {title} ===\n{group_context}")
        
    context_block = "\n\n".join(context_parts)
    
    if image_context:
        system_prompt = f"{system_prompt}\n\n{image_context}"

    messages = [
        {"role": "system", "content": f"{system_prompt}\n\nRETRIEVED CONTEXT:\n{context_block}"}
    ]
    
    if history:
        valid_history = [m for m in history if m.get('role') in ('user', 'assistant') and m.get('content')]
        messages.extend(valid_history[-12:])
        
    messages.append({"role": "user", "content": question})
    return messages

def generate_multi_answer(question: str, result: MultiSourceResult, history: list[dict] | None = None, image_context: str | None = None, is_legal: bool = False) -> GeneratedAnswer:
    if not result.all_chunks:
        return GeneratedAnswer(
            answer="I searched your knowledge base but found no relevant information. \n This usually means: (1) no documents have been ingested yet, \n (2) your question doesn't match any uploaded content, or \n (3) the FAISS index is empty. \n Please upload a PDF or website first, then try again.",
            citations=[],
            chunks=[]
        )

    # 1. Select builder
    if result.query_intent == "comparison":
        messages = build_comparison_prompt(question, result, history, image_context, is_legal=is_legal)
    elif result.query_intent == "synthesis":
        messages = build_synthesis_prompt(question, result, history, image_context, is_legal=is_legal)
    else:
        messages = build_single_source_prompt(question, result, history, image_context, is_legal=is_legal)
        
    # 2. Call Groq
    try:
        client = _get_groq_client()
        response = client.chat.completions.create(
            model=GROQ_MODEL,
            messages=messages,
            stream=False, # Instruction said collect tokens, but typically we return GeneratedAnswer non-streaming here
            timeout=GROQ_TIMEOUT
        )
        answer = response.choices[0].message.content
    except Exception as e:
        print(f"[MultiGenerator] Groq Error: {e}")
        answer = f"Error generating answer: {e}"

    # 3. Build citations
    citations = _build_citations(result.all_chunks)
    
    return GeneratedAnswer(
        answer=answer,
        citations=citations,
        chunks=result.all_chunks
    )