Spring AI RAG example: index source code with PgVector

A developer built a Spring AI retrieval-augmented generation (RAG) example that indexes dependency source code into a PgVector vector store, deduplicates files, splits them into embedding-sized chunks, and returns grounded answers with file citations drawn from retrieval metadata. The indexer is exposed as an MCP tool, and a companion /api/chat endpoint answers questions against the indexed code. To avoid duplicate rows on re-indexing, the approach combines stable text-based document IDs with checksum metadata and a JdbcTemplate pre-check that skips unchanged files before the embedding call.

This Spring AI RAG example builds a source-code index with a Spring AI vector store backed by PgVector. It deduplicates files, splits them into embedding-sized chunks, and returns grounded answers with file citations taken from retrieval metadata. Think of it as a local Context7-style index for the library sources your project actually uses, with indexing exposed over MCP. A model trained last year does not know the API you upgraded to last week. By the end, you will have a Spring AI RAG indexer for dependency sources and a chat client that answers from indexed code with file citations. Gradle downloads the source jars, the MCP tool builds the index, and /api/chat asks questions against it. If an outside MCP agent also needs to search or ask, wrap CodeAssistantService.ask ... in a separate MCP tool. The project from part one already has the pieces. It is a Spring Boot RAG application running Spring Boot 4.1 and Spring AI with PgVector. I added a Gradle task, downloadDependencySources , that resolves each runtime dependency’s source jar and unpacks it into sources/ . The flow looks like this: downloadDependencySources once to export and unpack your own dependency sources. sources/ , splits each file into chunks, embeds those chunks, and stores them in pgvector. index sources MCP tool builds the index, and the /api/chat endpoint asks questions against it. The first index uses two small services: SourcesIndexerService and SourcesQueryService . The MCP tool does not need a separate wrapper class; it lives as an annotated method on the indexer service. Indexing has a problem: you will run it more than once. The app restarts. The source export task runs again. Or you change the embedding model and rebuild the index. If every run inserts the same files again, retrieval quality drops. Top-K results fill with duplicate source files. The database also grows without adding value. Spring AI does not give you one generic “deduplicate this” switch. The VectorStore contract stays small by design. Deduplication belongs to your document identity model. Three options can help: hash checks, fixed ids with PgVector upsert, and a JDBC pre-check before embedding. This demo combines the last two. Fixed text ids keep rows stable. Checksum metadata detects changes. JdbcTemplate skips unchanged files before the embedding call. The first version calculated a checksum for each file, stored that checksum as metadata, and asked the vector store whether a document with that checksum already existed. The query wrapped the same VectorStore . It ran a metadata-filtered search with an empty query, topK 1 , and a zero threshold. That version works across vector stores, but it hides a cost: similaritySearch ... still asks the embedding model for a query vector before PgVector can run SQL. For an indexer you run often, the dedupe check should not call the embedding model at all. The demo solution rests on fixed document identity. A Document created with new Document text, metadata gets a random id. Every re-index then creates a new row. A Document created with new Document id, text, metadata keeps the same id every time. For source files, the id should come from the file identity, not just the file content. Say you key on the checksum alone. Then two different files with identical content collapse into one document. In the demo I use a readable text id built from the library and the source path: private String documentId String library, Path relative { String sourcePath = relative.getNameCount 1 ? relative.subpath 1, relative.getNameCount .toString : relative.toString ; return library + "/" + sourcePath; } The current indexer uses that method to build the parent id before it splits a source file into chunks. The exact string format does not matter. The same source file must produce the same id on every run. If the file changes later, the id stays stable. The next run can then replace the content and embedding for that row. One PgVector detail matters here. This Spring AI PgVector example uses readable ids, but Spring AI’s PgVector store defaults to UUID ids. A readable id like library/path needs a text id column: spring.ai.vectorstore.pgvector.id-type=TEXT If you keep the default UUID id type, use a deterministic UUID generator instead of a readable path id. Spring AI’s PgVector store inserts with ON CONFLICT id DO UPDATE , so fixed ids make repeat indexing update existing rows instead of piling up duplicates. Without chunking, fixed ids keep one row per source file. In the final chunked indexer, fixed chunk ids make repeat indexing update matching chunk rows. A fixed id alone does not skip embedding. Spring AI computes embeddings before the insert. The id protects the table from duplicate rows. It does not protect the embedding model from repeated work. For true “do nothing” behavior on unchanged files, make the decision before vectorStore.add ... . Spring AI’s PgVector store does not expose a switch from ON CONFLICT DO UPDATE to ON CONFLICT DO NOTHING . By the time PgVector runs its insert, Spring AI has already calculated the embedding. The practical version uses a direct id-and-checksum lookup with JdbcTemplate . If the stored row has the same id and checksum, return before vectorStore.add ... . If no row matches, or the checksum changed, call vectorStore.add ... and let PgVector upsert the row. @Service @RequiredArgsConstructor public class SourcesQueryService { private final JdbcTemplate jdbcTemplate; public boolean isDocumentStoredByIdAndCheckSum String id, Long checkSum { return Boolean.TRUE.equals jdbcTemplate.queryForObject """ select exists select 1 from public.vector store where id = ? and metadata - 'checksum' = ? """, Boolean.class, id, checkSum.toString ; } } Then the current indexer can skip unchanged files before it splits and embeds the source: Long checksum = crc32 file.toAbsolutePath ; String parentDocumentId = documentId library, relative ; String firstChunkId = chunkId parentDocumentId, 0 ; if queryService.isDocumentStoredByIdAndCheckSum firstChunkId, checksum { log.debug "Source document already stored: {}", relative ; return; } This trades portability for behavior this service needs. It depends on PgVector’s table and metadata layout: The id lookup already uses PgVector’s primary key. Add metadata expression indexes only for queries that scan by checksum or parent metadata without a specific id. For the demo service, use this version: a cheap JdbcTemplate pre-check plus PgVector upsert for changed files. That tradeoff pays off. One SQL read can save one model call. In the chunked index, each row maps to one chunk: no duplicate rows and no repeated embedding call for unchanged files. Deduplication decides whether to store a document. Chunking decides the document shape. That boundary matters with source files. Spring AI’s default vector-store batch strategy counts tokens before embedding. It can batch many small documents, but it will not split one oversize file. If a file exceeds the embedding model’s limit, vectorStore.add ... throws before the file reaches Postgres. You will see this failure: java.lang.IllegalArgumentException: Tokens in a single document exceeds the maximum number of allowed input tokens at org.springframework.ai.embedding.TokenCountBatchingStrategy.batch TokenCountBatchingStrategy.java:151 ~ spring-ai-model-2.0.0-RC2.jar:2.0.0-RC2 at org.springframework.ai.embedding.EmbeddingModel.embed EmbeddingModel.java:106 ~ spring-ai-model-2.0.0-RC2.jar:2.0.0-RC2 That error means the batching strategy rejected one Document . A bigger batch size will not help. The document itself has to shrink. The Spring AI ETL pipeline has a tool for this job: TokenTextSplitter . It implements DocumentTransformer . It takes a list of Document objects and returns a new list of smaller Document objects. The split runs before vectorStore.add ... , so the embedding model sees chunks instead of the original large file. Found this useful? Follow me on social media to stay updated. TokenTextSplitter runs locally. It does not call the embedding model or send text to an AI server. It uses jtokkit to map text to token ids, takes up to chunkSize tokens, decodes that window, then moves the cut back to the last configured punctuation mark after minChunkSizeChars . Finally, it trims the chunk and drops anything shorter than minChunkLengthToEmbed . Tokenization depends on the model. In this demo the embedding endpoint speaks an OpenAI-compatible API, but the model still runs text-embedding-nomic-embed-text-v1.5 . That API shape does not make the tokenizer OpenAI-compatible. Nomic examples use a BERT tokenizer. Spring AI’s TokenTextSplitter uses com.knuddels:jtokkit , a Java tokenizer for OpenAI-style encodings, and defaults to CL100K BASE . Spring AI can switch among supported OpenAI-style encodings, but none matches a true BERT tokenizer. When the target model uses another encoding supported by jtokkit , set the encoding type on the TokenTextSplitter builder. For Nomic, this does not give exact counts. A different EncodingType will not close the gap because Nomic uses a different tokenizer family. If exact token accounting matters, swap in a custom DocumentTransformer . Back it with a Nomic/BERT-compatible tokenizer. Until then, TokenTextSplitter still works as a safe local splitter. The local token count gives an estimate. I treat chunkSize as a safety budget, not the model’s real limit. The exact endpoint limit depends on the local embedding server configuration. Chunks around 1000 tokens leave room for batch overhead. Exact counts need a Nomic/BERT-aware splitter or token counter. For source code, I do not want chunks close to the endpoint limit: the batching strategy reserves space, Nomic token counts remain approximate here, and code has dense punctuation. A chunk size around 1000 tokens leaves room for the embedding request while keeping methods and classes searchable. Because the chunk limit belongs to the embedding model, configure the splitter as a shared bean rather than per-method setup. Spring AI will not apply this bean for you inside VectorStore.add ... . The app still has to call the splitter in its indexing pipeline. One shared bean gives the project one chunking policy. @Configuration public class SourcesIndexingConfiguration { @Bean public TokenTextSplitter tokenTextSplitter { return TokenTextSplitter.builder .withEncodingType EncodingType.CL100K BASE .withChunkSize 1000 .withMinChunkSizeChars 400 .withMinChunkLengthToEmbed 10 .withMaxNumChunks 10000 .withKeepSeparator true .withPunctuationMarks List.of '\n', ';', '}' .build ; } } This still does not parse Java. It can split inside a method, a class, or a multiline expression. The project still gets a token safety layer: it keeps line breaks, prefers line or statement boundaries, and keeps ordinary source files well under the local embedding limit. maxNumChunks does not provide safe truncation. If you set it too low for a generated file, the final chunk can still contain the rest of the file and still trip the batching error. To reject generated files, do it at the file boundary by size, extension, or path. Chunking changes document identity. Before chunking, the deterministic id belongs to the source file. From here I call it the parent id. After chunking, the vector store stores chunk rows, and each one must still point back to the same source file. String parentDocumentId = documentId library, relative ; Document sourceDocument = new Document parentDocumentId, text, Map.of LIBRARY FIELD, library, PATH FIELD, relative.toString , CHECK SUM FIELD, checksum ; List<Document chunks = tokenTextSplitter.split sourceDocument ; This code still lacks full dedupe behavior. Spring AI’s TextSplitter uses the id of sourceDocument as the parent id, copies the original metadata, and adds parent document id , chunk index , and total chunks to every chunk. We do not have to put parent document id into the source document metadata ourselves. TokenTextSplitter creates new Document instances instead of keeping the parent id as the chunk id. Those documents get random ids by default. That works for a one-off index, but it breaks deterministic PgVector upsert: a changed file inserts a fresh set of chunks and leaves the old chunks behind. The parent id only lives in metadata; PgVector’s ON CONFLICT id DO UPDATE sees only the row id. Once one source file becomes many vector-store rows, the deterministic id moves from “one id per file” to “one id per chunk”. The file id remains the parent id. Each stored row gets a deterministic chunk id: private String chunkId String parentDocumentId, int chunkIndex { return parentDocumentId + " chunk-" + chunkIndex; } After splitting, assign ids to the chunks. Each chunk needs a distinct deterministic id; otherwise PgVector would collapse the file to one row and each later chunk would overwrite the previous one. In the code below, chunkId parentDocumentId, index produces stable ids like library/Foo.java chunk-0 , library/Foo.java chunk-1 , and so on. Repeat indexing updates the matching chunk row. List<Document chunkDocuments = IntStream.range 0, chunks.size .mapToObj index - new Document chunkId parentDocumentId, index , chunks.get index .getText , chunks.get index .getMetadata .toList ; vectorStore.add chunkDocuments ; Rebuild each chunk with new Document id, text, metadata to preserve the metadata TextSplitter added, including parent document id . That metadata supports stale-chunk cleanup later. The other option, chunks.get index .mutate .id ... .build , carries the metadata for you. But the explicit constructor keeps the storage identity rule visible in one place. I prefer to assign the chunk id in the indexing service. It keeps the storage identity rule next to the checksum and upsert logic. The splitter does not hide it. The checksum story changes too, but only a little. The source file still owns the checksum; chunks inherit that checksum in metadata. A small edit near the top of a file can shift all later chunk boundaries. Chunk-level hashes look less stable than they seem. For this indexer, the file acts as the unit of freshness. The unchanged-file check can still happen before embedding. Do not check the old file row id. Check the first deterministic chunk id plus the file checksum: String firstChunkId = chunkId parentDocumentId, 0 ; if queryService.isDocumentStoredByIdAndCheckSum firstChunkId, checksum { log.debug "Source document already stored: {}", relative ; return; } Checking chunk 0 works for the demo, but it assumes indexing either stored every chunk or failed before storing the first chunk. If you need to tolerate partial indexing, store a manifest row per source file, or split first and compare the stored chunk count with total chunks . That turns the check into a stronger invariant: the index holds the whole source file at this checksum. If the checksum changed, split the file, assign deterministic chunk ids, and call vectorStore.add ... . PgVector will upsert chunks whose ids still exist. One cleanup case remains: a file used to produce ten chunks and now produces six. Upsert updates chunks 0..5 , but chunks 6..9 go stale. Delete old chunks for that parent before or after reindexing the changed file. Because the splitter already adds parent document id , the cleanup can target the source file with one more JdbcTemplate operation: public void deleteStaleChunks String parentDocumentId, Long checksum { jdbcTemplate.update """ delete from public.vector store where metadata - 'parent document id' = ? and metadata - 'checksum' is distinct from ? """, parentDocumentId, checksum.toString ; } This storage model makes a reliable Spring AI RAG indexer: deterministic identity, checksum metadata, a cheap JDBC pre-check, and PgVector upsert. Chunking only makes the identity model more exact. The file acts as the parent document. The vector-store rows hold chunk documents. For better source-code retrieval later, replace the local TokenTextSplitter setup with a code-aware splitter that understands classes, methods, and declarations. The same storage rules still hold: deterministic chunk ids, file-level checksum metadata, and stale-chunk cleanup when a file changes. With those rules in place, your Spring AI RAG pipeline stays idempotent as dependency sources grow and change. The index only covers half the system. Now we read from it. In the Spring AI beginner guide https://protsenko.dev/spring-ai-for-beginners-build-your-first-ai-app-in-java/ , a small @Service wrapped a ChatClient ; then I added a memory advisor and a QuestionAnswerAdvisor . Here we point that same shape at the VectorStore the indexer fills, so the code assistant answers from indexed sources instead of model training data. To make the answer machine-readable, we ask for structured output instead of free text. The reply carries the prose answer, the snippets it leaned on, and the files those snippets came from. public record CodeAssistantReply String answer, List<String evidences, List<String absoluteSourceFilePaths {} The first service version kept things minimal: wire the two advisors, set a system prompt, and map the model’s JSON onto the record with .entity ... . That version helps explain the failure, but the checked-in project now uses the grounded version shown below. A thin controller gives the conversation an id when the caller omits one, so memory keeps separate chats apart: @RestController @RequestMapping "/api/chat" @RequiredArgsConstructor public class ChatController { private final CodeAssistantService codeAssistantService; @PostMapping public GenerationResponse generation @RequestBody GenerationRequest request { var conversationId = request.conversationId ; if conversationId == null { conversationId = UUID.randomUUID .toString ; } return new GenerationResponse codeAssistantService.ask conversationId, request.userMessage ; } } Use the checked-in requests/chat-controller.http request unchanged. It asks the assistant about an MCP client: POST localhost:8080/api/chat Content-Type: application/json { "userMessage": "What's mcp client and how to use it?" } The answer looks useful, and the evidences look plausible. Then you check absoluteSourceFilePaths : the files do not exist. The model returned paths it never saw. To understand why, look at what QuestionAnswerAdvisor actually puts into the prompt. QuestionAnswerAdvisor packages the retrieval step of Spring AI RAG as an advisor. On every call it does three things before the model runs. First, it searches the vector store with the user’s question as the query: List<Document documents = this.vectorStore.similaritySearch searchRequestToUse ; Second, it turns those hits into a context string. This line matters for our bug: String documentContext = documents.stream .map Document::getText .collect Collectors.joining System.lineSeparator ; It joins Document.getText and nothing else. The chunk metadata, including the library , path , and checksum fields the indexer stored, never enters the string. Third, it renders a default template that wraps the question and that context together: {query} Context information is below, surrounded by --------------------- --------------------- {question answer context} --------------------- Given the context and provided history information and not prior knowledge, reply to the user comment. If the answer is not in the context, inform the user that you can't answer the question. The model sees code text and an instruction to answer from it, but no file paths. The retrieved documents still exist in Java: the advisor stores them under QuestionAnswerAdvisor.RETRIEVED DOCUMENTS and copies them to ChatResponse metadata after the call. The path data becomes available after the call, but not in the prompt. The bad output has three causes. The model never receives a path. The structured-output converter appends the JSON schema for CodeAssistantReply , so the model sees an absoluteSourceFilePaths field. And the original prompt tells it to fill that field with real file locations. With no path data in context, the model writes plausible strings that point at nothing. First, make the failure visible instead of guessing at it. Three checks, cheapest first. Turn on advisor logging while debugging. SimpleLoggerAdvisor can print the request the model receives, including the augmented user message, so you can confirm the prompt carries code but no paths. Inspect what retrieval returned. Response metadata contains the documents, so you can compare their real paths with the model’s claimed ones. When you add tests, pin the invariant: every cited path must resolve to a real file under the indexed root. If evidence must come from real quotes, add a second assertion that each evidence string appears in the joined retrieved text. The fix follows from the mechanics: stop asking the model for paths, and read them from retrieval instead. Let the model write the answer and quote evidence. Let the service take paths from the documents the advisor already retrieved. To reach both the parsed entity and the raw response, swap .entity ... for .responseEntity ... . That call returns both objects: the converted CodeAssistantReply and the full ChatResponse , whose metadata carries RETRIEVED DOCUMENTS . @Service public class CodeAssistantService { private final ChatClient chatClient; private static final String systemPrompt = """ You are a code assistant that explains how libraries and frameworks work. Answer ONLY from the context block provided below the user's question. That context is real source code retrieved from the user's own dependency sources. Rules: - Ground every statement in the provided context. Do not rely on prior knowledge of the library, and do not guess. - In evidences , copy short verbatim snippets from the context that back your answer. - If the context does not contain the answer, say so plainly in answer and return empty lists. - Leave absoluteSourceFilePaths empty. The system fills it from retrieval metadata. Never invent file paths. """; CodeAssistantService ChatClient.Builder builder, VectorStore vectorStore { var chatMemory = MessageWindowChatMemory.builder .build ; var memoryAdvisor = MessageChatMemoryAdvisor.builder chatMemory .build ; var questionAnswerAdvisor = QuestionAnswerAdvisor.builder vectorStore .build ; this.chatClient = builder .defaultSystem systemPrompt .defaultAdvisors memoryAdvisor, questionAnswerAdvisor .build ; } public CodeAssistantReply ask String conversationId, String question { ResponseEntity<ChatResponse, CodeAssistantReply result = chatClient .prompt .user question .advisors spec - spec.param ChatMemory.CONVERSATION ID, conversationId .call .responseEntity CodeAssistantReply.class ; CodeAssistantReply reply = result.entity ; if reply == null { return null; } List<String realPaths = retrievedPaths result.response ; // Real paths overwrite whatever the model produced. return new CodeAssistantReply reply.answer , reply.evidences , realPaths ; } private List<String retrievedPaths ChatResponse chatResponse { if chatResponse == null { return List.of ; } List<Document documents = chatResponse.getMetadata .get QuestionAnswerAdvisor.RETRIEVED DOCUMENTS ; if documents == null { return List.of ; } return documents.stream .map d - d.getMetadata .get MetadataFields.PATH FIELD .filter Objects::nonNull .map Object::toString .distinct .toList ; } } Two changes fix it. The system prompt tells the model to leave absoluteSourceFilePaths empty, which removes the temptation to invent. Then retrievedPaths ... reads RETRIEVED DOCUMENTS from response metadata, pulls the path field from each document, and overwrites the field the model returned. The index now supplies the paths as facts, not guesses. One caveat about the field name. The indexer stored PATH FIELD as relative.toString , so metadata holds a path relative to the indexed root, not an absolute one. If you want a full path, either store file.toAbsolutePath .toString at index time, or resolve the relative path against the known sources root when you build the reply. Use this pattern for grounded answers in Spring AI RAG: let the model reason over retrieved text, but take citable facts such as paths, ids, and versions from the retrieved documents themselves. The advisor exposes those documents in response metadata. The mistake comes from asking the model to repeat data the prompt never showed. Precision helps because three versions of the answer exist. answer , evidences , and whatever it guessed for absoluteSourceFilePaths . CodeAssistantReply . ask ... , your code replaces the paths and builds a new The model never sees step 3. It only produced the message in step 1. Your code creates the corrected reply after the model finishes and outside its view. That split has one consequence for memory. MessageChatMemoryAdvisor saves history in its after step, and it saves the model’s own output, not your record: assistantMessages = chatClientResponse.chatResponse .getResults .stream .map g - Message g.getOutput .toList ; this.chatMemory.add conversationId, assistantMessages ; Chat memory keeps the message from step 1. The user saw one version, but the next turn remembers another. The refined system prompt makes memory less risky because it tells the model to leave absoluteSourceFilePaths empty. The broader rule still matters. Paths, redactions, reranked evidence, and other post-call changes exist only in the reply you return. If later turns must reason over that corrected reply, store the corrected assistant message under the same conversation id. Keep the two roles apart. The model writes the answer from context. Your service owns the final reply and what gets stored for later turns. You now have a repeatable Spring AI RAG pipeline: export source jars, skip unchanged files, split the files that changed, store deterministic chunks in PgVector, and take citations from retrieval metadata instead of model output. For the engineering workflow around it, read AI Coding Workflow: TDD, Reviews, and Guardrails https://protsenko.dev/ai-driven-development-workflow-guard-rails/ , see where automated reviews still fall short in AI Code Review: 7 Limitations I Found in Production https://protsenko.dev/ai-code-review-7-limitations-i-found-in-production/ , or browse all artificial intelligence articles https://protsenko.dev/category/artificial-intelligence/ . New to Spring AI? Start with Spring AI for beginners: build your first AI app in Java https://protsenko.dev/spring-ai-for-beginners-build-your-first-ai-app-in-java/ .