Spaces:
Sleeping
Sleeping
| """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(), | |
| } | |