From 292c7e9db2c5c419adc02fec3436b878643c0ec0 Mon Sep 17 00:00:00 2001 From: sinnercode228 Date: Fri, 30 Oct 2026 19:58:37 +0700 Subject: [PATCH] Retrieve follow-up questions with the conversation, not the bare question stream_answer retrieved on the raw question only, so a follow-up like "what does it cost?" was embedded with no idea what "it" was and pulled back the wrong chunks; the history reached the LLM but never the retriever. Build the retrieval query from the recent user turns plus the current question and retrieve on that, while still passing the original question to the prompt and using it for the citation highlight. search() takes an optional retrieval_query so the plain /v1/search endpoint is unchanged. Fixes #8 --- backend/src/docmind/rag/service.py | 21 ++++++-- backend/tests/test_rag_service.py | 78 ++++++++++++++++++++++++++++++ 2 files changed, 96 insertions(+), 3 deletions(-) create mode 100644 backend/tests/test_rag_service.py diff --git a/backend/src/docmind/rag/service.py b/backend/src/docmind/rag/service.py index 50e8651..35a469d 100644 --- a/backend/src/docmind/rag/service.py +++ b/backend/src/docmind/rag/service.py @@ -91,13 +91,22 @@ async def search( *, top_k: int | None = None, document_ids: Sequence[str] | None = None, + retrieval_query: str | None = None, ) -> list[Citation]: hits = await self._retriever.retrieve( - tenant_id, query, top_k=top_k, document_ids=document_ids + tenant_id, retrieval_query or query, top_k=top_k, document_ids=document_ids ) titles = await self._titles(tenant_id, sorted({h.chunk.document_id for h in hits})) return to_citations(query, hits, titles) + @staticmethod + def _retrieval_query(question: str, history: Sequence[ChatTurn]) -> str: + """Prepend the recent user turns so a follow-up retrieves on its real subject.""" + prior = [t.content.strip() for t in history if t.role == "user" and t.content.strip()] + if not prior: + return question + return "\n".join([*prior, question]) + async def stream_answer( self, tenant_id: str, @@ -107,7 +116,14 @@ async def stream_answer( top_k: int | None = None, document_ids: Sequence[str] | None = None, ) -> AsyncIterator[RagEvent]: - citations = await self.search(tenant_id, question, top_k=top_k, document_ids=document_ids) + recent = list(history)[-self._history_turns * 2 :] if self._history_turns else [] + citations = await self.search( + tenant_id, + question, + top_k=top_k, + document_ids=document_ids, + retrieval_query=self._retrieval_query(question, recent), + ) yield SourcesEvent(citations) sources = [ @@ -119,7 +135,6 @@ async def stream_answer( ) for c in citations ] - recent = list(history)[-self._history_turns * 2 :] if self._history_turns else [] request = build_request(question, sources, recent, max_tokens=self._max_tokens) parts: list[str] = [] diff --git a/backend/tests/test_rag_service.py b/backend/tests/test_rag_service.py new file mode 100644 index 0000000..d54ea4d --- /dev/null +++ b/backend/tests/test_rag_service.py @@ -0,0 +1,78 @@ +"""Behaviour of RAGService.stream_answer around retrieval and conversation history.""" + +from __future__ import annotations + +from collections.abc import AsyncIterator, Sequence + +import numpy as np + +from docmind.domain import ChatTurn, Chunk, ScoredChunk +from docmind.llm.base import Completion, LLMRequest, TextDelta +from docmind.rag.service import RAGService, SourcesEvent + + +class RecordingRetriever: + """Captures the query string the service actually retrieves on.""" + + def __init__(self) -> None: + self.queries: list[str] = [] + + async def retrieve( + self, + tenant_id: str, + query: str, + *, + top_k: int | None = None, + document_ids: Sequence[str] | None = None, + ) -> list[ScoredChunk]: + self.queries.append(query) + chunk = Chunk( + id="d:1", document_id="d", ordinal=0, text="The Nimbus plan costs $9 a month." + ) + return [ScoredChunk(chunk, 0.9, np.ones(4, dtype=np.float32))] + + +class StubLLM: + name = "stub" + + async def stream(self, request: LLMRequest) -> AsyncIterator[TextDelta | Completion]: + yield TextDelta("It costs $9 a month [1].") + yield Completion(stop_reason="end_turn", model=self.name) + + async def aclose(self) -> None: + return None + + +async def _titles(tenant_id: str, document_ids: Sequence[str]) -> dict[str, str]: + return dict.fromkeys(document_ids, "Nimbus plan") + + +def _service(retriever: RecordingRetriever) -> RAGService: + return RAGService(retriever, StubLLM(), _titles) # type: ignore[arg-type] + + +async def test_followup_retrieval_includes_the_conversation() -> None: + retriever = RecordingRetriever() + history = [ + ChatTurn(role="user", content="Tell me about the Nimbus plan"), + ChatTurn(role="assistant", content="The Nimbus plan is our mid tier."), + ] + events = [ + event + async for event in _service(retriever).stream_answer( + "t", "what does it cost?", history=history + ) + ] + + assert any(isinstance(event, SourcesEvent) for event in events) + assert retriever.queries, "retriever was never called" + query = retriever.queries[0] + assert "what does it cost?" in query + # "it" only means something because the subject is in the history. + assert "Nimbus" in query + + +async def test_retrieval_query_is_the_bare_question_without_history() -> None: + retriever = RecordingRetriever() + _ = [event async for event in _service(retriever).stream_answer("t", "what does it cost?")] + assert retriever.queries == ["what does it cost?"]