Spring AI RAG example: index source code with PgVector
DEV Community

Spring AI RAG example: index source code with PgVector

How this Spring AI RAG example fits together

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.

Table of contents

  • How this Spring AI RAG example fits together
  • Deduplicate documents in the Spring AI vector store
  • Deduplication by document hash
  • Deduplication by deterministic id (default upsert)
  • Deduplication by deterministic id (via JdbcTemplate)
  • Chunk documents for Spring AI RAG retrieval
  • Nuances of tokenization and text splitting
  • Splitting text by tokens configuration
  • From index to answers: a grounded chat client
  • How Spring AI QuestionAnswerAdvisor grounds a RAG prompt
  • When the model invents file paths
  • Add file citations to a Spring AI RAG response
  • The client reply differs from the model’s message
  • Where to go next with Spring AI RAG

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:

  1. Run downloadDependencySources once to export and unpack your own dependency sources.
  2. An indexer walks sources/, splits each file into chunks, embeds those chunks, and stores them in pgvector.
  3. The 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.

Deduplicate documents in the Spring AI vector store

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.

Deduplication by document hash

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.

Deduplication by deterministic id (default upsert)

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.

Deduplication by deterministic id (via JdbcTemplate)

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:

  • same id and same checksum: do nothing;
  • same id and different checksum: embed and upsert;
  • missing id: embed and insert.

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.

Chunk documents for Spring AI RAG retrieval

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. Telegram · LinkedIn

Nuances of tokenization and text splitting

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.

Splitting text by tokens configuration

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 parent

Read on DEV Community ↗ ← Back to News

Comments

No comments yet. Start the discussion.