diff --git a/app.py b/app.py index 61597b6..ff1fcbb 100644 --- a/app.py +++ b/app.py @@ -5,7 +5,13 @@ import chromadb from src.config import ConfigError, load_config -from src.generation.answerer import answer, build_source_elements +from src.generation.answerer import ( + answer, + build_ref_map, + build_reference_list, + build_source_elements, + strip_llm_references, +) from src.generation.condenser import Condenser from src.health_check import check_models, check_ollama from src.ingestion.ingest import IngestResult, ingest_folder @@ -108,6 +114,7 @@ async def on_chat_start(): condenser = Condenser(model=config.models.llm) cl.user_session.set("condenser", condenser) cl.user_session.set("chat_history", []) + cl.user_session.set("ref_map", {}) # Set up Chainlit settings panel settings = await cl.ChatSettings( @@ -280,18 +287,32 @@ async def on_message(message: cl.Message): retrieval_step.output = f"Retrieved {len(results)} relevant chunks" + # Build ref_map for this answer, extending the session-wide map + existing_ref_map = cl.user_session.get("ref_map") + ref_map = build_ref_map(results, existing_ref_map) + cl.user_session.set("ref_map", ref_map) + + # Collect which reference numbers are used in this answer + used_nums = {ref_map[r.metadata.get("relative_path", r.metadata.get("filename", "unknown"))] for r in results} + # Stream the answer token by token msg = cl.Message(content="") async for token in answer( query, results, model=config.models.llm, + ref_map=ref_map, ): await msg.stream_token(token) + + # Strip any reference list the LLM generated, then append ours + msg.content = strip_llm_references(msg.content) + ref_list = build_reference_list(ref_map, only=used_nums) + msg.content += f"\n\n---\n**References:**\n{ref_list}" await msg.send() # Attach expandable source chunks below the answer - for el_data in build_source_elements(results): + for el_data in build_source_elements(results, ref_map=ref_map): element = cl.Text( name=el_data["name"], content=el_data["content"], diff --git a/config.yaml b/config.yaml index f947252..05ba9af 100644 --- a/config.yaml +++ b/config.yaml @@ -11,7 +11,7 @@ retrieval: semantic_top_k: 20 rrf_k: 60 rerank_top_k: 10 - reranker_model: "bge-reranker-v2-m3" + reranker_model: "BAAI/bge-reranker-v2-m3" paths: chroma_db: "~/.multi_doc_query/chroma_db/" diff --git a/src/config.py b/src/config.py index ce180c8..855a925 100644 --- a/src/config.py +++ b/src/config.py @@ -23,7 +23,7 @@ class RetrievalConfig(BaseModel): semantic_top_k: int = 20 rrf_k: int = 60 rerank_top_k: int = 10 - reranker_model: str = "bge-reranker-v2-m3" + reranker_model: str = "BAAI/bge-reranker-v2-m3" class PathsConfig(BaseModel): diff --git a/src/generation/answerer.py b/src/generation/answerer.py index b2a43be..9ea1418 100644 --- a/src/generation/answerer.py +++ b/src/generation/answerer.py @@ -1,8 +1,9 @@ +import re from collections.abc import AsyncIterator import ollama -from src.retrieval.vector_store import SearchResult +from src.models import SearchResult # --- Not via LangChain: prompt construction and LLM calls use direct # ollama.chat() for full control over streaming and prompt format. --- @@ -11,56 +12,146 @@ "You are a helpful assistant that answers questions based on the provided " "document sources. Follow these rules:\n" "1. Base your answer only on the provided sources.\n" - "2. Cite sources inline using the exact source label shown in the " - "headers, e.g. [filename.pdf, p. 12]. Always use the filename from " - "the source header, never invent labels like 'Excerpt 1'.\n" + "2. Cite sources inline using numeric references, e.g. [1, p. 12] or " + "[2, Section: Methods]. Use the reference numbers shown in the source " + "headers and the Reference List. If citing the same document on different " + "pages, reuse its number, e.g. [1, p. 5] and [1, p. 22]. Do NOT " + "reproduce the Reference List — it will be appended automatically.\n" "3. If sources from different documents conflict, highlight the " "discrepancy and cite both sources.\n" "4. If no sources are relevant to the question, say so clearly." ) -def build_prompt(question: str, results: list[SearchResult]) -> list[dict]: - """Build chat messages for Ollama from a question and search results.""" +def build_ref_map( + results: list[SearchResult], + existing: dict[str, int] | None = None, +) -> dict[str, int]: + """Assign a stable numeric ID to each unique document path. + + Returns a new dict mapping document path to its reference number (1-based). + If *existing* is provided, those mappings are preserved and new documents + get numbers starting after the current maximum. + """ + ref_map = dict(existing) if existing else {} + next_num = max(ref_map.values(), default=0) + 1 + for r in results: + path = _doc_path(r.metadata) + if path not in ref_map: + ref_map[path] = next_num + next_num += 1 + return ref_map + + +def build_reference_list( + ref_map: dict[str, int], + only: set[int] | None = None, +) -> str: + """Format a numbered reference list string. + + If *only* is given, include only the reference numbers in that set. + """ + lines = [] + for path, num in sorted(ref_map.items(), key=lambda x: x[1]): + if only is not None and num not in only: + continue + lines.append(f"[{num}] {path}") + return "\n".join(lines) + + +_REFS_TAIL_RE = re.compile( + r"\n*(?:---\n)?" # optional horizontal rule + r"\*{0,2}" # optional bold markers + r"[Rr]eferences?:?" # "Reference:", "References:", bold variants + r"\*{0,2}" # closing bold markers + r"\n.*", # everything after + re.DOTALL, +) + + +def strip_llm_references(text: str) -> str: + """Remove any trailing reference/references section the LLM generated.""" + return _REFS_TAIL_RE.sub("", text).rstrip() + + +def _doc_path(metadata: dict[str, str | int]) -> str: + """Get the document path from metadata, preferring relative_path.""" + return metadata.get("relative_path", metadata.get("filename", "unknown")) + + +def _source_label(metadata: dict[str, str | int], ref_num: int) -> str: + """Build the inline source label for a context header, e.g. [1, p. 5].""" + section = metadata.get("section_header") + if section: + return f"[{ref_num}, Section: {section}]" + page = metadata.get("page_number", "?") + return f"[{ref_num}, p. {page}]" + + +def _element_name(metadata: dict[str, str | int], ref_num: int) -> str: + """Build a display name for a source element, e.g. [1] path, p. 5.""" + path = _doc_path(metadata) + section = metadata.get("section_header") + if section: + return f"[{ref_num}] {path}, Section: {section}" + page = metadata.get("page_number", "?") + return f"[{ref_num}] {path}, p. {page}" + + +def build_prompt( + question: str, + results: list[SearchResult], + *, + ref_map: dict[str, int] | None = None, +) -> list[dict]: + """Build chat messages for Ollama from a question and search results. + + If *ref_map* is provided, it is used for numbering; otherwise a fresh + map is built from *results*. + """ + if ref_map is None: + ref_map = build_ref_map(results) + context_parts = [] for r in results: - context_parts.append( - f"--- {_source_name(r.metadata)} ---\n{r.text}" - ) + ref_num = ref_map[_doc_path(r.metadata)] + label = _source_label(r.metadata, ref_num) + context_parts.append(f"--- {label} ---\n{r.text}") context = "\n\n".join(context_parts) + # Include reference list in prompt so LLM knows the mapping + ref_list = build_reference_list(ref_map) + return [ {"role": "system", "content": SYSTEM_PROMPT}, { "role": "user", "content": ( f"Document excerpts:\n\n{context}\n\n" + f"Reference List:\n{ref_list}\n\n" f"Question: {question}" ), }, ] -def _source_name(metadata: dict[str, str | int]) -> str: - """Build a display name for a source from its metadata.""" - path = metadata.get("relative_path", metadata.get("filename", "unknown")) - section = metadata.get("section_header") - if section: - return f"Source: {path} | Section: {section}" - page = metadata.get("page_number", "?") - return f"Source: {path} | Page {page}" - - -def build_source_elements(results: list[SearchResult]) -> list[dict]: +def build_source_elements( + results: list[SearchResult], + *, + ref_map: dict[str, int] | None = None, +) -> list[dict]: """Build source element data for Chainlit display. Returns a list of dicts with keys: name, content, display. Order matches the input (relevance-ranked by caller). """ + if ref_map is None: + ref_map = build_ref_map(results) + return [ { - "name": _source_name(r.metadata), + "name": _element_name(r.metadata, ref_map[_doc_path(r.metadata)]), "content": r.text, "display": "side", } @@ -73,9 +164,10 @@ async def answer( results: list[SearchResult], *, model: str = "llama3.1:8b", + ref_map: dict[str, int] | None = None, ) -> AsyncIterator[str]: """Stream answer tokens from Ollama.""" - messages = build_prompt(question, results) + messages = build_prompt(question, results, ref_map=ref_map) stream = ollama.chat(model=model, messages=messages, stream=True) for chunk in stream: token = chunk["message"]["content"] diff --git a/tests/test_answerer.py b/tests/test_answerer.py index 5a2fa91..2c8fd6a 100644 --- a/tests/test_answerer.py +++ b/tests/test_answerer.py @@ -1,5 +1,11 @@ -from src.generation.answerer import build_prompt, build_source_elements -from src.retrieval.vector_store import SearchResult +from src.generation.answerer import ( + build_prompt, + build_ref_map, + build_reference_list, + build_source_elements, + strip_llm_references, +) +from src.models import SearchResult def _make_results(): @@ -17,6 +23,30 @@ def _make_results(): ] +def _make_results_same_doc(): + """Two chunks from the same document on different pages.""" + return [ + SearchResult( + text="Chapter 1 content.", + metadata={"filename": "report.pdf", "relative_path": "reports/annual.pdf", "doc_type": "pdf", "page_number": 5}, + distance=0.1, + ), + SearchResult( + text="Chapter 3 content.", + metadata={"filename": "report.pdf", "relative_path": "reports/annual.pdf", "doc_type": "pdf", "page_number": 22}, + distance=0.2, + ), + SearchResult( + text="Berlin is the capital of Germany.", + metadata={"filename": "europe.pdf", "relative_path": "countries/europe.pdf", "doc_type": "pdf", "page_number": 12}, + distance=0.3, + ), + ] + + +# --- build_prompt tests --- + + def test_build_prompt_includes_context(): """The prompt includes the text from search results.""" messages = build_prompt("What is the capital of France?", _make_results()) @@ -32,16 +62,44 @@ def test_build_prompt_includes_question(): assert "What is the capital of France?" in content -def test_build_prompt_source_format(): - """Context is formatted with source citations using relative path.""" +def test_build_prompt_numeric_source_headers(): + """Context headers use numeric reference format [N].""" messages = build_prompt("question", _make_results()) content = " ".join(m["content"] for m in messages) - assert "--- Source: countries/geo.pdf | Page 5 ---" in content - assert "--- Source: countries/europe.pdf | Page 12 ---" in content + assert "--- [1, p. 5] ---" in content + assert "--- [2, p. 12] ---" in content + + +def test_build_prompt_reference_list_appended(): + """A numbered reference list is appended after the context.""" + messages = build_prompt("question", _make_results()) + user_content = next(m["content"] for m in messages if m["role"] == "user") + assert "[1] countries/geo.pdf" in user_content + assert "[2] countries/europe.pdf" in user_content -def test_build_prompt_uses_relative_path(): - """build_prompt prefers relative_path over filename in source headers.""" +def test_build_prompt_same_doc_shares_number(): + """Two chunks from the same document share the same reference number.""" + messages = build_prompt("question", _make_results_same_doc()) + content = " ".join(m["content"] for m in messages) + # Both chunks from reports/annual.pdf should be [1] + assert "--- [1, p. 5] ---" in content + assert "--- [1, p. 22] ---" in content + # europe.pdf should be [2] + assert "--- [2, p. 12] ---" in content + + +def test_build_prompt_reference_list_deduplicates(): + """Reference list has one entry per unique document, not per chunk.""" + messages = build_prompt("question", _make_results_same_doc()) + user_content = next(m["content"] for m in messages if m["role"] == "user") + # Should have exactly one entry for reports/annual.pdf + assert user_content.count("[1] reports/annual.pdf") == 1 + assert user_content.count("[2] countries/europe.pdf") == 1 + + +def test_build_prompt_uses_relative_path_in_references(): + """Reference list uses relative_path, not filename.""" results = [ SearchResult( text="Some text.", @@ -50,13 +108,12 @@ def test_build_prompt_uses_relative_path(): ), ] messages = build_prompt("question", results) - content = " ".join(m["content"] for m in messages) - assert "--- Source: reports/annual.pdf | Page 3 ---" in content - assert "--- Source: report.pdf" not in content + user_content = next(m["content"] for m in messages if m["role"] == "user") + assert "[1] reports/annual.pdf" in user_content def test_build_prompt_falls_back_to_filename(): - """Without relative_path, build_prompt uses filename.""" + """Without relative_path, reference list uses filename.""" results = [ SearchResult( text="Some text.", @@ -65,8 +122,41 @@ def test_build_prompt_falls_back_to_filename(): ), ] messages = build_prompt("question", results) + user_content = next(m["content"] for m in messages if m["role"] == "user") + assert "[1] legacy.pdf" in user_content + + +def test_build_prompt_section_header_in_source(): + """Markdown sources show section header instead of page number in context header.""" + results = [ + SearchResult( + text="Install instructions.", + metadata={ + "filename": "README.md", + "relative_path": "README.md", + "doc_type": "md", + "page_number": 1, + "section_header": "Getting Started > Installation", + }, + distance=0.1, + ), + ] + messages = build_prompt("How to install?", results) content = " ".join(m["content"] for m in messages) - assert "--- Source: legacy.pdf | Page 1 ---" in content + assert "--- [1, Section: Getting Started > Installation] ---" in content + # Reference list should still show the document path + user_content = next(m["content"] for m in messages if m["role"] == "user") + assert "[1] README.md" in user_content + + +def test_build_prompt_system_prompt_numeric_format(): + """System prompt instructs the LLM to use [N, p. X] citation format.""" + messages = build_prompt("question", _make_results()) + system = next(m["content"] for m in messages if m["role"] == "system") + assert "[1, p. 12]" in system or "[N, p." in system + + +# --- build_source_elements tests --- def test_build_source_elements_returns_list(): @@ -79,17 +169,19 @@ def test_build_source_elements_returns_list(): assert "display" in el -def test_build_source_elements_name_format(): - """Element name follows 'Source: path | Page N' format.""" - results = [ - SearchResult( - text="Some text.", - metadata={"filename": "report.pdf", "relative_path": "reports/annual.pdf", "doc_type": "pdf", "page_number": 5}, - distance=0.1, - ), - ] - elements = build_source_elements(results) - assert elements[0]["name"] == "Source: reports/annual.pdf | Page 5" +def test_build_source_elements_numeric_names(): + """Element names include numeric reference number.""" + elements = build_source_elements(_make_results()) + assert elements[0]["name"] == "[1] countries/geo.pdf, p. 5" + assert elements[1]["name"] == "[2] countries/europe.pdf, p. 12" + + +def test_build_source_elements_same_doc_shares_number(): + """Elements from the same document share the same reference number.""" + elements = build_source_elements(_make_results_same_doc()) + assert elements[0]["name"] == "[1] reports/annual.pdf, p. 5" + assert elements[1]["name"] == "[1] reports/annual.pdf, p. 22" + assert elements[2]["name"] == "[2] countries/europe.pdf, p. 12" def test_build_source_elements_preserves_order(): @@ -109,11 +201,11 @@ def test_build_source_elements_falls_back_to_filename(): ), ] elements = build_source_elements(results) - assert elements[0]["name"] == "Source: old.pdf | Page 1" + assert elements[0]["name"] == "[1] old.pdf, p. 1" -def test_source_name_with_section_header(): - """Markdown sources show section header instead of page number.""" +def test_build_source_elements_section_header(): + """Markdown source elements show section header.""" results = [ SearchResult( text="Install instructions.", @@ -127,23 +219,163 @@ def test_source_name_with_section_header(): distance=0.1, ), ] - messages = build_prompt("How to install?", results) - content = " ".join(m["content"] for m in messages) - assert "--- Source: README.md | Section: Getting Started > Installation ---" in content - elements = build_source_elements(results) - assert elements[0]["name"] == "Source: README.md | Section: Getting Started > Installation" + assert elements[0]["name"] == "[1] README.md, Section: Getting Started > Installation" + + +# --- build_ref_map tests --- + + +def test_build_ref_map_assigns_numbers(): + """build_ref_map assigns 1-based numbers to unique documents.""" + ref_map = build_ref_map(_make_results()) + assert ref_map == {"countries/geo.pdf": 1, "countries/europe.pdf": 2} -def test_source_name_without_section_header(): - """Non-Markdown sources still show page number (unchanged behaviour).""" +def test_build_ref_map_deduplicates(): + """Same document gets the same number regardless of how many chunks.""" + ref_map = build_ref_map(_make_results_same_doc()) + assert ref_map == {"reports/annual.pdf": 1, "countries/europe.pdf": 2} + + +def test_build_ref_map_extends_existing(): + """Passing an existing ref_map preserves those numbers and continues.""" + existing = {"countries/geo.pdf": 1} + ref_map = build_ref_map(_make_results(), existing) + # geo.pdf keeps number 1, europe.pdf gets 2 + assert ref_map["countries/geo.pdf"] == 1 + assert ref_map["countries/europe.pdf"] == 2 + + +def test_build_ref_map_extends_with_new_docs(): + """New docs in second query get numbers after the existing max.""" + existing = {"old_doc.pdf": 1, "another.pdf": 2} results = [ SearchResult( - text="Some text.", - metadata={"filename": "report.pdf", "doc_type": "pdf", "page_number": 7}, + text="New content.", + metadata={"filename": "new.pdf", "relative_path": "new.pdf", "doc_type": "pdf", "page_number": 1}, distance=0.1, ), ] - messages = build_prompt("question", results) + ref_map = build_ref_map(results, existing) + assert ref_map["old_doc.pdf"] == 1 + assert ref_map["another.pdf"] == 2 + assert ref_map["new.pdf"] == 3 + + +def test_build_ref_map_no_mutation(): + """build_ref_map does not mutate the existing dict passed in.""" + existing = {"countries/geo.pdf": 1} + original = existing.copy() + build_ref_map(_make_results(), existing) + assert existing == original + + +# --- build_reference_list tests --- + + +def test_build_reference_list_format(): + """build_reference_list returns a formatted reference list string.""" + ref_map = {"countries/geo.pdf": 1, "countries/europe.pdf": 2} + text = build_reference_list(ref_map) + assert "[1] countries/geo.pdf" in text + assert "[2] countries/europe.pdf" in text + + +def test_build_reference_list_sorted_by_number(): + """References are listed in numeric order.""" + ref_map = {"b.pdf": 2, "a.pdf": 1, "c.pdf": 3} + text = build_reference_list(ref_map) + lines = text.strip().split("\n") + assert lines[0] == "[1] a.pdf" + assert lines[1] == "[2] b.pdf" + assert lines[2] == "[3] c.pdf" + + +def test_build_reference_list_only_for_given_numbers(): + """When filtering to specific numbers, only those appear.""" + ref_map = {"a.pdf": 1, "b.pdf": 2, "c.pdf": 3} + text = build_reference_list(ref_map, only={1, 3}) + assert "[1] a.pdf" in text + assert "[3] c.pdf" in text + assert "[2]" not in text + + +# --- build_prompt with explicit ref_map --- + + +def test_build_prompt_with_ref_map_uses_existing_numbers(): + """build_prompt respects a pre-existing ref_map for numbering.""" + ref_map = {"countries/geo.pdf": 5, "countries/europe.pdf": 6} + messages = build_prompt("question", _make_results(), ref_map=ref_map) content = " ".join(m["content"] for m in messages) - assert "--- Source: report.pdf | Page 7 ---" in content + assert "--- [5, p. 5] ---" in content + assert "--- [6, p. 12] ---" in content + + +# --- build_source_elements with explicit ref_map --- + + +def test_build_source_elements_with_ref_map(): + """build_source_elements respects a pre-existing ref_map.""" + ref_map = {"countries/geo.pdf": 5, "countries/europe.pdf": 6} + elements = build_source_elements(_make_results(), ref_map=ref_map) + assert elements[0]["name"] == "[5] countries/geo.pdf, p. 5" + assert elements[1]["name"] == "[6] countries/europe.pdf, p. 12" + + +# --- strip_llm_references tests --- + + +def test_strip_llm_references_removes_trailing_section(): + """Strips a trailing 'Reference:' or 'References:' section.""" + text = ( + "The answer is 42.\n\n" + "Reference:\n" + "[2, p. 1] and [1, p. 4]" + ) + assert strip_llm_references(text) == "The answer is 42." + + +def test_strip_llm_references_removes_with_list(): + """Strips a trailing references section with numbered entries.""" + text = ( + "FPCA has been applied in rowing.\n\n" + "References:\n" + "[1] some_doc.pdf\n" + "[2] other_doc.pdf" + ) + assert strip_llm_references(text) == "FPCA has been applied in rowing." + + +def test_strip_llm_references_preserves_mid_text(): + """Does not strip 'reference' that appears mid-text.""" + text = "See the reference list in [1, p. 5] for details." + assert strip_llm_references(text) == text + + +def test_strip_llm_references_no_references(): + """Returns text unchanged when there is no trailing references section.""" + text = "A clean answer with no references section." + assert strip_llm_references(text) == text + + +def test_strip_llm_references_handles_bold_heading(): + """Strips references section with markdown bold heading.""" + text = ( + "The answer.\n\n" + "**References:**\n" + "[1] doc.pdf" + ) + assert strip_llm_references(text) == "The answer." + + +def test_strip_llm_references_handles_hr_before(): + """Strips references section preceded by a horizontal rule.""" + text = ( + "The answer.\n\n" + "---\n" + "References:\n" + "[1] doc.pdf" + ) + assert strip_llm_references(text) == "The answer."