Upload research_synthesis.py

metadata_search/research_synthesis.py

@@ -1,38 +1,45 @@
 """
-Research-only multi-publication synthesis route for WaterScope-AI Search.
+WaterScope-AI Part 3 research endpoint — v3
 
 REPLACE:
     metadata_search/research_synthesis.py
 
-This module is registered by the existing metadata_search/research_routes.py:
-
-    from metadata_search.research_synthesis import register_research_synthesis_route
-    register_research_synthesis_route(research_bp)
-
-Design
-------
-The research route reuses the CURRENT production functions in
-metadata_search.chat_api:
-
-    focus_document_ids_for_question
-    build_context_for_documents
-    enhance_ordinal_question
-    _generate_answer
-    prefix_context
-    sanitize/model/agent path through _generate_answer
-
-It bypasses ONLY the normal selected-paper count check so that a research
-request can supply up to 20 ordered document IDs. The production /api/chat
-route and MAX_CHAT_PAPERS=4 remain unchanged.
-
-No monkey-patching of production globals is used.
+Research-only changes relative to the deployed production workflow
+------------------------------------------------------------------
+1. Up to 20 ordered selected papers.
+2. FULL stored summary for every selected paper:
+       - no per-paper truncation
+       - no research-side total input cap
+3. Output ceiling scales with the number of papers:
+       max_words = 300 * N
+   The word ceiling is stated explicitly in the prompt.
+4. Generation token allowance is sized conservatively at:
+       max_new_tokens = 2 * max_words = 600 * N
+   This is only a token safety allowance; the scientific output ceiling
+   remains 300 words per paper.
+5. Model = explicit request (normally "base").
+6. Research Task = explicit request (normally "auto"), so normal task
+   auto-detection is retained.
+7. Production /api/chat and its 4-paper limit are NOT changed.
+
+Why a cloned generation function is used
+----------------------------------------
+The deployed production _generate_answer() obtains WATERSCOPE_MAX_TOKENS
+through its agent path rather than accepting a request-level token argument.
+
+For the research endpoint only, v3 creates an isolated clone of the production
+generation function (and its imported agent_answer function when present) with
+a request-scoped WATERSCOPE_MAX_TOKENS override. It does NOT mutate os.environ,
+does NOT patch module globals, and therefore does not change production calls.
 """
 
 from __future__ import annotations
 
+import math
 import os
 import time
-from typing import Any
+import types
+from typing import Any, Callable
 
 from flask import jsonify, request
 
@@ -40,6 +47,8 @@ import metadata_search.chat_api as chat_api
 
 
 MAX_RESEARCH_PAPERS = 20
+DEFAULT_MAX_WORDS_PER_PAPER = 300
+DEFAULT_TOKEN_RESERVE_PER_WORD = 2.0
 
 
 def _clean(value: Any) -> str:
@@ -52,49 +61,58 @@ def _word_count(text: str) -> int:
     return len(_clean(text).split())
 
 
-def _split_context_excerpts(context: str) -> list[str]:
-    """
-    Extract the exact per-paper excerpt from the assembled production context.
+def _application_from_strategy(strategy: str) -> str:
+    strategy = _clean(strategy)
+    prefix = "waterscope:"
+    if strategy.startswith(prefix):
+        return strategy[len(prefix):]
+    return ""
 
-    metadata_search.mcp_retrieval builds each context block as:
 
-        Paper N: <title>
-        File: <file_name>
+def _build_full_context(
+    sources: list[dict[str, Any]],
+) -> str:
+    """
+    Build the model evidence context from the COMPLETE stored summary
+    contained in each production source record.
 
-        <excerpt>
+    The preceding v2 smoke test demonstrated that source["text"] contains
+    the complete stored summary while the normal production context contains
+    a shorter question-aware excerpt. v3 deliberately uses source["text"].
+    """
+    blocks: list[str] = []
 
-    with blocks separated by "\n\n---\n\n".
+    for index, source in enumerate(sources, start=1):
+        title = _clean(source.get("title")) or f"Paper {index}"
+        file_name = _clean(source.get("file_name"))
+        full_summary = _clean(source.get("text"))
 
-    Reading excerpts back from the final assembled context guarantees that the
-    trace records what was actually supplied to the model, rather than a
-    separately reconstructed approximation.
-    """
-    if not context:
-        return []
+        if not full_summary:
+            raise ValueError(
+                f"Paper {index} ({_clean(source.get('document_id'))}) "
+                "has an empty stored summary."
+            )
 
-    blocks = context.split("\n\n---\n\n")
-    excerpts: list[str] = []
+        header = f"Paper {index}: {title}"
+        if file_name:
+            header += f"\nFile: {file_name}"
 
-    for block in blocks:
-        if "\n\n" in block:
-            _, excerpt = block.split("\n\n", 1)
-            excerpts.append(excerpt.strip())
-        else:
-            excerpts.append("")
+        blocks.append(f"{header}\n\n{full_summary}")
 
-    return excerpts
+    return "\n\n---\n\n".join(blocks)
 
 
-def _evidence_trace(
+def _build_evidence_trace(
     sources: list[dict[str, Any]],
-    context: str,
 ) -> list[dict[str, Any]]:
-    excerpts = _split_context_excerpts(context)
+    """
+    In v3 the exact evidence supplied for each paper is the full stored
+    summary, so excerpt_used == full_summary by design.
+    """
     rows: list[dict[str, Any]] = []
 
     for index, source in enumerate(sources, start=1):
         full_summary = _clean(source.get("text"))
-        excerpt = excerpts[index - 1] if index - 1 < len(excerpts) else ""
 
         rows.append(
             {
@@ -105,26 +123,122 @@ def _evidence_trace(
                 "full_summary": full_summary,
                 "full_summary_chars": len(full_summary),
                 "full_summary_words": _word_count(full_summary),
-                "excerpt_used": excerpt,
-                "excerpt_chars": len(excerpt),
-                "excerpt_words": _word_count(excerpt),
-                "excerpt_truncated_vs_full": (
-                    bool(full_summary)
-                    and bool(excerpt)
-                    and len(excerpt) < len(full_summary)
-                ),
+                "excerpt_used": full_summary,
+                "excerpt_chars": len(full_summary),
+                "excerpt_words": _word_count(full_summary),
+                "excerpt_truncated_vs_full": False,
             }
         )
 
     return rows
 
 
-def _application_from_strategy(strategy: str) -> str:
-    strategy = _clean(strategy)
-    prefix = "waterscope:"
-    if strategy.startswith(prefix):
-        return strategy[len(prefix):]
-    return ""
+class _ResearchOSProxy:
+    """
+    Delegate normal os behavior to the real os module, but override only
+    WATERSCOPE_MAX_TOKENS for the cloned research call.
+    """
+
+    def __init__(self, max_new_tokens: int):
+        self._max_new_tokens = int(max_new_tokens)
+
+    def getenv(self, key: str, default: Any = None):
+        if key == "WATERSCOPE_MAX_TOKENS":
+            return str(self._max_new_tokens)
+        return os.getenv(key, default)
+
+    def __getattr__(self, name: str):
+        return getattr(os, name)
+
+
+def _copy_function(
+    fn: Callable,
+    globals_dict: dict[str, Any],
+) -> Callable:
+    cloned = types.FunctionType(
+        fn.__code__,
+        globals_dict,
+        name=fn.__name__,
+        argdefs=fn.__defaults__,
+        closure=fn.__closure__,
+    )
+    cloned.__kwdefaults__ = getattr(fn, "__kwdefaults__", None)
+    cloned.__annotations__ = dict(getattr(fn, "__annotations__", {}))
+    cloned.__doc__ = fn.__doc__
+    return cloned
+
+
+def _override_token_globals(
+    globals_dict: dict[str, Any],
+    max_new_tokens: int,
+) -> dict[str, Any]:
+    """
+    Create an isolated globals dictionary for a research-only cloned function.
+
+    Covers both common implementations:
+      - os.getenv("WATERSCOPE_MAX_TOKENS", ...)
+      - module-level MAX_*TOKEN* constants
+    """
+    new_globals = dict(globals_dict)
+    new_globals["os"] = _ResearchOSProxy(max_new_tokens)
+
+    for key, value in list(new_globals.items()):
+        upper = key.upper()
+        if "MAX" in upper and "TOKEN" in upper:
+            if isinstance(value, int):
+                new_globals[key] = int(max_new_tokens)
+            elif isinstance(value, str) and value.strip().isdigit():
+                new_globals[key] = str(int(max_new_tokens))
+
+    return new_globals
+
+
+def _research_generate_answer(
+    question: str,
+    context: str,
+    *,
+    application: str,
+    model: str | None,
+    max_new_tokens: int,
+):
+    """
+    Execute the SAME deployed production _generate_answer code with an
+    isolated, request-scoped token allowance.
+
+    No production module global or environment variable is changed.
+    """
+    original_generate = chat_api._generate_answer
+
+    generate_globals = _override_token_globals(
+        original_generate.__globals__,
+        max_new_tokens,
+    )
+
+    # _generate_answer normally calls an imported agent_answer function.
+    # Clone that function too so its os.getenv / token constants are isolated.
+    original_agent_answer = generate_globals.get("agent_answer")
+    if isinstance(original_agent_answer, types.FunctionType):
+        agent_globals = _override_token_globals(
+            original_agent_answer.__globals__,
+            max_new_tokens,
+        )
+        cloned_agent_answer = _copy_function(
+            original_agent_answer,
+            agent_globals,
+        )
+        generate_globals["agent_answer"] = cloned_agent_answer
+
+    cloned_generate = _copy_function(
+        original_generate,
+        generate_globals,
+    )
+
+    return cloned_generate(
+        question,
+        context,
+        application=application,
+        model=model,
+    )
 
 
 def register_research_synthesis_route(research_bp):
@@ -132,19 +246,14 @@ def register_research_synthesis_route(research_bp):
     Register:
         POST /api/research/synthesis
 
-    Expected JSON:
+    Example request:
     {
         "document_ids": ["...", "..."],
         "question": "...",
         "model": "base",
-        "application": "auto"
+        "application": "auto",
+        "max_words_per_paper": 300
     }
-
-    Optional Space secret:
-        WATERSCOPE_RESEARCH_TOKEN
-
-    If set, the caller must send:
-        X-Research-Token: <same value>
     """
 
     @research_bp.post("/synthesis")
@@ -157,22 +266,31 @@ def register_research_synthesis_route(research_bp):
 
         body = request.get_json(force=True, silent=True) or {}
 
-        question = _clean(body.get("question") or body.get("message"))
+        question_original = _clean(
+            body.get("question") or body.get("message")
+        )
         raw_ids = body.get("document_ids") or []
 
-        if not question:
-            return jsonify({"ok": False, "error": "question is required"}), 400
+        if not question_original:
+            return jsonify(
+                {"ok": False, "error": "question is required"}
+            ), 400
 
         if not isinstance(raw_ids, list):
             return jsonify(
                 {"ok": False, "error": "document_ids must be a list"}
             ), 400
 
-        document_ids = [_clean(x) for x in raw_ids if _clean(x)]
+        document_ids = [
+            _clean(x) for x in raw_ids if _clean(x)
+        ]
 
         if not document_ids:
             return jsonify(
-                {"ok": False, "error": "At least one document_id is required"}
+                {
+                    "ok": False,
+                    "error": "At least one document_id is required",
+                }
             ), 400
 
         if len(document_ids) > MAX_RESEARCH_PAPERS:
@@ -188,114 +306,235 @@ def register_research_synthesis_route(research_bp):
 
         if len(set(document_ids)) != len(document_ids):
             return jsonify(
-                {"ok": False, "error": "Duplicate document_ids are not allowed"}
+                {
+                    "ok": False,
+                    "error": "Duplicate document_ids are not allowed",
+                }
             ), 400
 
         application = _clean(body.get("application")) or "auto"
         model = _clean(body.get("model")) or None
 
+        try:
+            max_words_per_paper = int(
+                body.get(
+                    "max_words_per_paper",
+                    DEFAULT_MAX_WORDS_PER_PAPER,
+                )
+            )
+        except (TypeError, ValueError):
+            return jsonify(
+                {
+                    "ok": False,
+                    "error": "max_words_per_paper must be an integer",
+                }
+            ), 400
+
+        if not 1 <= max_words_per_paper <= 1000:
+            return jsonify(
+                {
+                    "ok": False,
+                    "error": (
+                        "max_words_per_paper must be between 1 and 1000"
+                    ),
+                }
+            ), 400
+
         total_start = time.perf_counter()
 
         try:
-            # Same production behavior for ordinal-paper questions.
-            # Part 3 prompts do not use "first paper", "second paper", etc.,
-            # so formal runs should retain all supplied IDs.
-            document_ids_used = chat_api.focus_document_ids_for_question(
-                question,
-                document_ids,
+            # Preserve current production behavior for questions that
+            # explicitly refer to an ordinal paper.
+            document_ids_used = (
+                chat_api.focus_document_ids_for_question(
+                    question_original,
+                    document_ids,
+                )
             )
 
-            context_start = time.perf_counter()
-            context, sources = chat_api.build_context_for_documents(
-                document_ids_used,
-                question=question,
+            # We use the production context builder ONLY to resolve/load the
+            # requested paper summaries and metadata. Its truncated context
+            # output is intentionally NOT sent to the model in this experiment.
+            load_start = time.perf_counter()
+            production_context, sources = (
+                chat_api.build_context_for_documents(
+                    document_ids_used,
+                    question=question_original,
+                )
+            )
+            source_load_elapsed_s = (
+                time.perf_counter() - load_start
             )
-            context_elapsed_s = time.perf_counter() - context_start
 
             if not sources:
-                raise ValueError("No summaries found for the selected papers.")
+                raise ValueError(
+                    "No summaries found for the selected papers."
+                )
+
+            if len(sources) != len(document_ids_used):
+                raise ValueError(
+                    f"Requested {len(document_ids_used)} papers but "
+                    f"loaded {len(sources)} complete summaries."
+                )
+
+            n_papers = len(sources)
+
+            # FULL-SUMMARY research context: no truncate_summary result is used.
+            full_context = _build_full_context(sources)
 
-            # Same production prompt enhancement.
-            question_to_model = chat_api.enhance_ordinal_question(
-                question,
+            # Preserve the platform's ordinal-question enhancement.
+            question_enhanced = chat_api.enhance_ordinal_question(
+                question_original,
                 sources,
             )
 
-            # This is the exact paper context after the same production
-            # prefix_context() transformation used inside _generate_answer().
-            model_context = chat_api.prefix_context(context)
+            max_output_words = (
+                max_words_per_paper * n_papers
+            )
 
-            model_start = time.perf_counter()
-            answer, strategy, diagnostics = chat_api._generate_answer(
-                question_to_model,
-                context,
-                application=application,
-                model=model,
+            # Word ceiling is the scientific constraint.
+            length_instruction = (
+                "\n\nRESPONSE LENGTH CEILING: "
+                f"Do not exceed {max_output_words:,} words total "
+                f"({max_words_per_paper} words per supplied paper). "
+                "This is a maximum, not a target; be substantially shorter "
+                "when the question can be answered completely in fewer words."
             )
-            model_elapsed_s = time.perf_counter() - model_start
 
-            total_elapsed_s = time.perf_counter() - total_start
+            question_to_model = (
+                question_enhanced + length_instruction
+            )
 
-            evidence = _evidence_trace(sources, context)
-            application_used = _application_from_strategy(strategy)
+            # Token allowance is deliberately more generous than the word
+            # ceiling so tokenization does not truncate a valid response.
+            max_new_tokens_requested = int(
+                math.ceil(
+                    max_output_words
+                    * DEFAULT_TOKEN_RESERVE_PER_WORD
+                )
+            )
 
-            # _generate_answer() does not accept a request-level token argument.
-            # Its agent path uses WATERSCOPE_MAX_TOKENS, defaulting to 1024.
-            effective_max_new_tokens = int(
-                os.getenv("WATERSCOPE_MAX_TOKENS", "1024")
+            # Record the exact prefixed evidence the production generation
+            # path will construct from this full context.
+            model_context = chat_api.prefix_context(
+                full_context
             )
 
+            model_start = time.perf_counter()
+            answer, strategy, diagnostics = (
+                _research_generate_answer(
+                    question_to_model,
+                    full_context,
+                    application=application,
+                    model=model,
+                    max_new_tokens=max_new_tokens_requested,
+                )
+            )
+            model_call_elapsed_s = (
+                time.perf_counter() - model_start
+            )
+
+            total_elapsed_s = (
+                time.perf_counter() - total_start
+            )
+
+            evidence = _build_evidence_trace(sources)
+            application_used = _application_from_strategy(
+                strategy
+            )
+            answer_words = _word_count(answer)
+
             response = {
                 "ok": True,
                 "answer": answer,
+                "answer_words": answer_words,
+                "answer_exceeded_word_ceiling": (
+                    answer_words > max_output_words
+                ),
                 "strategy": strategy,
                 "application_requested": application,
                 "application_used": application_used,
                 "model_requested": model or "",
-                "max_new_tokens_effective": effective_max_new_tokens,
+                "input_mode": "full_stored_summaries",
+                "research_input_truncation": False,
+                "max_words_per_paper": max_words_per_paper,
+                "max_output_words": max_output_words,
+                "token_reserve_per_word": (
+                    DEFAULT_TOKEN_RESERVE_PER_WORD
+                ),
+                "max_new_tokens_requested": (
+                    max_new_tokens_requested
+                ),
                 "document_ids_requested": document_ids,
                 "document_ids_used": document_ids_used,
                 "n_papers_requested": len(document_ids),
-                "n_papers_used": len(sources),
+                "n_papers_used": n_papers,
                 "sources": [
                     {
-                        "document_id": _clean(s.get("document_id")),
+                        "document_id": _clean(
+                            s.get("document_id")
+                        ),
                         "title": _clean(s.get("title")),
-                        "file_name": _clean(s.get("file_name")),
+                        "file_name": _clean(
+                            s.get("file_name")
+                        ),
                     }
                     for s in sources
                 ],
                 "evidence": evidence,
-                "context_used": context,
-                "context_used_chars": len(context),
-                "context_used_words": _word_count(context),
+                "full_context": full_context,
+                "full_context_chars": len(full_context),
+                "full_context_words": _word_count(
+                    full_context
+                ),
                 "model_context": model_context,
                 "model_context_chars": len(model_context),
-                "model_context_words": _word_count(model_context),
+                "model_context_words": _word_count(
+                    model_context
+                ),
+                "production_context_chars_not_used": len(
+                    production_context or ""
+                ),
+                "question_original": question_original,
                 "question_to_model": question_to_model,
-                "context_build_elapsed_s": context_elapsed_s,
-                "model_call_elapsed_s": model_elapsed_s,
+                "source_load_elapsed_s": source_load_elapsed_s,
+                "model_call_elapsed_s": model_call_elapsed_s,
                 "research_total_elapsed_s": total_elapsed_s,
-                "retrieval": "mcp" if chat_api.use_mcp() else "direct",
+                "retrieval": (
+                    "mcp" if chat_api.use_mcp() else "direct"
+                ),
                 "diagnostics": diagnostics,
                 "llm_backend": chat_api.llm_backend_name(),
                 "trace_status": {
-                    "captured_context": bool(context),
-                    "captured_model_context": bool(model_context),
-                    "captured_source_summaries": (
-                        len(evidence) == len(sources)
-                        and all(bool(row["full_summary"]) for row in evidence)
+                    "captured_full_context": bool(
+                        full_context
                     ),
-                    "captured_per_paper_excerpts": (
-                        len(evidence) == len(sources)
-                        and all(bool(row["excerpt_used"]) for row in evidence)
+                    "captured_model_context": bool(
+                        model_context
                     ),
-                    "source_count": len(sources),
-                    "excerpt_count": sum(
-                        bool(row["excerpt_used"]) for row in evidence
+                    "captured_source_summaries": (
+                        len(evidence) == n_papers
+                        and all(
+                            bool(row["full_summary"])
+                            for row in evidence
+                        )
+                    ),
+                    "all_papers_untruncated": all(
+                        row[
+                            "excerpt_truncated_vs_full"
+                        ] is False
+                        and row["excerpt_chars"]
+                        == row["full_summary_chars"]
+                        for row in evidence
                     ),
+                    "source_count": n_papers,
                     "production_chat_limit_unchanged": True,
-                    "research_paper_limit": MAX_RESEARCH_PAPERS,
+                    "research_paper_limit": (
+                        MAX_RESEARCH_PAPERS
+                    ),
+                    "token_override_scope": (
+                        "isolated_function_clone"
+                    ),
                 },
             }
 
@@ -314,7 +553,9 @@ def register_research_synthesis_route(research_bp):
             return jsonify(
                 {
                     "ok": False,
-                    "error": f"{type(exc).__name__}: {exc}",
+                    "error": (
+                        f"{type(exc).__name__}: {exc}"
+                    ),
                     "document_ids_requested": document_ids,
                 }
             ), 500