Retrieve follow-up questions with the conversation, not the bare question

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] = []

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?"]