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