| from __future__ import annotations |
|
|
| from researchmind.citations import ( |
| clean_model_answer, |
| format_context_block, |
| format_references, |
| ) |
| from researchmind.store import StoredChunk |
|
|
|
|
| def _chunk(chunk_id: str, doc_uri: str, text: str) -> StoredChunk: |
| return StoredChunk( |
| id=chunk_id, |
| doc_id="doc1", |
| ordinal=0, |
| text=text, |
| doc_title="AI Agents Review", |
| doc_uri=doc_uri, |
| metadata={}, |
| ) |
|
|
|
|
| def test_format_context_groups_chunks_by_document(): |
| chunks = [ |
| _chunk("c1", "https://example.com/paper", "First passage about agents."), |
| _chunk("c2", "https://example.com/paper", "Second passage about planning."), |
| ] |
| context, citations = format_context_block(chunks) |
| assert context.count("[1]") == 1 |
| assert "[2]" not in context |
| assert len(citations) == 1 |
| assert "First passage" in context |
| assert "Second passage" in context |
|
|
|
|
| def test_format_references_one_line_per_source(): |
| _, citations = format_context_block( |
| [ |
| _chunk("c1", "https://a.test", "alpha"), |
| _chunk("c2", "https://a.test", "beta"), |
| ] |
| ) |
| refs = format_references(citations) |
| assert refs.count("https://a.test") == 1 |
|
|
|
|
| def test_clean_passage_collapses_citation_runs(): |
| chunks = [_chunk("c1", "https://a.test", "[1] [2] [3] [4] [5] actual content")] |
| context, _ = format_context_block(chunks) |
| assert "[1] [2] [3] [4] [5]" not in context |
| assert "actual content" in context |
|
|
|
|
| def test_clean_model_answer_strips_reference_spam(): |
| raw = "Summary here [1][2][3][4][5].\n\n**References**\n- [1] dup" |
| cleaned = clean_model_answer(raw) |
| assert "**References**" not in cleaned |
| assert "[1][2][3]" not in cleaned |
| assert "Summary here" in cleaned |
|
|
|
|
| def test_clean_model_answer_strips_thinking_block(): |
| think_open = "<" + "think" + ">" |
| think_close = "</" + "think" + ">" |
| raw = f"{think_open}\nplan\n{think_close}\n\nAgents use tools and memory [1]." |
| cleaned = clean_model_answer(raw) |
| assert cleaned == "Agents use tools and memory [1]." |
|
|
|
|
| def test_clean_model_answer_rejects_unclosed_thinking(): |
| rt_open = "<" + "redacted_thinking" + ">" |
| raw = f"{rt_open}\nWe are given a context and need to plan the answer." |
| cleaned = clean_model_answer(raw) |
| assert "redacted_thinking" not in cleaned |
| assert "planning text without a final answer" in cleaned |
|
|
