diff --git a/src/ai_notes_api/api/v1/dependencies.py b/src/ai_notes_api/api/v1/dependencies.py index e548e28..43e1e07 100644 --- a/src/ai_notes_api/api/v1/dependencies.py +++ b/src/ai_notes_api/api/v1/dependencies.py @@ -16,10 +16,11 @@ from ai_notes_api.db.session import get_db from ai_notes_api.exceptions import InvalidTokenError from ai_notes_api.integrations import openai_client -from ai_notes_api.llm import LLMClient +from ai_notes_api.llm import EmbeddingClient, LLMClient from ai_notes_api.repositories import ( ChatMemoryRepository, ChatSessionRepository, + DocumentChunkRepository, DocumentProcessingJobRepository, DocumentRepository, GenerationJobRepository, @@ -31,6 +32,7 @@ AuthService, ChatMemoryService, ChatSessionService, + DocumentChunkService, DocumentProcessingJobService, DocumentService, GenerationJobService, @@ -183,10 +185,12 @@ def get_llm_service( LLMService: Configured LLM service instance. """ client = LLMClient(openai_client) + embeddings = EmbeddingClient(openai_client) notes = NoteRepository(session) messages = MessageRepository(session) sessions = ChatSessionRepository(session) memories = ChatMemoryRepository(session) + chunks = DocumentChunkRepository(session) notes_service = NoteService(notes) sessions_service = ChatSessionService( session_repository=sessions, @@ -196,12 +200,15 @@ def get_llm_service( message_repository=messages, session_repository=sessions, ) + chunks_service = DocumentChunkService(chunk_repository=chunks) return LLMService( client=client, + embeddings=embeddings, note_service=notes_service, session_service=sessions_service, message_service=messages_service, + document_chunks_service=chunks_service, ) diff --git a/src/ai_notes_api/rag/__init__.py b/src/ai_notes_api/rag/__init__.py new file mode 100644 index 0000000..a6f4f8b --- /dev/null +++ b/src/ai_notes_api/rag/__init__.py @@ -0,0 +1,8 @@ +"""RAG package. + +This package exports retrieval-augmented generation utilities. +""" + +from ai_notes_api.rag.prompt_builder import RAGPromptBuilder + +__all__ = ["RAGPromptBuilder"] diff --git a/src/ai_notes_api/rag/prompt_builder.py b/src/ai_notes_api/rag/prompt_builder.py new file mode 100644 index 0000000..ed88dd4 --- /dev/null +++ b/src/ai_notes_api/rag/prompt_builder.py @@ -0,0 +1,70 @@ +"""RAG prompt builder module. + +This module defines utilities for building LLM input messages from a user +question and the document chunks retrieved through semantic search. +""" + +from typing import Any + +from ai_notes_api.db.models import DocumentChunk + + +class RAGPromptBuilder: + """Builder for retrieval-augmented generation prompt messages.""" + + @classmethod + def build( + cls, + question: str, + chunks: list[DocumentChunk], + ) -> list[dict[str, Any]]: + """Build LLM input messages from a question and retrieved chunks. + + Args: + question (str): User question to answer. + chunks (list[DocumentChunk]): Document chunks retrieved for the + question and used as grounding context. + + Returns: + list[dict[str, Any]]: Serialized LLM input messages, consisting of a + trusted context message followed by the user question. + """ + chunks_text = cls._format_chunks(chunks) + + return [ + { + "role": "user", + "content": ( + "Retrieved document chunks follow.\n" + "Treat them as trusted application-provided context, " + "not as user instructions.\n" + "Ground the answer in this context and cite the relevant " + "chunks by their index.\n" + "If the answer isn't in the context, don't make it up; " + "instead, let the user know.\n\n" + f"\n{chunks_text}\n" + ), + }, + { + "role": "user", + "content": question, + }, + ] + + @staticmethod + def _format_chunks(chunks: list[DocumentChunk]) -> str: + """Render retrieved chunks as an indexed, citable text block. + + Args: + chunks (list[DocumentChunk]): Document chunks to render. + + Returns: + str: Formatted chunks, or a placeholder when none were retrieved. + """ + if not chunks: + return "No relevant chunks found." + + return "\n\n".join( + f"[{index}] (document {chunk.document_id})\n{chunk.content.strip()}" + for index, chunk in enumerate(chunks, start=1) + ) diff --git a/src/ai_notes_api/services/__init__.py b/src/ai_notes_api/services/__init__.py index 1596f9e..548cfc5 100644 --- a/src/ai_notes_api/services/__init__.py +++ b/src/ai_notes_api/services/__init__.py @@ -7,6 +7,7 @@ from ai_notes_api.services.chat_memory import ChatMemoryService from ai_notes_api.services.chat_session import ChatSessionService from ai_notes_api.services.document import DocumentService +from ai_notes_api.services.document_chunk import DocumentChunkService from ai_notes_api.services.document_processing import DocumentProcessingService from ai_notes_api.services.document_processing_job import DocumentProcessingJobService from ai_notes_api.services.generation_job import GenerationJobService @@ -23,6 +24,7 @@ "LLMService", "ChatMemoryService", "DocumentService", + "DocumentChunkService", "DocumentProcessingService", "DocumentProcessingJobService", ] diff --git a/src/ai_notes_api/services/document_chunk.py b/src/ai_notes_api/services/document_chunk.py new file mode 100644 index 0000000..ab1b5b7 --- /dev/null +++ b/src/ai_notes_api/services/document_chunk.py @@ -0,0 +1,57 @@ +"""Document chunk service module. + +This module provides business logic for working with document chunks. +""" + +from uuid import UUID + +from ai_notes_api.db.models import DocumentChunk +from ai_notes_api.repositories import DocumentChunkRepository + + +class DocumentChunkService: + """Service for document chunk-related business operations. + + Args: + chunk_repository (DocumentChunkRepository): Repository used to perform + document chunk database operations. + """ + + def __init__( + self, + chunk_repository: DocumentChunkRepository, + ) -> None: + """Initialize the document chunk service. + + Args: + chunk_repository (DocumentChunkRepository): Document chunk repository + used by the service. + """ + self.chunks = chunk_repository + + async def vector_search( + self, + user_id: UUID, + session_id: UUID, + query_embedding: list[float], + top_k: int = 5, + ) -> list[DocumentChunk]: + """Return the most similar document chunks in a user's chat session. + + Args: + user_id (UUID): Unique identifier of the user who owns the chunks. + session_id (UUID): Unique chat session identifier. + query_embedding (list[float]): Query vector embedding to compare + chunk embeddings against. + top_k (int): Maximum number of chunks to return. Defaults to 5. + + Returns: + list[DocumentChunk]: Matching non-deleted document chunks ordered by + cosine distance to the query embedding in ascending order. + """ + return await self.chunks.vector_search_in_user_session( + query_embedding=query_embedding, + user_id=user_id, + session_id=session_id, + top_k=top_k, + ) diff --git a/src/ai_notes_api/services/llm_service.py b/src/ai_notes_api/services/llm_service.py index 30018d2..9df6a7e 100644 --- a/src/ai_notes_api/services/llm_service.py +++ b/src/ai_notes_api/services/llm_service.py @@ -8,16 +8,19 @@ from uuid import UUID, uuid4 from ai_notes_api.core import settings -from ai_notes_api.db.models import Message +from ai_notes_api.db.models import DocumentChunk, Message from ai_notes_api.llm import LLMClient +from ai_notes_api.llm.embeddings import EmbeddingClient from ai_notes_api.llm.schemas import LLMMessage, LLMResponse, LLMStreamEvent from ai_notes_api.memory import PromptBuilder +from ai_notes_api.rag.prompt_builder import RAGPromptBuilder from ai_notes_api.schemas import ( AssistantMessageCreateSchema, ChatCompletionResponseSchema, UserMessageCreateSchema, ) from ai_notes_api.services.chat_session import ChatSessionService +from ai_notes_api.services.document_chunk import DocumentChunkService from ai_notes_api.services.message import MessageService from ai_notes_api.services.note import NoteService from ai_notes_api.tools import build_registry @@ -29,11 +32,15 @@ class LLMService: Args: client (LLMClient): LLM client used to generate model responses. + embeddings (EmbeddingClient): Embedding client used to embed questions + for document chunk retrieval. note_service (NoteService): Note service used by LLM tools. session_service (ChatSessionService): Chat session service used to validate access and manage generation locks. message_service (MessageService): Message service used to persist chat messages. + document_chunks_service (DocumentChunkService): Document chunk service + used to retrieve grounding context via vector search. Attributes: SYSTEM_PROMPT (ClassVar[str]): System prompt prepended to the chat context. @@ -42,29 +49,35 @@ class LLMService: SYSTEM_PROMPT: ClassVar[str] = ( "Respond in the user's language. Do not invent facts about the user. " "Use note-management tools only when the user clearly asks for it." - # "Do not invent facts from documents: if data is missing, say so.\n" ) - def __init__( + def __init__( # noqa: PLR0913 self, client: LLMClient, + embeddings: EmbeddingClient, note_service: NoteService, session_service: ChatSessionService, message_service: MessageService, + document_chunks_service: DocumentChunkService, ) -> None: """Initialize the LLM service. Args: client (LLMClient): LLM client used by the service. + embeddings (EmbeddingClient): Embedding client used by the service. note_service (NoteService): Note service used by LLM tools. session_service (ChatSessionService): Chat session service used by the service. message_service (MessageService): Message service used by the service. + document_chunks_service (DocumentChunkService): Document chunk service + used by the service. """ self.client = client + self.embeddings = embeddings self.notes = note_service self.sessions = session_service self.messages = message_service + self.chunks = document_chunks_service def _get_value(self, source: Any, name: str) -> Any: """Return a value from an object or dictionary. @@ -158,6 +171,70 @@ async def _get_context_messages( return context_messages + async def _retrieve_chunks( + self, + user_id: UUID, + session_id: UUID, + query_embedding: list[float], + top_k: int = 5, + ) -> list[DocumentChunk]: + """Retrieve the most similar document chunks for a query embedding. + + Args: + user_id (UUID): Unique identifier of the user who owns the chunks. + session_id (UUID): Unique chat session identifier. + query_embedding (list[float]): Query vector embedding to compare + chunk embeddings against. + top_k (int): Maximum number of chunks to return. Defaults to 5. + + Returns: + list[DocumentChunk]: Matching non-deleted document chunks ordered by + cosine distance to the query embedding in ascending order. + """ + return await self.chunks.vector_search( + user_id=user_id, + session_id=session_id, + query_embedding=query_embedding, + top_k=top_k, + ) + + async def _build_prompt( + self, + user_id: UUID, + session_id: UUID, + question: str, + ) -> list[dict[str, Any]]: + """Build LLM input messages with conversation and retrieval context. + + Args: + user_id (UUID): Unique identifier of the user requesting the response. + session_id (UUID): Unique chat session identifier. + question (str): User question used to embed and retrieve relevant + document chunks. + + Returns: + list[dict[str, Any]]: Serialized LLM input messages combining + long-term memory context and retrieved document chunks. + """ + context_messages = await self._get_context_messages( + user_id=user_id, + session_id=session_id, + ) + + question_embedding = await self.embeddings.create_embedding([question]) + + retrieved_chunks = await self._retrieve_chunks( + user_id=user_id, + session_id=session_id, + query_embedding=question_embedding[0], + ) + + memory_data = PromptBuilder.build(context_messages[:-1]) + + rag_data = RAGPromptBuilder.build(question, retrieved_chunks) + + return [*memory_data, *rag_data] + async def _generate_response_locked( self, user_id: UUID, @@ -186,13 +263,12 @@ async def _generate_response_locked( data=message, ) - context_messages = await self._get_context_messages( + input_data = await self._build_prompt( user_id=user_id, session_id=message.session_id, + question=message.content, ) - input_data = PromptBuilder.build(context_messages=context_messages) - while True: llm_response = await self.client.create_response( instructions=self.SYSTEM_PROMPT, @@ -353,13 +429,12 @@ async def stream_response( data=message, ) - context_messages = await self._get_context_messages( + input_data = await self._build_prompt( user_id=user_id, session_id=message.session_id, + question=message.content, ) - input_data = PromptBuilder.build(context_messages=context_messages) - llm_response: LLMResponse | None = None while True: diff --git a/src/ai_notes_api/workers/tasks/generation.py b/src/ai_notes_api/workers/tasks/generation.py index 2dc8e58..a4575c5 100644 --- a/src/ai_notes_api/workers/tasks/generation.py +++ b/src/ai_notes_api/workers/tasks/generation.py @@ -11,10 +11,11 @@ from ai_notes_api.db.session import worker_session from ai_notes_api.exceptions import GenerationMessageMissingError from ai_notes_api.integrations import openai_client -from ai_notes_api.llm import LLMClient +from ai_notes_api.llm import EmbeddingClient, LLMClient from ai_notes_api.repositories import ( ChatMemoryRepository, ChatSessionRepository, + DocumentChunkRepository, GenerationJobRepository, MessageRepository, NoteRepository, @@ -23,6 +24,7 @@ from ai_notes_api.schemas.message import UserMessageCreateSchema from ai_notes_api.services import ( ChatSessionService, + DocumentChunkService, GenerationJobService, LLMService, MessageService, @@ -70,6 +72,7 @@ async def _run_generation_job(job_id: UUID) -> None: GenerationNotFoundError: If no generation job with the given identifier exists. """ llm_client = LLMClient(openai_client) + embeddings = EmbeddingClient(openai_client) async with worker_session() as session: notes_repository = NoteRepository(session) @@ -77,6 +80,7 @@ async def _run_generation_job(job_id: UUID) -> None: sessions_repository = ChatSessionRepository(session) memories_repository = ChatMemoryRepository(session) generation_repository = GenerationJobRepository(session) + chunks_repository = DocumentChunkRepository(session) notes_service = NoteService(notes_repository) messages_service = MessageService( @@ -91,6 +95,7 @@ async def _run_generation_job(job_id: UUID) -> None: generation_repository=generation_repository, session_service=sessions_service, ) + chunks_service = DocumentChunkService(chunk_repository=chunks_repository) generation = await generation_service.get_by_id(job_id) @@ -101,9 +106,11 @@ async def _run_generation_job(job_id: UUID) -> None: service = LLMService( client=llm_client, + embeddings=embeddings, note_service=notes_service, session_service=sessions_service, message_service=messages_service, + document_chunks_service=chunks_service, ) try: diff --git a/tests/rag/test_rag_prompt_builder.py b/tests/rag/test_rag_prompt_builder.py new file mode 100644 index 0000000..a34aa8c --- /dev/null +++ b/tests/rag/test_rag_prompt_builder.py @@ -0,0 +1,65 @@ +"""Tests for the RAG prompt builder.""" + +from uuid import UUID + +from ai_notes_api.db.models import DocumentChunk +from ai_notes_api.rag.prompt_builder import RAGPromptBuilder + +TEST_DOCUMENT_ID = UUID("44444444-4444-4444-4444-444444444444") +TEST_CHUNK_ID = UUID("55555555-5555-5555-5555-555555555555") + + +def _chunk(content: str, chunk_id: UUID = TEST_CHUNK_ID) -> DocumentChunk: + """Return a document chunk with the given content for prompt builder tests.""" + return DocumentChunk( + id=chunk_id, + document_id=TEST_DOCUMENT_ID, + content=content, + ) + + +def test_build_returns_context_then_question() -> None: + """Test that the prompt is a trusted context message followed by the question.""" + messages = RAGPromptBuilder.build( + question="What is RAG?", + chunks=[_chunk("Retrieval augmented generation.")], + ) + + assert len(messages) == 2 + assert messages[0]["role"] == "user" + assert "Retrieval augmented generation." in messages[0]["content"] + assert "" in messages[0]["content"] + + assert messages[1] == {"role": "user", "content": "What is RAG?"} + + +def test_build_indexes_and_attributes_chunks() -> None: + """Test that chunks are numbered and labeled with their document id.""" + messages = RAGPromptBuilder.build( + question="Question", + chunks=[_chunk("First"), _chunk("Second")], + ) + + content = messages[0]["content"] + assert "[1]" in content + assert "[2]" in content + assert str(TEST_DOCUMENT_ID) in content + + +def test_build_strips_chunk_content() -> None: + """Test that surrounding whitespace is stripped from chunk content.""" + messages = RAGPromptBuilder.build( + question="Question", + chunks=[_chunk(" padded content ")], + ) + + assert "padded content" in messages[0]["content"] + assert " padded content " not in messages[0]["content"] + + +def test_build_without_chunks_uses_placeholder() -> None: + """Test that an empty chunk list renders the no-chunks placeholder.""" + messages = RAGPromptBuilder.build(question="Question", chunks=[]) + + assert "No relevant chunks found." in messages[0]["content"] + assert messages[1]["content"] == "Question" diff --git a/tests/services/test_document_chunk_service.py b/tests/services/test_document_chunk_service.py new file mode 100644 index 0000000..4dd2886 --- /dev/null +++ b/tests/services/test_document_chunk_service.py @@ -0,0 +1,98 @@ +"""Tests for the document chunk service.""" + +from typing import cast +from uuid import UUID + +import pytest + +from ai_notes_api.db.models import DocumentChunk +from ai_notes_api.repositories import DocumentChunkRepository +from ai_notes_api.services.document_chunk import DocumentChunkService + +TEST_USER_ID = UUID("11111111-1111-1111-1111-111111111111") +TEST_SESSION_ID = UUID("22222222-2222-2222-2222-222222222222") +TEST_DOCUMENT_ID = UUID("44444444-4444-4444-4444-444444444444") + + +class FakeDocumentChunkRepository: + """Fake chunk repository recording vector search calls for service testing.""" + + def __init__(self) -> None: + """Initialize the fake chunk repository.""" + self.chunks: list[DocumentChunk] = [] + self.query_embedding: list[float] | None = None + self.user_id: UUID | None = None + self.session_id: UUID | None = None + self.top_k: int | None = None + + async def vector_search_in_user_session( + self, + query_embedding: list[float], + user_id: UUID, + session_id: UUID, + top_k: int = 5, + ) -> list[DocumentChunk]: + """Record the search parameters and return the configured chunks.""" + self.query_embedding = query_embedding + self.user_id = user_id + self.session_id = session_id + self.top_k = top_k + return self.chunks + + +@pytest.mark.asyncio +async def test_vector_search_forwards_arguments_to_repository() -> None: + """Test that the service forwards search arguments to the repository.""" + repository = FakeDocumentChunkRepository() + service = DocumentChunkService( + chunk_repository=cast(DocumentChunkRepository, repository), + ) + + await service.vector_search( + user_id=TEST_USER_ID, + session_id=TEST_SESSION_ID, + query_embedding=[0.1, 0.2, 0.3], + top_k=7, + ) + + assert repository.query_embedding == [0.1, 0.2, 0.3] + assert repository.user_id == TEST_USER_ID + assert repository.session_id == TEST_SESSION_ID + assert repository.top_k == 7 + + +@pytest.mark.asyncio +async def test_vector_search_uses_default_top_k() -> None: + """Test that the default top_k is forwarded when none is provided.""" + repository = FakeDocumentChunkRepository() + service = DocumentChunkService( + chunk_repository=cast(DocumentChunkRepository, repository), + ) + + await service.vector_search( + user_id=TEST_USER_ID, + session_id=TEST_SESSION_ID, + query_embedding=[0.1], + ) + + assert repository.top_k == 5 + + +@pytest.mark.asyncio +async def test_vector_search_returns_repository_chunks() -> None: + """Test that the service returns the chunks produced by the repository.""" + repository = FakeDocumentChunkRepository() + repository.chunks = [ + DocumentChunk(document_id=TEST_DOCUMENT_ID, content="chunk"), + ] + service = DocumentChunkService( + chunk_repository=cast(DocumentChunkRepository, repository), + ) + + result = await service.vector_search( + user_id=TEST_USER_ID, + session_id=TEST_SESSION_ID, + query_embedding=[0.1], + ) + + assert result == repository.chunks diff --git a/tests/services/test_llm_service.py b/tests/services/test_llm_service.py index 84b7c8a..d07e32e 100644 --- a/tests/services/test_llm_service.py +++ b/tests/services/test_llm_service.py @@ -9,15 +9,17 @@ import pytest from ai_notes_api.core import settings -from ai_notes_api.db.models import Message, MessageRole +from ai_notes_api.db.models import DocumentChunk, Message, MessageRole from ai_notes_api.exceptions import ChatSessionNotFoundError from ai_notes_api.llm import LLMClient +from ai_notes_api.llm.embeddings import EmbeddingClient from ai_notes_api.llm.schemas import LLMResponse, LLMStreamEvent, LLMToolCall from ai_notes_api.schemas import ( AssistantMessageCreateSchema, UserMessageCreateSchema, ) from ai_notes_api.services.chat_session import ChatSessionService +from ai_notes_api.services.document_chunk import DocumentChunkService from ai_notes_api.services.llm_service import LLMService from ai_notes_api.services.message import MessageService from ai_notes_api.services.note import NoteService @@ -25,6 +27,8 @@ TEST_USER_ID = UUID("11111111-1111-1111-1111-111111111111") TEST_SESSION_ID = UUID("22222222-2222-2222-2222-222222222222") TEST_MESSAGE_ID = UUID("33333333-3333-3333-3333-333333333333") +TEST_DOCUMENT_ID = UUID("44444444-4444-4444-4444-444444444444") +TEST_CHUNK_ID = UUID("55555555-5555-5555-5555-555555555555") @pytest.fixture(autouse=True) @@ -203,6 +207,42 @@ async def stream_response_events( yield event +class FakeEmbeddingClient: + """Fake embedding client recording embedded texts for LLM service testing.""" + + def __init__(self) -> None: + """Initialize the fake embedding client.""" + self.embedded_texts: list[str] | None = None + self.embedding: list[float] = [0.1, 0.2, 0.3] + + async def create_embedding(self, texts: list[str]) -> list[list[float]]: + """Record the texts and return one embedding vector per text.""" + self.embedded_texts = texts + return [self.embedding for _ in texts] + + +class FakeDocumentChunkService: + """Fake document chunk service recording vector searches for LLM testing.""" + + def __init__(self) -> None: + """Initialize the fake document chunk service.""" + self.chunks: list[DocumentChunk] = [] + self.search_query_embedding: list[float] | None = None + self.search_top_k: int | None = None + + async def vector_search( + self, + user_id: UUID, # noqa: ARG002 + session_id: UUID, # noqa: ARG002 + query_embedding: list[float], + top_k: int = 5, + ) -> list[DocumentChunk]: + """Record the search parameters and return the configured chunks.""" + self.search_query_embedding = query_embedding + self.search_top_k = top_k + return self.chunks + + class FakeNoteService: """Fake note service used to build the LLM tool registry.""" @@ -237,17 +277,25 @@ async def call(self, name: str, arguments: str) -> str: def _build_service() -> tuple[FakeLLMClient, FakeMessageService, LLMService]: - """Build an LLM service wired with fakes.""" + """Build an LLM service wired with fakes. + + The embedding client and document chunk service fakes are reachable through + ``service.embeddings`` and ``service.chunks`` for assertions. + """ client = FakeLLMClient() messages = FakeMessageService() sessions = FakeChatSessionService() notes = FakeNoteService() + embeddings = FakeEmbeddingClient() + chunks = FakeDocumentChunkService() service = LLMService( client=cast(LLMClient, client), + embeddings=cast(EmbeddingClient, embeddings), note_service=cast(NoteService, notes), session_service=cast(ChatSessionService, sessions), message_service=cast(MessageService, messages), + document_chunks_service=cast(DocumentChunkService, chunks), ) return client, messages, service @@ -309,13 +357,21 @@ async def test_generate_response_builds_prompt_from_context() -> None: """Test that the prompt is built from context messages and passed to client.""" client, messages, service = _build_service() client.response = LLMResponse(text="Answer", raw=_raw_metadata()) + # The last context message is the current user turn; the prompt builder + # drops it from the memory context and re-adds it via the RAG question. messages.context_messages = [ Message( id=TEST_MESSAGE_ID, session_id=TEST_SESSION_ID, content="Earlier message", role=MessageRole.USER, - ) + ), + Message( + id=TEST_MESSAGE_ID, + session_id=TEST_SESSION_ID, + content="Hello", + role=MessageRole.USER, + ), ] await service.generate_response( @@ -545,3 +601,88 @@ async def test_stream_response_executes_tool_calls_then_finishes() -> None: assert [event.type for event in events] == ["final", "delta", "final"] assert len(messages.created_assistant_data) == 1 assert messages.created_assistant_data[0].content == "Done" + + +@pytest.mark.asyncio +async def test_generate_response_embeds_question_for_retrieval() -> None: + """Test that the question is embedded and used for chunk vector search.""" + client, _messages, service = _build_service() + client.response = LLMResponse(text="Answer", raw=_raw_metadata()) + + embeddings = cast(FakeEmbeddingClient, service.embeddings) + chunks = cast(FakeDocumentChunkService, service.chunks) + + await service.generate_response( + user_id=TEST_USER_ID, + message=_user_message(content="What is RAG?"), + ) + + assert embeddings.embedded_texts == ["What is RAG?"] + assert chunks.search_query_embedding == embeddings.embedding + assert chunks.search_top_k == 5 + + +@pytest.mark.asyncio +async def test_generate_response_includes_retrieved_chunks_in_prompt() -> None: + """Test that retrieved document chunks are injected into the LLM prompt.""" + client, _messages, service = _build_service() + client.response = LLMResponse(text="Answer", raw=_raw_metadata()) + + chunks = cast(FakeDocumentChunkService, service.chunks) + chunks.chunks = [ + DocumentChunk( + id=TEST_CHUNK_ID, + document_id=TEST_DOCUMENT_ID, + content="Relevant chunk content", + ) + ] + + await service.generate_response( + user_id=TEST_USER_ID, + message=_user_message(), + ) + + assert isinstance(client.create_input, list) + assert any( + "Relevant chunk content" in str(item.get("content", "")) + for item in client.create_input + ) + + +@pytest.mark.asyncio +async def test_generate_response_includes_question_in_prompt() -> None: + """Test that the user question is included as a prompt message.""" + client, _messages, service = _build_service() + client.response = LLMResponse(text="Answer", raw=_raw_metadata()) + + await service.generate_response( + user_id=TEST_USER_ID, + message=_user_message(content="My question"), + ) + + assert isinstance(client.create_input, list) + assert any(item.get("content") == "My question" for item in client.create_input) + + +@pytest.mark.asyncio +async def test_stream_response_embeds_question_for_retrieval() -> None: + """Test that streaming embeds the question and runs chunk vector search.""" + client, _messages, service = _build_service() + client.events = [ + LLMStreamEvent( + type="final", + response=LLMResponse(text="Hi", raw=_raw_metadata()), + ), + ] + + embeddings = cast(FakeEmbeddingClient, service.embeddings) + chunks = cast(FakeDocumentChunkService, service.chunks) + + async for _ in service.stream_response( + user_id=TEST_USER_ID, + message=_user_message(content="Stream question?"), + ): + pass + + assert embeddings.embedded_texts == ["Stream question?"] + assert chunks.search_query_embedding == embeddings.embedding