Spaces:
Sleeping
Sleeping
| """Unit tests for src/generate.py — gate, context packing, think-splitting, and the | |
| refusal/answer/stream paths with the LLM and retrieval monkeypatched out.""" | |
| import pytest | |
| import src.generate as gen | |
| # ---------- pure helpers ---------- | |
| def test_split_think_no_tags(): | |
| ans, think = gen._split_think("just an answer") | |
| assert ans == "just an answer" and think == "" | |
| def test_split_think_extracts_and_strips(): | |
| raw = "<think>step 1\nstep 2</think>The answer is $\\lambda$." | |
| ans, think = gen._split_think(raw) | |
| assert ans == "The answer is $\\lambda$." | |
| assert "step 1" in think and "step 2" in think | |
| def test_split_think_multiple_blocks(): | |
| raw = "<think>a</think>mid<think>b</think>end" | |
| ans, think = gen._split_think(raw) | |
| assert ans == "midend" | |
| assert think == "a\nb" | |
| def test_truncate_under_cap_unchanged(): | |
| assert gen._truncate("short", 100) == "short" | |
| def test_truncate_over_cap_marks_and_bounds(): | |
| text = "line\n" * 2000 | |
| out = gen._truncate(text, 100) | |
| assert out.endswith("... [passage truncated]") | |
| assert len(out) <= 100 * gen._CHARS_PER_TOK + len("\n... [passage truncated]") | |
| def _hit(i, *, page=None, score=0.5, text="x" * 200): | |
| return {"parent_text": text, "chunk_id": f"ck{i}", "source": "book", | |
| "chapter": f"Chapter {i}", "section": "Sec", "page": page, | |
| "rerank_score": score, "parent_id": f"p{i}"} | |
| def test_build_context_tags_and_metadata(): | |
| ctx, sources = gen._build_context([_hit(1, page=42), _hit(2)]) | |
| assert "[S1] (Chapter 1 / Sec)" in ctx and "[S2]" in ctx | |
| assert sources[0] == {"tag": "S1", "chunk_id": "ck1", "location": "Chapter 1 / Sec", | |
| "page": 42, "score": 0.5} | |
| assert sources[1]["page"] is None | |
| def test_build_context_location_falls_back_to_source(): | |
| h = _hit(1) | |
| h["chapter"] = h["section"] = None | |
| _, sources = gen._build_context([h]) | |
| assert sources[0]["location"] == "book" | |
| def test_build_context_stops_at_token_budget(): | |
| big = "y" * (gen.MAX_PARENT_TOKENS * gen._CHARS_PER_TOK) # each hit ~MAX_PARENT_TOKENS | |
| hits = [_hit(i, text=big) for i in range(10)] | |
| ctx, sources = gen._build_context(hits) | |
| assert 0 < len(sources) < 10 # budget cut it off | |
| est = sum(len(b) for b in ctx.split("\n\n")) // gen._CHARS_PER_TOK | |
| assert est <= gen.CTX_TOKEN_BUDGET + gen.MAX_PARENT_TOKENS # never wildly over | |
| # ---------- decompose gate ---------- | |
| def test_gate_matches_multi_part_questions(q): | |
| assert gen.DECOMPOSE_GATE.search(q) | |
| def test_gate_ignores_single_hop_questions(q): | |
| assert not gen.DECOMPOSE_GATE.search(q) | |
| def test_decompose_skips_llm_when_gate_misses(monkeypatch): | |
| def boom(messages, entry): | |
| raise AssertionError("LLM must not be called for single-hop questions") | |
| monkeypatch.setattr(gen, "_chat", boom) | |
| assert gen.decompose("What is the variance of the Poisson?") == [] | |
| def test_decompose_strips_bullets_and_caps_at_four(monkeypatch): | |
| raw = "1. mean of A\n- variance of A\n\n2) mean of B\nvariance of B\nextra query five\n" | |
| monkeypatch.setattr(gen, "_chat", lambda m, entry: raw) | |
| subs = gen.decompose("compare A and B") | |
| assert subs == ["mean of A", "variance of A", "mean of B", "variance of B"] | |
| def test_decompose_fails_open_on_llm_error(monkeypatch): | |
| def boom(messages, entry): | |
| raise OSError("ollama down") | |
| monkeypatch.setattr(gen, "_chat", boom) | |
| assert gen.decompose("compare A and B") == [] | |
| # ---------- answer / stream paths ---------- | |
| def test_answer_refuses_on_no_hits(monkeypatch): | |
| monkeypatch.setattr(gen, "_retrieve_for", lambda c, q, k: []) | |
| out = gen.answer(None, "q") | |
| assert out["refused"] and out["sources"] == [] | |
| assert out["answer"] == "I don't know based on the provided material." | |
| def test_answer_refuses_below_threshold(monkeypatch): | |
| monkeypatch.setattr(gen, "_retrieve_for", | |
| lambda c, q, k: [_hit(1, score=gen.MIN_RERANK_SCORE / 2)]) | |
| assert gen.answer(None, "q")["refused"] | |
| def test_answer_grounds_and_splits_thinking(monkeypatch): | |
| monkeypatch.setattr(gen, "_retrieve_for", lambda c, q, k: [_hit(1, page=7)]) | |
| monkeypatch.setattr(gen, "_chat", lambda m, entry: "<think>hmm</think>It is $\\mu$. [S1]") | |
| out = gen.answer(None, "q", show_thinking=True) | |
| assert not out["refused"] | |
| assert out["answer"] == "It is $\\mu$. [S1]" | |
| assert out["thinking"] == "hmm" | |
| assert out["sources"][0]["tag"] == "S1" | |
| # without the flag, thinking never leaks | |
| out2 = gen.answer(None, "q") | |
| assert "thinking" not in out2 | |
| def test_stream_event_order_and_final_answer(monkeypatch): | |
| monkeypatch.setattr(gen, "_retrieve_for", lambda c, q, k: [_hit(1)]) | |
| monkeypatch.setattr(gen, "_chat_stream", lambda m, entry: iter(["The ", "answer."])) | |
| events = list(gen.stream(None, "q")) | |
| assert [e["type"] for e in events] == ["sources", "token", "token", "done"] | |
| assert events[-1]["answer"] == "The answer." | |
| assert not events[-1]["refused"] | |
| def test_stream_refusal_is_single_done_event(monkeypatch): | |
| monkeypatch.setattr(gen, "_retrieve_for", lambda c, q, k: []) | |
| events = list(gen.stream(None, "q")) | |
| assert len(events) == 1 and events[0]["type"] == "done" and events[0]["refused"] | |
| def test_stream_style_selects_prompt(monkeypatch): | |
| seen = {} | |
| monkeypatch.setattr(gen, "_retrieve_for", lambda c, q, k: [_hit(1)]) | |
| def fake_stream(messages, entry): | |
| seen["system"] = messages[0]["content"] | |
| return iter(["ok"]) | |
| monkeypatch.setattr(gen, "_chat_stream", fake_stream) | |
| list(gen.stream(None, "q", style="eli5")) | |
| assert seen["system"] == gen.ELI5_SYSTEM | |
| list(gen.stream(None, "q", style="nonsense")) | |
| assert seen["system"] == gen.SYSTEM # unknown style falls back to default | |