diff --git a/README.md b/README.md index 4642318..28a402e 100644 --- a/README.md +++ b/README.md @@ -26,5 +26,6 @@ Runnable companion code for the Spring AI articles on [ankurm.com](https://ankur | [`langchain4j-interop/`](langchain4j-interop) | The same order assistant in Spring AI 2.0.1 and LangChain4j 1.22.0 in one Spring Boot 4.1.1 app, against a local fake OpenAI server: the bean collision when both OpenAI starters are present, the two requests compared on the wire, a throwing tool and a looping model (what each framework bounds), the meters each records, one MCP tool server used from both, and which lines of the 2025 version of the article still compile. Java 25. | [Spring AI vs LangChain4j in 2026: Is the Pairing Still Worth It?](https://ankurm.com/spring-ai-langchain4j-production-ai-microservices-java/) | | [`agent-skills/`](agent-skills) | `SKILL.md` folders loaded on demand by the `Skill` tool from `spring-ai-agent-utils` 0.13.0: what the model sees before it loads anything (name and description only), the two-call load round trip, unknown, duplicate and description-less skills, and listing size against inlining every body (character counts). Scripted chat model, no live model. Spring Boot 4.1.1, Spring AI 2.0.1, Java 25. | [Agent Skills in Spring AI: Packaging Reusable Capabilities for Your Agents](https://ankurm.com/agent-skills-spring-ai-packaging-reusable-capabilities/) | | [`rag-ingestion/`](rag-ingestion) | PDF, Word and HTML versions of one 12-section runbook through `PagePdfDocumentReader`, `TikaDocumentReader` and `JsoupDocumentReader`, then `TokenTextSplitter` at six sizes, a paragraph chunker, a sentence splitter and a section chunker; retrieval hit@1/3/5 on 37 fixed questions, metadata in the embedding, content-hash idempotent ingestion into pgvector and the re-embedding cost of one edit. Local all-MiniLM-L6-v2; a synthetic corpus, so differences of two or three questions are noise. Spring Boot 4.1.2, Spring AI 2.0.2, Java 25. | [Spring AI RAG Ingestion Pipeline: PDF, Word and HTML Readers with Chunking Strategies Compared](https://ankurm.com/spring-ai-rag-ingestion-pipeline-pdf-word-html-chunking/) | +| [`hybrid-search/`](hybrid-search) | Vector, `tsvector` and a Java Okapi BM25 over one PostgreSQL table (pgvector HNSW + GIN), fused with Reciprocal Rank Fusion in Java and in a single SQL statement, behind a `DocumentRetriever` for `RetrievalAugmentationAdvisor`. 60 questions in four kinds over 94 passages: hybrid helps exact error codes and hurts paraphrases; `k` tuned on the evaluation set did not hold up on held-out folds; plans and latency at 94 and 50,094 rows. Local all-MiniLM-L6-v2. Spring Boot 4.1.2, Spring AI 2.0.2, Java 25. | [Hybrid Search in Java: BM25 + Vector with PostgreSQL and Reciprocal Rank Fusion](https://ankurm.com/hybrid-search-java-pgvector-tsvector-reciprocal-rank-fusion/) | Upgrading from Spring AI 1.x: [migration guide](https://ankurm.com/spring-ai-1-to-2-migration-guide/). diff --git a/hybrid-search/.gitignore b/hybrid-search/.gitignore new file mode 100644 index 0000000..1ebf457 --- /dev/null +++ b/hybrid-search/.gitignore @@ -0,0 +1,3 @@ +target/ +*.iml +.idea/ diff --git a/hybrid-search/README.md b/hybrid-search/README.md new file mode 100644 index 0000000..63020a4 --- /dev/null +++ b/hybrid-search/README.md @@ -0,0 +1,7 @@ +# hybrid-search + +Companion code for [Hybrid Search in Java: BM25 + Vector with PostgreSQL and Reciprocal Rank Fusion](https://ankurm.com/hybrid-search-java-pgvector-tsvector-reciprocal-rank-fusion/). + +One PostgreSQL table holds a `vector(384)` column (pgvector, HNSW) and a generated `tsvector` column (GIN). Vector search, keyword search and Reciprocal Rank Fusion are measured on 94 passages and 60 questions in four kinds. Spring Boot 4.1.2, Spring AI 2.0.2, Java 25. + +Run `scripts/run-all.sh` (needs PostgreSQL 16 with pgvector; see the script header). Every number in the article is a line in `output/`. diff --git a/hybrid-search/output/01-rrf-by-hand.txt b/hybrid-search/output/01-rrf-by-hand.txt new file mode 100644 index 0000000..7171d5b --- /dev/null +++ b/hybrid-search/output/01-rrf-by-hand.txt @@ -0,0 +1,34 @@ +# Reciprocal Rank Fusion on two short rankings + +vector ranking : [A, B, C, D] +keyword ranking: [C, A, E, B] + +--- k = 1 --- +A score 0.83333 (vector rank 1, keyword rank 2) +C score 0.75000 (vector rank 3, keyword rank 1) +B score 0.53333 (vector rank 2, keyword rank 4) +E score 0.25000 (vector rank 0, keyword rank 3) +D score 0.20000 (vector rank 4, keyword rank 0) + +--- k = 60 --- +A score 0.03252 (vector rank 1, keyword rank 2) +C score 0.03227 (vector rank 3, keyword rank 1) +B score 0.03175 (vector rank 2, keyword rank 4) +E score 0.01587 (vector rank 0, keyword rank 3) +D score 0.01563 (vector rank 4, keyword rank 0) + +--- k = 1000 --- +A score 0.00200 (vector rank 1, keyword rank 2) +C score 0.00200 (vector rank 3, keyword rank 1) +B score 0.00199 (vector rank 2, keyword rank 4) +E score 0.00100 (vector rank 0, keyword rank 3) +D score 0.00100 (vector rank 4, keyword rank 0) + +--- the arithmetic for A at k = 60 --- +rank 1 in the vector list: 1/(60+1) = 0.01639 +rank 2 in the keyword list: 1/(60+2) = 0.01613 +sum = 0.03252 + +--- an item in one list only --- +E is rank 3 in the keyword list and absent from the vector list: 1/(60+3) = 0.01587 +D is rank 4 in the vector list and absent from the keyword list: 1/(60+4) = 0.01563 diff --git a/hybrid-search/output/02-quality-by-query-type.txt b/hybrid-search/output/02-quality-by-query-type.txt new file mode 100644 index 0000000..3e9f686 --- /dev/null +++ b/hybrid-search/output/02-quality-by-query-type.txt @@ -0,0 +1,73 @@ +# Six rankers on 60 labelled queries over 94 passages (all-MiniLM-L6-v2, PostgreSQL 16, pgvector 0.8.7) + +fusion: RRF with k=60 over the top 20 of each list. hit@1 and hit@5: the answer passage is first / in the first five. MRR@10: mean of 1/rank. + +--- exact error code (12 queries) --- +ranker hit@1 hit@5 MRR@10 +vector (pgvector cosine) 2/12 8/12 0.374 +keyword (tsvector, ts_rank_cd) 9/12 11/12 0.819 +keyword (BM25, in Java) 9/12 12/12 0.875 +hybrid: vector + tsvector, RRF 9/12 11/12 0.806 +hybrid: vector + BM25, RRF 10/12 11/12 0.883 +hybrid in one SQL statement 9/12 11/12 0.806 + +--- exact SKU or setting (12 queries) --- +ranker hit@1 hit@5 MRR@10 +vector (pgvector cosine) 12/12 12/12 1.000 +keyword (tsvector, ts_rank_cd) 12/12 12/12 1.000 +keyword (BM25, in Java) 11/12 12/12 0.958 +hybrid: vector + tsvector, RRF 12/12 12/12 1.000 +hybrid: vector + BM25, RRF 12/12 12/12 1.000 +hybrid in one SQL statement 12/12 12/12 1.000 + +--- paraphrase (24 queries) --- +ranker hit@1 hit@5 MRR@10 +vector (pgvector cosine) 15/24 23/24 0.767 +keyword (tsvector, ts_rank_cd) 5/24 11/24 0.296 +keyword (BM25, in Java) 4/24 13/24 0.290 +hybrid: vector + tsvector, RRF 13/24 20/24 0.663 +hybrid: vector + BM25, RRF 11/24 21/24 0.626 +hybrid in one SQL statement 13/24 20/24 0.663 + +--- code plus words (12 queries) --- +ranker hit@1 hit@5 MRR@10 +vector (pgvector cosine) 12/12 12/12 1.000 +keyword (tsvector, ts_rank_cd) 10/12 12/12 0.896 +keyword (BM25, in Java) 11/12 12/12 0.958 +hybrid: vector + tsvector, RRF 12/12 12/12 1.000 +hybrid: vector + BM25, RRF 12/12 12/12 1.000 +hybrid in one SQL statement 12/12 12/12 1.000 + +--- all 60 queries --- +ranker hit@1 hit@5 MRR@10 +vector (pgvector cosine) 41/60 55/60 0.782 +keyword (tsvector, ts_rank_cd) 36/60 46/60 0.662 +keyword (BM25, in Java) 35/60 49/60 0.674 +hybrid: vector + tsvector, RRF 46/60 55/60 0.826 +hybrid: vector + BM25, RRF 45/60 56/60 0.827 +hybrid in one SQL statement 46/60 55/60 0.826 + +--- the SQL version against the Java version --- +queries whose top 5 are identical, in the same order, from both implementations: 60 of 60 + +--- one query of each kind: the top 3 from each ranker --- +[exact error code] "What does error E-4106 mean?" -> answer is E-4106 + vector : E-4108, E-4105, E-4114 + tsvector : E-4106, E-4109, E-4139 + BM25 : E-4109, E-4106, E-4140 + hybrid : E-4108, E-4109, E-4111 +[exact SKU or setting] "Does OP-TERM-5520 work when the internet is down?" -> answer is OP-TERM-5520 + vector : OP-TERM-5520, OP-SUP-GOLD, G-30 + tsvector : OP-TERM-5520, OP-TERM-6100, OP-TERM-5530 + BM25 : OP-TERM-5520 + hybrid : OP-TERM-5520, OP-SUP-GOLD, OP-TERM-5530 +[paraphrase] "What happens when a customer's monthly payment does not go through?" -> answer is G-01 + vector : G-15, G-04, G-16 + tsvector : E-4110, E-4111, G-15 + BM25 : G-04, G-06, orbit.receipts.locale + hybrid : G-15, G-16, G-04 +[code plus words] "E-4106 keeps appearing in our Express webhook handler, is the body being parsed?" -> answer is E-4106 + vector : E-4106, E-4107, E-4116 + tsvector : E-4106, E-4102, E-4107 + BM25 : E-4106, E-4102, E-4107 + hybrid : E-4106, E-4107, E-4102 diff --git a/hybrid-search/output/03-rrf-parameters.txt b/hybrid-search/output/03-rrf-parameters.txt new file mode 100644 index 0000000..4d77560 --- /dev/null +++ b/hybrid-search/output/03-rrf-parameters.txt @@ -0,0 +1,26 @@ +# RRF parameters on the 60 queries: k, candidates per list, list weights + + +--- k, with 20 candidates per list --- +k hit@1 hit@5 MRR@10 +1 46/60 58/60 0.851 +10 45/60 55/60 0.826 +60 46/60 55/60 0.826 +200 46/60 55/60 0.825 +1000 46/60 55/60 0.825 + +--- candidates taken from each list, with k = 60 --- +cands hit@1 hit@5 MRR@10 +3 41/60 57/60 0.802 +5 43/60 59/60 0.828 +10 45/60 57/60 0.823 +20 46/60 55/60 0.826 +50 46/60 55/60 0.829 + +--- weights (vector : keyword), k = 60, 20 candidates --- +weights hit@1 hit@5 MRR@10 +1.0:0.0 41/60 55/60 0.782 +1.0:0.5 45/60 56/60 0.822 +1.0:1.0 46/60 55/60 0.826 +0.5:1.0 42/60 54/60 0.781 +0.0:1.0 37/60 50/60 0.692 diff --git a/hybrid-search/output/04-held-out-k.txt b/hybrid-search/output/04-held-out-k.txt new file mode 100644 index 0000000..bbb98dc --- /dev/null +++ b/hybrid-search/output/04-held-out-k.txt @@ -0,0 +1,15 @@ +# Choosing k on half of the queries and scoring it on the other half + +candidates per list: 20. The 60 queries are split by position into two folds of 30; k comes from {1, 10, 60, 200}. + +--- fold 1: train on 30, test on 30 --- +k chosen on the training half (best MRR@10 0.872): 1 +held-out half, k=1 : 22/30 29/30 0.831 +held-out half, k=60 : 24/30 28/30 0.859 +held-out half, vector only: 21/30 30/30 0.831 + +--- fold 2: train on 30, test on 30 --- +k chosen on the training half (best MRR@10 0.859): 60 +held-out half, k=60 : 22/30 27/30 0.794 +held-out half, k=60 : 22/30 27/30 0.794 +held-out half, vector only: 20/30 25/30 0.733 diff --git a/hybrid-search/output/05-retriever-and-advisor.txt b/hybrid-search/output/05-retriever-and-advisor.txt new file mode 100644 index 0000000..5222a74 --- /dev/null +++ b/hybrid-search/output/05-retriever-and-advisor.txt @@ -0,0 +1,43 @@ +# HybridDocumentRetriever inside RetrievalAugmentationAdvisor (top 3 passages per question) + + +--- what the retriever returns for "Why am I getting E-4116?" --- +score 0.03279 vector rank 1 keyword rank 1 Error E-4116 - Charge already captured. A capture was attemp +score 0.03125 vector rank 4 keyword rank 4 Error E-4111 - Authentication challenge failed. The customer +score 0.02988 vector rank 9 keyword rank 5 Error E-4118 - Payout account not verified. Payouts are bloc + +--- the prompt the model received --- +Context information is below. + +--------------------- +Error E-4116 - Charge already captured. A capture was attempted on a charge that has already been captured in full. Fetch the charge first and check its status instead of capturing blindly. +Error E-4111 - Authentication challenge failed. The customer abandoned or failed the 3-D Secure challenge. The payment cannot be completed with this card; offer another payment method. +Error E-4118 - Payout account not verified. Payouts are blocked because the bank account has not passed verification. Confirm the two micro-deposits shown in the dashboard to unlock transfers. +--------------------- + +Given the context information and no prior knowledge, answer the query. + +Follow these rules: + +1. If the answer is not in the context, just say that you don't know. +2. Avoid statements like "Based on the context..." or "The provided information...". + +Query: Why am I getting E-4116? + +Answer: + + +--- does the answer passage reach the prompt? the 12 exact-code questions, top 3 passages --- +E-4106 vector-only: false hybrid: false What does error E-4106 mean? +E-4116 vector-only: true hybrid: true Why am I getting E-4116? +E-4118 vector-only: false hybrid: true We see E-4118 in the logs, what should we do? +E-4120 vector-only: true hybrid: true Explain E-4120 +E-4122 vector-only: false hybrid: true Support told me to look up E-4122 +E-4126 vector-only: true hybrid: true What is E-4126? +E-4129 vector-only: false hybrid: true How do I fix E-4129? +E-4131 vector-only: false hybrid: true Our service started returning E-4131 this morning +E-4133 vector-only: false hybrid: true Meaning of E-4133 +E-4135 vector-only: true hybrid: true Is E-4135 our fault? +E-4138 vector-only: false hybrid: true What causes E-4138 on the shop floor? +E-4140 vector-only: true hybrid: true Customer reports E-4140, is the money safe? +answer passage in the top 3: vector-only 5 of 12, hybrid 11 of 12 diff --git a/hybrid-search/output/06-scale-plans-and-latency.txt b/hybrid-search/output/06-scale-plans-and-latency.txt new file mode 100644 index 0000000..31f11ca --- /dev/null +++ b/hybrid-search/output/06-scale-plans-and-latency.txt @@ -0,0 +1,45 @@ +# 94 rows against 50,094 rows: plans, index sizes, latency (PostgreSQL 16, pgvector 0.8.7, 2 vCPU) + +filler rows inserted: 50000 in 7819 ms. HNSW index built in 94632 ms, GIN index in 292 ms. +rows in hybrid_scale: 50094 +heap 98 MB, HNSW index 98 MB, GIN index 2248 kB + +--- plan for the vector query at 50,094 rows (scan nodes only) --- +Limit (actual rows=20 loops=1) + -> Index Scan using hybrid_scale_embedding_idx on hybrid_scale (actual rows=20 loops=1) + +--- plan for the keyword query at 50,094 rows (scan nodes only) --- +Limit (actual rows=20 loops=1) + -> Sort (actual rows=20 loops=1) + Sort Key: (ts_rank_cd(d.tsv, '''get'' | ''e'' | ''-4116'''::tsquery)) DESC + Sort Method: quicksort Memory: 26kB + -> Bitmap Heap Scan on hybrid_scale d (actual rows=40 loops=1) + -> Bitmap Index Scan on hybrid_scale_tsv_idx (actual rows=40 loops=1) + +--- plan for the same keyword query at 94 rows --- +Limit (actual rows=20 loops=1) + -> Sort (actual rows=20 loops=1) + Sort Key: (ts_rank_cd(d.tsv, '''get'' | ''e'' | ''-4116'''::tsquery)) DESC + Sort Method: quicksort Memory: 26kB + -> Seq Scan on hybrid_docs d (actual rows=40 loops=1) + Rows Removed by Filter: 54 + +--- rows the keyword query has to rank, by query kind (50,094 rows) --- +exact error code average rows matching the OR query: 1547 +exact SKU or setting average rows matching the OR query: 1970 +paraphrase average rows matching the OR query: 3781 +code plus words average rows matching the OR query: 5220 + +--- latency in milliseconds over 60 queries x 3 runs after 2 warm-up passes (query embedding excluded unless stated) --- + 94 p50 94 p95 50k p50 50k p95 +embedding the question (MiniLM) 7.4 10.4 (same for both table sizes) +vector, top 20 25.3 32.2 28.0 36.4 +keyword, top 20 27.0 39.3 85.5 177.1 +hybrid in Java (2 queries + fuse) 51.6 65.3 111.2 208.2 +hybrid in one SQL statement 31.5 40.1 132.3 226.4 + +--- keyword query latency at 50,094 rows by query kind (p50 / p95 ms) --- +exact error code 51.1 / 124.0 +exact SKU or setting 35.7 / 223.0 +paraphrase 85.9 / 169.6 +code plus words 121.2 / 214.9 diff --git a/hybrid-search/pom.xml b/hybrid-search/pom.xml new file mode 100644 index 0000000..b80c6ea --- /dev/null +++ b/hybrid-search/pom.xml @@ -0,0 +1,64 @@ + + + 4.0.0 + + + org.springframework.boot + spring-boot-starter-parent + 4.1.2 + + + + com.ankurm + hybrid-search + 1.0.0 + hybrid-search + Hybrid search on one PostgreSQL: pgvector and tsvector queries fused with Reciprocal Rank Fusion in Java, a real BM25 for comparison, and a Spring AI DocumentRetriever, measured on a labelled set of 60 queries. + + + 25 + 2.0.2 + + + + + + org.springframework.ai + spring-ai-bom + ${spring-ai.version} + pom + import + + + + + + + org.springframework.aispring-ai-pgvector-store + + org.springframework.aispring-ai-rag + + org.springframework.aispring-ai-transformers + org.postgresqlpostgresql + org.springframework.bootspring-boot-starter-jdbc + + org.springframework.boot + spring-boot-starter-test + test + + + + + + + org.apache.maven.plugins + maven-surefire-plugin + + -Duser.timezone=UTC -Dstdout.encoding=UTF-8 -Dfile.encoding=UTF-8 -Xmx3g + + + + + diff --git a/hybrid-search/scripts/run-all.sh b/hybrid-search/scripts/run-all.sh new file mode 100755 index 0000000..08d1833 --- /dev/null +++ b/hybrid-search/scripts/run-all.sh @@ -0,0 +1,10 @@ +#!/usr/bin/env bash +# Regenerates every file in output/ by running the test suite. +# Needs: JDK 25, Maven, internet for the first run (about 90 MB for the MiniLM model, cached in ~/.cache/minilm), +# and a PostgreSQL 16 with pgvector 0.8 on localhost:5432 (database ragdemo, user demo / demo). +# Set HYBRID_DB_URL to point somewhere else. ScaleTest inserts 50,000 rows and builds an HNSW index (about two minutes). +set -euo pipefail +cd "$(dirname "$0")/.." +rm -rf output +mvn -B -q test +ls output diff --git a/hybrid-search/src/main/java/com/ankurm/hybridsearch/Bm25.java b/hybrid-search/src/main/java/com/ankurm/hybridsearch/Bm25.java new file mode 100644 index 0000000..4994960 --- /dev/null +++ b/hybrid-search/src/main/java/com/ankurm/hybridsearch/Bm25.java @@ -0,0 +1,85 @@ +package com.ankurm.hybridsearch; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; +import java.util.regex.Matcher; +import java.util.regex.Pattern; + +/** + * Okapi BM25 over the passages, held in memory. PostgreSQL's ts_rank and ts_rank_cd are not BM25: they have no + * inverse document frequency, so a word found in every passage counts as much as a word found in one. This + * class exists to put a real BM25 next to them. Tokens are lower-case runs of letters and digits that may + * contain dots, underscores and hyphens, so "E-4116" and "orbit.retry.backoff_ms" each stay one token. + */ +public final class Bm25 { + + private static final Pattern TOKEN = Pattern.compile("[a-z0-9]+(?:[._-][a-z0-9]+)*"); + private static final Set STOP = Set.of("a", "an", "the", "of", "to", "in", "is", "are", "and", "or", "for", "on", "at", "be", "by", + "do", "does", "how", "what", "why", "can", "it", "its", "we", "our", "i", "my", "me", "with", "was", "this", "that", "from", "as", + "have", "has", "when", "which", "if", "so", "not", "no", "am"); + + private final double k1; + private final double b; + private final List ids = new ArrayList<>(); + private final List> termFrequencies = new ArrayList<>(); + private final List lengths = new ArrayList<>(); + private final Map documentFrequency = new HashMap<>(); + private final double averageLength; + + public Bm25(List passages, double k1, double b) { + this.k1 = k1; + this.b = b; + long total = 0; + for (KnowledgeBase.Passage p : passages) { + List tokens = tokens(p.text()); + Map tf = new LinkedHashMap<>(); + tokens.forEach(t -> tf.merge(t, 1, Integer::sum)); + ids.add(p.pid()); + termFrequencies.add(tf); + lengths.add(tokens.size()); + total += tokens.size(); + tf.keySet().forEach(t -> documentFrequency.merge(t, 1, Integer::sum)); + } + this.averageLength = (double) total / passages.size(); + } + + public List search(String query, int n) { + List terms = tokens(query); + List scored = new ArrayList<>(); + for (int d = 0; d < ids.size(); d++) { + double score = 0; + for (String term : terms) { + Integer f = termFrequencies.get(d).get(term); + if (f == null) { + continue; + } + int df = documentFrequency.get(term); + double idf = Math.log(1 + (ids.size() - df + 0.5) / (df + 0.5)); + score += idf * f * (k1 + 1) / (f + k1 * (1 - b + b * lengths.get(d) / averageLength)); + } + if (score > 0) { + scored.add(new double[] {d, score}); + } + } + scored.sort((x, y) -> Double.compare(y[1], x[1])); + return scored.stream().limit(n).map(s -> ids.get((int) s[0])).toList(); + } + + static List tokens(String text) { + List out = new ArrayList<>(); + Matcher m = TOKEN.matcher(text.toLowerCase(Locale.ROOT)); + while (m.find()) { + String t = m.group(); + if (STOP.contains(t)) { + continue; + } + out.add(t.length() > 3 && t.endsWith("s") && !t.contains("-") && !t.contains(".") ? t.substring(0, t.length() - 1) : t); + } + return out; + } +} diff --git a/hybrid-search/src/main/java/com/ankurm/hybridsearch/HybridDocumentRetriever.java b/hybrid-search/src/main/java/com/ankurm/hybridsearch/HybridDocumentRetriever.java new file mode 100644 index 0000000..88051c3 --- /dev/null +++ b/hybrid-search/src/main/java/com/ankurm/hybridsearch/HybridDocumentRetriever.java @@ -0,0 +1,35 @@ +package com.ankurm.hybridsearch; + +import java.util.List; +import java.util.Map; + +import org.springframework.ai.document.Document; +import org.springframework.ai.rag.Query; +import org.springframework.ai.rag.retrieval.search.DocumentRetriever; + +/** + * Spring AI's extension point for "where do the passages come from": give it to RetrievalAugmentationAdvisor + * and every question is answered from the fused ranking instead of from the vector store alone. + */ +public class HybridDocumentRetriever implements DocumentRetriever { + + private final HybridSearch search; + private final int topK; + private final int candidates; + private final int k; + + public HybridDocumentRetriever(HybridSearch search, int topK, int candidates, int k) { + this.search = search; + this.topK = topK; + this.candidates = candidates; + this.k = k; + } + + @Override + public List retrieve(Query query) { + return search.hybrid(query.text(), topK, candidates, k).stream() + .map(f -> Document.builder().text(f.content()).score(f.score()) + .metadata(Map.of("pid", f.pid(), "vector_rank", f.vectorRank(), "keyword_rank", f.keywordRank())).build()) + .toList(); + } +} diff --git a/hybrid-search/src/main/java/com/ankurm/hybridsearch/HybridSearch.java b/hybrid-search/src/main/java/com/ankurm/hybridsearch/HybridSearch.java new file mode 100644 index 0000000..5fd3ad9 --- /dev/null +++ b/hybrid-search/src/main/java/com/ankurm/hybridsearch/HybridSearch.java @@ -0,0 +1,155 @@ +package com.ankurm.hybridsearch; + +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.stream.Collectors; + +import org.springframework.ai.document.Document; +import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.vectorstore.pgvector.PgVectorStore; +import org.springframework.jdbc.core.JdbcTemplate; +import org.springframework.jdbc.core.namedparam.MapSqlParameterSource; +import org.springframework.jdbc.core.namedparam.NamedParameterJdbcTemplate; + +/** + * Keyword search, vector search and their fusion over one PostgreSQL table. + * + *

The table is the one Spring AI's PgVectorStore creates (id, content, metadata, embedding). The only thing + * added by hand is a generated tsvector column over {@code content} and a GIN index on it, so a passage is + * stored once and searched two ways. + */ +public class HybridSearch { + + public record Hit(String pid, String content) { + } + + /** One fused result: the passage, its RRF score and the rank it had in each list (0 when absent). */ + public record Fused(String pid, String content, double score, int vectorRank, int keywordRank) { + } + + /** The English tsquery of a question with its AND turned into OR, so one rare word is enough to match. */ + static final String OR_QUERY = "replace(plainto_tsquery('english', :text)::text, ' & ', ' | ')::tsquery"; + + private final JdbcTemplate jdbc; + private final NamedParameterJdbcTemplate named; + private final EmbeddingModel model; + private final String table; + + public HybridSearch(JdbcTemplate jdbc, EmbeddingModel model, String table) { + this.jdbc = jdbc; + this.named = new NamedParameterJdbcTemplate(jdbc); + this.model = model; + this.table = table; + } + + /** Creates the PgVectorStore table, loads the passages through it, then adds the tsvector column and its index. */ + public static HybridSearch load(JdbcTemplate jdbc, EmbeddingModel model, String table, List passages) throws Exception { + PgVectorStore store = PgVectorStore.builder(jdbc, model).vectorTableName(table).dimensions(384) + .removeExistingVectorStoreTable(true).initializeSchema(true).build(); + store.afterPropertiesSet(); + List docs = new ArrayList<>(); + for (KnowledgeBase.Passage p : passages) { + docs.add(Document.builder().id(UUID.nameUUIDFromBytes(p.pid().getBytes(StandardCharsets.UTF_8)).toString()) + .text(p.text()).metadata(Map.of("pid", p.pid(), "kind", p.kind())).build()); + } + store.add(docs); + addKeywordIndex(jdbc, table); + return new HybridSearch(jdbc, model, table); + } + + public static void addKeywordIndex(JdbcTemplate jdbc, String table) { + jdbc.execute("alter table " + table + " add column if not exists tsv tsvector generated always as (to_tsvector('english', content)) stored"); + jdbc.execute("create index if not exists " + table + "_tsv_idx on " + table + " using gin (tsv)"); + jdbc.execute("analyze " + table); + } + + public List vector(String text, int n) { + return vectorByLiteral(literal(model.embed(text)), n); + } + + /** The same query with the question already embedded, so a timing can leave the embedding model out. */ + public List vectorByLiteral(String vectorLiteral, int n) { + return named.query("select metadata->>'pid' as pid, content from " + table + + " order by embedding <=> cast(:vec as vector) limit :n", + new MapSqlParameterSource("vec", vectorLiteral).addValue("n", n), + (rs, i) -> new Hit(rs.getString("pid"), rs.getString("content"))); + } + + public List keyword(String text, int n) { + return named.query("with q as (select " + OR_QUERY + " as query) select d.metadata->>'pid' as pid, d.content, " + + "ts_rank_cd(d.tsv, q.query) as score from " + table + " d, q where d.tsv @@ q.query " + + "order by score desc, d.id limit :n", + new MapSqlParameterSource("text", text).addValue("n", n), + (rs, i) -> new Hit(rs.getString("pid"), rs.getString("content"))); + } + + /** Two queries, fused in Java: this is what a DocumentRetriever does. */ + public List hybrid(String text, int n, int candidates, int k) { + List vector = vector(text, candidates); + List keyword = keyword(text, candidates); + return fuse(vector, keyword, n, k); + } + + public static List fuse(List vector, List keyword, int n, int k) { + List v = vector.stream().map(Hit::pid).toList(); + List w = keyword.stream().map(Hit::pid).toList(); + Map content = new LinkedHashMap<>(); + vector.forEach(h -> content.putIfAbsent(h.pid(), h.content())); + keyword.forEach(h -> content.putIfAbsent(h.pid(), h.content())); + return ReciprocalRankFusion.fuse(List.of(v, w), k).stream().limit(n) + .map(f -> new Fused(f.id(), content.get(f.id()), f.score(), v.indexOf(f.id()) + 1, w.indexOf(f.id()) + 1)) + .collect(Collectors.toList()); + } + + /** + * The same fusion written as one SQL statement. Equal scores are common in RRF (rank 1 in one list only + * equals rank 1 in the other), so the tie-break is spelled out to match fuse(): vector-list order first, + * then keyword-only items in keyword order. + */ + public List hybridSql(String text, int n, int candidates, int k) { + return hybridSqlByLiteral(text, literal(model.embed(text)), n, candidates, k); + } + + public List hybridSqlByLiteral(String text, String vectorLiteral, int n, int candidates, int k) { + return named.query(""" + with q as (select %s as tsq, cast(:vec as vector) as v), + sem as (select id, row_number() over (order by dist) as r from + (select d.id, d.embedding <=> q.v as dist from %s d, q order by dist limit :cand) s), + kw as (select id, row_number() over (order by score desc, id) as r from + (select d.id, ts_rank_cd(d.tsv, q.tsq) as score from %s d, q where d.tsv @@ q.tsq + order by score desc, d.id limit :cand) s) + select d.metadata->>'pid' as pid, + coalesce(1.0 / (:k + sem.r), 0) + coalesce(1.0 / (:k + kw.r), 0) as score + from sem full outer join kw on sem.id = kw.id + join %s d on d.id = coalesce(sem.id, kw.id) + order by score desc, sem.r nulls last, kw.r nulls last + limit :n + """.formatted(OR_QUERY, table, table, table), + new MapSqlParameterSource("text", text).addValue("vec", vectorLiteral) + .addValue("cand", candidates).addValue("k", k).addValue("n", n), + (rs, i) -> rs.getString("pid")); + } + + public JdbcTemplate jdbc() { + return jdbc; + } + + public String table() { + return table; + } + + public static String literal(float[] v) { + StringBuilder sb = new StringBuilder("["); + for (int i = 0; i < v.length; i++) { + if (i > 0) { + sb.append(','); + } + sb.append(v[i]); + } + return sb.append(']').toString(); + } +} diff --git a/hybrid-search/src/main/java/com/ankurm/hybridsearch/KnowledgeBase.java b/hybrid-search/src/main/java/com/ankurm/hybridsearch/KnowledgeBase.java new file mode 100644 index 0000000..6b24944 --- /dev/null +++ b/hybrid-search/src/main/java/com/ankurm/hybridsearch/KnowledgeBase.java @@ -0,0 +1,222 @@ +package com.ankurm.hybridsearch; + +import java.util.ArrayList; +import java.util.List; + +/** + * A made-up support knowledge base for a payments API called Orbit: 94 passages and 60 labelled queries. + * Every query has exactly one passage that answers it, named by {@code pid}. + * + *

The passages come in three kinds, because the three kinds are where the two search methods differ: + * 40 error codes whose texts all look alike and differ in the identifier and one fact, 24 hardware and + * configuration identifiers (SKUs and setting names), and 30 how-to guides written in plain language. + * The corpus is synthetic and small. It shows how the methods behave, not how any of them scales. + */ +public final class KnowledgeBase { + + public record Passage(String pid, String kind, String text) { + } + + public record Query(String text, String pid, String category) { + } + + public static final String EXACT_CODE = "exact error code"; + public static final String EXACT_ID = "exact SKU or setting"; + public static final String PARAPHRASE = "paraphrase"; + public static final String MIXED = "code plus words"; + + public static final List PASSAGES = new ArrayList<>(); + public static final List QUERIES = new ArrayList<>(); + + private static void error(String code, String title, String body) { + PASSAGES.add(new Passage(code, "error", "Error " + code + " - " + title + ". " + body)); + } + + private static void item(String id, String kind, String text) { + PASSAGES.add(new Passage(id, kind, text)); + } + + private static void guide(String id, String title, String body) { + PASSAGES.add(new Passage(id, "guide", title + ". " + body)); + } + + private static void q(String category, String pid, String text) { + QUERIES.add(new Query(text, pid, category)); + } + + static { + // ---- 40 error codes ------------------------------------------------------------------------------- + error("E-4101", "Missing idempotency key", "Requests that create a charge must send an Idempotency-Key header. A request without one is rejected before any money moves. Generate a fresh UUID for each logical payment attempt."); + error("E-4102", "Idempotency key reused with a different body", "The same Idempotency-Key arrived with a different request body. Orbit replays the stored response only when the body is identical, and otherwise refuses the call."); + error("E-4103", "Unsupported currency", "The currency code is not enabled for this merchant account. Enable it in the dashboard under Settings, or convert the amount before calling the API."); + error("E-4104", "Amount below minimum", "The charge is smaller than the minimum for its currency, which is fifty cents for dollars. Orbit does not process smaller amounts."); + error("E-4105", "Amount above limit", "The amount is higher than the per-transaction ceiling set for the merchant. Ask support to raise the ceiling or split the order into several charges."); + error("E-4106", "Webhook signature mismatch", "The signature header on a webhook did not match the value computed with the endpoint secret. Verify against the raw request body, not the parsed JSON, because re-serialising changes the bytes."); + error("E-4107", "Webhook timestamp too old", "The timestamp inside the webhook signature is more than five minutes in the past. This guards against replays. Check for clock drift on the receiving server."); + error("E-4108", "Card expired", "The expiry date on the card has passed. Ask the customer for a different card or for the new expiry date before trying again."); + error("E-4109", "Card number invalid", "The card number fails the Luhn checksum, which usually means a typing mistake. Orbit rejects it locally without contacting the issuer."); + error("E-4110", "Strong customer authentication required", "The issuer demands a 3-D Secure challenge before it will approve the payment. Redirect the customer to the authentication page and confirm the payment afterwards."); + error("E-4111", "Authentication challenge failed", "The customer abandoned or failed the 3-D Secure challenge. The payment cannot be completed with this card; offer another payment method."); + error("E-4112", "Issuer declined without a reason", "The card issuer returned a generic do-not-honor response. Orbit cannot see the real cause. The customer should call their bank or use a different card."); + error("E-4113", "Insufficient funds", "The account behind the card does not have enough available balance for the amount. A later retry may succeed once the customer tops up."); + error("E-4114", "Refund larger than the captured amount", "The refund requested is more than what was captured on the charge, after earlier partial refunds are taken off. Lower the refund amount."); + error("E-4115", "Refund window closed", "Refunds to cards are possible for 180 days after capture. After that the money must be returned by bank transfer outside Orbit."); + error("E-4116", "Charge already captured", "A capture was attempted on a charge that has already been captured in full. Fetch the charge first and check its status instead of capturing blindly."); + error("E-4117", "Authorization expired", "An authorization that is not captured within seven days lapses and the hold is released. Create a new authorization if the goods ship later."); + error("E-4118", "Payout account not verified", "Payouts are blocked because the bank account has not passed verification. Confirm the two micro-deposits shown in the dashboard to unlock transfers."); + error("E-4119", "Payout below threshold", "The available balance is under ten units of currency, so the scheduled payout was skipped. It will be included in the next run once the balance grows."); + error("E-4120", "Dispute evidence deadline passed", "Evidence for a chargeback must be submitted before the deadline printed on the dispute. After it passes the dispute is lost automatically."); + error("E-4121", "Customer not found", "No customer exists with the identifier supplied. It may belong to the other mode, since test and live objects are kept apart."); + error("E-4122", "Payment method not attached", "The payment method is not attached to this customer. Attach it first, or pass it directly on a one-off charge."); + error("E-4123", "Duplicate customer email", "A customer with this email address already exists and the account is configured to enforce unique emails. Reuse the existing customer."); + error("E-4124", "Plan archived", "The subscription plan was archived and cannot be chosen by new subscribers. Existing subscribers keep it until they cancel."); + error("E-4125", "Proration not allowed during trial", "Prorated billing cannot be applied while the subscription is still in its free trial. Change the plan at the end of the trial instead."); + error("E-4126", "Coupon expired", "The coupon passed its redemption end date. Create a new coupon or extend the validity in the dashboard."); + error("E-4127", "Coupon already redeemed", "This single-use coupon was already redeemed by another order. Issue a fresh code for each customer."); + error("E-4128", "Tax identifier invalid", "The tax registration number does not match the format for its country. Check the prefix and the number of digits."); + error("E-4129", "API key lacks permission", "The key in use does not carry the scope needed for this endpoint. Create a restricted key that includes it."); + error("E-4130", "API key revoked", "The secret key was revoked, either by a team member or automatically after it was found in a public repository. Issue a new key and rotate it in your deployment."); + error("E-4131", "Rate limit exceeded", "More than one hundred requests per second arrived on this key. Back off with jitter and respect the Retry-After header."); + error("E-4132", "Request body too large", "The request payload is bigger than one megabyte. Upload large files through the file endpoint and reference them by identifier."); + error("E-4133", "Unknown parameter", "The request contains a field that this API version does not define. Remove it or pin a newer API version."); + error("E-4134", "Metadata key too long", "A metadata key is longer than forty characters. Shorten the key names."); + error("E-4135", "Too many metadata keys", "An object may carry at most fifty metadata keys. Store bulky data in your own database and keep only a reference."); + error("E-4136", "Test card used in live mode", "One of the published test card numbers was used with a live key. Switch to the test key, or use a real card."); + error("E-4137", "Live card used in test mode", "A real card number was sent with a test key. Orbit never charges real cards in test mode, so use a test card."); + error("E-4138", "Terminal offline", "The card terminal has not sent a heartbeat and cannot take payments. Check its network connection and power."); + error("E-4139", "Terminal firmware outdated", "The terminal runs firmware older than the minimum supported release. Connect it to power and wifi so it can update overnight."); + error("E-4140", "Receipt email bounced", "The receipt could not be delivered because the address rejected it. The payment itself is unaffected."); + + // ---- 12 hardware, plan and support identifiers ------------------------------------------------------ + item("OP-TERM-5520", "sku", "OP-TERM-5520 countertop card terminal with a 4G fallback modem. Accepts chip, tap and swipe, and keeps working for up to eight hours of sales if the shop network drops."); + item("OP-TERM-5530", "sku", "OP-TERM-5530 handheld card terminal with a built-in receipt printer and a six hour battery. Intended for table service in restaurants."); + item("OP-TERM-6100", "sku", "OP-TERM-6100 unattended card reader for kiosks and vending machines. Vandal-resistant housing, no display, powered over the serial port."); + item("OP-DOCK-210", "sku", "OP-DOCK-210 charging dock for the handheld terminals. Holds a single unit and charges it to full in about ninety minutes."); + item("OP-PRNT-77", "sku", "OP-PRNT-77 thermal paper rolls, 58 millimetres wide, sold in boxes of fifty. Fits the handheld terminal."); + item("OP-SIM-EU1", "sku", "OP-SIM-EU1 data SIM for terminals used across the European Union, with roaming included and no monthly minimum."); + item("OP-STND-45", "sku", "OP-STND-45 adjustable stand for countertop terminals. Tilts from zero to forty-five degrees and has a lockable base."); + item("OP-PLAN-STD", "plan", "OP-PLAN-STD standard pricing plan: 2.4 percent plus twenty cents per successful card charge, no monthly fee, payouts every two business days."); + item("OP-PLAN-PLUS", "plan", "OP-PLAN-PLUS pricing plan for growing merchants: 1.9 percent plus twenty cents per charge, a monthly fee, and next business day payouts."); + item("OP-PLAN-ENT", "plan", "OP-PLAN-ENT enterprise plan with negotiated rates, a named account manager and a custom payout calendar."); + item("OP-SUP-GOLD", "support", "OP-SUP-GOLD gold support tier: a human answers within one hour at any time of day, with a phone line for outages."); + item("OP-SUP-SILV", "support", "OP-SUP-SILV silver support tier: email answers within one business day, no phone line."); + + // ---- 12 configuration settings --------------------------------------------------------------------- + item("orbit.retry.max_attempts", "setting", "Setting orbit.retry.max_attempts is the number of times the client library repeats a failed request before giving up. The default is 5."); + item("orbit.retry.backoff_ms", "setting", "Setting orbit.retry.backoff_ms is the initial wait between retries in milliseconds. It doubles after each attempt. The default is 250."); + item("orbit.webhook.tolerance_seconds", "setting", "Setting orbit.webhook.tolerance_seconds is how old a webhook signature timestamp may be before verification fails. The default is 300."); + item("orbit.http.timeout_ms", "setting", "Setting orbit.http.timeout_ms is the total time the client waits for a response. The default is 30000."); + item("orbit.capture.mode", "setting", "Setting orbit.capture.mode chooses between automatic capture at authorization and manual capture later. The default is automatic."); + item("orbit.idempotency.ttl_hours", "setting", "Setting orbit.idempotency.ttl_hours is how long Orbit remembers an idempotency key and replays its response. The default is 24."); + item("orbit.radar.threshold", "setting", "Setting orbit.radar.threshold is the risk score from 0 to 100 above which a payment is blocked. The default is 75."); + item("orbit.payout.schedule", "setting", "Setting orbit.payout.schedule picks daily, weekly or monthly payouts. The default is daily."); + item("orbit.receipts.locale", "setting", "Setting orbit.receipts.locale is the language of emailed receipts. The default follows the customer's browser."); + item("orbit.terminal.heartbeat_seconds", "setting", "Setting orbit.terminal.heartbeat_seconds is how often a terminal reports in. After three missed beats it is shown as offline. The default is 60."); + item("orbit.log.level", "setting", "Setting orbit.log.level controls how much the client library writes to its log. The default is warn."); + item("orbit.cache.customer_ttl_s", "setting", "Setting orbit.cache.customer_ttl_s is how long customer objects are cached inside the client. The default is 120."); + + // ---- 30 guides in plain language -------------------------------------------------------------------- + guide("G-01", "Recovering failed subscription renewals", "When a renewal charge fails, Orbit runs a dunning schedule: it retries after one, three and five days, then marks the subscription past due and emails the customer a link to update their card."); + guide("G-02", "Paying several sellers from one order", "A marketplace can divide the proceeds of a single charge between sellers using connected accounts and transfers. Each seller is paid out on their own schedule and the platform keeps its fee."); + guide("G-03", "Authorise now, capture on dispatch", "Place a hold on the card when the order is made and capture it when the goods ship. The hold guarantees the funds for seven days and can be released without any fee."); + guide("G-04", "Responding to a chargeback", "When a cardholder disputes a payment their bank takes the money back at once. Gather proof of delivery and customer communication and submit it through the dispute page before the deadline."); + guide("G-05", "Trying the integration without real money", "Test mode mirrors live mode with separate keys and published card numbers that simulate approvals, declines and challenges. Nothing in test mode moves real funds."); + guide("G-06", "Remembering a shopper's card", "Attach a payment method to a customer object after the first purchase. Later orders can then reuse it, and the shopper never types the number again."); + guide("G-07", "Taking payment from people without cards", "Orbit supports bank redirects, wallets and instant bank transfers. Enable the methods in the dashboard and they appear in the checkout automatically for eligible countries."); + guide("G-08", "When refunded money arrives", "A refund is sent to the bank immediately but most issuers need five to ten business days to show it on the customer's statement."); + guide("G-09", "What shows on the cardholder's statement", "The statement descriptor is the short text a cardholder sees next to the charge. Choose one customers will recognise to avoid unnecessary disputes."); + guide("G-10", "Blocking suspicious orders", "Radar scores every payment for fraud risk using signals across the network. Rules let you block, review or allow payments based on the score, the country or the card type."); + guide("G-11", "When funds reach your bank account", "Settled funds are paid to your bank on a rolling schedule. New accounts wait longer for the first payout while Orbit builds a history."); + guide("G-12", "If your server is unreachable", "When your endpoint does not answer, Orbit redelivers the event with increasing gaps for up to three days. You can also list missed events through the events endpoint and replay them."); + guide("G-13", "Charging in one currency and settling in another", "Present prices in the shopper's currency and Orbit converts to your settlement currency at the daily rate, adding a small conversion fee to the payout."); + guide("G-14", "Adding sales tax to bills", "Orbit Tax calculates the rate from the customer's address and the product category, adds it to the invoice and keeps the reports you need for filing."); + guide("G-15", "Giving new customers a free period", "Set a trial length on the plan. The customer is not charged until the trial ends, and a reminder is emailed three days before the first payment."); + guide("G-16", "Putting a subscription on hold", "Pausing stops invoices while keeping the customer and their payment method. Billing resumes on the date you choose and the cycle is shifted accordingly."); + guide("G-17", "Handing transactions to the accountants", "Export every transaction, fee and payout for a date range as a spreadsheet. The report reconciles to the bank deposits line by line."); + guide("G-18", "Restricting what a team member can see", "Dashboard roles control who can view balances, issue refunds or read customer details. Give support staff the limited role so they cannot see payout accounts."); + guide("G-19", "Keeping European data in Europe", "Choose the EU region when you create the account and customer and card data are stored and processed only in European data centres."); + guide("G-20", "Staying out of card data rules", "Hosted fields and the checkout page collect card numbers on Orbit's servers, so they never touch yours and your compliance questionnaire stays short."); + guide("G-21", "Accepting taps on a phone", "The mobile SDK turns a supported phone into a contactless reader, with no extra hardware, for cards and wallets."); + guide("G-22", "Customising the emailed receipt", "Add your logo, brand colour and a footer message to receipts under Settings. Receipts are sent automatically after each successful payment."); + guide("G-23", "Resending a request that timed out", "If a call times out you cannot know whether it succeeded. Send it again with the same idempotency key and Orbit returns the original result instead of creating a second charge."); + guide("G-24", "Taking part of an authorised amount", "You may capture less than was authorised, for example when one item of an order is out of stock. The rest of the hold is released straight away."); + guide("G-25", "Letting customers add a tip", "Collect a gratuity on the terminal after the amount is shown. The tip is added to the same charge and reported separately."); + guide("G-26", "Showing local prices", "Define a price per currency on each product so shoppers see round local amounts instead of converted ones."); + guide("G-27", "Clearing out sandbox data", "Reset the test environment from the developer settings page to delete test customers, charges and subscriptions and start clean."); + guide("G-28", "Seeing who changed what", "The audit log records every dashboard action with the person, time and address, and can be exported for your security team."); + guide("G-29", "Signing in with your company directory", "Single sign-on lets the team log in to the dashboard through your identity provider, and access is removed automatically when someone leaves."); + guide("G-30", "Checking whether Orbit is having problems", "The status page lists current incidents and past outages per region, and you can subscribe to updates by email or webhook."); + + // ---- 60 queries -------------------------------------------------------------------------------------- + // 12 where the only useful words are an error code + q(EXACT_CODE, "E-4106", "What does error E-4106 mean?"); + q(EXACT_CODE, "E-4116", "Why am I getting E-4116?"); + q(EXACT_CODE, "E-4118", "We see E-4118 in the logs, what should we do?"); + q(EXACT_CODE, "E-4120", "Explain E-4120"); + q(EXACT_CODE, "E-4122", "Support told me to look up E-4122"); + q(EXACT_CODE, "E-4126", "What is E-4126?"); + q(EXACT_CODE, "E-4129", "How do I fix E-4129?"); + q(EXACT_CODE, "E-4131", "Our service started returning E-4131 this morning"); + q(EXACT_CODE, "E-4133", "Meaning of E-4133"); + q(EXACT_CODE, "E-4135", "Is E-4135 our fault?"); + q(EXACT_CODE, "E-4138", "What causes E-4138 on the shop floor?"); + q(EXACT_CODE, "E-4140", "Customer reports E-4140, is the money safe?"); + + // 12 where the identifier is a SKU, plan or setting name + q(EXACT_ID, "OP-TERM-5520", "Does OP-TERM-5520 work when the internet is down?"); + q(EXACT_ID, "OP-TERM-5530", "Tell me about OP-TERM-5530"); + q(EXACT_ID, "OP-TERM-6100", "Can OP-TERM-6100 be used in a vending machine?"); + q(EXACT_ID, "OP-DOCK-210", "How long does OP-DOCK-210 take to charge a unit?"); + q(EXACT_ID, "OP-PRNT-77", "What size is OP-PRNT-77?"); + q(EXACT_ID, "OP-PLAN-PLUS", "What are the rates on OP-PLAN-PLUS?"); + q(EXACT_ID, "OP-SUP-GOLD", "Response time for OP-SUP-GOLD"); + q(EXACT_ID, "orbit.retry.backoff_ms", "What is the default for orbit.retry.backoff_ms?"); + q(EXACT_ID, "orbit.radar.threshold", "Value of orbit.radar.threshold"); + q(EXACT_ID, "orbit.idempotency.ttl_hours", "What does orbit.idempotency.ttl_hours do?"); + q(EXACT_ID, "orbit.terminal.heartbeat_seconds", "Default of orbit.terminal.heartbeat_seconds"); + q(EXACT_ID, "orbit.cache.customer_ttl_s", "Explain orbit.cache.customer_ttl_s"); + + // 24 plain-language questions written without the passage's own distinctive words + q(PARAPHRASE, "G-01", "What happens when a customer's monthly payment does not go through?"); + q(PARAPHRASE, "G-02", "How can an online mall split one shopper's purchase among several vendors?"); + q(PARAPHRASE, "G-03", "Can I reserve money on a card and only take it once the parcel leaves?"); + q(PARAPHRASE, "G-04", "A buyer told their bank they never got the item. What now?"); + q(PARAPHRASE, "G-05", "How do I experiment safely before launching?"); + q(PARAPHRASE, "G-06", "Store a shopper's card so they do not re-enter it next time"); + q(PARAPHRASE, "G-07", "Accept money from people who do not own a credit card"); + q(PARAPHRASE, "G-08", "How long until the buyer sees their money back?"); + q(PARAPHRASE, "G-09", "Which words appear on the customer's bank statement?"); + q(PARAPHRASE, "G-10", "Stop dodgy purchases automatically"); + q(PARAPHRASE, "G-11", "When will the money land in our account?"); + q(PARAPHRASE, "G-12", "What if our website is down when you send us news?"); + q(PARAPHRASE, "G-13", "Show prices in euros but receive dollars"); + q(PARAPHRASE, "G-14", "Include VAT on our bills without doing sums"); + q(PARAPHRASE, "G-15", "Let new users try the product for two weeks at no cost"); + q(PARAPHRASE, "G-16", "A client wants to stop being billed for the summer"); + q(PARAPHRASE, "G-17", "Give our bookkeeper everything that came in last quarter"); + q(PARAPHRASE, "G-18", "Limit what a helpdesk employee can look at"); + q(PARAPHRASE, "G-19", "Our lawyers want shopper information to stay on this continent"); + q(PARAPHRASE, "G-20", "How do we avoid ever handling raw card numbers ourselves?"); + q(PARAPHRASE, "G-21", "Take contactless payments with just a smartphone"); + q(PARAPHRASE, "G-23", "Our call hung and we do not know if it worked. Is it safe to send it again?"); + q(PARAPHRASE, "G-24", "One item in the basket is sold out, can we bill only for the rest?"); + q(PARAPHRASE, "G-28", "Which colleague altered the settings last Tuesday?"); + + // 12 that carry a code and also say what it is about + q(MIXED, "E-4106", "E-4106 keeps appearing in our Express webhook handler, is the body being parsed?"); + q(MIXED, "E-4113", "Customer card gets E-4113 although they swear there is money on it"); + q(MIXED, "E-4117", "E-4117 on capture, the goods shipped ten days after the order"); + q(MIXED, "E-4125", "Plan change during trial returns E-4125"); + q(MIXED, "E-4130", "E-4130 on every call since the key leaked into our repository"); + q(MIXED, "E-4115", "Cannot refund a purchase from last year, E-4115"); + q(MIXED, "E-4110", "E-4110 and the customer is asked for a bank challenge"); + q(MIXED, "E-4127", "Promo code gives E-4127 for the second shopper"); + q(MIXED, "E-4139", "Terminal says E-4139 and needs an update"); + q(MIXED, "E-4101", "Getting E-4101 on the charge endpoint, our client does not set the idempotency header"); + q(MIXED, "E-4131", "E-4131 during the flash sale, too many requests per second"); + q(MIXED, "E-4119", "Why was our payout skipped, E-4119, the balance was small"); + } + + private KnowledgeBase() { + } +} diff --git a/hybrid-search/src/main/java/com/ankurm/hybridsearch/MiniLm.java b/hybrid-search/src/main/java/com/ankurm/hybridsearch/MiniLm.java new file mode 100644 index 0000000..1095ffa --- /dev/null +++ b/hybrid-search/src/main/java/com/ankurm/hybridsearch/MiniLm.java @@ -0,0 +1,44 @@ +package com.ankurm.hybridsearch; + +import java.util.Map; + +import org.springframework.ai.document.MetadataMode; +import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.ai.transformers.TransformersEmbeddingModel; + +/** + * all-MiniLM-L6-v2 (384 dimensions) running inside the JVM through ONNX Runtime. Both model files come from + * Hugging Face and are cached under the user's home directory, so the first run downloads about 90 MB and + * later runs are offline. No API key is involved, which is what lets anyone repeat the measurements. + */ +public final class MiniLm { + + private static final String HF = "https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/resolve/main/"; + private static EmbeddingModel shared; + + private MiniLm() { + } + + /** The shared instance, which embeds the text of a Document only (MetadataMode.NONE, the constructor default). */ + public static synchronized EmbeddingModel get() { + if (shared == null) { + shared = create(MetadataMode.NONE); + } + return shared; + } + + /** A fresh instance whose Document embedding content is built with the given metadata mode. */ + public static EmbeddingModel create(MetadataMode mode) { + try { + TransformersEmbeddingModel m = new TransformersEmbeddingModel(mode); + m.setTokenizerResource(HF + "tokenizer.json"); + m.setModelResource(HF + "onnx/model.onnx"); + m.setResourceCacheDirectory(System.getProperty("user.home") + "/.cache/minilm"); + m.setTokenizerOptions(Map.of("padding", "true", "truncation", "true", "maxLength", "256")); + m.afterPropertiesSet(); + return m; + } catch (Exception e) { + throw new IllegalStateException(e); + } + } +} diff --git a/hybrid-search/src/main/java/com/ankurm/hybridsearch/ReciprocalRankFusion.java b/hybrid-search/src/main/java/com/ankurm/hybridsearch/ReciprocalRankFusion.java new file mode 100644 index 0000000..a191f18 --- /dev/null +++ b/hybrid-search/src/main/java/com/ankurm/hybridsearch/ReciprocalRankFusion.java @@ -0,0 +1,43 @@ +package com.ankurm.hybridsearch; + +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +/** + * Reciprocal Rank Fusion: every ranking gives each item {@code weight / (k + rank)} (rank starts at 1) and the + * item's fused score is the sum over the rankings it appears in. Only ranks are used, never the scores + * the rankings were built from, which is why it can join a cosine similarity and a text rank that live on + * different scales without any normalisation. {@code k} (60 in the original paper) flattens the curve: the + * larger it is, the less the very top places matter. + */ +public final class ReciprocalRankFusion { + + public record Fused(String id, double score) { + } + + private ReciprocalRankFusion() { + } + + public static List fuse(List> rankings, int k) { + double[] ones = new double[rankings.size()]; + java.util.Arrays.fill(ones, 1.0); + return fuse(rankings, ones, k); + } + + public static List fuse(List> rankings, double[] weights, int k) { + Map scores = new LinkedHashMap<>(); + for (int r = 0; r < rankings.size(); r++) { + List ranking = rankings.get(r); + for (int i = 0; i < ranking.size(); i++) { + scores.merge(ranking.get(i), weights[r] / (k + i + 1), Double::sum); + } + } + List out = new ArrayList<>(); + scores.forEach((id, score) -> out.add(new Fused(id, score))); + // equal scores keep the order in which the ids first appeared, so the result is deterministic + out.sort((a, b) -> Double.compare(b.score(), a.score())); + return out; + } +} diff --git a/hybrid-search/src/test/java/com/ankurm/hybridsearch/KnowledgeBaseTest.java b/hybrid-search/src/test/java/com/ankurm/hybridsearch/KnowledgeBaseTest.java new file mode 100644 index 0000000..2df45f4 --- /dev/null +++ b/hybrid-search/src/test/java/com/ankurm/hybridsearch/KnowledgeBaseTest.java @@ -0,0 +1,24 @@ +package com.ankurm.hybridsearch; + +import static org.assertj.core.api.Assertions.assertThat; + +import java.util.HashSet; +import java.util.Set; +import java.util.stream.Collectors; + +import org.junit.jupiter.api.Test; + +class KnowledgeBaseTest { + + @Test + void sizesAndLabels() { + assertThat(KnowledgeBase.PASSAGES).hasSize(94); + assertThat(KnowledgeBase.QUERIES).hasSize(60); + Set pids = new HashSet<>(); + KnowledgeBase.PASSAGES.forEach(p -> assertThat(pids.add(p.pid())).as("unique pid %s", p.pid()).isTrue()); + KnowledgeBase.QUERIES.forEach(q -> assertThat(pids).as("answer of: %s", q.text()).contains(q.pid())); + assertThat(KnowledgeBase.QUERIES.stream().collect(Collectors.groupingBy(KnowledgeBase.Query::category, Collectors.counting()))) + .containsEntry(KnowledgeBase.EXACT_CODE, 12L).containsEntry(KnowledgeBase.EXACT_ID, 12L) + .containsEntry(KnowledgeBase.PARAPHRASE, 24L).containsEntry(KnowledgeBase.MIXED, 12L); + } +} diff --git a/hybrid-search/src/test/java/com/ankurm/hybridsearch/QualityTest.java b/hybrid-search/src/test/java/com/ankurm/hybridsearch/QualityTest.java new file mode 100644 index 0000000..d94e836 --- /dev/null +++ b/hybrid-search/src/test/java/com/ankurm/hybridsearch/QualityTest.java @@ -0,0 +1,135 @@ +package com.ankurm.hybridsearch; + +import java.util.ArrayList; +import java.util.List; +import java.util.function.Function; +import java.util.stream.Collectors; + +import com.ankurm.hybridsearch.support.Fixture; +import com.ankurm.hybridsearch.support.Metrics; +import com.ankurm.hybridsearch.support.Transcript; +import org.junit.jupiter.api.Test; + +/** Six rankers on 60 labelled queries, split by the kind of query, plus what k and the candidate count change. */ +class QualityTest { + + static final int CANDIDATES = 20; + + @Test + void rankersByCategory() { + HybridSearch hs = Fixture.search(); + Bm25 bm25 = Fixture.bm25(); + + Function> vector = q -> hs.vector(q, 10).stream().map(HybridSearch.Hit::pid).toList(); + Function> tsvector = q -> hs.keyword(q, 10).stream().map(HybridSearch.Hit::pid).toList(); + Function> bm25Only = q -> bm25.search(q, 10); + Function> rrfTs = q -> hs.hybrid(q, 10, CANDIDATES, 60).stream().map(HybridSearch.Fused::pid).toList(); + Function> rrfBm25 = q -> ReciprocalRankFusion.fuse(List.of( + hs.vector(q, CANDIDATES).stream().map(HybridSearch.Hit::pid).toList(), bm25.search(q, CANDIDATES)), 60) + .stream().limit(10).map(ReciprocalRankFusion.Fused::id).toList(); + Function> rrfSql = q -> hs.hybridSql(q, 10, CANDIDATES, 60); + + String[] names = {"vector (pgvector cosine)", "keyword (tsvector, ts_rank_cd)", "keyword (BM25, in Java)", + "hybrid: vector + tsvector, RRF", "hybrid: vector + BM25, RRF", "hybrid in one SQL statement"}; + List>> rankers = List.of(vector, tsvector, bm25Only, rrfTs, rrfBm25, rrfSql); + + try (Transcript t = new Transcript("02-quality-by-query-type.txt", "Six rankers on 60 labelled queries over 94 passages (all-MiniLM-L6-v2, PostgreSQL 16, pgvector 0.8.7)")) { + t.line("fusion: RRF with k=60 over the top %d of each list. hit@1 and hit@5: the answer passage is first / in the first five. MRR@10: mean of 1/rank.", CANDIDATES); + for (String category : List.of(KnowledgeBase.EXACT_CODE, KnowledgeBase.EXACT_ID, KnowledgeBase.PARAPHRASE, KnowledgeBase.MIXED)) { + t.section(category + " (" + Metrics.category(category).size() + " queries)"); + t.line("%-34s %s", "ranker", "hit@1 hit@5 MRR@10"); + for (int i = 0; i < names.length; i++) { + t.line("%-34s %s", names[i], Metrics.score(Metrics.category(category), rankers.get(i)).cells()); + } + } + t.section("all 60 queries"); + t.line("%-34s %s", "ranker", "hit@1 hit@5 MRR@10"); + for (int i = 0; i < names.length; i++) { + t.line("%-34s %s", names[i], Metrics.score(KnowledgeBase.QUERIES, rankers.get(i)).cells()); + } + + t.section("the SQL version against the Java version"); + int same = 0; + for (KnowledgeBase.Query q : KnowledgeBase.QUERIES) { + if (rrfTs.apply(q.text()).subList(0, 5).equals(rrfSql.apply(q.text()).subList(0, 5))) { + same++; + } + } + t.line("queries whose top 5 are identical, in the same order, from both implementations: %d of %d", same, KnowledgeBase.QUERIES.size()); + + t.section("one query of each kind: the top 3 from each ranker"); + for (String category : List.of(KnowledgeBase.EXACT_CODE, KnowledgeBase.EXACT_ID, KnowledgeBase.PARAPHRASE, KnowledgeBase.MIXED)) { + KnowledgeBase.Query q = Metrics.category(category).get(0); + t.line("[%s] \"%s\" -> answer is %s", category, q.text(), q.pid()); + t.line(" vector : %s", top(vector.apply(q.text()))); + t.line(" tsvector : %s", top(tsvector.apply(q.text()))); + t.line(" BM25 : %s", top(bm25Only.apply(q.text()))); + t.line(" hybrid : %s", top(rrfTs.apply(q.text()))); + } + } + } + + @Test + void whatKAndTheCandidateCountChange() { + HybridSearch hs = Fixture.search(); + try (Transcript t = new Transcript("03-rrf-parameters.txt", "RRF parameters on the 60 queries: k, candidates per list, list weights")) { + t.section("k, with 20 candidates per list"); + t.line("%-6s %s", "k", "hit@1 hit@5 MRR@10"); + for (int k : new int[] {1, 10, 60, 200, 1000}) { + t.line("%-6d %s", k, Metrics.score(KnowledgeBase.QUERIES, + q -> hs.hybrid(q, 10, 20, k).stream().map(HybridSearch.Fused::pid).toList()).cells()); + } + t.section("candidates taken from each list, with k = 60"); + t.line("%-6s %s", "cands", "hit@1 hit@5 MRR@10"); + for (int c : new int[] {3, 5, 10, 20, 50}) { + t.line("%-6d %s", c, Metrics.score(KnowledgeBase.QUERIES, + q -> hs.hybrid(q, 10, c, 60).stream().map(HybridSearch.Fused::pid).toList()).cells()); + } + t.section("weights (vector : keyword), k = 60, 20 candidates"); + t.line("%-10s %s", "weights", "hit@1 hit@5 MRR@10"); + for (double[] w : new double[][] {{1, 0}, {1, 0.5}, {1, 1}, {0.5, 1}, {0, 1}}) { + t.line("%-10s %s", w[0] + ":" + w[1], Metrics.score(KnowledgeBase.QUERIES, q -> { + List v = hs.vector(q, 20).stream().map(HybridSearch.Hit::pid).toList(); + List kw = hs.keyword(q, 20).stream().map(HybridSearch.Hit::pid).toList(); + return ReciprocalRankFusion.fuse(List.of(v, kw), w, 60).stream().limit(10).map(ReciprocalRankFusion.Fused::id).toList(); + }).cells()); + } + } + } + + @Test + void doesTheBestKSurviveHeldOutQueries() { + HybridSearch hs = Fixture.search(); + List even = new ArrayList<>(); + List odd = new ArrayList<>(); + for (int i = 0; i < KnowledgeBase.QUERIES.size(); i++) { + (i % 2 == 0 ? even : odd).add(KnowledgeBase.QUERIES.get(i)); + } + try (Transcript t = new Transcript("04-held-out-k.txt", "Choosing k on half of the queries and scoring it on the other half")) { + t.line("candidates per list: 20. The 60 queries are split by position into two folds of 30; k comes from {1, 10, 60, 200}."); + for (int fold = 0; fold < 2; fold++) { + List train = fold == 0 ? even : odd; + List test = fold == 0 ? odd : even; + int best = -1; + double bestMrr = -1; + for (int k : new int[] {1, 10, 60, 200}) { + double mrr = Metrics.score(train, q -> hs.hybrid(q, 10, 20, k).stream().map(HybridSearch.Fused::pid).toList()).mrr(); + if (mrr > bestMrr) { + bestMrr = mrr; + best = k; + } + } + final int chosen = best; + t.section("fold " + (fold + 1) + ": train on " + train.size() + ", test on " + test.size()); + t.line("k chosen on the training half (best MRR@10 %.3f): %d", bestMrr, chosen); + t.line("held-out half, k=%-3d : %s", chosen, Metrics.score(test, q -> hs.hybrid(q, 10, 20, chosen).stream().map(HybridSearch.Fused::pid).toList()).cells()); + t.line("held-out half, k=60 : %s", Metrics.score(test, q -> hs.hybrid(q, 10, 20, 60).stream().map(HybridSearch.Fused::pid).toList()).cells()); + t.line("held-out half, vector only: %s", Metrics.score(test, q -> hs.vector(q, 10).stream().map(HybridSearch.Hit::pid).toList()).cells()); + } + } + } + + private static String top(List ranking) { + return ranking.stream().limit(3).collect(Collectors.joining(", ")); + } +} diff --git a/hybrid-search/src/test/java/com/ankurm/hybridsearch/RetrieverTest.java b/hybrid-search/src/test/java/com/ankurm/hybridsearch/RetrieverTest.java new file mode 100644 index 0000000..46a0849 --- /dev/null +++ b/hybrid-search/src/test/java/com/ankurm/hybridsearch/RetrieverTest.java @@ -0,0 +1,64 @@ +package com.ankurm.hybridsearch; + +import static org.assertj.core.api.Assertions.assertThat; + +import java.util.List; + +import com.ankurm.hybridsearch.support.Fixture; +import com.ankurm.hybridsearch.support.Metrics; +import com.ankurm.hybridsearch.support.RecordingChatModel; +import com.ankurm.hybridsearch.support.Transcript; +import org.junit.jupiter.api.Test; +import org.springframework.ai.chat.client.ChatClient; +import org.springframework.ai.document.Document; +import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor; +import org.springframework.ai.rag.retrieval.search.DocumentRetriever; +import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever; +import org.springframework.ai.vectorstore.pgvector.PgVectorStore; + +/** The fused ranking behind Spring AI's RetrievalAugmentationAdvisor, and what reaches the prompt. */ +class RetrieverTest { + + @Test + void hybridRetrieverInTheAdvisor() throws Exception { + HybridSearch hs = Fixture.search(); + DocumentRetriever hybrid = new HybridDocumentRetriever(hs, 3, 20, 60); + PgVectorStore store = PgVectorStore.builder(Fixture.jdbc(), MiniLm.get()).vectorTableName(Fixture.TABLE).dimensions(384) + .initializeSchema(false).build(); + DocumentRetriever vectorOnly = VectorStoreDocumentRetriever.builder().vectorStore(store).topK(3).build(); + + try (Transcript t = new Transcript("05-retriever-and-advisor.txt", "HybridDocumentRetriever inside RetrievalAugmentationAdvisor (top 3 passages per question)")) { + String question = "Why am I getting E-4116?"; + t.section("what the retriever returns for \"" + question + "\""); + for (Document d : hybrid.retrieve(new org.springframework.ai.rag.Query(question))) { + t.line("score %.5f vector rank %-2s keyword rank %-2s %s", d.getScore(), + zero(d.getMetadata().get("vector_rank")), zero(d.getMetadata().get("keyword_rank")), d.getText().substring(0, 60)); + } + + RecordingChatModel chat = new RecordingChatModel(); + ChatClient client = ChatClient.builder(chat) + .defaultAdvisors(RetrievalAugmentationAdvisor.builder().documentRetriever(hybrid).build()).build(); + client.prompt().user(question).call().content(); + String prompt = chat.prompts().getFirst().getContents(); + t.section("the prompt the model received"); + t.line(prompt); + assertThat(prompt).contains("Error E-4116"); + + t.section("does the answer passage reach the prompt? the 12 exact-code questions, top 3 passages"); + int hybridHits = 0, vectorHits = 0; + for (KnowledgeBase.Query q : Metrics.category(KnowledgeBase.EXACT_CODE)) { + String answer = KnowledgeBase.PASSAGES.stream().filter(p -> p.pid().equals(q.pid())).findFirst().orElseThrow().text(); + boolean h = hybrid.retrieve(new org.springframework.ai.rag.Query(q.text())).stream().anyMatch(d -> d.getText().equals(answer)); + boolean v = vectorOnly.retrieve(new org.springframework.ai.rag.Query(q.text())).stream().anyMatch(d -> d.getText().equals(answer)); + hybridHits += h ? 1 : 0; + vectorHits += v ? 1 : 0; + t.line("%-7s vector-only: %-5s hybrid: %-5s %s", q.pid(), v, h, q.text()); + } + t.line("answer passage in the top 3: vector-only %d of 12, hybrid %d of 12", vectorHits, hybridHits); + } + } + + private static Object zero(Object rank) { + return ((Integer) rank) == 0 ? "-" : rank; + } +} diff --git a/hybrid-search/src/test/java/com/ankurm/hybridsearch/RrfTest.java b/hybrid-search/src/test/java/com/ankurm/hybridsearch/RrfTest.java new file mode 100644 index 0000000..69bc1ec --- /dev/null +++ b/hybrid-search/src/test/java/com/ankurm/hybridsearch/RrfTest.java @@ -0,0 +1,42 @@ +package com.ankurm.hybridsearch; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.within; + +import java.util.List; + +import com.ankurm.hybridsearch.support.Transcript; +import org.junit.jupiter.api.Test; + +/** Reciprocal Rank Fusion worked by hand and then by code, with the effect of k. */ +class RrfTest { + + @Test + void workedExample() { + List vector = List.of("A", "B", "C", "D"); + List keyword = List.of("C", "A", "E", "B"); + + try (Transcript t = new Transcript("01-rrf-by-hand.txt", "Reciprocal Rank Fusion on two short rankings")) { + t.line("vector ranking : %s", vector); + t.line("keyword ranking: %s", keyword); + for (int k : new int[] {1, 60, 1000}) { + t.section("k = " + k); + for (ReciprocalRankFusion.Fused f : ReciprocalRankFusion.fuse(List.of(vector, keyword), k)) { + t.line("%s score %.5f (vector rank %d, keyword rank %d)", f.id(), f.score(), vector.indexOf(f.id()) + 1, keyword.indexOf(f.id()) + 1); + } + } + t.section("the arithmetic for A at k = 60"); + t.line("rank 1 in the vector list: 1/(60+1) = %.5f", 1.0 / 61); + t.line("rank 2 in the keyword list: 1/(60+2) = %.5f", 1.0 / 62); + t.line("sum = %.5f", 1.0 / 61 + 1.0 / 62); + t.section("an item in one list only"); + t.line("E is rank 3 in the keyword list and absent from the vector list: 1/(60+3) = %.5f", 1.0 / 63); + t.line("D is rank 4 in the vector list and absent from the keyword list: 1/(60+4) = %.5f", 1.0 / 64); + } + + var fused = ReciprocalRankFusion.fuse(List.of(vector, keyword), 60); + assertThat(fused.get(0).id()).isEqualTo("A"); + assertThat(fused.get(0).score()).isCloseTo(1.0 / 61 + 1.0 / 62, within(1e-12)); + assertThat(ReciprocalRankFusion.fuse(List.of(vector, keyword), 1).get(0).id()).isEqualTo("A"); + } +} diff --git a/hybrid-search/src/test/java/com/ankurm/hybridsearch/ScaleTest.java b/hybrid-search/src/test/java/com/ankurm/hybridsearch/ScaleTest.java new file mode 100644 index 0000000..9918121 --- /dev/null +++ b/hybrid-search/src/test/java/com/ankurm/hybridsearch/ScaleTest.java @@ -0,0 +1,174 @@ +package com.ankurm.hybridsearch; + +import java.util.ArrayList; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.function.Consumer; + +import com.ankurm.hybridsearch.support.Fixture; +import com.ankurm.hybridsearch.support.Metrics; +import com.ankurm.hybridsearch.support.Transcript; +import org.junit.jupiter.api.Test; +import org.springframework.jdbc.core.JdbcTemplate; + +/** + * The same four queries at 94 rows and at 50,094: query plans, index sizes and latency. The 50,000 extra rows are + * generated in SQL from the knowledge base's own vocabulary (14 random words each) with random 384-dimension vectors, + * so keyword posting lists are long and the vector index has work to do. They say nothing about relevance. + */ +class ScaleTest { + + static final int FILLER = 50_000; + + @Test + void scale() throws Exception { + JdbcTemplate jdbc = Fixture.jdbc(); + Fixture.search(); // makes sure hybrid_docs exists + + Set vocabulary = new LinkedHashSet<>(); + KnowledgeBase.PASSAGES.forEach(p -> vocabulary.addAll(Bm25.tokens(p.text()))); + String array = vocabulary.stream().filter(w -> w.matches("[a-z0-9]+")).map(w -> "'" + w + "'") + .collect(java.util.stream.Collectors.joining(",", "array[", "]::text[]")); + + jdbc.execute("drop table if exists hybrid_scale"); + jdbc.execute("create table hybrid_scale (id uuid primary key, content text, metadata json, embedding vector(384), " + + "tsv tsvector generated always as (to_tsvector('english', content)) stored)"); + jdbc.execute("insert into hybrid_scale (id, content, metadata, embedding) select id, content, metadata, embedding from hybrid_docs"); + long t0 = System.nanoTime(); + jdbc.update("with v as (select " + array + " as w) insert into hybrid_scale (id, content, metadata, embedding) " + + "select gen_random_uuid(), c.t, '{\"pid\":\"filler\",\"kind\":\"filler\"}'::json, e.v " + + "from generate_series(1, " + FILLER + ") s(i) cross join v " + + "cross join lateral (select string_agg(v.w[1 + floor(random() * array_length(v.w, 1))::int + 0 * g], ' ') as t " + + "from generate_series(1 + 0 * s.i, 14) g) c " + + "cross join lateral (select array_agg(random() - 0.5 + 0 * g)::vector as v from generate_series(1 + 0 * s.i, 384) g) e"); + long loadMs = (System.nanoTime() - t0) / 1_000_000; + jdbc.execute("set maintenance_work_mem = '512MB'"); + // the PgVectorStore index: HNSW with cosine distance + HybridSearchIndexes small = new HybridSearchIndexes(jdbc); + t0 = System.nanoTime(); + jdbc.execute("create index hybrid_scale_embedding_idx on hybrid_scale using hnsw (embedding vector_cosine_ops)"); + long hnswMs = (System.nanoTime() - t0) / 1_000_000; + t0 = System.nanoTime(); + jdbc.execute("create index hybrid_scale_tsv_idx on hybrid_scale using gin (tsv)"); + long ginMs = (System.nanoTime() - t0) / 1_000_000; + jdbc.execute("analyze hybrid_scale"); + + HybridSearch tiny = Fixture.search(); + HybridSearch big = new HybridSearch(jdbc, MiniLm.get(), "hybrid_scale"); + + try (Transcript t = new Transcript("06-scale-plans-and-latency.txt", "94 rows against 50,094 rows: plans, index sizes, latency (PostgreSQL 16, pgvector 0.8.7, 2 vCPU)")) { + t.line("filler rows inserted: %d in %d ms. HNSW index built in %d ms, GIN index in %d ms.", FILLER, loadMs, hnswMs, ginMs); + t.line("rows in hybrid_scale: %d", jdbc.queryForObject("select count(*) from hybrid_scale", Integer.class)); + t.line("heap %s, HNSW index %s, GIN index %s", + jdbc.queryForObject("select pg_size_pretty(pg_table_size('hybrid_scale'))", String.class), + jdbc.queryForObject("select pg_size_pretty(pg_relation_size('hybrid_scale_embedding_idx'))", String.class), + jdbc.queryForObject("select pg_size_pretty(pg_relation_size('hybrid_scale_tsv_idx'))", String.class)); + + String sample = "Why am I getting E-4116?"; + String lit = HybridSearch.literal(MiniLm.get().embed(sample)); + t.section("plan for the vector query at 50,094 rows (scan nodes only)"); + plan(jdbc, t, "explain (analyze, costs off, timing off, summary off) select metadata->>'pid' from hybrid_scale order by embedding <=> '" + lit + "'::vector limit 20"); + t.section("plan for the keyword query at 50,094 rows (scan nodes only)"); + plan(jdbc, t, "explain (analyze, costs off, timing off, summary off) with q as (select replace(plainto_tsquery('english', '" + sample + + "')::text, ' & ', ' | ')::tsquery as query) select d.metadata->>'pid', ts_rank_cd(d.tsv, q.query) as s from hybrid_scale d, q " + + "where d.tsv @@ q.query order by s desc limit 20"); + t.section("plan for the same keyword query at 94 rows"); + plan(jdbc, t, "explain (analyze, costs off, timing off, summary off) with q as (select replace(plainto_tsquery('english', '" + sample + + "')::text, ' & ', ' | ')::tsquery as query) select d.metadata->>'pid', ts_rank_cd(d.tsv, q.query) as s from hybrid_docs d, q " + + "where d.tsv @@ q.query order by s desc limit 20"); + + t.section("rows the keyword query has to rank, by query kind (50,094 rows)"); + for (String category : List.of(KnowledgeBase.EXACT_CODE, KnowledgeBase.EXACT_ID, KnowledgeBase.PARAPHRASE, KnowledgeBase.MIXED)) { + double avg = Metrics.category(category).stream().mapToInt(q -> jdbc.queryForObject( + "select count(*) from hybrid_scale where tsv @@ replace(plainto_tsquery('english', ?)::text, ' & ', ' | ')::tsquery", + Integer.class, q.text())).average().orElse(0); + t.line("%-24s average rows matching the OR query: %8.0f", category, avg); + } + + List texts = KnowledgeBase.QUERIES.stream().map(KnowledgeBase.Query::text).toList(); + List lits = texts.stream().map(x -> HybridSearch.literal(MiniLm.get().embed(x))).toList(); + t.section("latency in milliseconds over 60 queries x 3 runs after 2 warm-up passes (query embedding excluded unless stated)"); + t.line("%-34s %9s %9s %9s %9s", "", "94 p50", "94 p95", "50k p50", "50k p95"); + t.line("%-34s %s", "embedding the question (MiniLM)", embedTiming(texts)); + row(t, "vector, top 20", texts, lits, (i) -> tiny.vectorByLiteral(lits.get(i), 20), (i) -> big.vectorByLiteral(lits.get(i), 20)); + row(t, "keyword, top 20", texts, lits, (i) -> tiny.keyword(texts.get(i), 20), (i) -> big.keyword(texts.get(i), 20)); + row(t, "hybrid in Java (2 queries + fuse)", texts, lits, + (i) -> HybridSearch.fuse(tiny.vectorByLiteral(lits.get(i), 20), tiny.keyword(texts.get(i), 20), 10, 60), + (i) -> HybridSearch.fuse(big.vectorByLiteral(lits.get(i), 20), big.keyword(texts.get(i), 20), 10, 60)); + row(t, "hybrid in one SQL statement", texts, lits, (i) -> tiny.hybridSqlByLiteral(texts.get(i), lits.get(i), 10, 20, 60), + (i) -> big.hybridSqlByLiteral(texts.get(i), lits.get(i), 10, 20, 60)); + + t.section("keyword query latency at 50,094 rows by query kind (p50 / p95 ms)"); + for (String category : List.of(KnowledgeBase.EXACT_CODE, KnowledgeBase.EXACT_ID, KnowledgeBase.PARAPHRASE, KnowledgeBase.MIXED)) { + List qs = Metrics.category(category).stream().map(KnowledgeBase.Query::text).toList(); + List samples = new ArrayList<>(); + for (int pass = 0; pass < 5; pass++) { + for (String q : qs) { + long s = System.nanoTime(); + big.keyword(q, 20); + if (pass >= 2) { + samples.add((System.nanoTime() - s) / 1e6); + } + } + } + t.line("%-24s %7.1f / %7.1f", category, percentile(samples, 50), percentile(samples, 95)); + } + } + jdbc.execute("drop table if exists hybrid_scale"); + } + + /** Placeholder holder so the maintenance setting above is visibly a separate step. */ + private record HybridSearchIndexes(JdbcTemplate jdbc) { + } + + private static void plan(JdbcTemplate jdbc, Transcript t, String sql) { + for (Map r : jdbc.queryForList(sql)) { + String line = r.values().iterator().next().toString(); + if (line.contains("Scan") || line.contains("Limit") || line.contains("Sort") || line.contains("Rows Removed")) { + t.line("%s", line.length() > 118 ? line.substring(0, 118) + "..." : line); + } + } + } + + private static String embedTiming(List texts) { + List samples = new ArrayList<>(); + for (int pass = 0; pass < 3; pass++) { + for (String q : texts) { + long s = System.nanoTime(); + MiniLm.get().embed(q); + if (pass >= 1) { + samples.add((System.nanoTime() - s) / 1e6); + } + } + } + return String.format("%9.1f %9.1f (same for both table sizes)", percentile(samples, 50), percentile(samples, 95)); + } + + private static void row(Transcript t, String name, List texts, List lits, java.util.function.IntConsumer small, + java.util.function.IntConsumer large) { + List a = time(texts.size(), small); + List b = time(texts.size(), large); + t.line("%-34s %9.1f %9.1f %9.1f %9.1f", name, percentile(a, 50), percentile(a, 95), percentile(b, 50), percentile(b, 95)); + } + + private static List time(int n, java.util.function.IntConsumer call) { + List samples = new ArrayList<>(); + for (int pass = 0; pass < 5; pass++) { + for (int i = 0; i < n; i++) { + long s = System.nanoTime(); + call.accept(i); + if (pass >= 2) { + samples.add((System.nanoTime() - s) / 1e6); + } + } + } + return samples; + } + + private static double percentile(List samples, int p) { + List sorted = samples.stream().sorted().toList(); + return sorted.get(Math.min(sorted.size() - 1, (int) Math.ceil(p / 100.0 * sorted.size()) - 1)); + } +} diff --git a/hybrid-search/src/test/java/com/ankurm/hybridsearch/support/Fixture.java b/hybrid-search/src/test/java/com/ankurm/hybridsearch/support/Fixture.java new file mode 100644 index 0000000..c48d4cf --- /dev/null +++ b/hybrid-search/src/test/java/com/ankurm/hybridsearch/support/Fixture.java @@ -0,0 +1,43 @@ +package com.ankurm.hybridsearch.support; + +import com.ankurm.hybridsearch.Bm25; +import com.ankurm.hybridsearch.HybridSearch; +import com.ankurm.hybridsearch.KnowledgeBase; +import com.ankurm.hybridsearch.MiniLm; +import org.springframework.jdbc.core.JdbcTemplate; +import org.springframework.jdbc.datasource.DriverManagerDataSource; + +/** One table, loaded once per test run: 94 passages embedded by all-MiniLM-L6-v2 and indexed for keyword search. */ +public final class Fixture { + + public static final String URL = System.getenv().getOrDefault("HYBRID_DB_URL", "jdbc:postgresql://localhost:5432/ragdemo"); + public static final String TABLE = "hybrid_docs"; + + private static HybridSearch search; + private static Bm25 bm25; + + private Fixture() { + } + + public static JdbcTemplate jdbc() { + return new JdbcTemplate(new DriverManagerDataSource(URL, "demo", "demo")); + } + + public static synchronized HybridSearch search() { + if (search == null) { + try { + search = HybridSearch.load(jdbc(), MiniLm.get(), TABLE, KnowledgeBase.PASSAGES); + } catch (Exception e) { + throw new IllegalStateException(e); + } + } + return search; + } + + public static synchronized Bm25 bm25() { + if (bm25 == null) { + bm25 = new Bm25(KnowledgeBase.PASSAGES, 1.2, 0.75); + } + return bm25; + } +} diff --git a/hybrid-search/src/test/java/com/ankurm/hybridsearch/support/Metrics.java b/hybrid-search/src/test/java/com/ankurm/hybridsearch/support/Metrics.java new file mode 100644 index 0000000..d9eee17 --- /dev/null +++ b/hybrid-search/src/test/java/com/ankurm/hybridsearch/support/Metrics.java @@ -0,0 +1,43 @@ +package com.ankurm.hybridsearch.support; + +import java.util.List; +import java.util.function.Function; + +import com.ankurm.hybridsearch.KnowledgeBase; + +/** hit@1, hit@5 and MRR@10 of a ranker over a set of labelled queries. */ +public final class Metrics { + + public record Score(int n, int hit1, int hit5, double mrr) { + + public String cells() { + return String.format("%2d/%-2d %2d/%-2d %.3f", hit1, n, hit5, n, mrr); + } + } + + private Metrics() { + } + + public static Score score(List queries, Function> ranker) { + int h1 = 0, h5 = 0; + double mrr = 0; + for (KnowledgeBase.Query q : queries) { + List ranking = ranker.apply(q.text()); + int rank = ranking.subList(0, Math.min(10, ranking.size())).indexOf(q.pid()) + 1; + if (rank == 1) { + h1++; + } + if (rank >= 1 && rank <= 5) { + h5++; + } + if (rank >= 1) { + mrr += 1.0 / rank; + } + } + return new Score(queries.size(), h1, h5, mrr / queries.size()); + } + + public static List category(String name) { + return KnowledgeBase.QUERIES.stream().filter(q -> q.category().equals(name)).toList(); + } +} diff --git a/hybrid-search/src/test/java/com/ankurm/hybridsearch/support/RecordingChatModel.java b/hybrid-search/src/test/java/com/ankurm/hybridsearch/support/RecordingChatModel.java new file mode 100644 index 0000000..27b27e5 --- /dev/null +++ b/hybrid-search/src/test/java/com/ankurm/hybridsearch/support/RecordingChatModel.java @@ -0,0 +1,26 @@ +package com.ankurm.hybridsearch.support; + +import java.util.List; +import java.util.concurrent.CopyOnWriteArrayList; + +import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.model.ChatModel; +import org.springframework.ai.chat.model.ChatResponse; +import org.springframework.ai.chat.model.Generation; +import org.springframework.ai.chat.prompt.Prompt; + +/** Records the prompts it is sent and answers "ok". What Spring AI sends is the part this module is responsible for. */ +public class RecordingChatModel implements ChatModel { + + private final List prompts = new CopyOnWriteArrayList<>(); + + public List prompts() { + return prompts; + } + + @Override + public ChatResponse call(Prompt prompt) { + prompts.add(prompt); + return new ChatResponse(List.of(new Generation(new AssistantMessage("ok")))); + } +} diff --git a/hybrid-search/src/test/java/com/ankurm/hybridsearch/support/Transcript.java b/hybrid-search/src/test/java/com/ankurm/hybridsearch/support/Transcript.java new file mode 100644 index 0000000..665bdea --- /dev/null +++ b/hybrid-search/src/test/java/com/ankurm/hybridsearch/support/Transcript.java @@ -0,0 +1,52 @@ +package com.ankurm.hybridsearch.support; + +import java.io.IOException; +import java.io.PrintWriter; +import java.io.StringWriter; +import java.nio.file.Files; +import java.nio.file.Path; + +/** + * Writes a numbered transcript under {@code output/} and echoes it to the console. Every console block + * quoted in the article comes out of one of these files verbatim. + */ +public final class Transcript implements AutoCloseable { + + private final Path path; + private final StringWriter buffer = new StringWriter(); + private final PrintWriter out = new PrintWriter(buffer); + + public Transcript(String fileName, String title) { + this.path = Path.of("output", fileName); + out.println("# " + title); + out.println(); + } + + public Transcript line(String format, Object... args) { + out.println(args.length == 0 ? format : String.format(format, args)); + return this; + } + + public Transcript blank() { + out.println(); + return this; + } + + public Transcript section(String heading) { + out.println(); + out.println("--- " + heading + " ---"); + return this; + } + + @Override + public void close() { + out.flush(); + try { + Files.createDirectories(path.getParent()); + Files.writeString(path, buffer.toString()); + } catch (IOException e) { + throw new IllegalStateException("could not write " + path, e); + } + System.out.print(buffer); + } +}