diff --git a/.env.example b/.env.example index 988b8d4..a8ef1f2 100644 --- a/.env.example +++ b/.env.example @@ -45,21 +45,50 @@ LLM_FALLBACK_MODEL=gemma-4-26b-a4b-it # Response generation parameters LLM_TEMPERATURE=0.3 -LLM_MAX_TOKENS=2048 +LLM_MAX_TOKENS=8192 AGENT_MAX_TOKENS=8192 -AGENT_TIMEOUT=120 +# Hard ceiling for one agentic run (seconds). Measured on a CPU-only box: +# ~45s retrieval + ~50s generation + ~25s evaluation. Anything near 120 leaves no +# room for AGENT_EVAL_RESERVE_S below, so verification gets skipped on every run. +AGENT_TIMEOUT=300 +# Per-request HTTP timeout for every LLM call (seconds). Without it the SDK +# defaults apply (OpenAI: 600s x 2 retries), so ONE stalled request outlives the +# whole agent budget. Sized for a legitimate call, not for the failover chain: +# agent answer generation is unary and measured 20-50s on CPU, so a materially +# lower value aborts real generations. Failover walks up to 3 (provider, model) +# attempts sequentially, so a fully-stalled chain can reach ~180s — past +# AGENT_REFLEXION_BUDGET_S, though still inside AGENT_TIMEOUT, which finalises +# the draft rather than 504. +LLM_REQUEST_TIMEOUT_S=60 +# Separate budget for STREAMED calls. For Gemini the HTTP timeout covers the whole +# stream rather than the gap between chunks, so reusing the 60s unary value above +# tore down long answers mid-generation (WinError 10054, truncated text). +LLM_STREAM_TIMEOUT_S=300 # Thinking budget: 0=off (cheapest), -1=dynamic, N=cap N thinking tokens. # Note: some models (e.g. gemini-3.6-flash) reject a literal 0 budget. The # Gemini backend detects that once per model and omits the field instead, so # 0 stays safe — it just means "no thinking config" on those models. AGENT_THINKING_BUDGET=0 +# Thinking LEVEL — the Gemini 3.x control, which SUPERSEDES thinking_budget above. +# gemini-3.6-flash rejects a budget outright and defaults to MEDIUM thinking when +# nothing is sent, so "budget 0" used to mean medium thinking in practice, billed +# and taken out of LLM_MAX_TOKENS. Values: minimal | low | medium | high, or empty +# to accept the model's own default. +LLM_THINKING_LEVEL=minimal +AGENT_THINKING_LEVEL=minimal # Max sub-queries per reflexion cycle (main agent-latency lever) AGENT_MAX_SUB_QUERIES=3 -# Wall-clock budget for the reflexion loop (seconds). Once exceeded, the +# Wall-clock budget for the reflexion LOOP (seconds). Once exceeded, the # evaluator finalizes the current draft instead of starting another -# retrieve -> generate -> verify cycle. Keep BELOW AGENT_TIMEOUT so it fires -# first and the user gets an answer rather than a 504 that discards the work. +# retrieve -> generate -> verify cycle. Applies from iteration 2 onwards only: +# on a CPU-only box the first pass alone runs past it, and gating iteration 1 +# here shipped every answer with no faithfulness score and no confidence. AGENT_REFLEXION_BUDGET_S=90 +# Wall-clock room that must remain under AGENT_TIMEOUT for the evaluator to +# attempt an evaluation at all — this one CAN skip iteration 1, but only when +# finishing would overrun the timeout and discard the draft entirely. Sized for +# one NLI pass plus one completeness call (measured ~30s + ~15s on CPU, doubled). +AGENT_EVAL_RESERVE_S=90 # Context caps for agentic mode. Wider than standard RAG because the agent pools # passages from several tools and the 12-chunk cut truncated equation chunks. AGENT_MAX_CONTEXT_CHUNKS=20 @@ -301,11 +330,32 @@ NLI_MODEL_NAME=MoritzLaurer/mDeBERTa-v3-base-xnli-multilingual-nli-2mil7 # cross-encoder/nli-deberta-v3-base: 0=contradiction, 1=entailment, 2=neutral NLI_ENTAILMENT_INDEX=0 -# A claim scoring below this is logged as ungrounded. -FAITHFULNESS_THRESHOLD=0.5 +# A claim scoring below this is logged as ungrounded. MUST be calibrated against +# the model above — this is not a "confidence percentage" you can set by taste. +# Measured for the default model: a sentence copied VERBATIM out of its own chunk +# scores median 0.226 / max 0.428, an unrelated paper's sentence median 0.099 / +# p90 0.158. The old 0.5 sat above every positive, so faithfulness read ~0 for +# every answer ever produced. 0.15 gives recall 0.70 at false-positive 0.10-0.15. +# Recalibrate (positives vs cross-paper negatives) whenever NLI_MODEL_NAME changes. +FAITHFULNESS_THRESHOLD=0.15 # What to do with ungrounded claims: warn | strip | regen FAITHFULNESS_ENFORCE=warn +# Fraction of an answer's claims that must be grounded for the reflexion loop to +# accept it and for the finalizer to trust it enough to abstain on completeness +# alone. Below 1.0 by design: per-claim recall is ~0.70, so even a fully grounded +# answer lands near 0.70 — a 0.75 bar could never fire. +AGENT_FAITHFULNESS_ACCEPT=0.6 + +# NLI cost is linear in pairs AND in premise length. Measured on a CPU-only box +# with the int8 ONNX model: 1.15s/pair at 512 tokens, 0.4s at 256. A 30-sentence +# answer citing multi-chunk papers hit ~275 pairs = 318s in ONE reflexion pass. +# These two knobs bound it (measured no quality cost: positive median 0.226 at +# 256 vs 0.221 at 512). Chunks arrive rerank-ordered, so the first are the best +# support. Drop the cap to 1 if the NLI pass still dominates your run. +NLI_MAX_SEQ_LENGTH=256 +NLI_MAX_CHUNKS_PER_CITATION=2 + # ============================================================================== # Caching (in-process, per worker) diff --git a/README.md b/README.md index bbe4022..30ebd43 100644 --- a/README.md +++ b/README.md @@ -32,7 +32,8 @@ Two pipelines ship side-by-side: **Standard RAG** (single-pass hybrid retrieval) ### v2.4 patch — correctness fixes * **Faithfulness verification restored** — `verify.check_claims` was passing `(index, text)` tuples to the NLI model instead of the chunk text, so every call raised and every answer reported `confidence: 0.0` with no evidence. -* **Generation restored on `gemini-3.6-flash`** — models that reject `thinking_budget=0` returned `400 INVALID_ARGUMENT`. The Gemini backend now remembers per-model rejections and retries once without `thinking_config`. +* **Generation restored on `gemini-3.6-flash`** — models that reject `thinking_budget=0` returned `400 INVALID_ARGUMENT`. The Gemini backend now remembers per-model rejections and retries once with the budget translated to a `thinking_level`. +* **Thinking is actually minimised, not defaulted** — the earlier fix dropped `thinking_config` on rejection, which means the model's *own* default (`medium` on `gemini-3.6-flash`), billed and taken out of `LLM_MAX_TOKENS`. New `LLM_THINKING_LEVEL` / `AGENT_THINKING_LEVEL` knobs (default `minimal`) set the level up front. * **Retrieval cache actually populates** — the cacheability check ran after the collection was materialized, so nothing was ever stored. Cached entries are now copied on both read and write so callers can't mutate them. * **Tags filter returns results** — tags are stored as one comma-joined metadata string and must be filtered in Python, which needs an over-fetch (`TAGS_OVERFETCH`); without it a tag query returned nothing. * **Cross-vendor failover is really cross-vendor** — the OpenRouter fallback was handed a bare Gemini model name, which OpenRouter rewrites to `google/`, routing back to the vendor that just failed. It now picks a `/`-shaped slug from `LLM_SELECTABLE_MODELS`. @@ -53,7 +54,7 @@ Two pipelines ship side-by-side: **Standard RAG** (single-pass hybrid retrieval) * **web_search** — Tavily web search for current events and non-academic info * **calculate** — numexpr math evaluation (identifier-whitelisted) * **execute_python** — process-isolated Python with AST-based validation (import whitelist, dunder + dangerous-builtin blocking) + 10s timeout -* **Reflexion loops with dual budget** — the evaluator checks faithfulness (NLI entailment, minimum across claims) and completeness (Gemini Flash). Below threshold it can regenerate, retrieve more, or reformulate — bounded by **both** an iteration cap (3) and a wall-clock budget (`AGENT_REFLEXION_BUDGET_S`). Stuck-loop detection auto-accepts when completeness stops improving. +* **Reflexion loops with layered budgets** — the evaluator checks faithfulness (fraction of claims grounded by NLI entailment) and completeness (Gemini Flash). Below threshold it can regenerate, retrieve more, or reformulate — bounded by an iteration cap (3), a loop budget (`AGENT_REFLEXION_BUDGET_S`, iteration 2+), a deadline reserve (`AGENT_EVAL_RESERVE_S`), and a per-call LLM timeout (`LLM_REQUEST_TIMEOUT_S`). Stuck-loop detection auto-accepts when completeness stops improving. * **Contradiction detection** — NLI-based cross-source contradiction flagging in the answer generator; surfaces both sides with citations when sources disagree * **Confidence & abstention** — finalizer surfaces a confidence score; low-confidence answers get an explicit abstention prefix with partial sourcing * **Multi-turn conversations** — session history threaded through `AgentState` so follow-ups resolve pronouns and references @@ -68,7 +69,8 @@ Two pipelines ship side-by-side: **Standard RAG** (single-pass hybrid retrieval) * **Dense + sparse** — BGE-M3 (1024d) fused with BM25 via Reciprocal Rank Fusion (RRF) * **Two-stage reranking** — `BAAI/bge-reranker-v2-m3` cross-encoder, with optional **ColBERT** multi-vector MaxSim rerank on the narrowed candidate set * **Optional HyDE** — generate a hypothetical answer, embed it, and retrieve against it for recall on sparse queries -* **Faithfulness verification** — `cross-encoder/nli-deberta-v3-base` scores entailment per claim; unsupported assertions flagged, stripped, or regenerated (`FAITHFULNESS_ENFORCE`) +* **Faithfulness verification** — a multilingual NLI cross-encoder (`NLI_MODEL_NAME`, int8 ONNX on CPU) scores entailment per claim against its cited chunks; unsupported assertions flagged, stripped, or regenerated (`FAITHFULNESS_ENFORCE`). The threshold is **model-specific and calibrated**, not a taste setting — see `FAITHFULNESS_THRESHOLD` +* **Citation integrity** — the answer's `[N]` markers are renumbered to a dense `1..M` matching the cited-only source panel, and markers are resolved against **only the chunks that reached the prompt**. The context is truncated by chunk count and by length, so numbering against everything retrieved let a marker the model invented resolve to a real paper it was never shown — a phantom citation that reads as legitimate. Unresolvable markers are dropped rather than left dangling * **HNSW tuning knobs** — `ef_search`, `ef_construction`, `M` all env-configurable ### 📥 Smart Ingestion @@ -89,7 +91,7 @@ Two pipelines ship side-by-side: **Standard RAG** (single-pass hybrid retrieval) ### 🛡️ Production-Ready Infrastructure * **SQLite session/job persistence** — restarts don't drop in-flight state (`SESSIONS_DB_PATH`) -* **SSE streaming** — token-by-token answers and live ingest progress +* **SSE streaming** — token-by-token answers and live ingest progress; the `done` event carries the citation-corrected answer, since chunks stream before numbering can be resolved * Thread-safe model init (double-checked locking on all singletons) * Startup warm-up via FastAPI lifespan (embeddings, vector store, reranker, BM25) — no cold first request * Request-ID correlation across log lines; Prometheus metrics @@ -160,7 +162,13 @@ TAVILY_API_KEY=your_tavily_key_here # Optional — higher token limit for agent answers (default 8192) AGENT_MAX_TOKENS=8192 -# Optional — agent thinking tokens: 0=off (cheapest), -1=dynamic, N=cap (default 0) +# Optional — thinking level (Gemini 3.x): minimal|low|medium|high, empty = model default +# Empty is not neutral: gemini-3.6-flash then thinks at medium, out of LLM_MAX_TOKENS +LLM_THINKING_LEVEL=minimal +AGENT_THINKING_LEVEL=minimal + +# Legacy — agent thinking tokens on models that still accept a budget: +# 0=off (cheapest), -1=dynamic, N=cap (default 0). Gemini 3.x translates it to a level. AGENT_THINKING_BUDGET=0 # Optional — retrieval quality boosters (off by default, cost more compute) @@ -234,6 +242,14 @@ with requests.post('http://localhost:8080/query/stream', print(line.decode()) # Server-Sent Events ``` +**Use the `done` event's `answer`, not the concatenated chunks.** Chunks are emitted as the model produces them, so they carry its raw `[N]` markers. The `done` event carries the *compacted* answer — citations renumbered to a dense `1..M` matching the source panel, and markers that resolve to nothing removed. Concatenating the chunks yields text whose numbering disagrees with `citations` (an answer drawing on papers 1 and 4 of 4 streams as `[1] … [4]` beside a two-entry panel). + +| SSE event | Fields | +|---|---| +| `chunk` | `text` — raw answer fragment, streamed live | +| `done` | `answer` (compacted), `citations`, `language`, `query_id`, `session_id` (chat only) | +| `error` | `message` | + #### Standard Chat — `POST /chat` ```python @@ -419,11 +435,16 @@ Key settings (all overridable via environment variables): | `LLM_MODEL_NAME` | `gemini-3.6-flash` | Gemini model for generation | | `LLM_FALLBACK_MODEL` | `gemma-4-26b-a4b-it` | Fallback when primary is overloaded (503/429) | | `LLM_SELECTABLE_MODELS` | `gemini-3.6-flash,gemini-3.5-flash,anthropic/claude-haiku,openai/gpt-5.4-nano` | Curated model dropdown (comma-separated; first entry is the default). `.env.example` ships a wider free-tier list. Bare name → Gemini, `/` slug → OpenRouter — **keep at least one `/` slug**, since cross-vendor failover picks the first one here | -| `LLM_MAX_TOKENS` | `2048` | Max tokens for standard RAG | +| `LLM_MAX_TOKENS` | `8192` | Max tokens for standard RAG (covers thinking + answer) | | `AGENT_MAX_TOKENS` | `8192` | Max tokens for agentic pipeline | -| `AGENT_TIMEOUT` | `120` | Agent pipeline timeout (seconds) → 504 | -| `AGENT_REFLEXION_BUDGET_S` | `90` | Wall-clock budget for reflexion loops | -| `AGENT_THINKING_BUDGET` | `0` | Agent thinking tokens: `0`=off, `-1`=dynamic, `N`=cap | +| `AGENT_TIMEOUT` | `300` | Agent pipeline timeout (seconds) → 504. Must leave room for `AGENT_EVAL_RESERVE_S`, or verification is skipped on every run | +| `LLM_REQUEST_TIMEOUT_S` | `60` | Per-request HTTP timeout for non-streaming LLM calls. Without it the SDK defaults apply (OpenAI: 600s × 2 retries) and one stalled request outlives the whole agent budget — multiplied by the 3-attempt failover chain | +| `LLM_STREAM_TIMEOUT_S` | `300` | Same, for **streamed** calls. Separate because the timeout covers the whole stream rather than the gap between chunks — sharing the unary value cuts long answers off mid-generation | +| `AGENT_REFLEXION_BUDGET_S` | `90` | Wall-clock budget for reflexion **loops** — blocks starting another cycle, from iteration 2 onwards | +| `AGENT_EVAL_RESERVE_S` | `90` | Room that must remain under `AGENT_TIMEOUT` to attempt an evaluation at all. Can skip iteration 1, but only when finishing would overrun the timeout and discard the draft | +| `AGENT_THINKING_BUDGET` | `0` | **Legacy.** Agent thinking tokens: `0`=off, `-1`=dynamic, `N`=cap. Gemini 3.x models reject this field; the backend translates it to a level | +| `LLM_THINKING_LEVEL` | `minimal` | Thinking level for standard RAG — the Gemini 3.x control. `minimal`, `low`, `medium`, `high`, or empty to accept the model default. **Not neutral:** `gemini-3.6-flash` defaults to `medium`, and those thought tokens come out of `LLM_MAX_TOKENS`, squeezing the answer | +| `AGENT_THINKING_LEVEL` | `minimal` | Same, for the agentic pipeline | | `AGENT_MAX_SUB_QUERIES` | `3` | Cap per-cycle retrievals to bound latency | | `CONTRADICTION_DETECT_ENABLE` | `false` | NLI-based cross-source contradiction flagging | | `CONTRADICTION_NLI_THRESHOLD` | `0.6` | NLI score threshold for contradiction detection | @@ -452,9 +473,12 @@ Key settings (all overridable via environment variables): | `HNSW_EF_CONSTRUCTION` | `100` | HNSW build-time breadth (index quality vs. ingest speed) | | `HNSW_M` | `16` | HNSW graph connectivity | | `FAITHFULNESS_ENFORCE` | `warn` | `warn`, `strip`, or `regen` | -| `FAITHFULNESS_THRESHOLD` | `0.5` | NLI support score threshold | +| `FAITHFULNESS_THRESHOLD` | `0.15` | Per-claim entailment probability above which a claim counts as grounded. **Model-specific — recalibrate if you change `NLI_MODEL_NAME`.** Measured for the default model: a sentence copied verbatim out of its own chunk scores median 0.226 / max 0.428; an unrelated paper's sentence median 0.099 / p90 0.158. A 0.5 bar sits above every positive, so faithfulness reads ~0 for every answer | +| `AGENT_FAITHFULNESS_ACCEPT` | `0.6` | Fraction of claims that must be grounded for reflexion to accept and for the finalizer to abstain on completeness alone. Below 1.0 by design: per-claim recall is ~0.70, so a fully grounded answer lands near 0.70 | | `NLI_MODEL_NAME` | `MoritzLaurer/mDeBERTa-v3-base-xnli-multilingual-nli-2mil7` | NLI model backing faithfulness + contradiction checks | | `NLI_ENTAILMENT_INDEX` | `0` | Index of the entailment label in that model's output — change if you swap models | +| `NLI_MAX_SEQ_LENGTH` | `256` | Premise truncation. NLI cost is linear in length: 1.15s/pair at 512 tokens vs 0.4s at 256 on CPU, with no measured quality cost | +| `NLI_MAX_CHUNKS_PER_CITATION` | `2` | Chunks scored per cited paper. Uncapped, one answer cost ~275 pairs = 318s in a single reflexion pass | | `GEMINI_CACHE_ENABLED` | `false` | Explicit Gemini prompt caching | | `GEMINI_CACHE_TTL` | `3600` | Prompt cache lifetime (seconds) | | `SESSIONS_DB_PATH` | `sessions.db` | SQLite path for session/job/watch persistence | @@ -535,7 +559,9 @@ flowchart TD | Guard | Behavior | |---|---| | **Iteration cap** | Max **3** reflexion cycles | -| **Wall-clock budget** | `AGENT_REFLEXION_BUDGET_S` | +| **Loop budget** | `AGENT_REFLEXION_BUDGET_S` — blocks starting another cycle, from iteration 2 onwards. Deliberately does *not* gate the first evaluation: on a CPU-only box the first pass alone runs past it, so gating iteration 1 would ship every answer with no faithfulness score and no confidence | +| **Deadline reserve** | `AGENT_EVAL_RESERVE_S` — skips the evaluation entirely when too little remains under `AGENT_TIMEOUT` to finish it, since being killed mid-evaluation discards the draft and 504s | +| **Per-call timeout** | `LLM_REQUEST_TIMEOUT_S` bounds each LLM request. Note it bounds one *attempt*, not the failover chain — three stalled attempts still take ~180s, which the deadline reserve absorbs | | **Stuck-loop detection** | Auto-accepts once completeness stops improving | ### ⚡ Standard RAG mode @@ -556,8 +582,10 @@ Typical query latency (on CPU): |------|---------|-------| | Standard RAG (Strategy A) | ~1–2s | Single-pass | | Standard RAG (Strategy B) | ~3–6s | Includes NLLB translation | -| Agentic RAG (1 reflexion) | ~15–30s | Multi-tool + evaluation (parallel tools) | -| Agentic RAG (max reflexions) | ~60–90s | Bounded by timeout + reflexion budget | +| Agentic RAG (1 reflexion) | ~2–4 min | Measured on a CPU-only box: ~45s retrieval (3 sub-queries, embed + rerank), ~50s generation, ~25s evaluation (NLI 20s / ~23 claims + completeness call). Tools run in parallel but contend for the same CPU | +| Agentic RAG (max reflexions) | bounded by `AGENT_TIMEOUT` | Loop budget stops further cycles; the draft is returned rather than discarded | + +On CPU the agent is dominated by the cross-encoders, not the LLM. The levers, cheapest first: lower `AGENT_MAX_SUB_QUERIES` (near-linear — the parallel retrievals thrash one CPU), add keys to `LLM_API_KEYS` so an exhausted model doesn't cost a failed call plus fallback on every step, and drop `NLI_MAX_CHUNKS_PER_CITATION` to `1`. A GPU removes most of this. Memory: base ~500MB · +BGE-M3 ~2.5GB · +reranker ~3.5GB · +NLLB (Strategy B) ~6GB. ColBERT rerank adds ~1GB when enabled. @@ -586,6 +614,10 @@ See [docs/evaluation.md](docs/evaluation.md) for methodology. **Agent answers truncated** — raise `AGENT_MAX_TOKENS` (e.g. `16384`) +**"Agent pipeline timed out"** — every node logs its wall time as `[Graph] took Ns`; read those before changing knobs. On CPU the usual culprit is retrieval or the NLI pass, not the LLM. If a single LLM call hangs, lower `LLM_REQUEST_TIMEOUT_S` so failover fires sooner — but not below ~60s, since agent answer generation is a unary call measured at 20–50s on CPU and would start aborting legitimately. Raising `AGENT_TIMEOUT` only delays the error + +**Faithfulness reads ~0 on every answer** — the threshold is calibrated per NLI model. If you changed `NLI_MODEL_NAME`, recalibrate `FAITHFULNESS_THRESHOLD`: score a sentence copied verbatim out of its own chunk (positive) against one from an unrelated paper (negative) over ~20 chunks, and put the threshold between the two distributions. A bar above every positive silently reports everything as ungrounded + **"Translation model gated"** — NLLB-200 needs no auth; first use downloads ~2.4GB automatically **Sessions lost on restart** — check `SESSIONS_DB_PATH` is writable; SQLite persistence is on by default diff --git a/agent/graph.py b/agent/graph.py index 6cec2e6..0d7aa52 100644 --- a/agent/graph.py +++ b/agent/graph.py @@ -1,3 +1,6 @@ +import logging +import time + from langgraph.graph import StateGraph, END from agent.state import AgentState @@ -9,6 +12,21 @@ from agent.nodes.finalizer import finalizer_node +logger = logging.getLogger(__name__) + + +def _timed(name, fn): + """Log each node's wall time so a run that hits AGENT_TIMEOUT says which node + ate the budget instead of leaving it to inference.""" + def wrapped(state): + t0 = time.monotonic() + try: + return fn(state) + finally: + logger.info("[Graph] %s took %.1fs", name, time.monotonic() - t0) + return wrapped + + def _route_reflexion(state: AgentState) -> str: count = state.get("reflexion_count", 0) if count >= MAX_REFLEXION or state.get("final_answer"): @@ -29,12 +47,15 @@ def _route_reflexion(state: AgentState) -> str: def build_agent_graph(): wf = StateGraph(AgentState) - wf.add_node("query_planner", query_planner_node) - wf.add_node("tool_selector", tool_selector_node) - wf.add_node("tool_executor", tool_executor_node) - wf.add_node("answer_generator", answer_generator_node) - wf.add_node("reflexion_evaluator", reflexion_evaluator_node) - wf.add_node("finalizer", finalizer_node) + for name, fn in ( + ("query_planner", query_planner_node), + ("tool_selector", tool_selector_node), + ("tool_executor", tool_executor_node), + ("answer_generator", answer_generator_node), + ("reflexion_evaluator", reflexion_evaluator_node), + ("finalizer", finalizer_node), + ): + wf.add_node(name, _timed(name, fn)) wf.set_entry_point("query_planner") wf.add_edge("query_planner", "tool_selector") diff --git a/agent/nodes/answer_generator.py b/agent/nodes/answer_generator.py index 0af2cf7..346bf36 100644 --- a/agent/nodes/answer_generator.py +++ b/agent/nodes/answer_generator.py @@ -4,6 +4,7 @@ import rag import config +import llm_client from agent.state import AgentState logger = logging.getLogger(__name__) @@ -63,7 +64,7 @@ def answer_generator_node(state: AgentState) -> dict: system_instruction=config.AGENT_SYSTEM_PROMPT, safety_settings=config.SAFETY_SETTINGS, # Thinking off by default so the full budget goes to the answer (config knob). - thinking_config=types.ThinkingConfig(thinking_budget=config.AGENT_THINKING_BUDGET), + thinking_config=llm_client.thinking_config_for("agent"), ) # The user's model choice reaches every other node (planner, selector, # evaluator) via state — the node that actually writes the answer must @@ -79,4 +80,7 @@ def answer_generator_node(state: AgentState) -> dict: return {"draft_answer": "The AI model is temporarily unavailable. Please try again."} logger.info(f"[AnswerGenerator] chunks_used={chunks_used}, ans_len={len(answer)}") - return {"draft_answer": answer} + # Surfaced so citation resolution can number only the chunks the model was + # actually shown — format_context truncates, and a marker invented past that + # point would otherwise resolve to a real paper that never reached the prompt. + return {"draft_answer": answer, "context_chunks_used": chunks_used} diff --git a/agent/nodes/finalizer.py b/agent/nodes/finalizer.py index f078d45..453e5bf 100644 --- a/agent/nodes/finalizer.py +++ b/agent/nodes/finalizer.py @@ -9,9 +9,11 @@ # citation-coverage denominator so they don't dilute the figure. _MIN_SENT_CHARS = 20 -# Faithfulness bar above which we trust what the answer *does* say. Matches the -# reflexion evaluator's "accept" faithfulness gate. -_ABSTAIN_FAITH_MIN = 0.75 +# Faithfulness bar above which we trust what the answer *does* say. Shares the +# reflexion evaluator's "accept" gate (config.AGENT_FAITHFULNESS_ACCEPT) so both +# move together when the NLI model or its threshold is recalibrated — the old +# hardcoded 0.75 sat above the score a fully grounded answer can reach, so +# abstention never fired. _ABSTAIN_PREFIX = ( "**Insufficient evidence in the corpus to fully answer this question.** " @@ -57,7 +59,7 @@ def finalizer_node(state: AgentState) -> dict: # Abstention: grounded but incomplete after the reflexion budget is spent — the # answer we have is trustworthy, the corpus just doesn't cover the rest. (Low # faithfulness is a different failure, already caveated by the reflexion node.) - if faith >= _ABSTAIN_FAITH_MIN and comp < config.ABSTAIN_COMPLETENESS_FLOOR: + if faith >= config.AGENT_FAITHFULNESS_ACCEPT and comp < config.ABSTAIN_COMPLETENESS_FLOOR: missing = last.get("missing_aspects") or ["parts of the question"] gaps = "\n".join(f"- {m}" for m in missing) answer = f"{_ABSTAIN_PREFIX}{base_answer}\n\n---\n*Not supported by the corpus:*\n{gaps}" diff --git a/agent/nodes/query_planner.py b/agent/nodes/query_planner.py index 89cedf5..80a8c90 100644 --- a/agent/nodes/query_planner.py +++ b/agent/nodes/query_planner.py @@ -102,11 +102,11 @@ def query_planner_node(state: AgentState) -> dict: max_output_tokens=1024, system_instruction=_DECOMPOSE_SYSTEM, # Structured JSON decomposition — thinking off by default (config knob). - thinking_config=types.ThinkingConfig(thinking_budget=config.AGENT_THINKING_BUDGET), + thinking_config=llm_client.thinking_config_for("agent"), ), provider=_provider, ) - raw_resp = resp.text or "" + raw_resp = rag.safe_extract_text(resp) active_provider = llm_client.resolve_provider(_model, _provider) @@ -115,11 +115,11 @@ def _gemini_retry(p, s): model=config.LLM_MODEL_NAME, contents=p, gen_config=types.GenerateContentConfig( temperature=0, max_output_tokens=1024, system_instruction=s, - thinking_config=types.ThinkingConfig(thinking_budget=config.AGENT_THINKING_BUDGET), + thinking_config=llm_client.thinking_config_for("agent"), ), provider="gemini", ) - return r.text or "" + return rag.safe_extract_text(r) parsed = extract_json_with_gemini_retry( raw_resp, active_provider, _gemini_retry, _prompt, _DECOMPOSE_SYSTEM, diff --git a/agent/nodes/reflexion_evaluator.py b/agent/nodes/reflexion_evaluator.py index b5eb1cf..3095ad1 100644 --- a/agent/nodes/reflexion_evaluator.py +++ b/agent/nodes/reflexion_evaluator.py @@ -81,21 +81,35 @@ def reflexion_evaluator_node(state: AgentState) -> dict: "reflexion_count": count, } - # Time budget: after at least one full evaluation, stop looping if we've spent - # the reflexion budget. Finalises the current draft rather than starting another - # retrieve→generate→verify cycle that would blow past AGENT_TIMEOUT and 504. + # Two time gates, because they protect different things. + # + # AGENT_REFLEXION_BUDGET_S stops LOOPING: past it, don't start another + # retrieve→generate→verify cycle. It must not block the FIRST evaluation — on a + # CPU-only box the first pass alone (retrieval + generation) runs past the budget, + # so gating iteration 1 on it means the answer ships with no faithfulness score, + # no completeness check and no confidence, every single time. + # + # AGENT_EVAL_RESERVE_S is the real deadline guard: only skip the evaluation when + # there isn't room left under AGENT_TIMEOUT to finish it, since being killed + # mid-evaluation discards the draft and 504s. start = state.get("start_time") - if start is not None and count >= 1: + draft = state.get("draft_answer") + if start is not None and draft: elapsed = time.monotonic() - start - if elapsed > config.AGENT_REFLEXION_BUDGET_S: + remaining = config.AGENT_TIMEOUT - elapsed + if remaining < config.AGENT_EVAL_RESERVE_S: + logger.info( + f"[Reflexion] iter={count + 1}/{MAX_REFLEXION} elapsed={elapsed:.0f}s, " + f"{remaining:.0f}s left < {config.AGENT_EVAL_RESERVE_S:.0f}s reserve " + f"→ finalising unverified draft" + ) + return {"final_answer": draft, "reflexion_count": count + 1} + if count >= 1 and elapsed > config.AGENT_REFLEXION_BUDGET_S: logger.info( f"[Reflexion] iter={count + 1}/{MAX_REFLEXION} elapsed={elapsed:.0f}s " f"> budget {config.AGENT_REFLEXION_BUDGET_S:.0f}s → finalising best draft" ) - return { - "final_answer": state.get("draft_answer", "Unable to produce a satisfactory answer."), - "reflexion_count": count + 1, - } + return {"final_answer": draft, "reflexion_count": count + 1} answer = state.get("draft_answer", "") _contexts = state.get("retrieved_contexts", []) @@ -105,8 +119,13 @@ def reflexion_evaluator_node(state: AgentState) -> dict: chunk_metas = [{"title": c.get("title", "Unknown"), "section": c.get("section", "body")} for c in _contexts] + _nli_t0 = time.monotonic() try: claims = verify.check_claims(answer, chunks, chunk_metas) + logger.info( + "[Reflexion] NLI scored %d claims in %.1fs", + len(claims), time.monotonic() - _nli_t0, + ) if claims: # Grounded fraction (RAGAS-style): min() collapsed to ~0 on any long # multi-claim answer because one synthesized/comparative sentence @@ -138,7 +157,7 @@ def reflexion_evaluator_node(state: AgentState) -> dict: temperature=0, max_output_tokens=1024, # JSON completeness verdict — thinking off by default (config knob). - thinking_config=types.ThinkingConfig(thinking_budget=config.AGENT_THINKING_BUDGET), + thinking_config=llm_client.thinking_config_for("agent"), ), provider=_provider, ) @@ -150,8 +169,8 @@ def _gemini_retry(p, s): r = rag.generate_with_failover( model=config.LLM_MODEL_NAME, contents=p, gen_config=types.GenerateContentConfig( - temperature=0, max_output_tokens=1024, - thinking_config=types.ThinkingConfig(thinking_budget=config.AGENT_THINKING_BUDGET), + temperature=0, max_output_tokens=1024, system_instruction=s, + thinking_config=llm_client.thinking_config_for("agent"), ), provider="gemini", ) @@ -165,7 +184,7 @@ def _gemini_retry(p, s): if not claims: action = parsed.get("action", "retrieve_more") - elif faithfulness_score >= 0.75 and completeness_score >= 0.75: + elif faithfulness_score >= config.AGENT_FAITHFULNESS_ACCEPT and completeness_score >= 0.75: action = "accept" else: action = parsed.get("action", "retrieve_more") @@ -176,7 +195,8 @@ def _gemini_retry(p, s): f"| raw={raw_text[:300]!r}" ) completeness_score, missing = 0.5, [] - action = "regenerate" if faithfulness_score >= 0.75 else "retrieve_more" + action = ("regenerate" if faithfulness_score >= config.AGENT_FAITHFULNESS_ACCEPT + else "retrieve_more") feedback = ReflexionFeedback( faithfulness_score=faithfulness_score, @@ -191,7 +211,7 @@ def _gemini_retry(p, s): if prev and action != "accept": prev_complete = prev[-1].get("completeness_score", 0.0) if completeness_score <= prev_complete + 0.05 and count >= 1: - if faithfulness_score < 0.75: + if faithfulness_score < config.AGENT_FAITHFULNESS_ACCEPT: missing_str = ", ".join(missing) or "the requested details" logger.info( f"[Reflexion] iter={count + 1}/{MAX_REFLEXION} " @@ -233,4 +253,23 @@ def _gemini_retry(p, s): "reflexion_history": history, } + # Post-evaluation budget check. The gate at the top of this node deliberately + # lets iteration 1 evaluate even when already over budget — but a retry verdict + # would then send the graph into a full retrieve→generate cycle (~95s on CPU) + # that the budget exists to prevent, and only the NEXT entry here would stop it. + # Returning without a final_answer is what allows that, so finalise instead. + if start is not None: + elapsed = time.monotonic() - start + if elapsed > config.AGENT_REFLEXION_BUDGET_S: + logger.info( + f"[Reflexion] iter={count + 1}/{MAX_REFLEXION} action={action} but " + f"elapsed={elapsed:.0f}s > budget " + f"{config.AGENT_REFLEXION_BUDGET_S:.0f}s → finalising instead of retrying" + ) + return { + "final_answer": answer, + "reflexion_count": count + 1, + "reflexion_history": history, + } + return {"reflexion_count": count + 1, "reflexion_history": history} diff --git a/agent/nodes/tool_selector.py b/agent/nodes/tool_selector.py index c3e4cfb..f914f5f 100644 --- a/agent/nodes/tool_selector.py +++ b/agent/nodes/tool_selector.py @@ -64,11 +64,11 @@ def _gate_model(state) -> tuple[str, str]: indicrag_retrieval with arxiv_search or open_access_search. RETRY RULES: -7. retrieve_more: Craft SHARPER queries using missing_aspects from the evaluator. \ +5. retrieve_more: Craft SHARPER queries using missing_aspects from the evaluator. \ Never repeat the original query verbatim. Re-use year_from from state. -8. reformulate: The query was misunderstood — build a corrected query \ +6. reformulate: The query was misunderstood — build a corrected query \ from missing_aspects before selecting tools. -9. regenerate: Context is adequate; answer needs rewriting. \ +7. regenerate: Context is adequate; answer needs rewriting. \ Return an EMPTY tool list so the answer generator runs without re-retrieval.\ """ @@ -146,7 +146,7 @@ def tool_selector_node(state: AgentState) -> dict: function_calling_config=types.FunctionCallingConfig(mode="AUTO") ), # Rule-based tool routing — thinking off by default (config knob). - thinking_config=types.ThinkingConfig(thinking_budget=config.AGENT_THINKING_BUDGET), + thinking_config=llm_client.thinking_config_for("agent"), ) gate_provider, gate_model = _gate_model(state) diff --git a/agent/state.py b/agent/state.py index 84f4d0d..e7d2968 100644 --- a/agent/state.py +++ b/agent/state.py @@ -21,6 +21,12 @@ class AgentState(TypedDict): draft_answer: Optional[str] final_answer: Optional[str] + # How many of retrieved_contexts actually reached the answer prompt — + # format_context truncates by chunk count and by length. Citation resolution + # numbers only this slice, so a marker invented past it dangles (and is + # dropped) instead of resolving to a paper the model never saw. + context_chunks_used: Optional[int] + reflexion_count: int reflexion_history: List[ReflexionFeedback] diff --git a/agent/tool_executor.py b/agent/tool_executor.py index 17b4b03..7dbfbbb 100644 --- a/agent/tool_executor.py +++ b/agent/tool_executor.py @@ -17,6 +17,7 @@ import rag import config +import llm_client logger = logging.getLogger(__name__) _tavily = None @@ -55,7 +56,7 @@ def _expand_query_variants(query: str) -> list[str]: max_output_tokens=256, system_instruction="Generate alternative search phrasings that preserve the original query's semantic meaning. Do not add new topics or narrow the scope.", # Short JSON list of paraphrases — thinking off by default (config knob). - thinking_config=types.ThinkingConfig(thinking_budget=config.AGENT_THINKING_BUDGET), + thinking_config=llm_client.thinking_config_for("agent"), ), ) clean = re.sub(r"```(?:json)?|```", "", resp.text or "").strip() diff --git a/config.py b/config.py index b5e2da1..d3a4b7f 100644 --- a/config.py +++ b/config.py @@ -190,7 +190,19 @@ def ensure_directories(): # ============================================================================ # Faithfulness Verification # ============================================================================ -FAITHFULNESS_THRESHOLD = float(os.getenv("FAITHFULNESS_THRESHOLD", "0.5")) +# Per-claim entailment probability above which a claim counts as grounded. +# Calibrated on this corpus with the shipped int8 mDeBERTa-xnli model: a sentence +# copied VERBATIM out of its own chunk scores median 0.226 / max 0.428, while a +# sentence from an unrelated paper scores p90 0.158. The old 0.5 was unreachable — +# recall 0.00 on guaranteed-grounded claims, so faithfulness read ~0 for every +# answer ever produced. 0.15 gives recall 0.70 at false-positive 0.10. +# Recalibrate (positives vs cross-paper negatives) if NLI_MODEL_NAME changes. +FAITHFULNESS_THRESHOLD = float(os.getenv("FAITHFULNESS_THRESHOLD", "0.15")) +# Fraction of an answer's claims that must be grounded for the reflexion loop to +# accept it, and for the finalizer to trust the answer enough to abstain on +# completeness alone. Sits below 1.0 by design: per-claim recall is 0.70, so even a +# fully grounded answer lands near 0.70 — the old hardcoded 0.75 could never fire. +AGENT_FAITHFULNESS_ACCEPT = float(os.getenv("AGENT_FAITHFULNESS_ACCEPT", "0.6")) FAITHFULNESS_ENFORCE = os.getenv("FAITHFULNESS_ENFORCE", "warn") # warn | strip | regen # NLI model for claim faithfulness. Default is MULTILINGUAL so Indic-language @@ -211,6 +223,18 @@ def ensure_directories(): # mDeBERTa-xnli (default): 2=contradiction # cross-encoder/nli-deberta-v3-base: 0=contradiction NLI_CONTRADICTION_INDEX = int(os.getenv("NLI_CONTRADICTION_INDEX", "2")) +# NLI cost is linear in pairs and in premise length — measured on this CPU box with +# the int8 ONNX model: 1.15s/pair at 512 tokens, 0.4s/pair at 256. A 30-sentence +# answer citing multi-chunk papers hit ~275 pairs = 318s in one reflexion pass. +# These two knobs bound that: shorter premise, and at most N chunks per cited paper +# (chunks arrive rerank-ordered, so the first ones are the best support anyway). +NLI_MAX_SEQ_LENGTH = int(os.getenv("NLI_MAX_SEQ_LENGTH", "256")) +# Clamped to >=1 because this value FAILS OPEN: a 0 (or negative) cap slices away +# every cited chunk, check_claims() then returns no claims, and the reflexion +# evaluator reads an empty claim list as "no citable claims != hallucination" — +# faithfulness 1.0, answer accepted with zero grounding. A typo in .env would +# silently disable verification while reporting perfect scores. +NLI_MAX_CHUNKS_PER_CITATION = max(1, int(os.getenv("NLI_MAX_CHUNKS_PER_CITATION", "2"))) # ============================================================================ # Vector Store @@ -257,7 +281,11 @@ def ensure_directories(): # LLM Configuration # ============================================================================ # Google Gemini API configuration -LLM_MAX_TOKENS = int(os.getenv("LLM_MAX_TOKENS", "2048")) # maximum tokens to generate +# Caps thinking + answer together, not just the answer. gemini-3.6-flash rejects +# thinking_budget=0 and spends 0-4856 thought tokens on identical prompts, so a +# 2048 cap left as little as 80 tokens for the answer and truncated mid-sentence. +# Measured: answer <=2100, worst thinking+answer 6926. +LLM_MAX_TOKENS = int(os.getenv("LLM_MAX_TOKENS", "8192")) # maximum tokens to generate AGENT_MAX_TOKENS = int(os.getenv("AGENT_MAX_TOKENS", "8192")) # higher limit for agentic pipeline # Thinking budget for ALL agentic-mode LLM calls (query planner, tool routing, # query expansion, answer generation, reflexion judge). Gemini semantics: @@ -265,13 +293,56 @@ def ensure_directories(): # -1 = DYNAMIC (model decides how much to think) # N = cap thinking to N tokens (billed; higher = smarter routing/judging, pricier) # Raise this only if agent answer/routing quality is the bottleneck, not the bill. +# LEGACY on Gemini 3.x: those models reject thinking_budget outright (see the level +# knobs below, which supersede it). Still honoured by models that accept budgets. AGENT_THINKING_BUDGET = int(os.getenv("AGENT_THINKING_BUDGET", "0")) -AGENT_TIMEOUT = int(os.getenv("AGENT_TIMEOUT", "120")) # seconds; CPU embedding can take 45s+ +# Thinking LEVEL — the Gemini 3.x control, replacing thinking_budget. Google's docs +# list minimal | low | medium | high for gemini-3.6-flash, defaulting to MEDIUM when +# nothing is sent. That default is the trap: sending the legacy thinking_budget=0 gets +# a 400, the backend used to drop the field entirely, and the model then thought at +# MEDIUM — the opposite of the "thinking off" that was asked for, with those thought +# tokens coming out of LLM_MAX_TOKENS and squeezing the answer. +# minimal = least thinking (default here; closest to the old budget=0 intent) +# low | medium | high = progressively more (slower, pricier, sometimes better) +# "" = send nothing, i.e. accept the model's own default +LLM_THINKING_LEVEL = os.getenv("LLM_THINKING_LEVEL", "minimal").strip().lower() +AGENT_THINKING_LEVEL = os.getenv("AGENT_THINKING_LEVEL", "minimal").strip().lower() +# Seconds. Measured end-to-end on a CPU-only box: ~45s retrieval + ~50s generation +# + ~25s evaluation. The old 120s default left under 30s of room once +# AGENT_EVAL_RESERVE_S was set aside, so the evaluator skipped verification on +# every stock-config run — the answer shipped with no faithfulness score at all. +AGENT_TIMEOUT = int(os.getenv("AGENT_TIMEOUT", "300")) +# Per-request HTTP timeout for non-streaming LLM calls. Without it the SDK defaults +# apply (OpenAI: 600s x 2 retries) and ONE stalled request outlasts the whole agent +# budget. 60s is sized for a legitimate call, not for the failover chain: agent +# answer generation is unary and measured 20-50s on CPU, so a materially lower +# value would abort real generations rather than stalled ones. +# +# The chain is therefore NOT bounded by AGENT_REFLEXION_BUDGET_S: generate_with_failover +# walks up to 3 (provider, model) attempts sequentially, so a fully-stalled chain can +# reach ~180s — past the 90s reflexion budget. What actually bounds it is AGENT_TIMEOUT +# plus AGENT_EVAL_RESERVE_S, which finalise the draft rather than 504. In practice the +# per-(provider, model) circuit breaker in llm_client skips recently-dead paths, so +# three consecutive full stalls are rare. If that worst case ever matters more than +# generation headroom, make the timeout deadline-aware (remaining budget / attempts +# left) instead of just lowering it. +LLM_REQUEST_TIMEOUT_S = int(os.getenv("LLM_REQUEST_TIMEOUT_S", "60")) +# Streaming needs its own, much larger budget: for Gemini the HTTP timeout covers +# the WHOLE stream, not the gap between chunks, so reusing the 60s unary value tore +# down long answers mid-generation (WinError 10054, truncated text). This still +# bounds a genuinely stuck stream without capping legitimate long generations. +LLM_STREAM_TIMEOUT_S = int(os.getenv("LLM_STREAM_TIMEOUT_S", "300")) # Wall-clock budget for the reflexion loop. Once exceeded, the evaluator finalizes # the current best draft instead of starting another retrieve→generate→verify cycle, # so the user gets an answer rather than a hard AGENT_TIMEOUT 504 that discards all # work. Keep below AGENT_TIMEOUT so it fires first. AGENT_REFLEXION_BUDGET_S = float(os.getenv("AGENT_REFLEXION_BUDGET_S", "90")) +# Wall-clock room that must remain under AGENT_TIMEOUT for the evaluator to attempt +# an evaluation at all. Unlike the budget above (which only stops FURTHER loops), +# this one can skip iteration 1 — but only when finishing would overrun the timeout +# and discard the draft entirely. Sized for one NLI pass plus one completeness LLM +# call: measured ~30s + ~15s on CPU, doubled for headroom. +AGENT_EVAL_RESERVE_S = float(os.getenv("AGENT_EVAL_RESERVE_S", "90")) # Max sub-queries the planner emits (and tools run per cycle). Each sub-query does a # retrieve + CPU reranker pass (~15 pairs); those passes are CPU-bound so N concurrent # ones thrash rather than parallelize. Over a small corpus, 3 covers most queries at a diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 262f331..ccc600c 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -417,10 +417,21 @@ vendor that just failed. `_fallback_model_for` therefore selects the first `/`-shaped slug in `LLM_SELECTABLE_MODELS`; keep at least one such slug in that list. -Some Gemini models reject `thinking_config` with `thinking_budget=0` and return -`400 INVALID_ARGUMENT`. `GeminiBackend` records the rejection per model and -retries once without `thinking_config`; streaming only retries if nothing has -been emitted yet. +Gemini 3.x models reject the legacy `thinking_budget` field with +`400 INVALID_ARGUMENT` — it is superseded by `thinking_level`. `GeminiBackend` +records the rejection per model and retries once with the budget translated to +the equivalent level (`0` → `MINIMAL`, `<=1024` → `LOW`, higher → `MEDIUM`, +`-1` → field dropped, i.e. the model decides); streaming only retries if nothing +has been emitted yet. + +Dropping `thinking_config` is *not* a neutral fallback, which is why the retry +translates instead: omitting the field means the model's own default, and +`gemini-3.6-flash` defaults to `MEDIUM`. Those thought tokens are billed and +come out of `max_output_tokens`, so "thinking off" silently became medium +thinking with a squeezed answer. `llm_client.thinking_config_for(scope)` builds +the config up front from `LLM_THINKING_LEVEL` / `AGENT_THINKING_LEVEL` +(default `minimal`), falling back to a budget only when the level is empty or +the installed `google-genai` has no `ThinkingLevel` enum. ### Memory Usage diff --git a/docs/GEMINI_SETUP.md b/docs/GEMINI_SETUP.md index 94a18d9..7c5bfda 100644 --- a/docs/GEMINI_SETUP.md +++ b/docs/GEMINI_SETUP.md @@ -99,13 +99,27 @@ Rough guidance rather than a fixed list, since model names change often: the cost and latency. - **Non-Gemini via OpenRouter**: useful as a genuinely independent failover leg. -### Thinking budget +### Thinking level (and the legacy budget) -Gemini "thinking" models accept `thinking_budget`: `0` disables it, `-1` lets -the model decide, a positive integer caps it. Some models reject -`thinking_budget=0` with a 400 — `providers/gemini.py` detects that and retries -once without `thinking_config` (streaming only retries if nothing was emitted -yet). +Gemini 3.x models take a **thinking level**: `minimal`, `low`, `medium`, or +`high`. Set it with `LLM_THINKING_LEVEL` (standard RAG) and +`AGENT_THINKING_LEVEL` (agentic pipeline); both default to `minimal`. + +```bash +LLM_THINKING_LEVEL=minimal # minimal | low | medium | high, or empty for the model default +AGENT_THINKING_LEVEL=minimal +``` + +Leaving these empty is a real choice, not a neutral one: with no level sent, +`gemini-3.6-flash` thinks at `medium`. Those thought tokens are billed and are +drawn from `LLM_MAX_TOKENS`, so the answer gets less room. + +The older `AGENT_THINKING_BUDGET` (`0` off, `-1` model decides, `N` caps) still +works on models that accept budgets. Gemini 3.x rejects it with a 400 — +`providers/gemini.py` detects that per model and retries once with the budget +translated to the closest level (`0` → `MINIMAL`, `<=1024` → `LOW`, higher → +`MEDIUM`, `-1` → omit the field so the model decides). Streaming only retries if +nothing was emitted yet. --- @@ -130,7 +144,7 @@ Set these in `.env` (defaults live in `config.py`): LLM_MODEL_NAME=gemini-3.6-flash LLM_SELECTABLE_MODELS=gemini-3.6-flash,gemini-3.5-flash,anthropic/claude-haiku -LLM_MAX_TOKENS=2048 # Maximum response length +LLM_MAX_TOKENS=8192 # Maximum response length (thinking + answer share this) LLM_TEMPERATURE=0.3 # Lower = more factual, higher = more creative ``` diff --git a/figure_captioner.py b/figure_captioner.py index 84335e8..0f1a441 100644 --- a/figure_captioner.py +++ b/figure_captioner.py @@ -139,6 +139,10 @@ def _caption_one(region: Dict[str, Any]) -> str: import rag gen_config = types.GenerateContentConfig( + system_instruction=( + "You are a scientific figure and table analyst: describe what the image " + "concretely shows, factually and without preamble." + ), temperature=0.0, max_output_tokens=512, safety_settings=config.SAFETY_SETTINGS, diff --git a/llm_client.py b/llm_client.py index e5417ad..8152bc5 100644 --- a/llm_client.py +++ b/llm_client.py @@ -156,6 +156,31 @@ def generate_with_failover(model: str, contents, gen_config, provider: str | Non raise last_exc # type: ignore[misc] +def thinking_config_for(scope: str = "standard"): + """ThinkingConfig for a call scope ("standard" or "agent"), or None to send nothing. + + Prefers the Gemini 3.x thinking_level knob and falls back to the legacy + thinking_budget when the configured level is empty or the installed SDK has no + ThinkingLevel enum. Returning None means "omit the field", which lets the model + apply its own default — MEDIUM on gemini-3.6-flash, so it is a real choice, not + a neutral one. + """ + from google.genai import types + + level_name = (_config.AGENT_THINKING_LEVEL if scope == "agent" + else _config.LLM_THINKING_LEVEL) + if level_name: + level = get_backend("gemini")._thinking_level(level_name) + if level is not None: + return types.ThinkingConfig(thinking_level=level) + logger.warning( + "Unknown thinking level %r for scope %s — falling back to thinking_budget", + level_name, scope, + ) + budget = _config.AGENT_THINKING_BUDGET if scope == "agent" else 0 + return types.ThinkingConfig(thinking_budget=budget) + + def _build_gemini_stream_config(model, max_tokens, system_instruction): from google.genai import types kwargs = dict( @@ -165,7 +190,7 @@ def _build_gemini_stream_config(model, max_tokens, system_instruction): system_instruction=system_instruction or _config.SYSTEM_PROMPT, ) if get_backend("gemini").supports_thinking(model): - kwargs["thinking_config"] = types.ThinkingConfig(thinking_budget=0) + kwargs["thinking_config"] = thinking_config_for("standard") return types.GenerateContentConfig(**kwargs) @@ -200,15 +225,27 @@ def llm_generate_stream(prompt: str, max_tokens: int = None, system_instruction: gen_config = _build_openrouter_stream_config(max_tokens, system_instruction) any_attempted = True emitted = False + chars = 0 + started = time.monotonic() try: for chunk in backend.generate_stream(mdl, prompt, gen_config): emitted = True + chars += len(chunk) yield chunk _circuit_clear(key) return except Exception as exc: last_exc = exc if emitted: + # A mid-stream death can't be retried (the client already holds the + # prefix), so log what tells the causes apart: elapsed near + # LLM_STREAM_TIMEOUT_S means our own timeout cut it; elapsed well + # under it means the provider dropped the connection. + logger.warning( + "[stream] %s:%s died after %.0fs and %d chars (limit %ds) — %s: %s", + prov, mdl, time.monotonic() - started, chars, + _config.LLM_STREAM_TIMEOUT_S, type(exc).__name__, str(exc)[:200], + ) raise # committed to this stream if backend.is_permanent(exc): raise diff --git a/providers/base.py b/providers/base.py index a8908a4..bf7a967 100644 --- a/providers/base.py +++ b/providers/base.py @@ -8,6 +8,14 @@ from abc import ABC, abstractmethod from typing import Iterator +# Appended to a stream that stopped because it ran out of output tokens. Backends +# end such streams normally, so every layer above sees a clean finish and the text +# is the only channel left to tell the user the answer is incomplete. +TRUNCATION_NOTE = ( + "\n\n*[Answer truncated — output token limit reached. " + "Ask a narrower question or raise `LLM_MAX_TOKENS`.]*" +) + class _FunctionCall: def __init__(self, name: str, args: dict): diff --git a/providers/gemini.py b/providers/gemini.py index 3c2066e..823d420 100644 --- a/providers/gemini.py +++ b/providers/gemini.py @@ -12,7 +12,8 @@ import config from google import genai -from providers.base import LLMBackend +from google.genai import types as genai_types +from providers.base import LLMBackend, TRUNCATION_NOTE logger = logging.getLogger(__name__) @@ -20,6 +21,7 @@ class GeminiBackend(LLMBackend): def __init__(self): self._pool: list[genai.Client] = [] + self._stream_pool: list[genai.Client] = [] self._lock = threading.Lock() self._index = itertools.cycle([]) @@ -34,7 +36,22 @@ def _init_pool(self) -> None: "Google Gemini API key not configured. " "Set LLM_API_KEY (single) or LLM_API_KEYS (comma-separated) in .env." ) - self._pool = [genai.Client(api_key=k) for k in config.LLM_API_KEY_POOL] + # Explicit per-request timeout (google-genai takes milliseconds). The SDK + # default is generous enough that one stalled call outlives the agent's whole + # budget and surfaces as "Agent pipeline timed out" instead of failing over. + # + # Streaming gets a SEPARATE, much larger budget because this timeout covers + # the whole stream rather than the gap between chunks — sharing the unary + # value tore down long answers mid-generation (WinError 10054, truncated text). + def _client(key, seconds): + return genai.Client( + api_key=key, + http_options=genai_types.HttpOptions(timeout=seconds * 1000), + ) + + self._pool = [_client(k, config.LLM_REQUEST_TIMEOUT_S) for k in config.LLM_API_KEY_POOL] + self._stream_pool = [_client(k, config.LLM_STREAM_TIMEOUT_S) + for k in config.LLM_API_KEY_POOL] self._index = itertools.cycle(range(len(self._pool))) def _ensure_pool(self) -> None: @@ -48,6 +65,12 @@ def next_client_idx(self) -> int: with self._lock: return next(self._index) + @property + def stream_pool(self) -> list: + """Clients whose HTTP timeout fits a whole streamed generation.""" + self._ensure_pool() + return self._stream_pool + @property def pool(self) -> list: self._ensure_pool() @@ -73,17 +96,47 @@ def supports_thinking(self, model: str) -> bool: # Models that reject thinking_budget=0 outright (gemini-3.6-flash returns # 400 INVALID_ARGUMENT; gemini-3.5-flash accepts it). Learned at runtime and # remembered per model, so a new model generation doesn't need a hardcoded - # list here — omitting thinking_config is the closest thing to "no thinking" - # those models accept. + # list here. # Class-level on purpose (the learning is about the model, not the instance), # so concurrent requests must not race on the check-then-add. _zero_budget_rejected: set = set() _zero_budget_lock = threading.Lock() + # Budget-to-level translation for models on the newer thinking API. Dropping + # thinking_config was the old fallback, but "send nothing" means the model's + # OWN default — MEDIUM on gemini-3.6-flash — so asking for no thinking produced + # medium thinking, billed and taken out of the answer's token budget. + _LOW_BUDGET_CEILING = 1024 # above this a positive budget reads as MEDIUM + + @staticmethod + def _thinking_level(name: str): + """SDK ThinkingLevel by name, or None if unsupported/unknown. + + Guarded so a pinned older google-genai without the enum keeps working. + """ + enum = getattr(genai_types, "ThinkingLevel", None) + if enum is None or not name: + return None + return getattr(enum, name.upper(), None) + + @classmethod + def _level_for_budget(cls, budget): + """Legacy thinking_budget → the equivalent ThinkingLevel, or None to omit. + + 0 is "off", and MINIMAL is the least thinking the new API offers. -1 is + "model decides", which is exactly what omitting the field already means. + """ + if budget is None or budget < 0: + return None + if budget == 0: + return cls._thinking_level("MINIMAL") + return cls._thinking_level("LOW" if budget <= cls._LOW_BUDGET_CEILING else "MEDIUM") + @staticmethod - def _has_zero_thinking_budget(gen_config) -> bool: + def _has_thinking_budget(gen_config) -> bool: + """True if the config carries a legacy thinking_budget of any value.""" tc = getattr(gen_config, "thinking_config", None) - return tc is not None and getattr(tc, "thinking_budget", None) == 0 + return tc is not None and getattr(tc, "thinking_budget", None) is not None @staticmethod def _is_invalid_argument(exc: Exception) -> bool: @@ -107,31 +160,50 @@ def _with_cache(self, client, model: str, gen_config): logger.debug("Gemini context caching skipped: %s", exc) return gen_config + def _translate_thinking(self, call_config): + """Re-express a thinking_budget this model rejects as a thinking_level. + + Falls back to dropping thinking_config when no level applies (budget -1, + or an SDK too old to expose the enum) — that means "model decides". + """ + tc = getattr(call_config, "thinking_config", None) + if tc is None: + return call_config + level = self._level_for_budget(getattr(tc, "thinking_budget", None)) + if level is None: + return call_config.model_copy(update={"thinking_config": None}) + return call_config.model_copy(update={ + "thinking_config": genai_types.ThinkingConfig(thinking_level=level), + }) + def _prep_config(self, client, model, gen_config): call_config = self._with_cache(client, model, gen_config) + if not self.supports_thinking(model): + if getattr(call_config, "thinking_config", None) is not None: + call_config = call_config.model_copy(update={"thinking_config": None}) + return call_config + with self._zero_budget_lock: known_rejected = model in self._zero_budget_rejected - drop_thinking = ( - not self.supports_thinking(model) - or (known_rejected and self._has_zero_thinking_budget(call_config)) - ) - if drop_thinking and getattr(call_config, "thinking_config", None) is not None: - call_config = call_config.model_copy(update={"thinking_config": None}) + if known_rejected and self._has_thinking_budget(call_config): + call_config = self._translate_thinking(call_config) return call_config def _remember_zero_budget_rejection(self, model: str, gen_config, exc: Exception) -> bool: - """True if `exc` is this model refusing thinking_budget=0, so the caller - can retry once without the field. Narrow on purpose: any other 400 is a - genuine bad request and must keep propagating.""" - if not (self._has_zero_thinking_budget(gen_config) and self._is_invalid_argument(exc)): + """True if `exc` is this model refusing a legacy thinking_budget, so the + caller can retry once with the translated thinking_level. Narrow on + purpose: any other 400 is a genuine bad request and must keep + propagating.""" + if not (self._has_thinking_budget(gen_config) and self._is_invalid_argument(exc)): return False with self._zero_budget_lock: first_time = model not in self._zero_budget_rejected self._zero_budget_rejected.add(model) if first_time: logger.warning( - "Model %s rejects thinking_budget=0; retrying without thinking_config " - "and omitting it for this model from now on.", model, + "Model %s rejects thinking_budget (superseded by thinking_level on " + "Gemini 3.x); translating and using a level for this model from now on.", + model, ) return True @@ -144,23 +216,35 @@ def generate(self, model: str, contents, gen_config, client=None): except Exception as exc: if not self._remember_zero_budget_rejection(model, call_config, exc): raise - retry_config = call_config.model_copy(update={"thinking_config": None}) - return client.models.generate_content(model=model, contents=contents, config=retry_config) + return client.models.generate_content( + model=model, contents=contents, + config=self._translate_thinking(call_config), + ) def generate_stream(self, model: str, contents, gen_config, client=None) -> Iterator[str]: - client = client or self.pool[self.next_client_idx()] + client = client or self.stream_pool[self.next_client_idx()] call_config = self._prep_config(client, model, gen_config) emitted = False def _iter(cfg): nonlocal emitted + finish = None for chunk in client.models.generate_content_stream(model=model, contents=contents, config=cfg): + cands = getattr(chunk, "candidates", None) + if cands and getattr(cands[0], "finish_reason", None): + finish = cands[0].finish_reason try: if chunk.text: emitted = True yield chunk.text except (ValueError, AttributeError) as exc: logger.debug("Skipping non-text Gemini stream chunk: %s", exc) + # max_output_tokens caps thinking + answer together, and this model's + # thinking budget is dynamic, so the cut is unpredictable. The stream + # still ends normally — say so in the text or nobody ever finds out. + if emitted and "MAX_TOKENS" in str(finish): + logger.warning("Gemini stream hit max_output_tokens for %s — answer truncated", model) + yield TRUNCATION_NOTE try: yield from _iter(call_config) @@ -170,7 +254,7 @@ def _iter(cfg): # error is something else and must propagate. if emitted or not self._remember_zero_budget_rejection(model, call_config, exc): raise - yield from _iter(call_config.model_copy(update={"thinking_config": None})) + yield from _iter(self._translate_thinking(call_config)) if not emitted: raise RuntimeError("No text generated from Gemini stream") diff --git a/providers/openrouter.py b/providers/openrouter.py index 4acb309..9147528 100644 --- a/providers/openrouter.py +++ b/providers/openrouter.py @@ -11,7 +11,7 @@ from typing import Iterator import config -from providers.base import LLMBackend, ShimResponse +from providers.base import LLMBackend, ShimResponse, TRUNCATION_NOTE logger = logging.getLogger(__name__) @@ -95,9 +95,13 @@ def _get_client(self): "OPENROUTER_API_KEY not configured. Set it in .env to use OpenRouter." ) from openai import OpenAI + # SDK defaults are 600s per request with 2 retries — one stalled + # call would outlive the agent budget and 504 the whole run. self._client = OpenAI( api_key=config.OPENROUTER_API_KEY, base_url=config.OPENROUTER_BASE_URL, + timeout=config.LLM_REQUEST_TIMEOUT_S, + max_retries=1, ) return self._client @@ -107,6 +111,11 @@ def _params(self, model, contents, gen_config, stream: bool) -> dict: "messages": _to_messages(contents, gen_config), "stream": stream, } + if stream: + # Per-request override: the streamed request stays open for the whole + # generation, so the client's unary timeout would cut long answers off + # mid-stream. Client default still applies to non-streaming calls. + params["timeout"] = config.LLM_STREAM_TIMEOUT_S temp = getattr(gen_config, "temperature", None) if temp is not None: params["temperature"] = temp @@ -136,9 +145,11 @@ def generate(self, model: str, contents, gen_config): def generate_stream(self, model: str, contents, gen_config) -> Iterator[str]: client = self._get_client() emitted = False + finish = None for chunk in client.chat.completions.create(**self._params(model, contents, gen_config, stream=True)): if not chunk.choices: continue + finish = getattr(chunk.choices[0], "finish_reason", None) or finish delta = chunk.choices[0].delta text = getattr(delta, "content", None) if text: @@ -146,6 +157,11 @@ def generate_stream(self, model: str, contents, gen_config) -> Iterator[str]: yield text if not emitted: raise RuntimeError("No text generated from OpenRouter stream") + # OpenAI's "length" is Gemini's MAX_TOKENS: a clean stream end that hides + # a cut-off answer. Same sentinel so both providers behave alike. + if finish == "length": + logger.warning("OpenRouter stream hit max_tokens for %s — answer truncated", model) + yield TRUNCATION_NOTE def is_transient(self, exc: Exception) -> bool: status = getattr(exc, "status_code", None) or getattr(exc, "code", None) diff --git a/rag.py b/rag.py index 2dafa31..22cdc32 100644 --- a/rag.py +++ b/rag.py @@ -4,6 +4,7 @@ from typing import Dict, List, Optional, Any import logging +import re import config import embeddings import vector_store @@ -50,7 +51,8 @@ def _crop_url(crop_path: str) -> Optional[str]: return None -def extract_citations(answer: str, metadatas: List[Dict], chunks: List[str] = None) -> List[Dict]: +def extract_citations(answer: str, metadatas: List[Dict], chunks: List[str] = None, + visible_chunks: int = None) -> List[Dict]: """ Extract [Cite:N] citations from answer text and resolve them to papers. @@ -61,12 +63,25 @@ def extract_citations(answer: str, metadatas: List[Dict], chunks: List[str] = No answer: Generated answer text containing citations metadatas: List of metadata dictionaries from retrieved chunks chunks: Unused; kept for backwards-compatible call sites + visible_chunks: How many leading chunks actually reached the prompt + (``format_context``'s ``chunks_used``). Callers hold the FULL + retrieved metadata, but format_context truncates by chunk count and + by total length — so without this, a number the model invented past + the truncation point resolves to a real paper it was never shown, + and the answer carries a citation that looks legitimate. Numbering + the visible slice only makes such a marker dangle, and dangling + markers are dropped. ``None`` means "truncation unknown", use all. Returns: List of citation dictionaries with number, title, and section """ import re + if visible_chunks is not None: + metadatas = metadatas[:visible_chunks] + if chunks: + chunks = chunks[:visible_chunks] + seen_nums = set() # Match [N] and comma-separated [N, N, ...] citation markers. # Still ignores ranges like [10-15] mg (no comma, so no match). @@ -116,6 +131,57 @@ def extract_citations(answer: str, metadatas: List[Dict], chunks: List[str] = No return citations +# Inline citation markers: [3], [1, 3, 5], [2,4]. Digit-only, so [NOT FOUND: ...] +# and ranges like [10-15] never match. Leading whitespace is captured separately so +# a fully-dangling marker is dropped together with the space in front of it — and +# the optional newline is included because a marker alone on its own line otherwise +# left the newline behind, which markdown renders as a paragraph break. +_CITE_MARKER_RE = re.compile(r'([ \t]*\n?[ \t]*)\[(\d+(?:\s*,\s*\d+)*)\]') + + +def compact_citations(answer: str, metadatas: List[Dict], chunks: List[str] = None, + visible_chunks: int = None): + """extract_citations, then renumber the survivors to a dense 1..M sequence. + + format_context numbers EVERY retrieved paper, but only the papers the answer + actually cites reach the citation panel — so an answer drawing on papers 1 + and 4 of 4 read "[1] ... [4]" beside a two-entry panel. Renumber the cited + papers in context order and rewrite the answer's markers to match. Markers + that resolve to no paper (the model over-numbered) are dropped rather than + left dangling, mirroring report_runner._remap_markers. + + Returns ``(rewritten answer, citations)`` — the citations carry the new + dense numbers, so callers must use the returned answer, not the original. + """ + citations = extract_citations(answer, metadatas, chunks, visible_chunks) + old_to_new = {int(c['number']): i for i, c in enumerate(citations, 1)} + + def _repl(m: "re.Match") -> str: + mapped: List[int] = [] + for part in m.group(2).split(','): + try: + new = old_to_new.get(int(part.strip())) + except ValueError: + new = None + if new is not None and new not in mapped: + mapped.append(new) + if not mapped: # every number in this marker was dangling + # Drop the preceding space too — no double/trailing space. When the + # marker sat alone on its line, the line's own trailing newline + # survives, so the captured leading newline must go with the marker + # or a blank line is left behind (markdown reads it as a paragraph + # break). Otherwise put it back, or the lines splice together. + if '\n' not in m.group(1): + return '' + rest = m.string[m.end():] + return '' if (rest == '' or rest[0] == '\n') else '\n' + return m.group(1) + '[' + ', '.join(str(n) for n in mapped) + ']' + + for i, c in enumerate(citations, 1): + c['number'] = str(i) + return _CITE_MARKER_RE.sub(_repl, answer), citations + + def _hyde_embedding(user_query: str): """Draft a hypothetical answer and embed it, for HyDE retrieval. @@ -129,7 +195,7 @@ def _hyde_embedding(user_query: str): max_output_tokens=256, safety_settings=config.SAFETY_SETTINGS, # Throwaway hypothetical draft for embedding — thinking is wasted spend. - thinking_config=types.ThinkingConfig(thinking_budget=0), + thinking_config=llm_client.thinking_config_for("standard"), ) response = llm_client.generate_with_failover( config.LLM_MODEL_NAME, @@ -585,8 +651,10 @@ def llm_generate(prompt: str, max_tokens: int = None, max_output_tokens=max_tokens, safety_settings=config.SAFETY_SETTINGS, system_instruction=system_instruction or config.SYSTEM_PROMPT, - # Disable thinking so the full token budget goes to the answer, not thoughts. - thinking_config=types.ThinkingConfig(thinking_budget=0), + # Minimise thinking so the token budget goes to the answer, not thoughts. + # Sending nothing here would mean the model's own default (MEDIUM on + # gemini-3.6-flash), whose thoughts come out of max_output_tokens. + thinking_config=llm_client.thinking_config_for("standard"), ) try: @@ -841,8 +909,11 @@ def answer_question_strategy_a( logger.info("Generating answer...") answer = llm_generate(prompt, model=model, provider=provider) - # Extract citations using robust parser - citations = extract_citations(answer, context_data['metadatas'], context_data.get('chunks')) + # Extract citations using robust parser, compacting [1],[4] → [1],[2] so the + # answer's markers match the cited-only panel. + answer, citations = compact_citations( + answer, context_data['metadatas'], context_data.get('chunks'), + visible_chunks=context_data.get('chunks_used')) result = { 'answer': answer, @@ -931,8 +1002,11 @@ def answer_question_strategy_b( logger.info("Generating answer in English...") english_answer = llm_generate(prompt, model=model, provider=provider) - # Extract citations from ENGLISH answer (before translation) using robust parser - citations = extract_citations(english_answer, context_data['metadatas'], context_data.get('chunks')) + # Extract citations from ENGLISH answer (before translation) using robust parser. + # Compacting here means the translated answer inherits the dense numbering. + english_answer, citations = compact_citations( + english_answer, context_data['metadatas'], context_data.get('chunks'), + visible_chunks=context_data.get('chunks_used')) # Translate answer to target language if needed if detected_lang != "en" and lang_utils.is_indic_language(detected_lang): @@ -1078,7 +1152,10 @@ def answer_with_history( prompt = f"## Conversation History\n{history_str}\n\n---\n\n{prompt}" english_answer = llm_generate(prompt, model=model, provider=provider) - citations = extract_citations(english_answer, context_data["metadatas"], context_data.get("chunks")) + # Compact before any translation so the translated answer carries the same numbers. + english_answer, citations = compact_citations( + english_answer, context_data["metadatas"], context_data.get("chunks"), + visible_chunks=context_data.get("chunks_used")) if strategy == "B" and detected_lang != "en" and lang_utils.is_indic_language(detected_lang): try: diff --git a/report_runner.py b/report_runner.py index 4611651..ea573b9 100644 --- a/report_runner.py +++ b/report_runner.py @@ -14,6 +14,7 @@ import re import config +import lang_utils import rag logger = logging.getLogger(__name__) @@ -32,12 +33,17 @@ def plan_sections(topic: str, language: str = "en", max_sections: int = None) -> """ if max_sections is None: max_sections = config.REPORT_MAX_SECTIONS + lang_name = lang_utils.get_language_name(language) + lang_rule = f"The titles themselves MUST be written in {lang_name}." + if language != "en": + lang_rule += " Do not write them in English." prompt = ( f"You are planning a literature-review report on: {topic!r}.\n" f"Propose at most {max_sections} section titles (e.g. background, methods " - f"comparison, key findings, open gaps). Write the section titles in the " - f"language with code {language!r}. Reply with ONLY a JSON array of " - f'short title strings, e.g. ["Background", "Methods", "Findings"].' + f"comparison, key findings, open gaps). Reply with ONLY a JSON array of " + f"short title strings. This example shows the required FORMAT ONLY — its " + f'shape, not its language: ["Background", "Methods", "Findings"].\n' + + lang_rule ) try: raw = rag.llm_generate(prompt, max_tokens=_PLAN_MAX_TOKENS) @@ -46,7 +52,7 @@ def plan_sections(topic: str, language: str = "en", max_sections: int = None) -> logger.warning(f"[Report] section planning failed, using default outline: {e}") sections = [] if not sections: - sections = list(_DEFAULT_SECTIONS) + sections = _default_sections(language) # de-dupe (case-insensitive) preserving order, then cap seen, out = set(), [] for s in sections: @@ -57,6 +63,33 @@ def plan_sections(topic: str, language: str = "en", max_sections: int = None) -> return out[:max_sections] +def _default_sections(language: str) -> list[str]: + """The fallback outline, in the requested language. + + plan_sections promises titles in `language`; returning the English skeleton + on the fallback path broke that promise precisely when the planner had + already failed. One short translation call is cheap here — this path only + runs when the planner returned nothing parseable — and English remains the + last resort if it fails too. + """ + if language == "en": + return list(_DEFAULT_SECTIONS) + try: + lang_name = lang_utils.get_language_name(language) + raw = rag.llm_generate( + f"Translate each title into {lang_name}. Reply with ONLY a JSON array " + f"of {len(_DEFAULT_SECTIONS)} strings, same order, no English: " + f"{json.dumps(_DEFAULT_SECTIONS)}", + max_tokens=_PLAN_MAX_TOKENS, + ) + translated = _parse_sections(raw) + if len(translated) == len(_DEFAULT_SECTIONS): + return translated + except Exception as e: + logger.warning(f"[Report] default-outline translation failed: {e}") + return list(_DEFAULT_SECTIONS) + + def _parse_sections(raw: str) -> list[str]: """Pull a JSON array of strings out of an LLM reply, tolerating code fences/prose.""" m = re.search(r"\[.*\]", raw, re.DOTALL) @@ -71,7 +104,10 @@ def _parse_sections(raw: str) -> list[str]: # Inline citation markers: [3], [1, 3, 5], [2,4]. Excludes [NOT FOUND: ...] # and other non-numeric brackets (the digit-only pattern won't match them). -_MARKER_RE = re.compile(r"\[(\d+(?:\s*,\s*\d+)*)\]") +# Leading whitespace is captured separately so a fully-dangling marker is dropped +# together with the space in front of it, including a newline when the marker sits +# alone on its own line (mirrors rag._CITE_MARKER_RE). +_MARKER_RE = re.compile(r"([ \t]*\n?[ \t]*)\[(\d+(?:\s*,\s*\d+)*)\]") def _remap_markers(body: str, cites: list[dict], registry: dict) -> str: @@ -94,13 +130,20 @@ def _remap_markers(body: str, cites: list[dict], registry: dict) -> str: def _repl(m: re.Match) -> str: mapped: list[int] = [] - for part in m.group(1).split(","): + for part in m.group(2).split(","): g = local_to_global.get(part.strip()) if g is not None and g not in mapped: mapped.append(g) if not mapped: # every number in this marker was dangling - return "" - return "[" + ", ".join(str(g) for g in mapped) + "]" + # Drops the preceding space too — no double/trailing space. Mirrors + # rag.compact_citations: a marker alone on its line takes the leading + # newline with it (the line's own newline survives), otherwise the + # newline is restored so neighbouring lines don't splice together. + if "\n" not in m.group(1): + return "" + rest = m.string[m.end():] + return "" if (rest == "" or rest[0] == "\n") else "\n" + return m.group(1) + "[" + ", ".join(str(g) for g in mapped) + "]" return _MARKER_RE.sub(_repl, body) @@ -181,6 +224,7 @@ def run_report(topic: str, language: str = "en", progress_cb=None) -> dict: assert "alpha [1]" in md, "P should be global [1] in section 1" assert "beta [2] gamma [1] delta" in md, f"remap/drop wrong: {md!r}" assert "delta [" not in md, "dangling [5] not dropped" + assert "delta " not in md, f"dropped marker left a stray space: {md!r}" assert "## References" in md assert "- [1] P" in md and "- [2] Q" in md # planner fallback when LLM returns junk diff --git a/routes/agent.py b/routes/agent.py index 5441ea5..751ddd8 100644 --- a/routes/agent.py +++ b/routes/agent.py @@ -157,14 +157,6 @@ async def agent_query( detail={"error": "Agent pipeline failed. Please try again later.", "code": "AGENT_ERROR"}, ) - _append_session_messages(session_id, body.question, result["final_answer"]) - processing_time = time.time() - start_time - - logger.info( - f"Agent query: lang={result['detected_language']} " - f"reflexion={result['reflexion_count']} time={processing_time:.2f}s" - ) - all_contexts = result.get("retrieved_contexts", []) final_answer = result["final_answer"] cited_titles: set[str] = set() @@ -172,15 +164,27 @@ async def agent_query( try: metas = [{"title": c.get("title", "Unknown"), "section": c.get("section", "body")} for c in all_contexts] - # Same per-paper numbering the LLM saw in the context, so the source - # panel's [N] matches [Cite:N] in the answer text. - for num, meta in rag.citation_number_map(metas).items(): - title_to_num[(meta.get("title") or "Unknown").strip()] = num - for cit in rag.extract_citations(final_answer, metas): - cited_titles.add(cit["title"].strip()) + # The context numbers every retrieved paper, but the panel below keeps + # only the cited ones — so an answer citing papers 1 and 4 of 4 read + # "[1] … [4]" beside a two-entry panel. compact_citations renumbers the + # answer's markers and the citations together to a dense 1..M. + final_answer, cits = rag.compact_citations( + final_answer, metas, visible_chunks=result.get("context_chunks_used")) + for cit in cits: + title = cit["title"].strip() + cited_titles.add(title) + title_to_num[title] = int(cit["number"]) except Exception: pass # fall through to dedup-only logic below + _append_session_messages(session_id, body.question, final_answer) + processing_time = time.time() - start_time + + logger.info( + f"Agent query: lang={result['detected_language']} " + f"reflexion={result['reflexion_count']} time={processing_time:.2f}s" + ) + seen_titles: set[str] = set() sources = [] for ctx in all_contexts: @@ -206,18 +210,18 @@ async def agent_query( query_id = str(uuid.uuid4()) try: persistence.log_query( - query_id=query_id, question=body.question, answer=result["final_answer"], + query_id=query_id, question=body.question, answer=final_answer, mode=f"agent_{body.strategy}", model=body.model or "default", language=result.get("detected_language", "en"), confidence=result.get("answer_confidence") or 0.0, - coverage=citation_coverage(result["final_answer"]), + coverage=citation_coverage(final_answer), created_at=datetime.now(timezone.utc).isoformat(), ) except Exception: logger.warning("Failed to log query for feedback correlation", exc_info=True) return AgentQueryResponse( - answer=result["final_answer"], + answer=final_answer, session_id=session_id, language=result.get("detected_language", "en"), reflexion_iterations=result.get("reflexion_count", 0), @@ -287,7 +291,6 @@ async def _run_and_stream(): return processing_time = time.time() - start_time - _append_session_messages(session_id, body.question, result["final_answer"]) # --- Phase 2: stream tool calls as thinking events --- for tc in result.get("tool_calls_log", []): @@ -297,27 +300,34 @@ async def _run_and_stream(): }) yield f"data: {tool_msg}\n\n" - # --- Phase 3: stream the final answer in chunks --- + # --- Phase 3: resolve citations BEFORE streaming --- + # The markers have to be compacted before the first chunk goes out, or + # the client renders [1],[4] against a two-entry panel. final_answer = result["final_answer"] or "" - chunk_size = 80 # characters per SSE chunk - for i in range(0, len(final_answer), chunk_size): - chunk = final_answer[i:i + chunk_size] - yield f"data: {json.dumps({'type': 'chunk', 'text': chunk})}\n\n" - - # --- Phase 4: build sources list --- all_contexts = result.get("retrieved_contexts", []) cited_titles: set = set() title_to_num: dict = {} try: metas = [{"title": c.get("title", "Unknown"), "section": c.get("section", "body")} for c in all_contexts] - for num, meta in rag.citation_number_map(metas).items(): - title_to_num[(meta.get("title") or "Unknown").strip()] = num - for cit in rag.extract_citations(final_answer, metas): - cited_titles.add(cit["title"].strip()) + final_answer, cits = rag.compact_citations( + final_answer, metas, visible_chunks=result.get("context_chunks_used")) + for cit in cits: + title = cit["title"].strip() + cited_titles.add(title) + title_to_num[title] = int(cit["number"]) except Exception: pass + _append_session_messages(session_id, body.question, final_answer) + + # --- Phase 3b: stream the final answer in chunks --- + chunk_size = 80 # characters per SSE chunk + for i in range(0, len(final_answer), chunk_size): + chunk = final_answer[i:i + chunk_size] + yield f"data: {json.dumps({'type': 'chunk', 'text': chunk})}\n\n" + + # --- Phase 4: build sources list --- seen_titles: set = set() sources = [] for ctx in all_contexts: diff --git a/routes/chat.py b/routes/chat.py index 2c60569..8e2f6c8 100644 --- a/routes/chat.py +++ b/routes/chat.py @@ -155,15 +155,18 @@ async def _no_docs(): async def _stream_and_save(): full_answer: list[str] = [] + final_answer: str | None = None # compacted answer from the done event hit_error = False async for event in sse_stream(prepared["prompt"], prepared["metadatas"], prepared["detected_lang"], strategy=body.strategy, query_id=query_id, - model=body.model, provider=body.provider): + model=body.model, provider=body.provider, + visible_chunks=prepared["chunks_used"]): if event.startswith('data: {"type": "error"'): hit_error = True if event.startswith('data: {"type": "done"'): payload = json.loads(event[6:]) payload["session_id"] = session_id + final_answer = payload.get("answer") yield f"data: {json.dumps(payload)}\n\n" else: if event.startswith('data: {"type": "chunk"'): @@ -173,7 +176,10 @@ async def _stream_and_save(): pass yield event if not hit_error: - _append_session_messages(session_id, body.message, "".join(full_answer)) + # Persist the compacted answer, not the raw streamed chunks — otherwise + # the follow-up turns inherit gapped and dangling [N] markers. + _append_session_messages( + session_id, body.message, final_answer or "".join(full_answer)) return StreamingResponse(_stream_and_save(), media_type="text/event-stream") diff --git a/routes/ingest.py b/routes/ingest.py index 2747b75..7346669 100644 --- a/routes/ingest.py +++ b/routes/ingest.py @@ -561,32 +561,19 @@ def _run_batch_url_ingest(job_id: str, urls_to_ingest: List[dict]): _inflight_paper_ids.discard(paper_id) failed += 1 continue - # Save a permanent copy to the papers directory + # Save a permanent copy to the papers directory. The filename must be + # `{paper_id}.pdf`: /ingest/health derives paper_id from the file stem, + # so a title-derived name made the paper list with 0 chunks and "Need + # re-index" despite indexing fine, and DELETE /papers/{id} 404'd on it. saved_path = None try: import shutil - title = item.get("title", "") - if title: - safe_name = _re.sub(r'[<>:"/\\|?*]', '_', title)[:120].strip(' ._') - else: - safe_name = paper_id - if not safe_name.endswith('.pdf'): - safe_name += '.pdf' - dest = config.PAPERS_DIR / safe_name + dest = config.PAPERS_DIR / f"{paper_id}.pdf" shutil.copy2(path, dest) saved_path = dest logger.info(f"[Ingest] Saved paper to {dest} (paper_id={paper_id})") except Exception as e: logger.warning(f"[Ingest] Failed to save paper to disk: {e}", exc_info=True) - # Fallback: save with paper_id as filename - try: - import shutil - fallback = config.PAPERS_DIR / (paper_id + ".pdf") - shutil.copy2(path, fallback) - saved_path = fallback - logger.info(f"[Ingest] Saved paper (fallback) to {fallback}") - except Exception as e2: - logger.warning(f"[Ingest] Fallback save also failed: {e2}", exc_info=True) try: ingest_path = str(saved_path) if saved_path else path logger.info(f"[Ingest] Ingesting from {ingest_path} with paper_id={paper_id}") diff --git a/routes/query.py b/routes/query.py index 7a73267..e805454 100644 --- a/routes/query.py +++ b/routes/query.py @@ -315,7 +315,8 @@ async def _no_docs(): return StreamingResponse( sse_stream(prepared["prompt"], prepared["metadatas"], prepared["detected_lang"], strategy=body.strategy, query_id=query_id, - model=body.model, provider=body.provider), + model=body.model, provider=body.provider, + visible_chunks=prepared["chunks_used"]), media_type="text/event-stream", ) diff --git a/sse_utils.py b/sse_utils.py index a546bf9..3738322 100644 --- a/sse_utils.py +++ b/sse_utils.py @@ -11,10 +11,19 @@ import rag import translation +# Distinct from providers.base.TRUNCATION_NOTE, which means "hit the token limit". +# This one means the connection died mid-generation, so the answer stops wherever +# the last chunk landed — usually mid-sentence. +INTERRUPTED_NOTE = ( + "\n\n*[Answer incomplete — the connection to the model dropped mid-response. " + "The sources below cover only what was generated.]*" +) + async def sse_stream(prompt: str, metadatas: list, language: str, strategy: str = "A", max_tokens: int = None, query_id: str = None, - model: str = None, provider: str = None): + model: str = None, provider: str = None, + visible_chunks: int = None): """Async SSE generator: bridges sync llm_generate_stream via asyncio.Queue. Strategy B + Indic target language: buffers all chunks, translates the full @@ -47,6 +56,7 @@ def _run(): # ponytail: buffer when translation needed, stream otherwise needs_translation = strategy == "B" and language != "en" and lang_utils.is_indic_language(language) full_answer: list[str] = [] + interrupted = False # stream died partway, but there is text worth keeping try: while True: kind, data = await q.get() @@ -56,21 +66,45 @@ def _run(): yield f"data: {json.dumps({'type': 'chunk', 'text': data})}\n\n" elif kind == "error": yield f"data: {json.dumps({'type': 'error', 'message': data})}\n\n" - yield "data: [DONE]\n\n" - return + # Don't discard what already streamed. A stream that dies partway + # (dropped connection, provider hiccup) used to return here, so the + # user kept the partial answer on screen but lost every citation + # with it. Fall through to the done event when there is text left + # to attribute; only a completely empty answer stops here. + if not full_answer: + yield "data: [DONE]\n\n" + return + # Say so in the answer itself. A stream cut off mid-sentence + # otherwise reads as a complete answer once the error toast is + # gone — and it arrives with citations, which makes it look + # more finished than it is. + interrupted = True + break else: # done break assembled = "".join(full_answer) - if needs_translation and assembled: + # Compact BEFORE translating, so the translated answer inherits the dense + # numbering (same order rag.answer_question uses). Chunks already streamed + # carry the raw markers — an answer citing papers 1 and 4 of 4 renders + # "[1] … [4]" against a two-entry panel, and a marker past visible_chunks + # has no source at all — so the done event carries the corrected answer and + # the client re-renders from it. visible_chunks keeps a marker invented past + # the prompt's truncation point from resolving to a paper never shown. + compacted, citations = rag.compact_citations( + assembled, metadatas, visible_chunks=visible_chunks) + if interrupted: + compacted += INTERRUPTED_NOTE + + if needs_translation and compacted: try: - translated = await run_in_threadpool(translation.translate_from_english, assembled, language) - yield f"data: {json.dumps({'type': 'chunk', 'text': translated})}\n\n" + compacted = await run_in_threadpool( + translation.translate_from_english, compacted, language) except Exception: - yield f"data: {json.dumps({'type': 'chunk', 'text': assembled})}\n\n" # fall back to English + pass # fall back to English + yield f"data: {json.dumps({'type': 'chunk', 'text': compacted})}\n\n" - citations = rag.extract_citations(assembled, metadatas) - yield f"data: {json.dumps({'type': 'done', 'citations': citations, 'language': language, 'query_id': query_id})}\n\n" + yield f"data: {json.dumps({'type': 'done', 'answer': compacted, 'citations': citations, 'language': language, 'query_id': query_id})}\n\n" yield "data: [DONE]\n\n" finally: stop_event.set() diff --git a/static/index.html b/static/index.html index f8bacd2..b451455 100644 --- a/static/index.html +++ b/static/index.html @@ -1419,8 +1419,12 @@

Ask anything

} else if (evt.type === 'done') { const cites = evt.citations || []; if (evt.session_id) sessionId = evt.session_id; - $('#streamBody').innerHTML = renderContent(acc); - lastAnswerText = acc; lastCites = cites; + // Streamed chunks can carry the model's raw [N] markers; when the done + // event carries a compacted answer (dense numbering, dangling markers + // dropped) it matches the citation panel, so re-render from that. + const finalText = evt.answer || acc; + $('#streamBody').innerHTML = renderContent(finalText); + lastAnswerText = finalText; lastCites = cites; turnDiv.querySelector('.asst-meta').innerHTML = 'Agentic' + '' + esc(evt.language || 'en') + '' + @@ -1469,8 +1473,12 @@

Ask anything

} else if (evt.type === 'done') { const cites = evt.citations || []; if (evt.session_id) sessionId = evt.session_id; - $('#streamBody').innerHTML = renderContent(acc); - lastAnswerText = acc; lastCites = cites; + // Streamed chunks can carry the model's raw [N] markers; when the done + // event carries a compacted answer (dense numbering, dangling markers + // dropped) it matches the citation panel, so re-render from that. + const finalText = evt.answer || acc; + $('#streamBody').innerHTML = renderContent(finalText); + lastAnswerText = finalText; lastCites = cites; turnDiv.querySelector('.asst-meta').innerHTML = '' + esc(evt.language_name || evt.language || 'en') + '' + (cites.length ? '' + cites.length + ' source' + (cites.length > 1 ? 's' : '') + '' : '') + diff --git a/tests/test_agent.py b/tests/test_agent.py index 3c8c4b9..3ee2c0e 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -713,6 +713,25 @@ def test_safe_stop_preserves_draft_answer(): assert draft in result["final_answer"], "draft answer must survive safe_stop" +def test_query_planner_survives_safety_blocked_response(): + """A safety-blocked reply whose .text raises must fall back, not crash the node.""" + from agent.nodes.query_planner import query_planner_node + + class _BlockedResp: + candidates = None + + @property + def text(self): + raise ValueError("Response has no valid Part") + + # NOTE: rag.safe_extract_text is the code under test — do not patch it. + with patch("rag.generate_with_failover", return_value=_BlockedResp()): + result = query_planner_node({"original_query": "antenna design for IoT"}) + + assert result["query_plan"] == ["antenna design for IoT"] + assert result["detected_language"] + + def test_year_filter_builds_chromadb_where_clause(): """Year range filter must produce valid ChromaDB where-clauses and ignore junk.""" from agent.tool_executor import _year_filter @@ -791,6 +810,98 @@ def test_reflexion_time_budget_finalises_draft(): assert result["final_answer"] == draft +def test_first_evaluation_runs_even_past_the_loop_budget(): + """The reflexion budget stops FURTHER loops, not the first evaluation. + + On a CPU-only box the first pass alone runs past the budget, so gating + iteration 1 on it shipped every answer unverified — no faithfulness score, no + completeness, no confidence. + """ + import time as _time + import config + from agent.nodes.reflexion_evaluator import reflexion_evaluator_node + + budget = 90.0 + state = _eval_state( + draft_answer="Substantive answer. [1]", + reflexion_count=0, + start_time=_time.monotonic() - (budget + 10), + ) + # Pin all three knobs: a developer .env (or its absence in CI) otherwise decides + # whether the reserve gate fires, and this test is about the loop budget. + with patch.object(config, "AGENT_REFLEXION_BUDGET_S", budget), \ + patch.object(config, "AGENT_TIMEOUT", 600), \ + patch.object(config, "AGENT_EVAL_RESERVE_S", 90.0), \ + patch("verify.check_claims", return_value=[]) as cc, \ + patch("rag.generate_with_failover", return_value=_mk_eval_resp(0.9, "accept")), \ + patch("rag.safe_extract_text", side_effect=lambda r: r.text): + reflexion_evaluator_node(state) + + assert cc.called, "first evaluation must not be skipped by the loop budget" + + +def test_over_budget_retry_verdict_finalises_instead_of_looping(): + """A retry verdict past the loop budget must not start another cycle. + + The top-of-node gate lets iteration 1 evaluate even when already over budget. + Without a post-evaluation check, a `retrieve_more` verdict there returns no + final_answer, so the graph runs a full retrieve→generate cycle (~95s on CPU) + that the budget exists to prevent. + """ + import time as _time + import config + from agent.nodes.reflexion_evaluator import reflexion_evaluator_node + + draft = "Substantive answer. [1]" + claims = [{"claim": "c", "support": 0.01, "grounded": False}] + state = _eval_state( + draft_answer=draft, + reflexion_count=0, + start_time=_time.monotonic() - (config.AGENT_REFLEXION_BUDGET_S + 10), + ) + with patch.object(config, "AGENT_TIMEOUT", 600), \ + patch.object(config, "AGENT_EVAL_RESERVE_S", 90.0), \ + patch("verify.check_claims", return_value=claims), \ + patch("rag.generate_with_failover", return_value=_mk_eval_resp(0.2, "retrieve_more")), \ + patch("rag.safe_extract_text", side_effect=lambda r: r.text): + result = reflexion_evaluator_node(state) + + assert result["final_answer"] == draft, ( + "over-budget retry must finalise; without final_answer the graph loops again" + ) + + +def test_nli_chunk_cap_cannot_disable_verification(): + """A 0 cap would slice away every cited chunk, and an empty claim list reads as + faithfulness 1.0 — accepting an ungrounded answer. config must clamp it to >=1.""" + import config + + assert config.NLI_MAX_CHUNKS_PER_CITATION >= 1 + + +def test_evaluation_skipped_when_timeout_reserve_is_gone(): + """With less than the reserve left under AGENT_TIMEOUT, return the draft without + paying for NLI or the completeness LLM call — finishing would 504 and lose it.""" + import time as _time + import config + from agent.nodes.reflexion_evaluator import reflexion_evaluator_node + + def _boom(*a, **kw): + raise AssertionError("must not be called with no time left") + + draft = "Best-effort answer so far." + state = _eval_state( + draft_answer=draft, + reflexion_count=0, + start_time=_time.monotonic() - (config.AGENT_TIMEOUT - config.AGENT_EVAL_RESERVE_S + 5), + ) + with patch("verify.check_claims", side_effect=_boom), \ + patch("rag.generate_with_failover", side_effect=_boom): + result = reflexion_evaluator_node(state) + + assert result["final_answer"] == draft + + def test_evaluator_sees_full_long_answer(): """Completeness evaluator must not judge only the first 4000 chars of a long answer.""" from agent.nodes.reflexion_evaluator import reflexion_evaluator_node diff --git a/tests/test_providers_gemini.py b/tests/test_providers_gemini.py index beadc6c..7c8319e 100644 --- a/tests/test_providers_gemini.py +++ b/tests/test_providers_gemini.py @@ -8,15 +8,19 @@ class _FakeModels: """Rejects thinking_budget=0 the way gemini-3.6-flash does (400 INVALID_ARGUMENT).""" - def __init__(self, stream_chunks=("hi",)): + def __init__(self, stream_chunks=("hi",), finish_reason=None): self.calls = [] + self.levels = [] # thinking_level per call, parallel to calls self._stream_chunks = stream_chunks + self.finish_reason = finish_reason def _check(self, config): tc = getattr(config, "thinking_config", None) budget = getattr(tc, "thinking_budget", None) if tc is not None else None self.calls.append(budget) - if budget == 0: + self.levels.append(getattr(tc, "thinking_level", None) if tc is not None else None) + if budget is not None: + # Gemini 3.x rejects the legacy budget field outright, whatever its value. raise Exception("400 INVALID_ARGUMENT. Request contains an invalid argument.") def generate_content(self, model, contents, config): @@ -25,7 +29,17 @@ def generate_content(self, model, contents, config): def generate_content_stream(self, model, contents, config): self._check(config) - return iter([SimpleNamespace(text=c) for c in self._stream_chunks]) + return iter(_stream(self._stream_chunks, self.finish_reason)) + + +def _stream(chunks, finish_reason=None): + """Gemini reports finish_reason on the last chunk's candidate, not the chunk.""" + out = [] + for i, c in enumerate(chunks): + last = i == len(chunks) - 1 + cands = [SimpleNamespace(finish_reason=finish_reason)] if (last and finish_reason) else None + out.append(SimpleNamespace(text=c, candidates=cands)) + return out def _backend_with(models): @@ -43,22 +57,48 @@ def zero_budget_config(): ) -def test_generate_retries_without_thinking_when_zero_budget_rejected(zero_budget_config): +def test_generate_retries_with_a_level_when_budget_rejected(zero_budget_config): + """The retry must ask for MINIMAL thinking, not simply drop the field. + + Omitting thinking_config means the model's own default — MEDIUM on + gemini-3.6-flash — so "thinking off" used to produce medium thinking, billed + and taken out of the answer's share of max_output_tokens. + """ + from google.genai import types + models = _FakeModels() b = _backend_with(models) resp = b.generate("gemini-3.6-flash", "q", zero_budget_config, client=b._pool[0]) assert resp.text == "ok" - assert models.calls == [0, None] # rejected, then retried without + assert models.calls == [0, None] # budget rejected, then not resent + assert models.levels == [None, types.ThinkingLevel.MINIMAL] assert "gemini-3.6-flash" in b._zero_budget_rejected -def test_second_call_skips_thinking_config_entirely(zero_budget_config): +def test_second_call_sends_the_level_without_retrying(zero_budget_config): + from google.genai import types + models = _FakeModels() b = _backend_with(models) client = b._pool[0] b.generate("gemini-3.6-flash", "q", zero_budget_config, client=client) b.generate("gemini-3.6-flash", "q2", zero_budget_config, client=client) assert models.calls == [0, None, None] # no repeat of the failing call + assert models.levels[-1] == types.ThinkingLevel.MINIMAL + + +def test_dynamic_budget_translates_to_no_thinking_config(): + """-1 means "model decides", which is exactly what omitting the field does.""" + from google.genai import types + + models = _FakeModels() + b = _backend_with(models) + cfg = types.GenerateContentConfig( + max_output_tokens=16, thinking_config=types.ThinkingConfig(thinking_budget=-1) + ) + b.generate("gemini-3.6-flash", "q", cfg, client=b._pool[0]) + assert models.calls == [-1, None] + assert models.levels == [None, None] def test_stream_retries_before_any_token_emitted(zero_budget_config): @@ -69,6 +109,22 @@ def test_stream_retries_before_any_token_emitted(zero_budget_config): assert models.calls == [0, None] +def test_stream_appends_truncation_note_on_max_tokens(zero_budget_config): + from providers.base import TRUNCATION_NOTE + models = _FakeModels(stream_chunks=("half an ans", "wer that stops"), + finish_reason="FinishReason.MAX_TOKENS") + b = _backend_with(models) + out = "".join(b.generate_stream("gemini-3.5-flash", "q", zero_budget_config, client=b._pool[0])) + assert out == "half an answer that stops" + TRUNCATION_NOTE + + +def test_stream_stays_clean_when_model_finishes(zero_budget_config): + models = _FakeModels(stream_chunks=("all ", "done"), finish_reason="FinishReason.STOP") + b = _backend_with(models) + out = "".join(b.generate_stream("gemini-3.5-flash", "q", zero_budget_config, client=b._pool[0])) + assert out == "all done" + + def test_other_invalid_argument_errors_still_propagate(): from google.genai import types diff --git a/tests/test_providers_openrouter.py b/tests/test_providers_openrouter.py index bb0bb8d..aab0c49 100644 --- a/tests/test_providers_openrouter.py +++ b/tests/test_providers_openrouter.py @@ -42,6 +42,36 @@ def test_generate_returns_shim_with_text(): assert resp.text == "the answer" +def _chunks(texts, finish_reason=None): + """OpenAI reports finish_reason on the final choice, alongside an empty delta.""" + out = [SimpleNamespace(choices=[SimpleNamespace( + delta=SimpleNamespace(content=t), finish_reason=None)]) for t in texts] + if finish_reason: + out.append(SimpleNamespace(choices=[SimpleNamespace( + delta=SimpleNamespace(content=None), finish_reason=finish_reason)])) + return out + + +def _stream_backend(texts, finish_reason=None): + b = OpenRouterBackend() + b._client = MagicMock() + b._client.chat.completions.create.return_value = _chunks(texts, finish_reason) + return b + + +def test_stream_appends_truncation_note_on_length(): + from providers.base import TRUNCATION_NOTE + b = _stream_backend(["half an ans", "wer that stops"], "length") + out = "".join(b.generate_stream("anthropic/claude-haiku", "q", _cfg())) + assert out == "half an answer that stops" + TRUNCATION_NOTE + + +def test_stream_stays_clean_when_model_finishes(): + b = _stream_backend(["all ", "done"], "stop") + out = "".join(b.generate_stream("anthropic/claude-haiku", "q", _cfg())) + assert out == "all done" + + def test_generate_returns_shim_with_function_call(): import json b = OpenRouterBackend() diff --git a/tests/test_rag.py b/tests/test_rag.py index 33ef3db..b562bce 100644 --- a/tests/test_rag.py +++ b/tests/test_rag.py @@ -88,6 +88,73 @@ def test_extract_citations_cite_format(): assert result[1]["title"] == "Paper B" +def test_compact_citations_closes_numbering_gaps(): + """Citing papers 1 and 4 of 4 must render as [1],[2] against a 2-entry panel.""" + from rag import compact_citations + + metadatas = [ + {"title": "Paper A", "section": "intro"}, + {"title": "Paper B", "section": "body"}, + {"title": "Paper C", "section": "body"}, + {"title": "Paper D", "section": "results"}, + ] + answer, cites = compact_citations("reward [1] and bandwidth [4], both [1, 4]", metadatas) + + assert answer == "reward [1] and bandwidth [2], both [1, 2]" + assert [c["number"] for c in cites] == ["1", "2"] + assert [c["title"] for c in cites] == ["Paper A", "Paper D"] + + +def test_dangling_marker_on_its_own_line_leaves_no_blank_line(): + """A marker alone on a line must take its newline with it, not leave a blank + line that markdown renders as a paragraph break — and must not splice the + surrounding lines together either.""" + from rag import compact_citations + + metadatas = [{"title": "Paper A", "section": "intro"}] + answer, _ = compact_citations("first line\n[99]\nsecond line", metadatas) + + assert answer == "first line\nsecond line" + + +def test_citations_ignore_papers_the_prompt_never_showed(): + """A marker past the prompt's truncation point must not resolve to a real paper. + + format_context truncates by chunk count and by length, but callers hold the + FULL retrieved metadata. Without visible_chunks, an invented [3] resolves to + Paper C — a paper the model was never shown — producing a citation that looks + legitimate. Numbering only the visible slice makes it dangle, so it is dropped. + """ + from rag import compact_citations, extract_citations + + metadatas = [ + {"title": "Paper A", "section": "intro"}, + {"title": "Paper B", "section": "body"}, + {"title": "Paper C", "section": "results"}, # truncated out of the prompt + ] + answer = "grounded [1] and invented [3]" + + # Only the first two chunks reached the prompt. + compacted, cites = compact_citations(answer, metadatas, visible_chunks=2) + assert [c["title"] for c in cites] == ["Paper A"] + assert compacted == "grounded [1] and invented" + + # Same call without the slice is what produced the phantom citation. + assert [c["title"] for c in extract_citations(answer, metadatas)] == ["Paper A", "Paper C"] + + +def test_compact_citations_drops_dangling_marker(): + """A number the model invented resolves to no paper — drop it, don't renumber around it.""" + from rag import compact_citations + + metadatas = [{"title": "Paper A", "section": "intro"}] + answer, cites = compact_citations("grounded [1] invented [7] and mid [7] sentence", metadatas) + + # dropping a marker must not leave a trailing or doubled space behind + assert answer == "grounded [1] invented and mid sentence" + assert [c["number"] for c in cites] == ["1"] + + def test_extract_citations_no_false_match(): from rag import extract_citations diff --git a/tests/test_report_routes.py b/tests/test_report_routes.py index 3e8acb2..3e9ccb5 100644 --- a/tests/test_report_routes.py +++ b/tests/test_report_routes.py @@ -62,6 +62,32 @@ def test_report_end_to_end(client): assert "# Literature Review: graphene sensors" in dl.text +def test_plan_sections_prompt_names_language_natively(): + import config + import rag + import report_runner + + seen = {} + + def _capture(prompt, **kwargs): + seen["prompt"] = prompt + return '["x"]' + + with patch.object(rag, "llm_generate", _capture): + report_runner.plan_sections("some topic", "hi") + + assert config.LANGUAGE_NAMES["hi"] in seen["prompt"] + assert "'hi'" not in seen["prompt"] + assert "Do not write them in English." in seen["prompt"] + + with patch.object(rag, "llm_generate", _capture): + report_runner.plan_sections("some topic", "en") + + # English must not get the self-contradictory "in English, not in English" + assert "Do not write them in English." not in seen["prompt"] + assert "MUST be written in English." in seen["prompt"] + + def test_report_rejects_empty_topic(client): with patch("config.REPORT_ENABLE", True): assert client.post("/report", json={"topic": ""}).status_code == 422 diff --git a/tests/test_sse_utils.py b/tests/test_sse_utils.py new file mode 100644 index 0000000..29f919f --- /dev/null +++ b/tests/test_sse_utils.py @@ -0,0 +1,104 @@ +"""SSE streaming path: the done event must carry a citation-corrected answer. + +Chunks go out as the model produces them, so they carry its raw [N] markers. +An answer citing papers 1 and 4 of 4 streams as "[1] ... [4]" while the panel +holds two entries, and a marker past the prompt's truncation point has no source +behind it at all. The done event therefore carries the compacted answer, and the +client re-renders from that. +""" + +import asyncio +import json +from unittest.mock import patch + +import sse_utils + + +def _drive(**kwargs): + """Run sse_stream to completion, returning the parsed events.""" + async def _collect(): + out = [] + async for raw in sse_utils.sse_stream(**kwargs): + payload = raw[len("data: "):].strip() + if payload != "[DONE]": + out.append(json.loads(payload)) + return out + + return asyncio.run(_collect()) + + +_METADATAS = [ + {"title": "Paper A", "section": "intro"}, + {"title": "Paper B", "section": "body"}, + {"title": "Paper C", "section": "results"}, +] + + +def test_done_event_carries_compacted_answer(): + answer = "reward [1] and bandwidth [3], plus invented [9]" + + with patch("llm_client.llm_generate_stream", return_value=iter([answer])): + events = _drive(prompt="p", metadatas=_METADATAS, language="en", + query_id="q1", visible_chunks=3) + + chunks = [e for e in events if e["type"] == "chunk"] + done = next(e for e in events if e["type"] == "done") + + # Streaming itself is untouched — the raw text still goes out live. + assert "".join(c["text"] for c in chunks) == answer + + # ...but the done event carries the corrected answer: [3] renumbered to [2] + # (Paper B was never cited) and the dangling [9] dropped. + assert done["answer"] == "reward [1] and bandwidth [2], plus invented" + assert [c["title"] for c in done["citations"]] == ["Paper A", "Paper C"] + assert [c["number"] for c in done["citations"]] == ["1", "2"] + + +def test_partial_answer_keeps_its_citations_after_a_stream_error(): + """A stream that dies partway must still emit done. + + Observed with a dropped upstream connection (WinError 10054): the user kept + the partial answer on screen but every citation vanished with it, because the + error path returned before the done event. + """ + def _die_midway(): + yield "grounded [1] and more" + raise RuntimeError("[WinError 10054] connection forcibly closed") + + with patch("llm_client.llm_generate_stream", return_value=_die_midway()): + events = _drive(prompt="p", metadatas=_METADATAS, language="en", + query_id="q3", visible_chunks=3) + + assert any(e["type"] == "error" for e in events), "the failure must still surface" + done = next(e for e in events if e["type"] == "done") + assert done["answer"].startswith("grounded [1] and more") + assert [c["title"] for c in done["citations"]] == ["Paper A"] + # ...and the answer must admit it is incomplete: with the error toast gone, + # a mid-sentence answer arriving with citations reads as a finished one. + assert done["answer"].endswith(sse_utils.INTERRUPTED_NOTE) + + +def test_empty_stream_error_stops_without_a_done_event(): + """Nothing salvageable — no answer text, so there is nothing to attribute.""" + def _die_immediately(): + raise RuntimeError("upstream refused") + yield # pragma: no cover - makes this a generator + + with patch("llm_client.llm_generate_stream", return_value=_die_immediately()): + events = _drive(prompt="p", metadatas=_METADATAS, language="en", + query_id="q4", visible_chunks=3) + + assert [e["type"] for e in events] == ["error"] + + +def test_marker_past_visible_chunks_resolves_to_nothing(): + """A number invented past the prompt's truncation point must not resolve to a + real paper — Paper C never reached the prompt when only 2 chunks were used.""" + with patch("llm_client.llm_generate_stream", + return_value=iter(["grounded [1] invented [3]"])): + events = _drive(prompt="p", metadatas=_METADATAS, language="en", + query_id="q2", visible_chunks=2) + + done = next(e for e in events if e["type"] == "done") + assert [c["title"] for c in done["citations"]] == ["Paper A"] + assert done["answer"] == "grounded [1] invented" diff --git a/tests/test_verify.py b/tests/test_verify.py index e89c464..e192ea3 100644 --- a/tests/test_verify.py +++ b/tests/test_verify.py @@ -57,6 +57,41 @@ def test_citation_marker_stripped_from_nli_hypothesis(): assert "[" not in premise +def test_faithfulness_threshold_is_reachable(): + """The grounded bar must sit inside the range the NLI model actually produces. + + Measured on this corpus with the shipped int8 mDeBERTa-xnli: a sentence copied + verbatim out of its own chunk scores median 0.226 / max 0.428; an unrelated + paper's sentence scores median 0.099, p90 0.158 — the two distributions overlap, + so 0.15 trades a ~0.10-0.15 false-positive rate for 0.70 recall. The old 0.5 was + above every positive, so `grounded` was always False and faithfulness read ~0 for + every answer. Guards against a future model/threshold change reintroducing that. + """ + import config + + assert 0.099 < config.FAITHFULNESS_THRESHOLD < 0.428, ( + "threshold outside the measured positive/negative separation — recalibrate " + "before changing NLI_MODEL_NAME or FAITHFULNESS_THRESHOLD" + ) + + +def test_chunks_per_citation_are_capped(): + """One paper contributing many chunks must not cost one NLI pair each — the + per-citation cap is what keeps the faithfulness pass inside the agent budget.""" + import config + + answer = "The framework uses deep Q-networks for optimization. [1]" + chunks = [f"chunk {i} about deep Q-networks" for i in range(8)] + metas = [{"title": "One Paper"} for _ in chunks] # all 8 chunks = citation [1] + + fake_model = _fake_model() + with patch("verify._load", return_value=fake_model): + verify.check_claims(answer, chunks, metas) + + called_pairs = fake_model.predict.call_args[0][0] + assert len(called_pairs) == config.NLI_MAX_CHUNKS_PER_CITATION + + def test_high_entailment_logit_yields_high_grounded_score(): answer = "The framework uses deep Q-networks for antenna optimization. [1]" chunks = ["The proposed framework uses deep Q-networks to optimize antenna parameters."] diff --git a/tests/test_watch_run.py b/tests/test_watch_run.py index 345a2bd..b583047 100644 --- a/tests/test_watch_run.py +++ b/tests/test_watch_run.py @@ -94,6 +94,60 @@ def test_abstract_only_when_no_pdf(monkeypatch): assert persistence.get_watch("w1")["seen_ids"] == ["2401.003"] +def test_downloaded_pdf_is_kept_in_papers_dir(monkeypatch, tmp_path): + """The library lists PAPERS_DIR and derives paper_id from the file stem, so a + watch that deleted its download left indexed-but-invisible orphan chunks.""" + import config + + papers = tmp_path / "papers" + monkeypatch.setattr(config, "PAPERS_DIR", papers) + + src = tmp_path / "download.pdf" + src.write_bytes(b"%PDF-1.4 fake") + + seen = {} + + def _ingest(path, paper_id=None, metadata=None): + seen["path"], seen["paper_id"] = path, paper_id + return 5, "Ingested Title" + + _save("w1") + _patch(monkeypatch, [_hit("2401.001")], ingest=_ingest) + monkeypatch.setattr(watch_runner, "_download_pdf", lambda url: str(src)) + + watch_runner.run_watch("w1") + + kept = papers / "2401_001.pdf" + assert kept.exists(), "PDF must survive in PAPERS_DIR or the paper never lists" + assert kept.stem == seen["paper_id"], "stem drives /ingest/health's paper_id lookup" + assert seen["path"] == str(kept) # ingested from the kept copy, not the temp file + assert not src.exists() # temp moved, not left behind + + +def test_failed_save_still_ingests_from_temp(monkeypatch, tmp_path): + """A save failure must degrade to the old behavior, not drop the paper.""" + import config + + monkeypatch.setattr(config, "PAPERS_DIR", tmp_path / "papers") + monkeypatch.setattr(watch_runner.shutil, "move", + lambda *a, **k: (_ for _ in ()).throw(OSError("disk full"))) + + seen = {} + + def _ingest(path, paper_id=None, metadata=None): + seen["path"] = path + return 5, "Ingested Title" + + _save("w1") + _patch(monkeypatch, [_hit("2401.001")], ingest=_ingest) + monkeypatch.setattr(watch_runner, "_download_pdf", lambda url: "/tmp/fake.pdf") + + res = watch_runner.run_watch("w1") + + assert res["new_count"] == 1 + assert seen["path"] == "/tmp/fake.pdf" + + def test_missing_watch_raises_keyerror(monkeypatch): _patch(monkeypatch, []) with pytest.raises(KeyError): diff --git a/verify.py b/verify.py index 12a999e..e481592 100644 --- a/verify.py +++ b/verify.py @@ -34,6 +34,10 @@ def _load(): # claims are actually verified, not scored by an English-only model. _model = CrossEncoder(config.NLI_MODEL_NAME, device=device, cache_folder=str(config.MODELS_CACHE_DIR)) + # Truncate the premise: cost is linear in sequence length (measured + # 1.15s/pair at 512 tokens vs 0.4s at 256 on this CPU box) and a + # chunk's support for a one-sentence claim is in its head, not tail. + _model.max_seq_length = config.NLI_MAX_SEQ_LENGTH return _model @@ -95,7 +99,15 @@ def check_claims(answer: str, chunks: List[str], metadatas=None) -> List[dict]: results = [] for sent in sentences: cited_nums = {int(n) for n in re.findall(r'\[(\d+)\]', sent)} - cited_chunks = [c for n in cited_nums for c in num_to_chunks.get(n, [])] + # Cap chunks per cited paper: a paper contributing 8 chunks used to cost 8 + # NLI pairs for ONE sentence, and best() over them is dominated by the top + # rerank-ordered chunks anyway. This is the difference between a 300s and a + # 30s faithfulness pass on CPU. + cited_chunks = [ + c + for n in sorted(cited_nums) + for c in num_to_chunks.get(n, [])[:config.NLI_MAX_CHUNKS_PER_CITATION] + ] if not cited_chunks: continue # Strip citation/not-found markers before scoring — leaving literal diff --git a/watch_runner.py b/watch_runner.py index 6e3817f..a997501 100644 --- a/watch_runner.py +++ b/watch_runner.py @@ -9,8 +9,10 @@ import asyncio import logging import os +import shutil import config +import lang_utils import persistence import rag from agent.tool_executor import execute_arxiv_search, execute_open_access_search @@ -30,7 +32,8 @@ def _summarize(topic: str, papers: list[dict], language: str) -> str: listing = "\n\n".join(f"[{p['arxiv_id']}] {p['title']}\n{p['text']}" for p in papers) prompt = ( f"You are compiling a research digest on the topic: {topic!r}.\n" - f"Below are newly published papers. Write a concise digest in {language} " + f"Below are newly published papers. Write a concise digest in " + f"{lang_utils.get_language_name(language)} " f"summarizing what is new. Cite each paper inline by its id in square " f"brackets, e.g. [2401.12345]. Use only the papers provided.\n\n" f"{listing}" @@ -38,6 +41,29 @@ def _summarize(topic: str, papers: list[dict], language: str) -> str: return rag.llm_generate(prompt, max_tokens=_DIGEST_MAX_TOKENS) +def _keep_pdf(tmp_path: str, paper_id: str) -> str: + """Move a downloaded PDF into PAPERS_DIR and return the path to ingest from. + + Watches used to ingest straight from the temp file and delete it, so the + paper never appeared in /papers or /ingest/health (both enumerate + PAPERS_DIR) and its chunks counted as orphans in the library panel. + + The filename must be `{paper_id}.pdf`: /ingest/health derives paper_id from + the file stem, so any other name shows the paper with 0 chunks. + Falls back to the temp path so a failed move degrades to the old behavior + (indexed but invisible) rather than losing the paper entirely. + """ + try: + config.PAPERS_DIR.mkdir(parents=True, exist_ok=True) + dest = config.PAPERS_DIR / f"{paper_id}.pdf" + shutil.move(tmp_path, dest) + logger.info("[Watch] saved %s", dest) + return str(dest) + except OSError as exc: + logger.warning("[Watch] could not save PDF for %s (%s) — indexing from temp", paper_id, exc) + return tmp_path + + def run_watch(watch_id: str) -> dict: """Search the watch topic, ingest genuinely-new papers, store a cited digest. @@ -70,16 +96,19 @@ def run_watch(watch_id: str) -> dict: if pdf_url: path = _download_pdf(pdf_url) if path: + paper_id = _bibtex_safe_id(arxiv_id) + ingest_path = _keep_pdf(path, paper_id) try: n_chunks, title = ingest_pdf( - path, paper_id=_bibtex_safe_id(arxiv_id), + ingest_path, paper_id=paper_id, metadata={"title": p.get("title", ""), "source": p.get("source", "")}, ) finally: - try: - os.remove(path) - except OSError: - pass + if ingest_path == path: # still the temp file — nothing kept + try: + os.remove(path) + except OSError: + pass if n_chunks > 0: # 0 = duplicate/unchanged in corpus → seen, but not "new" ingested.append({"arxiv_id": arxiv_id, "title": title or p.get("title", ""), "text": p.get("text", "")}) indexed = True