chat.py

"""
Core RAG chat service — pure turn-by-turn answer logic.

answer_one_turn() takes a question and returns a structured AnswerResult so
the CLI (or any other frontend) can display results however it chooses.
This separation makes the logic testable without needing stdin/stdout.
"""
import time
import logging
from dataclasses import dataclass, field

from langchain_core.documents import Document
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

from config.runtime import (
    get_retriever_k,
    get_confidence_threshold,
    get_image_top_n,
)
from config.settings import RAW_IMAGE_PATH
from retrieval.query import (
    rewrite_query_for_retrieval,
    is_image_intent,
    is_procedural_question,
)
from retrieval.fusion import (
    merge_retrieved_docs,
    prioritize_non_image_docs,
    search_non_image_context,
    rank_text_docs,
    apply_source_focus,
)
from retrieval.confidence import (
    build_vector_score_map,
    compute_text_confidence,
    compute_image_confidence,
)
from retrieval.images import (
    create_image_candidate_doc,
    apply_image_source_focus,
    get_image_candidates,
    score_image_path,
    rank_image_hits,
    open_images_in_browser,
)
from retrieval.extractive import build_extractive_answer
from retrieval.source_summary import (
    collect_source_summary,
    format_source_summary,
    should_show_sources,
)

logger = logging.getLogger(__name__)

_ANSWER_TEMPLATE = """Answer the question based on the documents below. If the answer is not found in the documents, say "I cannot find relevant information in the documents."

Documents:
{context}

Question: {question}

Answer:"""

_LLM_CANNOT_ANSWER = "I cannot find relevant information in the documents."


# ---------------------------------------------------------------------------
# Result dataclass
# ---------------------------------------------------------------------------

@dataclass
class AnswerResult:
    question: str
    response: str = ""
    image_hits: list[str] = field(default_factory=list)
    retrieved_docs: list[Document] = field(default_factory=list)
    confidence: float = 0.0
    fallback_mode: str = ""          # "llm" | "extractive" | "snippet" | "image" | "no_docs"
    provider: str = ""
    elapsed: float = 0.0
    is_image_query: bool = False


# ---------------------------------------------------------------------------
# Document formatting helper
# ---------------------------------------------------------------------------

def _format_docs(docs: list[Document]) -> str:
    if not docs:
        return ""
    formatted = []
    for i, doc in enumerate(docs, 1):
        source = doc.metadata.get("source", "unknown")
        page = doc.metadata.get("page")
        page_label = f", page={page}" if page is not None else ""
        formatted.append(f"[{i}] source={source}{page_label}\n{doc.page_content}")
    return "\n\n".join(formatted)


# ---------------------------------------------------------------------------
# Vector results
# ---------------------------------------------------------------------------

def _get_vector_results(vectorstore, question: str, limit: int):
    try:
        return vectorstore.similarity_search_with_relevance_scores(question, k=limit)
    except Exception:
        docs = vectorstore.similarity_search(question, k=limit)
        return [(doc, 0.0) for doc in docs]


# ---------------------------------------------------------------------------
# Core turn logic
# ---------------------------------------------------------------------------

def answer_one_turn(
    vectorstore,
    keyword_retriever,
    llm,
    llm_provider: str,
    question: str,
) -> AnswerResult:
    """
    Execute one question → answer turn.

    Returns an AnswerResult with all retrieval details so the caller can log,
    display, or test the outcome without parsing printed text.
    """
    top_k = get_retriever_k()
    confidence_threshold = get_confidence_threshold()
    started = time.time()

    result = AnswerResult(
        question=question,
        provider=llm_provider,
    )

    # ------------------------------------------------------------------
    # Retrieval
    # ------------------------------------------------------------------
    query = rewrite_query_for_retrieval(question)
    result.is_image_query = is_image_intent(question)

    vector_results = _get_vector_results(vectorstore, query, top_k)
    vector_docs = [doc for doc, _ in vector_results]
    vector_score_map = build_vector_score_map(vector_results)
    keyword_docs = keyword_retriever.invoke(query) if keyword_retriever else []

    if result.is_image_query:
        candidate_docs = [
            create_image_candidate_doc(path)
            for path in get_image_candidates(query, limit=max(top_k, 6))
        ]
        image_docs = [d for d in (keyword_docs or []) if d.metadata.get("source_type") == "image"]
        other_kw_docs = [d for d in (keyword_docs or []) if d.metadata.get("source_type") != "image"]
        docs = merge_retrieved_docs(
            candidate_docs + image_docs + other_kw_docs,
            vector_docs,
            max(top_k, 8),
        )
        docs = apply_image_source_focus(query, docs)
    else:
        docs = merge_retrieved_docs(vector_docs, keyword_docs, top_k)
        docs = prioritize_non_image_docs(docs, top_k)
        # If all returned docs are images for a text query, go wider
        if docs and all(d.metadata.get("source_type") == "image" for d in docs):
            expanded = vectorstore.similarity_search(query, k=max(12, top_k * 4))
            expanded_non_image = [d for d in expanded if d.metadata.get("source_type") != "image"]
            if expanded_non_image:
                docs = expanded_non_image[:top_k]
        if not docs or all(d.metadata.get("source_type") == "image" for d in docs):
            forced_non_image = search_non_image_context(vectorstore, query, top_k)
            if forced_non_image:
                docs = forced_non_image
        docs = rank_text_docs(query, docs)
        docs = apply_source_focus(vectorstore, query, docs, top_k)

    result.retrieved_docs = docs

    # ------------------------------------------------------------------
    # No docs
    # ------------------------------------------------------------------
    if not docs:
        result.response = _LLM_CANNOT_ANSWER
        result.fallback_mode = "no_docs"
        result.elapsed = time.time() - started
        return result

    # ------------------------------------------------------------------
    # Source summary logging
    # ------------------------------------------------------------------
    sources = collect_source_summary(docs)
    if should_show_sources():
        print(f"[INFO] Sources matched: {format_source_summary(sources)}")

    # ------------------------------------------------------------------
    # Image query path
    # ------------------------------------------------------------------
    if result.is_image_query:
        image_hits = []
        for doc in docs:
            src = doc.metadata.get("source", "unknown")
            if doc.metadata.get("source_type") == "image":
                img_path = doc.metadata.get("image_path") or f"{RAW_IMAGE_PATH}/{src}"
                if img_path not in image_hits:
                    image_hits.append(img_path)

        if image_hits:
            image_hits = rank_image_hits(query, image_hits)
            positive_hits = [p for p in image_hits if score_image_path(query, p) > 0]
            if positive_hits:
                image_hits = positive_hits

            image_confidence = compute_image_confidence(query, image_hits)
            result.confidence = image_confidence
            print(f"[INFO] Image confidence: {image_confidence:.2f}")

            if image_confidence < confidence_threshold:
                result.response = "No image found with sufficient confidence."
                result.fallback_mode = "image_low_confidence"
                result.elapsed = time.time() - started
                print(f"[WARN] Image confidence below threshold {confidence_threshold:.2f}")
                return result

            image_hits = image_hits[:get_image_top_n()]
            result.image_hits = image_hits
            result.fallback_mode = "image"
            result.response = "I found the following image(s):\n" + "\n".join(
                f"[IMAGE {i}] {p}" for i, p in enumerate(image_hits, 1)
            )
            result.elapsed = time.time() - started
            open_images_in_browser(image_hits)
            return result

    # ------------------------------------------------------------------
    # Text query path — confidence gate
    # ------------------------------------------------------------------
    text_confidence = compute_text_confidence(query, docs, vector_score_map)
    result.confidence = text_confidence
    print(f"[INFO] Retrieval confidence: {text_confidence:.2f}")

    if text_confidence < confidence_threshold:
        snippets = [" ".join(d.page_content.split())[:220] for d in docs[:2]]
        result.response = "; ".join(snippets) if snippets else "I cannot find sufficiently confident information in the documents."
        result.fallback_mode = "snippet"
        result.elapsed = time.time() - started
        print(f"[WARN] Retrieval confidence below threshold {confidence_threshold:.2f}")
        return result

    # ------------------------------------------------------------------
    # Procedural extractive path
    # ------------------------------------------------------------------
    if is_procedural_question(question):
        extractive = build_extractive_answer(question, docs)
        if extractive:
            result.response = extractive
            result.fallback_mode = "extractive"
            result.elapsed = time.time() - started
            return result

    # ------------------------------------------------------------------
    # LLM generative path (streaming)
    # ------------------------------------------------------------------
    prompt = ChatPromptTemplate.from_template(_ANSWER_TEMPLATE)
    answer_chain = prompt | llm | StrOutputParser()
    context = _format_docs(docs)

    response_parts: list[str] = []
    print("\nAnswer: ", end="", flush=True)
    try:
        for chunk in answer_chain.stream({"context": context, "question": question}):
            from io.voice import to_console_text
            print(to_console_text(chunk), end="", flush=True)
            response_parts.append(chunk)
    except Exception as e:
        logger.error("LLM generation failed: %s", e)
        print(f"\n[ERROR] Failed to generate answer: {e}")
        raise

    response = "".join(response_parts)

    # If LLM says it doesn't know but we have docs, try extractive fallback
    if (not response.strip()) or (_LLM_CANNOT_ANSWER in response and docs):
        extractive = build_extractive_answer(question, docs) if is_procedural_question(question) else ""
        if extractive:
            print(f"\n[INFO] Switched to extractive answer:\n{extractive}", end="")
            result.response = extractive
            result.fallback_mode = "extractive"
        else:
            snippets = [" ".join(d.page_content.split())[:220] for d in docs[:2]]
            fallback = "; ".join(snippets) if snippets else _LLM_CANNOT_ANSWER
            print(f"\n[INFO] Switched to document snippet: {fallback}", end="")
            result.response = fallback
            result.fallback_mode = "snippet"
    else:
        result.response = response
        result.fallback_mode = "llm"

    result.elapsed = time.time() - started
    return result