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