DEV Community

Cover image for Building Production AI Agents on the JVM: A Practical Guide to Java + Spring AI in 2026
Harshvardhan Singh
Harshvardhan Singh

Posted on

Building Production AI Agents on the JVM: A Practical Guide to Java + Spring AI in 2026

Let's skip the philosophical debate. This isn't about whether Java is better than Python for AI, Python owns model research and that's not changing. This is about what happens after you have a model endpoint and need to build a reliable system around it: retrieval, tool calling, multi-agent orchestration, observability, and enough throughput to not fall over on launch day.

This guide walks through building a production-grade RAG + agent service in Java 21, using Spring Boot 3.x, Spring AI, virtual threads, and structured concurrency — with real code, real tradeoffs, and a couple of diagrams so you're not just taking my word for it.


Why Backend Engineers Are Ending Up on AI Teams

AI products in 2026 aren't model calls. They're distributed systems: retrieval pipelines, tool-calling agents, multi-model routing for cost and reliability, semantic caching, and observability stacks that need to trace a request across five different services. That's backend engineering with an LLM bolted onto one node of the graph.

flowchart LR
    U[User Request] --> GW[AI Gateway]
    GW --> RT[Model Router]
    RT --> M1[GPT / Claude / Gemini]
    RT --> M2[Local vLLM Cluster]
    GW --> RAG[RAG Pipeline]
    RAG --> VDB[(Vector DB: pgvector / Qdrant)]
    GW --> TOOLS[MCP Tool Layer]
    TOOLS --> EXT[External APIs]
    GW --> CACHE[(Redis Semantic Cache)]
    GW --> OTEL[OpenTelemetry Collector]
Enter fullscreen mode Exit fullscreen mode

Every box in that diagram except "GPT / Claude / Gemini" is a job for backend infrastructure — and it's where Java's ecosystem (Spring Boot, resilience libraries, mature observability tooling) genuinely earns its keep.


Setting Up the Project

curl https://start.spring.io/starter.zip \
  -d dependencies=web,spring-ai-openai,spring-ai-vectorstore-pgvector,actuator \
  -d javaVersion=21 \
  -d bootVersion=3.4.0 \
  -o ai-service.zip
Enter fullscreen mode Exit fullscreen mode

build.gradle.kts dependencies:

dependencies {
    implementation("org.springframework.boot:spring-boot-starter-web")
    implementation("org.springframework.ai:spring-ai-openai-spring-boot-starter")
    implementation("org.springframework.ai:spring-ai-pgvector-store-spring-boot-starter")
    implementation("org.springframework.boot:spring-boot-starter-actuator")
    implementation("io.micrometer:micrometer-registry-prometheus")
    runtimeOnly("org.postgresql:postgresql")
}
Enter fullscreen mode Exit fullscreen mode

Enable virtual threads in application.yml — this one property changes your entire concurrency model:

spring:
  threads:
    virtual:
      enabled: true
  ai:
    openai:
      api-key: ${OPENAI_API_KEY}
      chat:
        options:
          model: gpt-4.1
Enter fullscreen mode Exit fullscreen mode

Why Virtual Threads Matter Here Specifically

AI request handling is almost entirely I/O wait: call the model, call the vector store, call a tool, call the model again. Platform threads made this expensive — each blocked thread held an OS thread hostage, so you'd tune thread pools and hope. Virtual threads decouple the two. You write normal, blocking, readable code, and the JVM schedules thousands of them onto a small number of carrier threads.

@RestController
@RequestMapping("/api/chat")
public class ChatController {

    private final ChatClient chatClient;
    private final VectorStore vectorStore;

    public ChatController(ChatClient.Builder builder, VectorStore vectorStore) {
        this.chatClient = builder.build();
        this.vectorStore = vectorStore;
    }

    @PostMapping(produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> chat(@RequestBody ChatRequest request) {
        List<Document> relevant = vectorStore.similaritySearch(
            SearchRequest.query(request.message()).withTopK(5)
        );

        String context = relevant.stream()
            .map(Document::getContent)
            .collect(Collectors.joining("\n---\n"));

        return chatClient.prompt()
            .system("Answer using only this context:\n" + context)
            .user(request.message())
            .stream()
            .content();
    }
}
Enter fullscreen mode Exit fullscreen mode

No manual thread pool tuning. No reactive boilerplate required to get high concurrency — though WebFlux remains a valid choice when you need true backpressure control over streaming pipelines, not just high thread counts. Virtual threads and reactive streams solve overlapping but distinct problems: virtual threads give you cheap concurrency with imperative code; WebFlux gives you backpressure and non-blocking composition. Pick reactive when you're streaming large volumes with flow control requirements; pick virtual threads when you want blocking-style simplicity at scale. Many production systems use both.


Structured Concurrency for Multi-Tool Agent Calls

Agents frequently need to fan out to multiple tools or retrieval sources in parallel and combine results. Java's structured concurrency (finalized as of JDK 21+ preview, stabilized in later releases) gives you a clean way to do this without leaking threads on partial failure:

public AgentResult gatherContext(String query) throws Exception {
    try (var scope = new StructuredTaskScope.ShutdownOnFailure()) {
        var vectorResults = scope.fork(() -> vectorStore.similaritySearch(query));
        var toolResults = scope.fork(() -> mcpToolClient.callTool("search_invoices", query));
        var cacheCheck = scope.fork(() -> semanticCache.lookup(query));

        scope.join();
        scope.throwIfFailed();

        return new AgentResult(
            vectorResults.get(),
            toolResults.get(),
            cacheCheck.get()
        );
    }
}
Enter fullscreen mode Exit fullscreen mode

If any subtask fails, the scope cancels the others automatically. No orphaned threads, no manual CompletableFuture cleanup, no silent leaks — a real problem in long-running agent services where a forgotten future eventually shows up as a memory profile nobody wants to explain.


RAG Architecture: pgvector vs. Qdrant vs. Milvus

A quick, honest comparison for teams choosing a vector store for a Java-based RAG pipeline:

Store Best fit Tradeoff
PostgreSQL + pgvector You already run Postgres; want transactional consistency between metadata and vectors Slower at very high dimensional scale (10M+ vectors) without careful indexing (HNSW tuning)
Qdrant Dedicated vector workloads, filtering-heavy queries, strong Rust-based performance One more system to operate and monitor
Milvus Very large scale (100M+ vectors), distributed deployments Higher operational complexity, heavier resource footprint

For most mid-scale production systems, pgvector wins on simplicity: one less database to run, one less failure domain, and Spring AI's PgVectorStore integrates directly with Spring's transaction management, which matters when you need vector writes to stay consistent with the rest of your data model.

@Bean
public VectorStore vectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel embeddingModel) {
    return PgVectorStore.builder(jdbcTemplate, embeddingModel)
        .withDimensions(1536)
        .withDistanceType(PgDistanceType.COSINE_DISTANCE)
        .withIndexType(PgIndexType.HNSW)
        .build();
}
Enter fullscreen mode Exit fullscreen mode

Multi-Model Routing and Cost Control

Production systems rarely call a single model for every request. Cheap, fast models handle simple queries; larger models handle complex reasoning; local models via Ollama or vLLM handle anything sensitive enough to keep off third-party infrastructure. LiteLLM or a custom Spring AI ChatModel router can sit in front of all of them:

@Component
public class ModelRouter {

    private final Map<String, ChatModel> models;

    public ChatResponse route(String query, ComplexityScore score) {
        ChatModel selected = switch (score) {
            case SIMPLE -> models.get("gpt-4o-mini");
            case MODERATE -> models.get("claude-sonnet");
            case COMPLEX -> models.get("gpt-4.1");
            case SENSITIVE -> models.get("local-vllm-llama3");
        };
        return selected.call(new Prompt(query));
    }
}
Enter fullscreen mode Exit fullscreen mode

This is where the "backend engineering" framing pays off directly: routing logic, cost attribution per customer, and failover between providers are ordinary service-layer concerns, not AI-specific magic. You're applying the same patterns you'd use for any multi-vendor integration — circuit breakers, retries with backoff, health checks — just pointed at model APIs instead of payment processors.


Semantic Caching with Redis

Re-answering identical or near-identical questions is pure waste. A semantic cache checks embedding similarity before hitting the model:

@Service
public class SemanticCacheService {

    private final RedisTemplate<String, String> redis;
    private final EmbeddingModel embeddingModel;
    private static final double SIMILARITY_THRESHOLD = 0.95;

    public Optional<String> lookup(String query) {
        float[] queryEmbedding = embeddingModel.embed(query);
        // Redis with a vector-capable index (RediSearch) performs the ANN lookup
        return redis.opsForValue()
            .get(hashKeyIfCloseEnough(queryEmbedding, SIMILARITY_THRESHOLD));
    }
}
Enter fullscreen mode Exit fullscreen mode

In practice this cuts model spend meaningfully on high-traffic, low-variance query patterns — support bots, FAQ-style assistants, internal tooling — where the same question shows up in a dozen phrasings.


Agent Orchestration and MCP

Model Context Protocol standardizes how agents discover and invoke tools, instead of every team hand-rolling a bespoke function-calling schema. Spring AI's MCP support lets you register a Java service as an MCP server that any compliant agent — regardless of what language or framework built it — can call:

@McpServer(name = "invoice-tools")
public class InvoiceToolServer {

    @McpTool(description = "Fetch the last N invoices for a customer")
    public List<Invoice> getRecentInvoices(String customerId, int count) {
        return invoiceRepository.findRecentByCustomer(customerId, count);
    }
}
Enter fullscreen mode Exit fullscreen mode

For multi-step agent workflows with branching logic, teams often pair Spring AI's tool-calling primitives with graph-based orchestration patterns similar to LangGraph — modeling the agent as an explicit state machine rather than an implicit chain, which makes debugging failed runs dramatically easier:

stateDiagram-v2
    [*] --> Retrieve
    Retrieve --> Decide
    Decide --> ToolCall: needs external data
    Decide --> Synthesize: has enough context
    ToolCall --> Decide
    Synthesize --> [*]
Enter fullscreen mode Exit fullscreen mode

Structured Outputs

Downstream systems need structured data, not prose. Spring AI's structured output converters map model responses directly to Java records, with schema validation baked in:

public record InvoiceSummary(
    String customerId,
    double totalAmount,
    List<String> flaggedItems
) {}

InvoiceSummary summary = chatClient.prompt()
    .user("Summarize these invoices: " + invoiceData)
    .call()
    .entity(InvoiceSummary.class);
Enter fullscreen mode Exit fullscreen mode

This eliminates an entire category of bugs — hand-parsed JSON from free-text model output — that plagues early-stage AI integrations.


Observability: You Can't Fix What You Can't See

AI systems fail in unusual ways: silent hallucination, slow degradation under load, cost spikes from a prompt template change nobody reviewed. Standard observability tooling still applies, and Java's ecosystem here is mature:

@Bean
public ObservationRegistry observationRegistry(MeterRegistry meterRegistry) {
    return ObservationRegistry.create()
        .observationConfig()
        .observationHandler(new DefaultMeterObservationHandler(meterRegistry))
        .observationRegistry();
}
Enter fullscreen mode Exit fullscreen mode

Spring AI auto-instruments chat calls, embedding calls, and vector store operations through Micrometer, which flows into Prometheus and OpenTelemetry with no extra glue code. In production, track at minimum: token usage per request, latency percentiles per model, cache hit rate, and tool-call failure rate. That last one catches problems weeks before users notice.


Deployment: Containers and Startup Time

A common objection to Java in latency-sensitive or serverless-adjacent AI deployments is JVM startup time. Two mitigations worth knowing:

  • CDS (Class Data Sharing) / AppCDS — reduces cold start meaningfully for containerized deployments by caching class metadata.
  • GraalVM native image for services where cold start truly matters (scale-to-zero Kubernetes deployments), compiling to a native binary can bring startup from seconds to milliseconds, at the cost of longer build times and some reflection-heavy library incompatibilities (check Spring AI's native support status for your specific starters before committing).
FROM eclipse-temurin:21-jre-alpine
COPY target/ai-service.jar app.jar
ENTRYPOINT ["java", "-XX:+UseZGC", "-jar", "app.jar"]
Enter fullscreen mode Exit fullscreen mode

ZGC (Z Garbage Collector) is worth defaulting to for AI services specifically, sub-millisecond pause times matter when you're streaming tokens over SSE and a GC pause shows up as a visible stutter in the response.


What This Adds Up To

None of this replaces Python for model development, evaluation, or research, that's not the goal. What it gives you is a production-grade layer for everything a model needs around it: retrieval, tool orchestration, multi-model routing, caching, structured outputs, and observability, built on a runtime with decades of tuning for exactly this class of problem high-concurrency, long-running, I/O-bound services that need to not fall over.

If you're currently running your AI orchestration layer as a Python script that grew into a service by accident, the migration path here doesn't require throwing anything away. Keep your model layer in Python. Put a Java service in front of it for routing, caching, and orchestration. Measure the difference in your p99 latency and your on-call load after a month.

Top comments (0)