Blame for metadata_search/chat_api.py

"""Chat API for the WaterScope research agent (HF Space answers)."""

from __future__ import annotations

import os
import logging

_logger = logging.getLogger(__name__)
from backend.chat_settings import MAX_CHAT_PAPERS, enhance_ordinal_question, focus_document_ids_for_question
from backend.mcp_sync import use_mcp
from backend.model import llm_backend_name
from backend.paper_voice import prefix_context, sanitize_answer
from backend.waterscope_agent import agent_answer
from metadata_search.mcp_retrieval import build_context_for_documents, list_papers_catalog, list_saved_papers

def _generate_answer(
    message: str,
    context: str,
    *,
    application: str | None = None,
    model: str | None = None,
    general_chat: bool = False,
    max_new_tokens: int | None = None,
) -> tuple[str, str, str]:
    from backend.research_fallback import (
        clear_llama_circuit, deepseek_research_answer, is_infrastructure_error,
        is_gradio_transport_json_error,
        llama_circuit_open, llama_space_paused, mark_llama_unavailable,
    )
    from backend.waterscope_agent import infer_application
    requested_app = application or os.getenv("WATERSCOPE_APPLICATION", "auto")
    app_used = infer_application(message, explicit=requested_app if requested_app != "auto" else None)
    fallback_enabled = os.getenv("WATERSCOPE_ENABLE_RESEARCH_FALLBACK", "0") == "1"
    # No publication evidence for free-form chat. Never prepend paper-only
    # grounding instructions to this path.
    model_context = "" if general_chat else prefix_context(context)
    def _fallback():
        return deepseek_research_answer(
            model_context, message, app_used,
            general_chat=general_chat, max_output_tokens=max_new_tokens,
        )
    if fallback_enabled and llama_circuit_open():
        answer, fallback_model = _fallback()
        return sanitize_answer(answer), f"deepseek_fallback:{app_used}", f"provider=deepseek;model={fallback_model};reason=llama_circuit_open"
    if fallback_enabled and llama_space_paused():
        _logger.info("Llama Space is PAUSED; skipping Llama and using DeepSeek fallback")
        answer, fallback_model = _fallback()
        return sanitize_answer(answer), f"deepseek_fallback:{app_used}", f"provider=deepseek;model={fallback_model};reason=llama_space_paused"
    try:
        answer, app_used, diagnostics = agent_answer(
            model_context,
            message,
            application=requested_app,
            model=model or os.getenv("WATERSCOPE_MODEL"),
            max_new_tokens=max_new_tokens,
        )
        clear_llama_circuit()
        return sanitize_answer(answer), f"waterscope:{app_used}", diagnostics
    except Exception as exc:
        eligible = is_infrastructure_error(exc) or is_gradio_transport_json_error(exc)
        if not fallback_enabled or not eligible:
            raise
        _logger.warning(
            "Llama request failed (%s); switching to configured research fallback",
            type(exc).__name__,
        )
        mark_llama_unavailable()
        answer, fallback_model = _fallback()
        return sanitize_answer(answer), f"deepseek_fallback:{app_used}", f"provider=deepseek;model={fallback_model};reason=llama_unavailable"

def get_papers_payload() -> dict:
    papers = list_papers_catalog()
    saved_ids = {paper.get("document_id") for paper in list_saved_papers() if paper.get("document_id")}
    for paper in papers:
        paper["is_saved"] = paper.get("document_id") in saved_ids
    return {"papers": papers, "max_select": MAX_CHAT_PAPERS}

def chat_payload(body: dict) -> dict:
    message = (body.get("message") or "").strip()
    document_ids = body.get("document_ids") or []
    application = body.get("application")
    agent_mode = bool(body.get("agent_mode"))

    if not message:
        raise ValueError("message is required")

    if not document_ids and body.get("general_chat") is True:
        # Opt-in direct LLM route for short no-paper conversations. This does
        # not run publication search or the agent orchestrator.
        history = body.get("history")
        if not isinstance(history, list):
            history = []
        transcript = []
        for turn in history[-6:]:
            if not isinstance(turn, dict) or turn.get("role") not in ("user", "assistant"):
                continue
            content = str(turn.get("content") or "").strip()[:1500]
            if content:
                transcript.append(f"{turn['role'].capitalize()}: {content}")
        prompt = (
            "Previous conversation (context only):\n" + "\n".join(transcript)
            + "\n\nCurrent user message:\n" + message
            if transcript else message
        )
        try:
            requested_max = int(body.get("max_new_tokens", 300))
        except (TypeError, ValueError):
            requested_max = 300
        requested_max = max(32, min(300, requested_max))
        answer, strategy, diagnostics = _generate_answer(
            prompt, "", application="general_qa", model=body.get("model"),
            general_chat=True, max_new_tokens=requested_max,
        )
        return {
            "answer": answer,
            "strategy": strategy,
            "diagnostics": diagnostics,
            "general_chat": True,
            "sources": [],
            "chunks": [],
            "retrieval": "none",
            "llm_backend": "deepseek_fallback" if strategy.startswith("deepseek_fallback:") else llm_backend_name(),
        }
    if agent_mode and not document_ids:
        from metadata_search.agent_orchestrator import agent_chat_payload

        return agent_chat_payload(body)
    if not document_ids:
        raise ValueError("Select at least one paper, or enable agent mode.")
    if len(document_ids) > MAX_CHAT_PAPERS:
        raise ValueError(f"Select at most {MAX_CHAT_PAPERS} papers.")

    document_ids = focus_document_ids_for_question(message, document_ids)
    context, sources = build_context_for_documents(document_ids, question=message)
    message = enhance_ordinal_question(message, sources)
    answer, strategy, diagnostics = _generate_answer(
        message,
        context,
        application=application,
        model=body.get("model"),
    )
    return {
        "answer": answer,
        "strategy": strategy,
        "include_excerpts": False,
        "sources": [
            {
                "document_id": source["document_id"],
                "title": source["title"],
                "file_name": source["file_name"],
            }
            for source in sources
        ],
        "chunks": [],
        "retrieval": "mcp" if use_mcp() else "direct",
        "diagnostics": diagnostics,
        "llm_backend": "deepseek_fallback" if strategy.startswith("deepseek_fallback:") else llm_backend_name(),
    }