answer.py

"""
Takes the top chunks from retriever, passes them to LLM along with query and conversation history, it returns answer with citations
"""

import os
from dotenv import load_dotenv
from google import genai
from config import LLM_MODEL
from memory import format_history_for_prompt
from logger import get_logger

load_dotenv()
logger = get_logger("answer")

llm = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))

def answer(query:str, chunks:list, history:dict=None) -> str:
    """
    Builds a prompt from top chunks + conversation memory and asks llm to answer

    Args:
        query: the user question
        chunks: top ranked points from retrieve()
        history: memory dict, pass None or new_memory() for first time
    """

    if not chunks:
        return "I couldn't find relevant information in the stored papers."
    
    try:
        # build context for chunks
        context_parts = []
        for i, chunk in enumerate(chunks):
            p = chunk.payload
            context_parts.append(
                f"[Chunk No. {i+1} | Paper: {p["paper_name"]} | Page: {p["page_number"]} | Section: {p["section"]}]\n\n{p["text"]}"
            )
        context = "\n\n---\n\n".join(context_parts)


        history_block = ""
        if history and (history.get("summary") or history.get("recent")):
            history_text = format_history_for_prompt(history)
            history_block = f"""Previous conversation (for context only — do NOT use it to answer if the chunks don't support it):
            {history_text}"""



        prompt = f"""You are a research assistant. Answer the question using ONLY the provided context chunks.
        For each piece of information you use for answer, cite the source at end like: (Paper: paper_name, Page: X, Section: Y)
        If the answer is not in context, say so clearly.
        {history_block}
        Context:
        {context}

        Question: {query}

        Answer:"""

        response = llm.models.generate_content(
            model=LLM_MODEL,
            contents=prompt
        )

        logger.info("answer: generated response for query '%s'", query)
        return response.text.strip()
    
    except Exception as e:
        logger.error("answer: LLM call failed for query '%s': %s", query, e)
        return "Sorry, I couldn't generate an answer right now. Please try again."