Knowledgebase-RagChat Module Design and Implementation

This note records the design and API implementation of the RagChat module in the interview-guide project. This module belongs to the knowledgebase capability. Its focus is connecting multi-knowledgebase session management, message persistence, RAG streaming answers, and historical context into a continuous conversational Q&A system.

Module Capability Overview

• Multi-knowledgebase sessions: a session can be bound to multiple knowledgebases, and later Q&A only retrieves within the knowledge scope associated with the current session.
• Session management: supports session listing, detail retrieval, renaming, pinning, deletion, and updating associated knowledgebases.
• Message persistence: user questions and AI answers are stored separately, and complete conversation history can be restored by messageOrder.
• Streaming answers: uses SSE to return generated content so the frontend can display output while it is being generated.
• RAG reuse: reuses Knowledgebase module’s queryService.answerQuestionStream(...), keeping vector retrieval, prompt construction, and fallback logic unified.
• Context memory: currently supports recent messages as short-term context, and can later be extended with session summaries, semantic recall, and long-term memory.

Core State Flow

sequenceDiagram
  participant FE as Frontend KnowledgeBaseQueryPage
  participant API as ragChatApi.sendMessageStream
  participant C as RagChatController
  participant S as RagChatSessionService
  participant Q as KnowledgeBaseQueryService
  participant V as KnowledgeBaseVectorService
  participant LLM as ChatClient Stream
  participant DB as DB

  FE->>API: Submit question
  API->>C: POST /messages/stream
  C->>S: prepareStreamMessage(sessionId, question)
  S->>DB: Save user message completed=true
  S->>DB: Create AI placeholder message completed=false
  S-->>C: Return assistantMessageId

  C->>S: getStreamAnswer(sessionId, question)
  S->>DB: Query session and knowledgebases
  S->>DB: Load completed historical messages
  S->>Q: answerQuestionStream(kbIds, question, history)

  Q->>Q: query rewrite
  Q->>V: similaritySearch(...)
  V-->>Q: Return relevant document chunks
  Q->>Q: Build context + prompt
  Q->>LLM: stream().content()
  LLM-->>C: Flux chunk

  C-->>API: SSE data chunk
  API-->>FE: onMessage(chunk)
  FE->>FE: Append to the last assistant message

  C->>S: doOnComplete completeStreamMessage(...)
  S->>DB: Update AI message content + completed=true

Key API Design

POST /api/rag-chat/sessions Create a RAG Chat Session

Returns:

• Result<SessionDTO>

Call chain:

sessionService.createSession(request);
knowledgeBaseRepository.findAllById(request.knowledgeBaseIds());
sessionRepository.save(session);
ragChatMapper.toSessionDTO(session);

Flow:

1. Controller receives CreateSessionRequest and uses @Valid to ensure knowledgeBaseIds is not empty.
2. Service queries knowledgebases by knowledgeBaseIds.
3. It checks whether the number of queried knowledgebases matches the requested number. If not, it throws:

BusinessException(ErrorCode.NOT_FOUND, "Some knowledgebases do not exist")

4. It creates a RagChatSessionEntity.
5. It sets the session title:
   • If the request provides a non-empty title, use it.
   • If no title is provided, call generateTitle(knowledgeBases).
6. It binds knowledgebases through session.setKnowledgeBases(new HashSet<>(knowledgeBases)).
7. It saves the session and converts it to SessionDTO.

GET /api/rag-chat/sessions Get Session List

Returns:

• Result<List<SessionListItemDTO>>

Call chain:

sessionService.listSessions();
sessionRepository.findAllOrderByPinnedAndUpdatedAtDesc();

Key points:

• Sessions are sorted by pinned status and updated time in descending order.
• The API returns lightweight SessionListItemDTO objects for the left-side session list or history entry.

GET /api/rag-chat/sessions/{sessionId} Get Session Detail

Returns:

• Result<SessionDetailDTO>

Call chain:

sessionService.getSessionDetail(sessionId);
sessionRepository.findByIdWithKnowledgeBases(sessionId);
messageRepository.findBySessionIdOrderByMessageOrderAsc(sessionId);
ragChatMapper.toSessionDetailDTO(session, messages, kbDTOs);

Flow:

1. Query the session and its associated knowledgebases by sessionId.
2. If the session does not exist, throw:

BusinessException(ErrorCode.NOT_FOUND, "Session does not exist")

3. Query all messages under the session and sort them by messageOrder ASC.
4. Convert associated KnowledgeBaseEntity objects into KnowledgeBaseListItemDTO.
5. Assemble SessionDetailDTO, including session info, knowledgebase list, and message history.

PUT /api/rag-chat/sessions/{sessionId}/title Update Session Title

Returns:

• Result<Void>

Call chain:

sessionService.updateSessionTitle(sessionId, request.title());
sessionRepository.findById(sessionId);
sessionRepository.save(session);

Key points:

• UpdateTitleRequest uses @Valid, and title cannot be empty.
• If the session does not exist, it throws BusinessException(ErrorCode.NOT_FOUND, "Session does not exist").
• After updating session.title, saving the entity relies on @PreUpdate to refresh updatedAt.

PUT /api/rag-chat/sessions/{sessionId}/pin Toggle Session Pin Status

Returns:

• Result<Void>

Call chain:

sessionService.togglePin(sessionId);
sessionRepository.findById(sessionId);
sessionRepository.save(session);

Logic:

Boolean currentPinned = session.getIsPinned() != null ? session.getIsPinned() : false;
session.setIsPinned(!currentPinned);

Notes:

• When isPinned = null, it is treated as false, then toggled to true.
• After saving, @PreUpdate refreshes the updated time.

PUT /api/rag-chat/sessions/{sessionId}/knowledge-bases Update Associated Knowledgebases

Returns:

• Result<Void>

Call chain:

sessionService.updateSessionKnowledgeBases(sessionId, request.knowledgeBaseIds());
sessionRepository.findById(sessionId);
knowledgeBaseRepository.findAllById(knowledgeBaseIds);
session.setKnowledgeBases(new HashSet<>(knowledgeBases));
sessionRepository.save(session);

Key points:

• The service re-queries knowledgebase entities using the requested knowledgeBaseIds.
• It replaces the current session associations with a new HashSet.
• This is suitable when the user switches or expands the knowledge scope in the same chat window.

DELETE /api/rag-chat/sessions/{sessionId} Delete Session

Returns:

• Result<Void>

Call chain:

sessionService.deleteSession(sessionId);
sessionRepository.existsById(sessionId);
sessionRepository.deleteById(sessionId);

Key points:

• Deletion is executed in a transactional method.
• The service first checks whether the session exists, then deletes the session record.
• Whether associated messages are deleted by cascade depends on the entity mapping and repository implementation.

POST /api/rag-chat/sessions/{sessionId}/messages/stream Send a Question and Stream the Answer

Returns:

• Flux<ServerSentEvent<String>>

Call chain:

sessionService.prepareStreamMessage(sessionId, request.question());
sessionService.getStreamAnswer(sessionId, request.question());
queryService.answerQuestionStream(kbIds, question, history);
sessionService.completeStreamMessage(messageId, fullContent.toString());

Flow:

1. Controller receives SendMessageRequest and validates the question with @Valid.
2. It calls prepareStreamMessage(...) before streaming:
   • Query the session and associated knowledgebases.
   • Save a USER message with completed status.
   • Create an ASSISTANT placeholder message with empty content and incomplete status.
   • Update the session messageCount and save it.
3. Controller records the placeholder assistant messageId.
4. It creates StringBuilder fullContent to collect the complete AI answer.
5. It calls getStreamAnswer(...):
   • Query the session and associated knowledgebases again.
   • Read the knowledgeBaseIds bound to the current session.
   • If historical context is enabled, load recent completed messages as multi-turn context.
   • Call queryService.answerQuestionStream(kbIds, question, history).
6. For each received chunk, append it to fullContent, then wrap it as an SSE event:

ServerSentEvent.<String>builder()
    .data(chunk.replace("\n", "\\n").replace("\r", "\\r"))
    .build();

7. After streaming completes, call the following in doOnComplete:

sessionService.completeStreamMessage(messageId, fullContent.toString());

This writes the complete AI answer back to the assistant placeholder message and marks it as completed.

8. If streaming fails, doOnError saves the partial content. If no content was generated, it saves an error message.

RAG Streaming Q&A Chain

RagChat does not reimplement vector retrieval. It reuses the knowledgebase query service:

queryService.answerQuestionStream(kbIds, question, history);

Core steps:

1. Restrict retrieval scope by the knowledgebase IDs bound to the current session.
2. Optionally include historical context as multi-turn input.
3. Build QueryContext, perform query normalization, query rewrite, and dynamic topK/minScore setup.
4. Call vectorService.similaritySearch(...) to retrieve relevant document chunks.
5. Concatenate matched documents into context.
6. Build the system prompt and user prompt, including prompt-injection constraints.
7. Call chatClient.prompt().stream().content() to output a token stream.
8. Normalize streaming output through normalizeStreamOutput(...).
9. Return unified fallback text streams for empty input, no retrieval hits, or exceptions.

Current Issues and Optimization Directions

The current context strategy is still mainly short-term memory. The main issues are:

• Context can grow too long: raw historical messages are directly inserted into the prompt, and longer AI answers increase later token cost.
• Long-term memory is easy to lose: only recent messages are included, so early key information becomes invisible after it leaves the window.
• Historical retrieval is not smart enough: current logic retrieves recent messages by time, not semantically relevant history for the current question.
• AI answers also consume context slots: maxMessages = 10 limits message count, not conversation turns, so long answers can waste context capacity.
• No long-term summary: there is no session-level summary, user profile, or preference memory yet.

Future improvement can follow a layered memory design:

• Short-term memory: cache the most recent rounds in Redis for fast context recovery.
• Long-term memory: store all raw messages, summaries, and states in PostgreSQL for reliable traceability.
• Recallable memory: write historical message summaries or session summaries into pgvector and retrieve them semantically by the current question.
• Session summary: after each completed Q&A round, asynchronously trigger a summarization agent to update the session summary.
• External knowledge: continue using knowledgebase RAG retrieval results as factual sources.

Summary

The Knowledgebase-RagChat module extends the knowledgebase from a “one-shot Q&A API” into a “manageable, recoverable, continuous conversation” chat system. Its key value is clear separation of concerns: sessions organize user context and knowledgebase scope, while the RAG query service handles retrieval and generation. This boundary makes it easier to gradually evolve toward long-term memory, session summaries, and semantic historical recall.