Skip to content

rag.md

Latest commit

 

History

History
146 lines (113 loc) · 5.83 KB

rag.md

File metadata and controls

146 lines (113 loc) · 5.83 KB
Raw

RAG research assistant walkthrough

Flink-runtime showcase, exercises Flink-only capabilities (Flink-state keyed HNSW vector memory). The portable cores ship their own in-process HNSW index, so RAG itself runs everywhere (see banking-rag.yaml and the banking agent on every runtime); what's Flink-specific here is binding the index to Flink's checkpointed state.

Source: src/main/java/org/agentic/flink/example/rag/RagResearchExample.java Inline README: src/main/java/org/agentic/flink/example/rag/README.md

Running it

bash examples-bin/run-ollama.sh      # Ollama in Podman on 127.0.0.1:11434, pulls qwen2.5:3b
bash examples-bin/run-rag.sh

Prerequisites: JDK 21, the Maven wrapper, Podman (for Ollama), and outbound internet on the first run: DJL downloads sentence-transformers/all-MiniLM-L6-v2 (embedder), cross-encoder/mmarco-mMiniLMv2-L12-H384-v1 (reranker) and the PyTorch CPU native runtime (about 500 MB) into its cache directory. No API key. The script installs ports/jagentic-core when missing, resolves the provided-scope Flink dependencies and runs the example in a forked JVM in an embedded MiniCluster.

The input is a fixed list of (topic, question) queries; each prints one Answer[...] record with the topic, the question, the grounded answer with [1] [2] [3] citations and the top rerank score. bash examples-bin/run-rag-stack.sh brings up the larger Podman stack (Ollama, Postgres, Redis, Fluss) that the RAG notebooks use; run-rag.sh itself needs only Ollama.

Why this shape

A retrieval-augmented assistant has four levers that all matter:

  1. Embedder quality: bad embeddings ≡ retrieving the wrong passages.
  2. Vector store latency: every query embeds + searches; cost compounds.
  3. Rerank precision: embedding similarity is a coarse signal. A cross-encoder over the top-k catches "looks similar, is wrong" cases (the example uses cross-encoder/mmarco-mMiniLMv2-L12-H384-v1).
  4. Citations: the LLM has to ground its answer or the system is just a confident summarizer.

The example wires all four with the framework's primitives:

Query
  │
  ▼  keyBy(topic)       - each topic gets its own KB scope
  │
  ▼  EmbeddingClient.embed(question)
  │
  ▼  VectorMemory.search(queryVec, k=4)   - brute-force over MapState
  │
  ▼  Scorer.scorePair(passage, question)  - cross-encoder rerank
  │
  ▼  LLM answers using top-3 passages, cites [1] [2] [3]

Why Flink-state vector memory?

The framework ships brute-force KNN backed by MapState<String, VectorEntry>. For a few thousand vectors per key (the typical "conversation-local recall" use case), exact search at d=384 is sub-millisecond. The state itself rides in Flink checkpoints, so a job restart picks up where it left off, no "rebuild the index" boot sequence.

When you outgrow that (10⁵+ vectors per key), drop in an HNSW-backed VectorMemorySpec via the SPI. The agent code doesn't change.

Why per-topic keying?

Topics are independent corpora. Keying by topic gives each topic its own VectorMemory scope, which means:

  • A noisy topic doesn't drown a clean one in the same brute-force scan.
  • Different topics can use different embedders or dimensions if you want (one operator per topic, different specs).
  • Eviction policies (TTL, max-items) apply per topic.

Seeding the corpus

The demo seeds 4 documents per topic on the first event using a ValueState<Boolean> flag. In production you'd either:

  • Hydrate from a LongTermMemoryStore on cold-start, the framework's standard pattern.
  • Stream new documents in via a Channel<KeyedContextItem>, KafkaContextChannel, PostgresChangeChannel, RedisPubSubChannel, or any custom transport. The channel produces KeyedContextItem records that the agent operator union-connects with its main input.

Both approaches are covered in docs/memory.md. For the unified channel SPI see docs/channels.md. For a more ambitious example with two inputs sharing a corpus, see live-research.

Adding an MCP tool source

The demo keeps the runtime self-contained; to let the LLM call out to file or web tools through Model Context Protocol, register an McpServerSpec and discover its tools at job startup:

McpServerSpec everything = McpServerSpec.stdio(
    "everything", "npx", "-y", "@modelcontextprotocol/server-everything");

Agent agent = Agent.builder()
    .withId("rag")
    .withChatConnection(chat)
    .withChatSetup(chatSetup)
    .withMcpServer(everything)
    .build();

// In the operator's open():
try (McpClient client = new McpClient(everything)) {
    client.initialize();
    for (McpToolExecutor t : McpToolRegistry.discover(everything, client)) {
        toolRegistry.register(t.getToolId(), t);
    }
}

Each discovered tool surfaces as a regular ToolExecutor, the LLM doesn't know or care that it's MCP under the hood.

Performance shape

On a recent laptop with the suggested models:

  • Embed: ~12 ms (MiniLM-L6-v2, batch 1, CPU)
  • Brute-force search (k=4 over 4 vectors): ~0.1 ms
  • Rerank top-4: ~80 ms (cross-encoder, batch 4)
  • LLM answer (qwen2.5:3b): ~600 ms

The LLM dominates; everything else is overhead. Scale the embedder + vector search up to 10k vectors per key before the search side starts to matter.

Failure modes

  • Cold first run: first call downloads the two Hugging Face models plus the PyTorch native runtime (about 500 MB). Subsequent runs are cache-warm.
  • Dimension mismatch: FlinkStateVectorMemory.spec(384) must match the embedder's actual dimension. The framework throws IllegalArgumentException at insert time if they disagree.
  • No-op rerank when k=1: with a single hit, the cross-encoder still runs. Skip it conditionally if you care about the savings.