Add graphrag-neo4j module: structured-output entity and relation extraction into a typed Neo4j graph with a vector index, multi-hop retrieval by span linking and variable-length Cypher paths, compared with plain vector RAG over the same 35 chunks and 9 questions (complete evidence 7/9 vs 5/9 at k=8); Testcontainers on neo4j:2026.10.1 and 5.26. The extraction step is a rule-based stand-in, not a language model.
Co-Authored-By: Claude Sonnet 5.5 <[email protected]> Claude-Session: https://claude.ai/code/session_01CifR2BE7PBvRy5hT8TS9oz
This commit is contained in:
Claude Sonnet 5.5
@@ -29,5 +29,6 @@ Runnable companion code for the Spring AI articles on [ankurm.com](https://ankur
|
||||
| [`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/) |
|
||||
| [`rerank/`](rerank) | A bi-encoder (all-MiniLM-L6-v2) retrieves candidates, a cross-encoder (ms-marco-MiniLM-L-6-v2) rescores them, and a `DocumentPostProcessor` keeps the top 5 for `RetrievalAugmentationAdvisor`. Both models run in-process through ONNX Runtime 1.31.0 with DJL tokenizers 0.39.0 (no database, no API key). 60 questions over 94 passages: hit@1 41 to 49 and MRR 0.78 to 0.87, almost all of it on exact error codes, with paraphrases slightly worse; pool size, batch size, threads and int8 against latency on 2 vCPUs; and five traps reproduced (missing `token_type_ids`, swapped pair order, logits are not probabilities, truncation, a pool that is too small). Spring Boot 4.1.2, Spring AI 2.0.2, Java 25. | [Reranking RAG Results in Java: Cross-Encoder Reranker with ONNX Runtime Before the LLM](https://ankurm.com/reranking-rag-results-java-cross-encoder-onnx-runtime-spring-ai/) |
|
||||
| [`audio/`](audio) | A meeting-minutes generator: `OpenAiAudioTranscriptionModel` -> `ChatClient.entity(Minutes.class)` -> `OpenAiAudioSpeechModel`. Splitting a 30 MB recording at a silence under the 25 MB upload limit (and what a blind split does), SDK retries that re-upload the file, the 4096-character speech limit, three replies that fail to parse, and a cost model with formulas. The audio endpoints are a local mock and the chat model a stub; no real model was called. Spring Boot 4.1.2, Spring AI 2.0.2, Java 25. | [Speech-to-Text and Text-to-Speech with Spring AI: Build a Meeting-Minutes Generator](https://ankurm.com/spring-ai-speech-to-text-text-to-speech-meeting-minutes/) |
|
||||
| [`graphrag-neo4j/`](graphrag-neo4j) | Entities and relations extracted with `ChatClient.entity(...)` into a typed Neo4j graph with a vector index on entity names (Testcontainers, `neo4j:2026.10.1-community`; also run on 5.26), then multi-hop retrieval by span linking, variable-length Cypher paths and path ranking, against plain vector RAG over the same 35 chunks and 9 questions: complete evidence 7/9 for the graph, 5/9 for vector k=8, 2/9 for k=4, with the traps reproduced (duplicate entity under an alias, an embedding merge that joins the wrong incidents, a reply with a sentence before the JSON fence, a hostile relation type, hop explosion). **The extraction step is a rule-based stand-in, not a language model**; Neo4j, embeddings and Cypher are real. Spring Boot 4.1.2, Spring AI 2.0.2, Java 25. | [GraphRAG in Java with Neo4j and Spring AI](https://ankurm.com/graphrag-java-neo4j-spring-ai-knowledge-graph/) |
|
||||
|
||||
Upgrading from Spring AI 1.x: [migration guide](https://ankurm.com/spring-ai-1-to-2-migration-guide/).
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
target/
|
||||
@@ -0,0 +1,58 @@
|
||||
# graphrag-neo4j
|
||||
|
||||
GraphRAG on Spring AI 2.0 and Neo4j, next to plain vector RAG over the same documents. Companion code for [GraphRAG in Java with Neo4j and Spring AI: Turning Documents into a Knowledge Graph](https://ankurm.com/graphrag-java-neo4j-spring-ai-knowledge-graph/).
|
||||
|
||||
**Read this first: no language model is called anywhere in this module.** The extraction step runs through the real `ChatClient.entity(...)` against `RuleBasedExtractorModel`, a stand-in that finds eight fixed sentence patterns with regular expressions and answers with the JSON a good model would. Everything downstream of it (Neo4j, the vector indexes, MiniLM embeddings, the Cypher, the retrieval numbers) is real. Nothing here says how well a real model extracts from real prose. Answer generation is not exercised at all; the tests print the prompt that would be sent.
|
||||
|
||||
## Versions
|
||||
|
||||
| | |
|
||||
|---|---|
|
||||
| Spring Boot | 4.1.2 |
|
||||
| Spring AI | 2.0.2 (`spring-ai-client-chat`, `spring-ai-neo4j-store`, `spring-ai-transformers`) |
|
||||
| Neo4j server | `neo4j:2026.10.1-community` (Docker); also run on `neo4j:5-community`, see `output/03b-graph-build-neo4j-5.txt` |
|
||||
| Neo4j Java driver | 6.1.0 (managed by Boot) |
|
||||
| Testcontainers | 2.0.5 (`testcontainers-neo4j`) |
|
||||
| Embeddings | all-MiniLM-L6-v2 (384 dims) in-process through ONNX Runtime, as in `rag-ingestion` |
|
||||
| Java | 25 (Temurin 25.0.5+8) |
|
||||
|
||||
## Run it
|
||||
|
||||
```bash
|
||||
./scripts/run-all.sh # regenerates output/ (needs Docker; first run downloads MiniLM and the Neo4j image)
|
||||
NEO4J_IMAGE=neo4j:5-community ./scripts/run-all.sh
|
||||
```
|
||||
|
||||
A full run takes a few minutes: the Neo4j container needs about a minute to start, and the sweep test runs a few hundred retrievals.
|
||||
|
||||
## Layout
|
||||
|
||||
| File | What it does |
|
||||
|---|---|
|
||||
| `Extraction`, `EntityType`, `RelationType` | The record the model must fill; the enums are the closed vocabulary and the allow-list for relationship types put into Cypher |
|
||||
| `Extractor` | `ChatClient.entity(Extraction.class)` per chunk, then drops relations whose endpoints were never listed |
|
||||
| `stub/RuleBasedExtractorModel` | The stand-in for the model, with failure modes (fenced JSON, truncated JSON, unknown and hostile relation types, dangling relation) |
|
||||
| `EntityResolver` | Merges names: exact key, alias table, optional embedding-similarity threshold |
|
||||
| `GraphStore` | Every Cypher statement: schema, vector index, writes, variable-length paths |
|
||||
| `GraphRetriever` | Span linking through the entity vector index, path walk, path ranking |
|
||||
| `Pipeline` | Chunks into `Neo4jVectorStore`, extraction, resolution, graph write; the plain vector search used as the baseline |
|
||||
| `Answerer` | Builds the final prompt from selected paths and source passages |
|
||||
| `Question` | The 9 evaluation questions and, for each, the documents that together hold the facts |
|
||||
| `src/main/resources/corpus/` | 35 short invented documents: 21 with facts, 14 distractors that share the vocabulary |
|
||||
|
||||
## Captured output
|
||||
|
||||
| File | Shows |
|
||||
|---|---|
|
||||
| `output/01-extraction.txt` | The exact user message (with the JSON schema Spring AI appends), the parsed record, call and size totals |
|
||||
| `output/02-extraction-failures.txt` | What each bad reply does |
|
||||
| `output/03-graph-build.txt` | Server version, the two vector indexes, the duplicate-entity problem and what each merge threshold does |
|
||||
| `output/03b-graph-build-neo4j-5.txt` | The same test against `neo4j:5-community` |
|
||||
| `output/04-vector-vs-graph.txt` | Plain vector RAG against GraphRAG, per question |
|
||||
| `output/05-graph-sweeps.txt` | Hops, paths kept and link threshold against the result; path counts by hop |
|
||||
| `output/06-prompt.txt` | The prompt GraphRAG builds for one question |
|
||||
| `output/07-fixed-shape-cypher.txt` | Hand-written Cypher for questions whose shape is known |
|
||||
|
||||
## What the numbers are and are not
|
||||
|
||||
The corpus is small and invented, the nine questions were written by hand after the corpus, and the default retriever settings (3 hops, 6 paths, link threshold 0.80) were fixed before the first result. The sweep in `05` was run afterwards, on the same questions, so a best cell in that table is tuned on the test set. "Complete" means every supporting document is in the context; it does not mean a model answered correctly.
|
||||
@@ -0,0 +1,76 @@
|
||||
# Entity and relation extraction through ChatClient.entity() (STUB model, not an LLM)
|
||||
|
||||
|
||||
--- the user message the model receives for svc-checkout.txt ---
|
||||
Extract the entities and relations from the document.
|
||||
|
||||
DOCUMENT_START
|
||||
Checkout Service. The Checkout Service takes a customer's cart and turns it into an order. It depends on the Payments Gateway to take card payments and on the Inventory Service to reserve stock. The Checkout Service is owned by the Storefront Team. Its latency dashboard lives in the shared Grafana folder.
|
||||
DOCUMENT_END
|
||||
|
||||
Your response should be in JSON format.
|
||||
Do not include any explanations, only provide a RFC8259 compliant JSON response following this format without deviation.
|
||||
Do not include markdown code blocks in your response.
|
||||
Remove the ```json markdown from the output.
|
||||
Here is the JSON Schema instance your output must adhere to:
|
||||
```{
|
||||
"$schema" : "https://json-schema.org/draft/2020-12/schema",
|
||||
"type" : "object",
|
||||
"properties" : {
|
||||
"entities" : {
|
||||
"type" : "array",
|
||||
"items" : {
|
||||
"type" : "object",
|
||||
"properties" : {
|
||||
"name" : {
|
||||
"type" : "string"
|
||||
},
|
||||
"type" : {
|
||||
"type" : "string",
|
||||
"enum" : [ "SYSTEM", "TEAM", "PERSON", "VENDOR", "INCIDENT", "PLACE" ]
|
||||
}
|
||||
},
|
||||
"required" : [ "name", "type" ],
|
||||
"additionalProperties" : false
|
||||
}
|
||||
},
|
||||
"relations" : {
|
||||
"type" : "array",
|
||||
"items" : {
|
||||
"type" : "object",
|
||||
"properties" : {
|
||||
"source" : {
|
||||
"type" : "string"
|
||||
},
|
||||
"target" : {
|
||||
"type" : "string"
|
||||
},
|
||||
"type" : {
|
||||
"type" : "string",
|
||||
"enum" : [ "DEPENDS_ON", "OWNED_BY", "LED_BY", "HOSTED_ON", "AFFECTED", "CAUSED_BY", "REPORTS_TO", "LOCATED_IN" ]
|
||||
}
|
||||
},
|
||||
"required" : [ "source", "target", "type" ],
|
||||
"additionalProperties" : false
|
||||
}
|
||||
}
|
||||
},
|
||||
"required" : [ "entities", "relations" ],
|
||||
"additionalProperties" : false
|
||||
}```
|
||||
|
||||
--- the parsed Extraction record ---
|
||||
entity SYSTEM Checkout Service
|
||||
entity SYSTEM Payments Gateway
|
||||
entity SYSTEM Inventory Service
|
||||
entity TEAM Storefront Team
|
||||
relation Checkout Service -[DEPENDS_ON]-> Payments Gateway
|
||||
relation Checkout Service -[DEPENDS_ON]-> Inventory Service
|
||||
relation Checkout Service -[OWNED_BY]-> Storefront Team
|
||||
|
||||
--- all 35 chunks ---
|
||||
model calls : 35 (one per chunk)
|
||||
characters sent : 72437 (about 18109 tokens at 4 characters per token, an estimate)
|
||||
entity mentions (before merging): 54
|
||||
relations : 33
|
||||
chunks with no relation at all : 14
|
||||
@@ -0,0 +1,20 @@
|
||||
# What a bad model reply does to extraction (replies produced by the STUB's failure modes)
|
||||
|
||||
|
||||
--- JSON wrapped in a markdown fence (reply starts with the fence) ---
|
||||
parsed OK, 2 relations
|
||||
|
||||
--- the same fence, after one sentence of chat ---
|
||||
tools.jackson.core.exc.StreamReadException: Unrecognized token 'Here': was expecting (JSON String, Number, Array, Object or token 'null', 'true' or 'false')
|
||||
|
||||
--- JSON cut off two thirds of the way through ---
|
||||
tools.jackson.core.exc.UnexpectedEndOfInputException: Unexpected end-of-input in property name
|
||||
|
||||
--- a relation type outside the allowed list (IS_RESPONSIBLE_FOR) ---
|
||||
tools.jackson.databind.exc.InvalidFormatException: Cannot deserialize value of type `com.ankurm.graphrag.RelationType` from String "IS_RESPONSIBLE_FOR": not one of the values accepted for Enum class: [OWNED_BY, LOCATED_IN, AFFECTED, HOSTED_ON, LED_BY, DEPENDS_ON, REPORTS...
|
||||
|
||||
--- a relation type that is a Cypher fragment ---
|
||||
tools.jackson.databind.exc.InvalidFormatException: Cannot deserialize value of type `com.ankurm.graphrag.RelationType` from String "OWNED_BY]->(x) DETACH DELETE x //": not one of the values accepted for Enum class: [OWNED_BY, LOCATED_IN, AFFECTED, HOSTED_ON, LED_BY, DEPE...
|
||||
|
||||
--- a relation to an entity that was never listed ---
|
||||
relations kept 2, dropped 1
|
||||
@@ -0,0 +1,45 @@
|
||||
# Building the graph in Neo4j; names that are the same thing (STUB extraction)
|
||||
|
||||
image : neo4j:2026.10.1-community
|
||||
server : Neo4j Kernel 2026.10.1 community
|
||||
|
||||
--- vector indexes (one from Neo4jVectorStore for chunks, one of ours for entity names) ---
|
||||
chunk_embedding on (:Chunk).embedding ONLINE
|
||||
entity_name_embedding on (:Entity).embedding ONLINE
|
||||
|
||||
--- A. exact-name matching only (no alias table, similarity merge off) ---
|
||||
entities: 28
|
||||
Payments Gateway : -[DEPENDS_ON]-> Ledger Service; -[HOSTED_ON]-> Arcadia Cloud; -[OWNED_BY]-> Payments Team; <-[DEPENDS_ON]- Checkout Service
|
||||
Card Gateway : <-[AFFECTED]- INC-2060
|
||||
|
||||
--- same-type name pairs by cosine of their embeddings, top 8 of 72 ---
|
||||
0.806 INCIDENT INC-2057 <-> INC-2060
|
||||
0.797 INCIDENT INC-2041 <-> INC-2057
|
||||
0.758 INCIDENT INC-2041 <-> INC-2060
|
||||
0.681 SYSTEM Card Gateway <-> Payments Gateway
|
||||
0.540 SYSTEM Ledger Service <-> Inventory Service
|
||||
0.535 TEAM Storefront Team <-> Platform Team
|
||||
0.530 TEAM Core Data Team <-> Platform Team
|
||||
0.519 SYSTEM Inventory Service <-> Checkout Service
|
||||
-- Card Gateway <-> Payments Gateway: cosine 0.681, rank 4
|
||||
|
||||
--- what a similarity-only merge would do at each threshold ---
|
||||
threshold 0.90 -> 0 merges:
|
||||
threshold 0.80 -> 1 merges: INC-2057 + INC-2060
|
||||
threshold 0.70 -> 3 merges: INC-2057 + INC-2060; INC-2041 + INC-2057; INC-2041 + INC-2060
|
||||
threshold 0.60 -> 4 merges: INC-2057 + INC-2060; INC-2041 + INC-2057; INC-2041 + INC-2060; Card Gateway + Payments Gateway
|
||||
|
||||
--- B. plus a one-line alias table (Card Gateway = Payments Gateway) ---
|
||||
entities: 27
|
||||
Payments Gateway : -[DEPENDS_ON]-> Ledger Service; -[HOSTED_ON]-> Arcadia Cloud; -[OWNED_BY]-> Payments Team; <-[AFFECTED]- INC-2060; <-[DEPENDS_ON]- Checkout Service
|
||||
chunks: 35, MENTIONS edges: 54, dropped relations: 0
|
||||
|
||||
--- relationships by type ---
|
||||
DEPENDS_ON 6
|
||||
OWNED_BY 6
|
||||
HOSTED_ON 5
|
||||
LED_BY 5
|
||||
REPORTS_TO 4
|
||||
AFFECTED 3
|
||||
CAUSED_BY 2
|
||||
LOCATED_IN 2
|
||||
@@ -0,0 +1,45 @@
|
||||
# Building the graph in Neo4j; names that are the same thing (STUB extraction)
|
||||
|
||||
image : neo4j:5-community
|
||||
server : Neo4j Kernel 5.26.32 community
|
||||
|
||||
--- vector indexes (one from Neo4jVectorStore for chunks, one of ours for entity names) ---
|
||||
chunk_embedding on (:Chunk).embedding ONLINE
|
||||
entity_name_embedding on (:Entity).embedding ONLINE
|
||||
|
||||
--- A. exact-name matching only (no alias table, similarity merge off) ---
|
||||
entities: 28
|
||||
Payments Gateway : -[DEPENDS_ON]-> Ledger Service; -[HOSTED_ON]-> Arcadia Cloud; -[OWNED_BY]-> Payments Team; <-[DEPENDS_ON]- Checkout Service
|
||||
Card Gateway : <-[AFFECTED]- INC-2060
|
||||
|
||||
--- same-type name pairs by cosine of their embeddings, top 8 of 72 ---
|
||||
0.806 INCIDENT INC-2057 <-> INC-2060
|
||||
0.797 INCIDENT INC-2041 <-> INC-2057
|
||||
0.758 INCIDENT INC-2041 <-> INC-2060
|
||||
0.681 SYSTEM Card Gateway <-> Payments Gateway
|
||||
0.540 SYSTEM Ledger Service <-> Inventory Service
|
||||
0.535 TEAM Storefront Team <-> Platform Team
|
||||
0.530 TEAM Core Data Team <-> Platform Team
|
||||
0.519 SYSTEM Inventory Service <-> Checkout Service
|
||||
-- Card Gateway <-> Payments Gateway: cosine 0.681, rank 4
|
||||
|
||||
--- what a similarity-only merge would do at each threshold ---
|
||||
threshold 0.90 -> 0 merges:
|
||||
threshold 0.80 -> 1 merges: INC-2057 + INC-2060
|
||||
threshold 0.70 -> 3 merges: INC-2057 + INC-2060; INC-2041 + INC-2057; INC-2041 + INC-2060
|
||||
threshold 0.60 -> 4 merges: INC-2057 + INC-2060; INC-2041 + INC-2057; INC-2041 + INC-2060; Card Gateway + Payments Gateway
|
||||
|
||||
--- B. plus a one-line alias table (Card Gateway = Payments Gateway) ---
|
||||
entities: 27
|
||||
Payments Gateway : -[DEPENDS_ON]-> Ledger Service; -[HOSTED_ON]-> Arcadia Cloud; -[OWNED_BY]-> Payments Team; <-[AFFECTED]- INC-2060; <-[DEPENDS_ON]- Checkout Service
|
||||
chunks: 35, MENTIONS edges: 54, dropped relations: 0
|
||||
|
||||
--- relationships by type ---
|
||||
DEPENDS_ON 6
|
||||
OWNED_BY 6
|
||||
HOSTED_ON 5
|
||||
LED_BY 5
|
||||
REPORTS_TO 4
|
||||
AFFECTED 3
|
||||
CAUSED_BY 2
|
||||
LOCATED_IN 2
|
||||
@@ -0,0 +1,31 @@
|
||||
# Plain vector RAG against GraphRAG: is every supporting document in the context?
|
||||
|
||||
corpus: 35 chunks (one document each); 9 questions; evidence = the documents that together hold every fact needed
|
||||
vector: Neo4jVectorStore, top-k chunks. graph: link spans (cosine >= 0.80), walk <= 3 hops, keep the 6 best-ranked paths, take the chunks that stated their edges
|
||||
|
||||
id hops | vector k4 vector k8 | graph chunks | question
|
||||
Q1 1 | yes yes | yes 6 | Who leads the Fulfilment Team?
|
||||
Q2 2 | yes yes | yes 5 | Who leads the team that owns the Payments Gateway?
|
||||
Q3 2 | NO NO | yes 4 | Which team owns the root cause of INC-2041?
|
||||
Q4 3 | NO yes | yes 5 | Who does the leader of the team that owns the Payments Gateway report to?
|
||||
Q5 3 | NO yes | yes 4 | Where is the company that hosts the database behind the Ledger Service located?
|
||||
Q6 3 | NO yes | NO 5 | Who leads the teams that own the services the Checkout Service depends on?
|
||||
Q7 3 | NO NO | NO 5 | Which services stop working if the Orders Database goes down?
|
||||
Q8 2 | NO NO | yes 6 | Which incidents affected services that the Checkout Service depends on directly?
|
||||
Q9 3 | NO NO | yes 4 | Who leads the team that owns the service INC-2057 affected?
|
||||
|
||||
complete evidence: vector k=4 2/9, vector k=8 5/9, graph 7/9 (graph context averaged 4.9 chunks; vector k=8 is always 8)
|
||||
1-hop questions (1): vector k=4 1, vector k=8 1, graph 1
|
||||
2-hop questions (3): vector k=4 1, vector k=8 1, graph 3
|
||||
3-hop questions (5): vector k=4 0, vector k=8 3, graph 3
|
||||
|
||||
--- what each retriever missed ---
|
||||
Q1 vector k=8 missing [] graph missing []
|
||||
Q2 vector k=8 missing [] graph missing []
|
||||
Q3 vector k=8 missing [db-orders] graph missing []
|
||||
Q4 vector k=8 missing [] graph missing []
|
||||
Q5 vector k=8 missing [] graph missing []
|
||||
Q6 vector k=8 missing [] graph missing [team-payments]
|
||||
Q7 vector k=8 missing [svc-payments-gateway] graph missing [svc-checkout, svc-payments-gateway]
|
||||
Q8 vector k=8 missing [inc-2057] graph missing []
|
||||
Q9 vector k=8 missing [team-fulfilment, svc-inventory] graph missing []
|
||||
@@ -0,0 +1,32 @@
|
||||
# How the graph retriever's own settings move the result (complete evidence out of 9)
|
||||
|
||||
|
||||
--- hops walked x paths kept (link threshold 0.80) ---
|
||||
max hops | 6 paths 10 paths 20 paths 40 paths
|
||||
1 | 3/9 ( 3.0 ch) 3/9 ( 3.4 ch) 3/9 ( 3.4 ch) 3/9 ( 3.4 ch)
|
||||
2 | 4/9 ( 3.9 ch) 4/9 ( 5.1 ch) 7/9 ( 7.8 ch) 7/9 ( 8.4 ch)
|
||||
3 | 7/9 ( 4.9 ch) 8/9 ( 6.8 ch) 9/9 (10.3 ch) 9/9 (13.6 ch)
|
||||
|
||||
--- entity-link threshold (3 hops, 6 paths) ---
|
||||
threshold | complete seeds per question (avg)
|
||||
0.70 | 7/9 3.9
|
||||
0.80 | 7/9 1.9
|
||||
0.85 | 8/9 1.2
|
||||
0.90 | 8/9 1.0
|
||||
0.95 | 8/9 1.0
|
||||
|
||||
--- the seeds found for three questions at 0.80 (span -> entity, cosine) ---
|
||||
Q2 Who leads the team that owns the Payments Gateway?
|
||||
'Payments Gateway' -> Payments Gateway (1.00)
|
||||
'Payments' -> Payments Team (0.88)
|
||||
'team' -> Fulfilment Team (0.83)
|
||||
Q5 Where is the company that hosts the database behind the Ledger Service located?
|
||||
'Ledger Service' -> Ledger Service (1.00)
|
||||
'database' -> Orders Database (0.82)
|
||||
Q7 Which services stop working if the Orders Database goes down?
|
||||
'Orders Database' -> Orders Database (1.00)
|
||||
|
||||
--- candidate paths from one seed as the hop limit grows ---
|
||||
Payments Gateway 1 hop : 5 2 hops: 14 3 hops: 32 4 hops: 54
|
||||
Orders Database 1 hop : 4 2 hops: 12 3 hops: 24 4 hops: 46
|
||||
Checkout Service 1 hop : 3 2 hops: 12 3 hops: 27 4 hops: 49
|
||||
@@ -0,0 +1,20 @@
|
||||
# The prompt GraphRAG assembles for Q4 (no model was called; nothing below is a model answer)
|
||||
|
||||
QUESTION: Who does the leader of the team that owns the Payments Gateway report to?
|
||||
|
||||
FACTS FROM THE GRAPH:
|
||||
- Payments Gateway is owned by Payments Team
|
||||
- Payments Gateway is owned by Payments Team; Payments Team is led by Priya Raman; Priya Raman reports to Dana Okafor
|
||||
- Payments Gateway is owned by Payments Team; Payments Team is led by Priya Raman
|
||||
- Payments Gateway is owned by Payments Team; INC-2060 affected Payments Gateway
|
||||
- Payments Gateway is owned by Payments Team; Checkout Service depends on Payments Gateway
|
||||
- Payments Gateway is owned by Payments Team; Payments Gateway is hosted on Arcadia Cloud
|
||||
|
||||
SOURCE PASSAGES:
|
||||
[svc-payments-gateway] Payments Gateway. The Payments Gateway talks to the card networks on behalf of every other service. The Payments Gateway depends on the Ledger Service to record each authorised charge. The Payments Gateway is owned by the Payments Team. It is hosted on Arcadia Cloud.
|
||||
[team-payments] Payments Team. The Payments Team looks after card acceptance and refunds. The Payments Team is led by Priya Raman. The Payments Team is based in Pune. The team holds a weekly review on Thursdays.
|
||||
[person-priya] Priya Raman. Priya Raman joined the company in 2021 as a payments engineer. Priya Raman reports to Dana Okafor. She is the named contact for card network certificate renewals.
|
||||
[inc-2060] INC-2060. INC-2060 was a 12-minute slowdown on card payments. INC-2060 affected the Card Gateway. The Card Gateway is the on-call channel's name for the Payments Gateway. A retry storm after a network blip was the trigger. The write-up recommends capping retries.
|
||||
[svc-checkout] Checkout Service. The Checkout Service takes a customer's cart and turns it into an order. It depends on the Payments Gateway to take card payments and on the Inventory Service to reserve stock. The Checkout Service is owned by the Storefront Team. Its latency dashboard lives in the shared Grafana folder.
|
||||
|
||||
characters in this prompt: 1922; characters in the 8 chunks a k=8 vector search would add instead: 1820
|
||||
@@ -0,0 +1,23 @@
|
||||
# Questions whose shape is known: the Cypher is the retrieval (hand-written, parameterised, run on the real graph)
|
||||
|
||||
|
||||
--- Q7 Which services stop working if the Orders Database goes down? ---
|
||||
MATCH (db:Entity {name: $name})<-[:DEPENDS_ON*1..5]-(s:Entity)
|
||||
RETURN DISTINCT s.name AS service ORDER BY service
|
||||
-> Checkout Service
|
||||
-> Ledger Service
|
||||
-> Payments Gateway
|
||||
|
||||
--- Q6 Who leads the teams that own the services the Checkout Service depends on? ---
|
||||
MATCH (c:Entity {name: $name})-[:DEPENDS_ON]->(d:Entity)-[:OWNED_BY]->(team:Entity)-[:LED_BY]->(lead:Entity)
|
||||
RETURN d.name AS dependency, team.name AS team, lead.name AS lead ORDER BY dependency
|
||||
-> Inventory Service / Fulfilment Team / Tomasz Nowak
|
||||
-> Payments Gateway / Payments Team / Priya Raman
|
||||
|
||||
--- Q8 entity linking and the traversal in ONE statement: the vector index finds the seed, MATCH walks from it ---
|
||||
CALL db.index.vector.queryNodes('entity_name_embedding', 1, $vector) YIELD node AS c, score
|
||||
MATCH (c)-[:DEPENDS_ON]->(d:Entity)<-[:AFFECTED]-(i:Entity)
|
||||
RETURN c.name AS seed, score, d.name AS dependency, i.name AS incident ORDER BY incident
|
||||
$vector = embedding of the span "Checkout Service"
|
||||
-> seed Checkout Service (1.00): INC-2057 affected Inventory Service
|
||||
-> seed Checkout Service (1.00): INC-2060 affected Payments Gateway
|
||||
@@ -0,0 +1,80 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<project xmlns="http://maven.apache.org/POM/4.0.0"
|
||||
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
|
||||
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 https://maven.apache.org/xsd/maven-4.0.0.xsd">
|
||||
<modelVersion>4.0.0</modelVersion>
|
||||
|
||||
<parent>
|
||||
<groupId>org.springframework.boot</groupId>
|
||||
<artifactId>spring-boot-starter-parent</artifactId>
|
||||
<version>4.1.2</version>
|
||||
<relativePath/>
|
||||
</parent>
|
||||
|
||||
<groupId>com.ankurm</groupId>
|
||||
<artifactId>graphrag-neo4j</artifactId>
|
||||
<version>1.0.0</version>
|
||||
<name>graphrag-neo4j</name>
|
||||
<description>GraphRAG on Spring AI 2.0 and Neo4j: structured-output entity and relation extraction, a knowledge graph with a vector index, multi-hop retrieval in Cypher, compared with plain vector RAG over the same corpus.</description>
|
||||
|
||||
<properties>
|
||||
<java.version>25</java.version>
|
||||
<spring-ai.version>2.0.2</spring-ai.version>
|
||||
</properties>
|
||||
|
||||
<dependencyManagement>
|
||||
<dependencies>
|
||||
<dependency>
|
||||
<groupId>org.springframework.ai</groupId>
|
||||
<artifactId>spring-ai-bom</artifactId>
|
||||
<version>${spring-ai.version}</version>
|
||||
<type>pom</type>
|
||||
<scope>import</scope>
|
||||
</dependency>
|
||||
</dependencies>
|
||||
</dependencyManagement>
|
||||
|
||||
<dependencies>
|
||||
<!-- ChatClient, BeanOutputConverter, ChatModel (core). -->
|
||||
<dependency><groupId>org.springframework.ai</groupId><artifactId>spring-ai-client-chat</artifactId></dependency>
|
||||
<!-- Store module, not the starter: the store is built by hand so every setting is visible. -->
|
||||
<dependency><groupId>org.springframework.ai</groupId><artifactId>spring-ai-neo4j-store</artifactId></dependency>
|
||||
<!-- all-MiniLM-L6-v2 in-process through ONNX Runtime: a real embedding model, no API key. -->
|
||||
<dependency><groupId>org.springframework.ai</groupId><artifactId>spring-ai-transformers</artifactId></dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.springframework.boot</groupId>
|
||||
<artifactId>spring-boot-starter-test</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.testcontainers</groupId>
|
||||
<artifactId>testcontainers-neo4j</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.testcontainers</groupId>
|
||||
<artifactId>testcontainers-junit-jupiter</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
<!-- Testcontainers 2.x reads its configuration through commons-lang3 but does not bring it in: without this
|
||||
the first `new Neo4jContainer(...)` fails with NoClassDefFoundError. -->
|
||||
<dependency>
|
||||
<groupId>org.apache.commons</groupId>
|
||||
<artifactId>commons-lang3</artifactId>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
</dependencies>
|
||||
|
||||
<build>
|
||||
<plugins>
|
||||
<plugin>
|
||||
<groupId>org.apache.maven.plugins</groupId>
|
||||
<artifactId>maven-surefire-plugin</artifactId>
|
||||
<configuration>
|
||||
<argLine>-Duser.timezone=UTC -Dstdout.encoding=UTF-8 -Dfile.encoding=UTF-8 -Xmx3g</argLine>
|
||||
</configuration>
|
||||
</plugin>
|
||||
</plugins>
|
||||
</build>
|
||||
</project>
|
||||
Executable
+10
@@ -0,0 +1,10 @@
|
||||
#!/usr/bin/env bash
|
||||
# Regenerates every file in output/ by running the test suite.
|
||||
# Needs: JDK 25, Maven, a running Docker daemon (Testcontainers starts Neo4j itself), internet for the first run
|
||||
# (about 90 MB for the MiniLM model, cached in ~/.cache/minilm, and the Neo4j image, about 400 MB).
|
||||
# NEO4J_IMAGE overrides the image, e.g. NEO4J_IMAGE=neo4j:5-community ./scripts/run-all.sh
|
||||
set -euo pipefail
|
||||
cd "$(dirname "$0")/.."
|
||||
rm -rf output
|
||||
mvn -B -q test
|
||||
ls output
|
||||
@@ -0,0 +1,37 @@
|
||||
package com.ankurm.graphrag;
|
||||
|
||||
import java.util.Map;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
import org.springframework.ai.chat.client.ChatClient;
|
||||
import org.springframework.ai.chat.model.ChatModel;
|
||||
|
||||
/**
|
||||
* The last step of GraphRAG, and the only one that needs a real language model: put the selected paths and the
|
||||
* chunks that stated them in front of the model and ask. Nothing in this module's tests calls it with a real
|
||||
* model; {@link #prompt} is public so the tests can print exactly what would be sent.
|
||||
*/
|
||||
public class Answerer {
|
||||
|
||||
static final String SYSTEM = """
|
||||
Answer the question using only the facts and passages below. If they do not contain the answer,
|
||||
say so. Name the facts you used.
|
||||
""";
|
||||
|
||||
private final ChatClient client;
|
||||
|
||||
public Answerer(ChatModel model) {
|
||||
this.client = ChatClient.builder(model).defaultSystem(SYSTEM).build();
|
||||
}
|
||||
|
||||
public String answer(GraphRetriever.Retrieval retrieval, Map<String, String> chunkTexts) {
|
||||
return client.prompt().user(prompt(retrieval, chunkTexts)).call().content();
|
||||
}
|
||||
|
||||
public static String prompt(GraphRetriever.Retrieval r, Map<String, String> chunkTexts) {
|
||||
String facts = r.selected().stream().map(p -> "- " + p.path().text()).collect(Collectors.joining("\n"));
|
||||
String passages = r.chunkIds().stream().filter(chunkTexts::containsKey)
|
||||
.map(id -> "[" + id + "] " + chunkTexts.get(id)).collect(Collectors.joining("\n"));
|
||||
return "QUESTION: " + r.question() + "\n\nFACTS FROM THE GRAPH:\n" + facts + "\n\nSOURCE PASSAGES:\n" + passages;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
package com.ankurm.graphrag;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.util.ArrayList;
|
||||
import java.util.Comparator;
|
||||
import java.util.List;
|
||||
|
||||
import org.springframework.core.io.Resource;
|
||||
import org.springframework.core.io.support.PathMatchingResourcePatternResolver;
|
||||
|
||||
/** The 35 short documents under {@code src/main/resources/corpus}: file name = id, first line = title. */
|
||||
public final class Corpus {
|
||||
|
||||
public record Doc(String id, String title, String body) {
|
||||
/** Title and body as one chunk. These documents are short enough that one document is one chunk. */
|
||||
public String text() {
|
||||
return title + ". " + body;
|
||||
}
|
||||
}
|
||||
|
||||
private Corpus() {
|
||||
}
|
||||
|
||||
public static List<Doc> load() {
|
||||
try {
|
||||
List<Doc> docs = new ArrayList<>();
|
||||
for (Resource r : new PathMatchingResourcePatternResolver().getResources("classpath:corpus/*.txt")) {
|
||||
String id = r.getFilename().replace(".txt", "");
|
||||
String[] lines = new String(r.getInputStream().readAllBytes(), StandardCharsets.UTF_8).strip().split("\n", 2);
|
||||
docs.add(new Doc(id, lines[0].strip(), lines[1].strip()));
|
||||
}
|
||||
docs.sort(Comparator.comparing(Doc::id));
|
||||
return docs;
|
||||
} catch (IOException e) {
|
||||
throw new IllegalStateException(e);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,126 @@
|
||||
package com.ankurm.graphrag;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.Comparator;
|
||||
import java.util.HashMap;
|
||||
import java.util.LinkedHashMap;
|
||||
import java.util.List;
|
||||
import java.util.Locale;
|
||||
import java.util.Map;
|
||||
|
||||
import org.springframework.ai.embedding.EmbeddingModel;
|
||||
|
||||
/**
|
||||
* Step 2: decide which names are the same thing. A model that reads each chunk on its own writes "Card
|
||||
* Gateway" in one post-mortem and "Payments Gateway" in a service page; unless something merges them they
|
||||
* become two nodes and every path through one is invisible from the other.
|
||||
*
|
||||
* <p>Three layers, applied in order: names that differ only in case, punctuation or spacing are the same key;
|
||||
* a hand-kept alias table; and, optionally, cosine similarity of the name embeddings between entities of the
|
||||
* same type above a threshold.
|
||||
*/
|
||||
public final class EntityResolver {
|
||||
|
||||
/** Two distinct names of the same type and how close their embeddings are. */
|
||||
public record Pair(String a, String b, EntityType type, double cosine) {
|
||||
}
|
||||
|
||||
/** What every mention maps to, and what was merged on the way. */
|
||||
public record Resolution(Map<String, String> canonicalByName, Map<String, EntityType> typeByCanonical,
|
||||
List<Pair> merged) {
|
||||
}
|
||||
|
||||
private final EmbeddingModel embeddings;
|
||||
|
||||
public EntityResolver(EmbeddingModel embeddings) {
|
||||
this.embeddings = embeddings;
|
||||
}
|
||||
|
||||
public static String key(String name) {
|
||||
return name.toLowerCase(Locale.ROOT).replaceAll("[^a-z0-9]", "");
|
||||
}
|
||||
|
||||
/**
|
||||
* @param mentions every entity mention from every chunk, repeats included (the most frequent spelling wins)
|
||||
* @param aliases lower-case-insensitive alias to canonical name
|
||||
* @param threshold merge same-type names whose cosine is at least this; pass a value above 1 to switch it off
|
||||
*/
|
||||
public Resolution resolve(List<Extraction.Entity> mentions, Map<String, String> aliases, double threshold) {
|
||||
Map<String, Integer> counts = new HashMap<>();
|
||||
Map<String, EntityType> types = new LinkedHashMap<>();
|
||||
Map<String, String> spelling = new LinkedHashMap<>();
|
||||
for (Extraction.Entity e : mentions) {
|
||||
counts.merge(key(e.name()), 1, Integer::sum);
|
||||
types.putIfAbsent(key(e.name()), e.type());
|
||||
spelling.putIfAbsent(key(e.name()), e.name());
|
||||
}
|
||||
Map<String, String> aliasByKey = new HashMap<>();
|
||||
aliases.forEach((alias, canonical) -> aliasByKey.put(key(alias), canonical));
|
||||
|
||||
// key -> canonical key after the alias table
|
||||
Map<String, String> parent = new LinkedHashMap<>();
|
||||
for (String k : types.keySet()) {
|
||||
String canonical = aliasByKey.get(k);
|
||||
parent.put(k, canonical == null ? k : key(canonical));
|
||||
if (canonical != null) {
|
||||
types.putIfAbsent(key(canonical), types.get(k));
|
||||
spelling.putIfAbsent(key(canonical), canonical);
|
||||
}
|
||||
}
|
||||
|
||||
List<Pair> merged = new ArrayList<>();
|
||||
List<String> roots = parent.values().stream().distinct().toList();
|
||||
if (threshold <= 1.0) {
|
||||
for (Pair p : similarPairs(roots, types, spelling)) {
|
||||
if (p.cosine() >= threshold) {
|
||||
String ka = key(p.a());
|
||||
String kb = key(p.b());
|
||||
String winner = counts.getOrDefault(ka, 0) >= counts.getOrDefault(kb, 0) ? ka : kb;
|
||||
String loser = winner.equals(ka) ? kb : ka;
|
||||
final String w = winner;
|
||||
final String l = loser;
|
||||
parent.replaceAll((k, v) -> v.equals(l) ? w : v);
|
||||
merged.add(p);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Map<String, String> canonicalByName = new LinkedHashMap<>();
|
||||
Map<String, EntityType> typeByCanonical = new LinkedHashMap<>();
|
||||
for (Extraction.Entity e : mentions) {
|
||||
String root = parent.get(key(e.name()));
|
||||
String canonical = spelling.get(root);
|
||||
canonicalByName.put(e.name(), canonical);
|
||||
typeByCanonical.putIfAbsent(canonical, types.get(root));
|
||||
}
|
||||
return new Resolution(canonicalByName, typeByCanonical, merged);
|
||||
}
|
||||
|
||||
/** All pairs of distinct same-type names, most similar first. Used to merge and to print the evidence. */
|
||||
public List<Pair> similarPairs(List<String> keys, Map<String, EntityType> types, Map<String, String> spelling) {
|
||||
List<float[]> vectors = embeddings.embed(keys.stream().map(spelling::get).toList());
|
||||
List<Pair> out = new ArrayList<>();
|
||||
for (int i = 0; i < keys.size(); i++) {
|
||||
for (int j = i + 1; j < keys.size(); j++) {
|
||||
if (types.get(keys.get(i)) == types.get(keys.get(j))) {
|
||||
out.add(new Pair(spelling.get(keys.get(i)), spelling.get(keys.get(j)), types.get(keys.get(i)),
|
||||
cosine(vectors.get(i), vectors.get(j))));
|
||||
}
|
||||
}
|
||||
}
|
||||
out.sort(Comparator.comparingDouble(Pair::cosine).reversed());
|
||||
return out;
|
||||
}
|
||||
|
||||
public static double cosine(float[] a, float[] b) {
|
||||
double dot = 0;
|
||||
double na = 0;
|
||||
double nb = 0;
|
||||
for (int i = 0; i < a.length; i++) {
|
||||
dot += a[i] * b[i];
|
||||
na += a[i] * a[i];
|
||||
nb += b[i] * b[i];
|
||||
}
|
||||
return dot / (Math.sqrt(na) * Math.sqrt(nb));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
package com.ankurm.graphrag;
|
||||
|
||||
/** The node kinds the extraction schema allows. A closed set, so the model cannot invent "Concept" or "Thing". */
|
||||
public enum EntityType {
|
||||
SYSTEM, TEAM, PERSON, VENDOR, INCIDENT, PLACE
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
package com.ankurm.graphrag;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* What the model must return for one chunk. {@code ChatClient.entity(Extraction.class)} turns this record into a
|
||||
* JSON schema, appends it to the prompt, and parses the reply back into the record.
|
||||
*/
|
||||
public record Extraction(List<Entity> entities, List<Relation> relations) {
|
||||
|
||||
public record Entity(String name, EntityType type) {
|
||||
}
|
||||
|
||||
public record Relation(String source, RelationType type, String target) {
|
||||
}
|
||||
|
||||
public static Extraction empty() {
|
||||
return new Extraction(List.of(), List.of());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
package com.ankurm.graphrag;
|
||||
|
||||
import java.util.HashSet;
|
||||
import java.util.List;
|
||||
import java.util.Set;
|
||||
|
||||
import org.springframework.ai.chat.client.ChatClient;
|
||||
import org.springframework.ai.chat.model.ChatModel;
|
||||
|
||||
/**
|
||||
* Step 1 of GraphRAG: one chunk of text in, entities and relations out, through
|
||||
* {@code ChatClient.entity(Extraction.class)}. The model sees the chunk and a JSON schema built from the
|
||||
* {@link Extraction} record; the enums in that record are what constrain it to a fixed vocabulary.
|
||||
*/
|
||||
public class Extractor {
|
||||
|
||||
static final String SYSTEM = """
|
||||
You turn short technical documents into a knowledge graph.
|
||||
List every named system, team, person, vendor, incident and place the document mentions as an entity,
|
||||
using the name exactly as the document writes it.
|
||||
List a relation only when the document states it. Never infer one.
|
||||
""";
|
||||
|
||||
static final String USER = """
|
||||
Extract the entities and relations from the document.
|
||||
|
||||
DOCUMENT_START
|
||||
{document}
|
||||
DOCUMENT_END
|
||||
""";
|
||||
|
||||
/** The cleaned extraction plus how many relations were dropped for naming an entity that was never listed. */
|
||||
public record Result(Extraction extraction, int droppedRelations) {
|
||||
}
|
||||
|
||||
private final ChatClient client;
|
||||
|
||||
public Extractor(ChatModel model) {
|
||||
this.client = ChatClient.builder(model).defaultSystem(SYSTEM).build();
|
||||
}
|
||||
|
||||
public Result extract(Corpus.Doc doc) {
|
||||
Extraction raw = client.prompt()
|
||||
.user(u -> u.text(USER).param("document", doc.text()))
|
||||
.call()
|
||||
.entity(Extraction.class);
|
||||
return clean(raw);
|
||||
}
|
||||
|
||||
/**
|
||||
* A relation is kept only if both endpoints appear in the entity list. A model that writes a relation to a
|
||||
* name it never listed would otherwise make the graph writer MATCH nothing and silently create no edge.
|
||||
*/
|
||||
static Result clean(Extraction raw) {
|
||||
if (raw == null || raw.entities() == null || raw.relations() == null) {
|
||||
return new Result(Extraction.empty(), 0);
|
||||
}
|
||||
Set<String> names = new HashSet<>();
|
||||
raw.entities().forEach(e -> names.add(e.name()));
|
||||
List<Extraction.Relation> kept = raw.relations().stream()
|
||||
.filter(r -> names.contains(r.source()) && names.contains(r.target()))
|
||||
.toList();
|
||||
return new Result(new Extraction(raw.entities(), kept), raw.relations().size() - kept.size());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
package com.ankurm.graphrag;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.Comparator;
|
||||
import java.util.LinkedHashMap;
|
||||
import java.util.LinkedHashSet;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.Set;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
import org.springframework.ai.embedding.EmbeddingModel;
|
||||
|
||||
/**
|
||||
* Multi-hop retrieval in three steps, with no language model anywhere in it:
|
||||
* <ol>
|
||||
* <li><b>Link</b>: cut the question into every 1 to 4 word span, embed them, ask the entity vector index for
|
||||
* the nearest entity of each, and keep the ones that clear {@link #linkThreshold}.</li>
|
||||
* <li><b>Walk</b>: from each seed entity, collect every simple path of up to {@code depth} relationships.</li>
|
||||
* <li><b>Rank</b>: render each path as a sentence, embed it, and keep the {@code topPaths} closest to the
|
||||
* question. The chunks that stated the relationships on those paths are the context.</li>
|
||||
* </ol>
|
||||
*/
|
||||
public class GraphRetriever {
|
||||
|
||||
private static final Set<String> STOP = Set.of("a", "an", "the", "of", "to", "in", "on", "is", "are", "who", "what",
|
||||
"which", "where", "that", "does", "do", "for", "by", "if", "and", "or", "it", "its", "this", "these", "those",
|
||||
"with", "from", "at", "as", "be", "goes", "down", "stop", "working", "whom", "behind");
|
||||
|
||||
public record Seed(String span, String name, double score) {
|
||||
}
|
||||
|
||||
public record RankedPath(GraphStore.GraphPath path, double score) {
|
||||
}
|
||||
|
||||
public record Retrieval(String question, List<Seed> seeds, int candidatePaths, List<RankedPath> selected,
|
||||
Set<String> chunkIds) {
|
||||
}
|
||||
|
||||
private final GraphStore graph;
|
||||
private final EmbeddingModel embeddings;
|
||||
private final double linkThreshold;
|
||||
|
||||
public GraphRetriever(GraphStore graph, EmbeddingModel embeddings, double linkThreshold) {
|
||||
this.graph = graph;
|
||||
this.embeddings = embeddings;
|
||||
this.linkThreshold = linkThreshold;
|
||||
}
|
||||
|
||||
public List<Seed> link(String question) {
|
||||
String[] words = question.replaceAll("[?,.]", "").split("\\s+");
|
||||
List<String> spans = new ArrayList<>();
|
||||
for (int i = 0; i < words.length; i++) {
|
||||
for (int n = 1; n <= 4 && i + n <= words.length; n++) {
|
||||
String first = words[i].toLowerCase();
|
||||
String last = words[i + n - 1].toLowerCase();
|
||||
if (STOP.contains(first) || STOP.contains(last)) {
|
||||
continue;
|
||||
}
|
||||
spans.add(String.join(" ", java.util.Arrays.copyOfRange(words, i, i + n)));
|
||||
}
|
||||
}
|
||||
List<float[]> vectors = embeddings.embed(spans);
|
||||
Map<String, Seed> best = new LinkedHashMap<>();
|
||||
for (int i = 0; i < spans.size(); i++) {
|
||||
for (GraphStore.EntityHit hit : graph.nearestEntities(vectors.get(i), 1)) {
|
||||
if (hit.score() >= linkThreshold) {
|
||||
Seed seed = new Seed(spans.get(i), hit.name(), hit.score());
|
||||
best.merge(hit.name(), seed, (a, b) -> a.score() >= b.score() ? a : b);
|
||||
}
|
||||
}
|
||||
}
|
||||
return best.values().stream().sorted(Comparator.comparingDouble(Seed::score).reversed()).toList();
|
||||
}
|
||||
|
||||
public Retrieval retrieve(String question, int depth, int topPaths) {
|
||||
List<Seed> seeds = link(question);
|
||||
// one walk per seed; a path found from two seeds in opposite directions counts once
|
||||
Map<Set<String>, GraphStore.GraphPath> unique = new LinkedHashMap<>();
|
||||
for (Seed s : seeds) {
|
||||
for (GraphStore.GraphPath p : graph.paths(EntityResolver.key(s.name()), depth, 5000)) {
|
||||
Set<String> edgeSet = p.edges().stream().map(GraphStore.Edge::sentence).collect(Collectors.toSet());
|
||||
unique.putIfAbsent(edgeSet, p);
|
||||
}
|
||||
}
|
||||
List<GraphStore.GraphPath> candidates = new ArrayList<>(unique.values());
|
||||
List<RankedPath> ranked = new ArrayList<>();
|
||||
if (!candidates.isEmpty()) {
|
||||
float[] q = embeddings.embed(question);
|
||||
List<float[]> vectors = embeddings.embed(candidates.stream().map(GraphStore.GraphPath::text).toList());
|
||||
for (int i = 0; i < candidates.size(); i++) {
|
||||
ranked.add(new RankedPath(candidates.get(i), EntityResolver.cosine(q, vectors.get(i))));
|
||||
}
|
||||
ranked.sort(Comparator.comparingDouble(RankedPath::score).reversed());
|
||||
}
|
||||
List<RankedPath> selected = ranked.subList(0, Math.min(topPaths, ranked.size()));
|
||||
Set<String> chunks = new LinkedHashSet<>();
|
||||
selected.forEach(rp -> rp.path().edges().forEach(e -> chunks.addAll(e.sources())));
|
||||
return new Retrieval(question, seeds, candidates.size(), List.copyOf(selected), chunks);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,164 @@
|
||||
package com.ankurm.graphrag;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.Arrays;
|
||||
import java.util.Collection;
|
||||
import java.util.LinkedHashMap;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
import org.neo4j.driver.Driver;
|
||||
import org.neo4j.driver.Record;
|
||||
import org.neo4j.driver.Value;
|
||||
import org.springframework.ai.embedding.EmbeddingModel;
|
||||
|
||||
/**
|
||||
* Everything that talks Cypher. Entities are {@code (:Entity {key, name, type, embedding})}; each relation is a
|
||||
* real typed relationship (:DEPENDS_ON, :OWNED_BY, ...) carrying the ids of the chunks that stated it in a
|
||||
* {@code sources} list; chunks are {@code (:Chunk)} nodes written by Spring AI's Neo4jVectorStore, and a
|
||||
* {@code (:Chunk)-[:MENTIONS]->(:Entity)} edge records where every entity was seen.
|
||||
*/
|
||||
public class GraphStore {
|
||||
|
||||
public static final String ENTITY_INDEX = "entity_name_embedding";
|
||||
|
||||
/** One stated relation, in its real direction, with the chunks that stated it. */
|
||||
public record Edge(String from, RelationType type, String to, List<String> sources) {
|
||||
public String sentence() {
|
||||
return from + " " + type.phrase() + " " + to;
|
||||
}
|
||||
}
|
||||
|
||||
/** A walk through the graph from a seed entity. {@code edges} are in walking order, each in its true direction. */
|
||||
public record GraphPath(String seed, List<Edge> edges) {
|
||||
public String text() {
|
||||
return edges.stream().map(Edge::sentence).collect(Collectors.joining("; "));
|
||||
}
|
||||
}
|
||||
|
||||
public record EntityHit(String name, String key, EntityType type, double score) {
|
||||
}
|
||||
|
||||
private final Driver driver;
|
||||
private final EmbeddingModel embeddings;
|
||||
|
||||
public GraphStore(Driver driver, EmbeddingModel embeddings) {
|
||||
this.driver = driver;
|
||||
this.embeddings = embeddings;
|
||||
}
|
||||
|
||||
public void initSchema() {
|
||||
run("CREATE CONSTRAINT entity_key IF NOT EXISTS FOR (e:Entity) REQUIRE e.key IS UNIQUE", Map.of());
|
||||
run("CREATE VECTOR INDEX " + ENTITY_INDEX + " IF NOT EXISTS FOR (e:Entity) ON (e.embedding) "
|
||||
+ "OPTIONS {indexConfig: {`vector.dimensions`: " + embeddings.dimensions()
|
||||
+ ", `vector.similarity_function`: 'cosine'}}", Map.of());
|
||||
run("CALL db.awaitIndexes(60)", Map.of());
|
||||
}
|
||||
|
||||
public void clearEntities() {
|
||||
run("MATCH (e:Entity) DETACH DELETE e", Map.of());
|
||||
}
|
||||
|
||||
/**
|
||||
* Writes one resolved extraction per chunk. Relationship types are concatenated into the Cypher text because
|
||||
* Cypher has no parameter for them; {@link RelationType} is the allow-list that makes that safe.
|
||||
*/
|
||||
public void write(Map<String, Extraction> extractionByChunk, EntityResolver.Resolution resolution) {
|
||||
List<Map<String, Object>> entityRows = new ArrayList<>();
|
||||
for (Map.Entry<String, EntityType> e : resolution.typeByCanonical().entrySet()) {
|
||||
Map<String, Object> row = new LinkedHashMap<>();
|
||||
row.put("key", EntityResolver.key(e.getKey()));
|
||||
row.put("name", e.getKey());
|
||||
row.put("type", e.getValue().name());
|
||||
row.put("embedding", embeddings.embed(e.getKey()));
|
||||
entityRows.add(row);
|
||||
}
|
||||
run("""
|
||||
UNWIND $rows AS row
|
||||
MERGE (e:Entity {key: row.key})
|
||||
SET e.name = row.name, e.type = row.type
|
||||
WITH e, row
|
||||
CALL db.create.setNodeVectorProperty(e, 'embedding', row.embedding)
|
||||
""", Map.of("rows", entityRows));
|
||||
|
||||
List<Map<String, Object>> mentions = new ArrayList<>();
|
||||
Map<RelationType, List<Map<String, Object>>> byType = new LinkedHashMap<>();
|
||||
extractionByChunk.forEach((chunkId, extraction) -> {
|
||||
for (Extraction.Entity e : extraction.entities()) {
|
||||
mentions.add(Map.of("chunk", chunkId, "key", EntityResolver.key(resolution.canonicalByName().get(e.name()))));
|
||||
}
|
||||
for (Extraction.Relation r : extraction.relations()) {
|
||||
byType.computeIfAbsent(r.type(), t -> new ArrayList<>()).add(Map.of("chunk", chunkId,
|
||||
"from", EntityResolver.key(resolution.canonicalByName().get(r.source())),
|
||||
"to", EntityResolver.key(resolution.canonicalByName().get(r.target()))));
|
||||
}
|
||||
});
|
||||
run("""
|
||||
UNWIND $rows AS row
|
||||
MATCH (c:Chunk {id: row.chunk}), (e:Entity {key: row.key})
|
||||
MERGE (c)-[:MENTIONS]->(e)
|
||||
""", Map.of("rows", mentions));
|
||||
byType.forEach((type, rows) -> run("""
|
||||
UNWIND $rows AS row
|
||||
MATCH (a:Entity {key: row.from}), (b:Entity {key: row.to})
|
||||
MERGE (a)-[r:%s]->(b)
|
||||
SET r.sources = CASE WHEN row.chunk IN coalesce(r.sources, []) THEN r.sources
|
||||
ELSE coalesce(r.sources, []) + row.chunk END
|
||||
""".formatted(type.name()), Map.of("rows", rows)));
|
||||
}
|
||||
|
||||
/** The k entities whose name embedding is nearest to the vector, through the vector index. */
|
||||
public List<EntityHit> nearestEntities(float[] vector, int k) {
|
||||
return query("""
|
||||
CALL db.index.vector.queryNodes($index, $k, $vector) YIELD node, score
|
||||
RETURN node.name AS name, node.key AS key, node.type AS type, score
|
||||
ORDER BY score DESC
|
||||
""", Map.of("index", ENTITY_INDEX, "k", k, "vector", vector)).stream()
|
||||
.map(r -> new EntityHit(r.get("name").asString(), r.get("key").asString(),
|
||||
EntityType.valueOf(r.get("type").asString()), r.get("score").asDouble()))
|
||||
.toList();
|
||||
}
|
||||
|
||||
/**
|
||||
* Every simple path of 1..maxHops relationships that starts at the seed, in either direction, over the
|
||||
* eight allowed relationship types. Chunk nodes are not walked through: the type list excludes MENTIONS.
|
||||
*/
|
||||
public List<GraphPath> paths(String seedKey, int maxHops, int limit) {
|
||||
String types = Arrays.stream(RelationType.values()).map(Enum::name).collect(Collectors.joining("|"));
|
||||
List<Record> rows = query("""
|
||||
MATCH p = (s:Entity {key: $key})-[:%s*1..%d]-(t:Entity)
|
||||
WHERE all(n IN nodes(p) WHERE single(m IN nodes(p) WHERE m = n))
|
||||
RETURN [r IN relationships(p) | {from: startNode(r).name, type: type(r), to: endNode(r).name,
|
||||
sources: r.sources}] AS edges
|
||||
LIMIT $limit
|
||||
""".formatted(types, maxHops), Map.of("key", seedKey, "limit", limit));
|
||||
List<GraphPath> out = new ArrayList<>();
|
||||
for (Record row : rows) {
|
||||
List<Edge> edges = new ArrayList<>();
|
||||
for (Value e : row.get("edges").values()) {
|
||||
edges.add(new Edge(e.get("from").asString(), RelationType.valueOf(e.get("type").asString()),
|
||||
e.get("to").asString(), e.get("sources").asList(Value::asString)));
|
||||
}
|
||||
out.add(new GraphPath(seedKey, edges));
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
public Map<String, String> chunkTexts(Collection<String> ids) {
|
||||
Map<String, String> out = new LinkedHashMap<>();
|
||||
query("MATCH (c:Chunk) WHERE c.id IN $ids RETURN c.id AS id, c.text AS text ORDER BY id", Map.of("ids", List.copyOf(ids)))
|
||||
.forEach(r -> out.put(r.get("id").asString(), r.get("text").asString()));
|
||||
return out;
|
||||
}
|
||||
|
||||
// ---- small helpers ---------------------------------------------------------------------------
|
||||
|
||||
public List<Record> query(String cypher, Map<String, Object> params) {
|
||||
return driver.executableQuery(cypher).withParameters(params).execute().records();
|
||||
}
|
||||
|
||||
public void run(String cypher, Map<String, Object> params) {
|
||||
driver.executableQuery(cypher).withParameters(params).execute();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
package com.ankurm.graphrag;
|
||||
|
||||
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. Same approach as the rag-ingestion module in this repository.
|
||||
*/
|
||||
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() {
|
||||
}
|
||||
|
||||
public static synchronized EmbeddingModel get() {
|
||||
if (shared == null) {
|
||||
try {
|
||||
TransformersEmbeddingModel m = new TransformersEmbeddingModel(MetadataMode.NONE);
|
||||
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();
|
||||
shared = m;
|
||||
} catch (Exception e) {
|
||||
throw new IllegalStateException(e);
|
||||
}
|
||||
}
|
||||
return shared;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,79 @@
|
||||
package com.ankurm.graphrag;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.LinkedHashMap;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
import org.neo4j.driver.Driver;
|
||||
import org.springframework.ai.chat.model.ChatModel;
|
||||
import org.springframework.ai.document.Document;
|
||||
import org.springframework.ai.embedding.EmbeddingModel;
|
||||
import org.springframework.ai.vectorstore.SearchRequest;
|
||||
import org.springframework.ai.vectorstore.neo4j.Neo4jVectorStore;
|
||||
|
||||
/**
|
||||
* The whole indexing side in one place: chunks into a Neo4j vector index (what plain vector RAG needs), then
|
||||
* extraction, entity resolution and graph writes (what GraphRAG adds on top).
|
||||
*/
|
||||
public class Pipeline {
|
||||
|
||||
/** What building the graph produced, for printing and for tests. */
|
||||
public record Built(Map<String, Extraction> extractionByChunk, EntityResolver.Resolution resolution,
|
||||
int droppedRelations, List<Extraction.Entity> allMentions) {
|
||||
}
|
||||
|
||||
private final GraphStore graph;
|
||||
private final EmbeddingModel embeddings;
|
||||
private final Extractor extractor;
|
||||
private final Neo4jVectorStore chunks;
|
||||
|
||||
public Pipeline(Driver driver, EmbeddingModel embeddings, ChatModel chatModel) {
|
||||
this.embeddings = embeddings;
|
||||
this.graph = new GraphStore(driver, embeddings);
|
||||
this.extractor = new Extractor(chatModel);
|
||||
this.chunks = Neo4jVectorStore.builder(driver, embeddings)
|
||||
.label("Chunk")
|
||||
.indexName("chunk_embedding")
|
||||
.embeddingDimension(embeddings.dimensions())
|
||||
.initializeSchema(true)
|
||||
.build();
|
||||
this.chunks.afterPropertiesSet();
|
||||
this.graph.initSchema();
|
||||
}
|
||||
|
||||
public GraphStore graph() {
|
||||
return graph;
|
||||
}
|
||||
|
||||
public Neo4jVectorStore chunks() {
|
||||
return chunks;
|
||||
}
|
||||
|
||||
public void loadChunks(List<Corpus.Doc> docs) {
|
||||
chunks.add(docs.stream().map(d -> new Document(d.id(), d.text(), Map.of("title", d.title()))).toList());
|
||||
}
|
||||
|
||||
/** Plain vector RAG: the ids of the k chunks nearest to the question. */
|
||||
public List<String> vectorSearch(String question, int k) {
|
||||
return chunks.similaritySearch(SearchRequest.builder().query(question).topK(k).build())
|
||||
.stream().map(Document::getId).toList();
|
||||
}
|
||||
|
||||
/** Extract from every chunk, resolve names, and (re)write the entity graph. */
|
||||
public Built buildGraph(List<Corpus.Doc> docs, Map<String, String> aliases, double mergeThreshold) {
|
||||
Map<String, Extraction> byChunk = new LinkedHashMap<>();
|
||||
List<Extraction.Entity> mentions = new ArrayList<>();
|
||||
int dropped = 0;
|
||||
for (Corpus.Doc d : docs) {
|
||||
Extractor.Result r = extractor.extract(d);
|
||||
byChunk.put(d.id(), r.extraction());
|
||||
mentions.addAll(r.extraction().entities());
|
||||
dropped += r.droppedRelations();
|
||||
}
|
||||
EntityResolver.Resolution resolution = new EntityResolver(embeddings).resolve(mentions, aliases, mergeThreshold);
|
||||
graph.clearEntities();
|
||||
graph.write(byChunk, resolution);
|
||||
return new Built(byChunk, resolution, dropped, mentions);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
package com.ankurm.graphrag;
|
||||
|
||||
import java.util.List;
|
||||
import java.util.Set;
|
||||
|
||||
/**
|
||||
* One evaluation question. {@code evidence} is the set of documents that together contain every fact needed to
|
||||
* answer it, worked out by hand from the corpus; retrieval is "complete" only when all of them are in the context.
|
||||
*/
|
||||
public record Question(String id, int hops, String text, Set<String> evidence, List<String> answers) {
|
||||
|
||||
public static List<Question> all() {
|
||||
return List.of(
|
||||
new Question("Q1", 1, "Who leads the Fulfilment Team?",
|
||||
Set.of("team-fulfilment"), List.of("Tomasz Nowak")),
|
||||
new Question("Q2", 2, "Who leads the team that owns the Payments Gateway?",
|
||||
Set.of("svc-payments-gateway", "team-payments"), List.of("Priya Raman")),
|
||||
new Question("Q3", 2, "Which team owns the root cause of INC-2041?",
|
||||
Set.of("inc-2041", "db-orders"), List.of("Core Data Team")),
|
||||
new Question("Q4", 3, "Who does the leader of the team that owns the Payments Gateway report to?",
|
||||
Set.of("svc-payments-gateway", "team-payments", "person-priya"), List.of("Dana Okafor")),
|
||||
new Question("Q5", 3, "Where is the company that hosts the database behind the Ledger Service located?",
|
||||
Set.of("svc-ledger", "db-orders", "vendor-borealis"), List.of("Frankfurt")),
|
||||
new Question("Q6", 3, "Who leads the teams that own the services the Checkout Service depends on?",
|
||||
Set.of("svc-checkout", "svc-payments-gateway", "svc-inventory", "team-payments", "team-fulfilment"),
|
||||
List.of("Priya Raman", "Tomasz Nowak")),
|
||||
new Question("Q7", 3, "Which services stop working if the Orders Database goes down?",
|
||||
Set.of("svc-ledger", "svc-payments-gateway", "svc-checkout"),
|
||||
List.of("Ledger Service", "Payments Gateway", "Checkout Service")),
|
||||
new Question("Q8", 2, "Which incidents affected services that the Checkout Service depends on directly?",
|
||||
Set.of("svc-checkout", "inc-2057", "inc-2060"), List.of("INC-2057", "INC-2060")),
|
||||
new Question("Q9", 3, "Who leads the team that owns the service INC-2057 affected?",
|
||||
Set.of("inc-2057", "svc-inventory", "team-fulfilment"), List.of("Tomasz Nowak")));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
package com.ankurm.graphrag;
|
||||
|
||||
/**
|
||||
* The relationship kinds the extraction schema allows. This enum is also the allow-list that makes it safe to
|
||||
* put a model-chosen relationship type into a Cypher string: Cypher cannot take a relationship type as a
|
||||
* parameter, so the type is concatenated into the query, and only these constants can ever get there.
|
||||
*/
|
||||
public enum RelationType {
|
||||
DEPENDS_ON("depends on"),
|
||||
OWNED_BY("is owned by"),
|
||||
LED_BY("is led by"),
|
||||
HOSTED_ON("is hosted on"),
|
||||
AFFECTED("affected"),
|
||||
CAUSED_BY("was caused by"),
|
||||
REPORTS_TO("reports to"),
|
||||
LOCATED_IN("is located in");
|
||||
|
||||
private final String phrase;
|
||||
|
||||
RelationType(String phrase) {
|
||||
this.phrase = phrase;
|
||||
}
|
||||
|
||||
/** The relation as an English phrase, used to render a path as a sentence. */
|
||||
public String phrase() {
|
||||
return phrase;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,202 @@
|
||||
package com.ankurm.graphrag.stub;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.LinkedHashMap;
|
||||
import java.util.LinkedHashSet;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.Set;
|
||||
import java.util.concurrent.atomic.AtomicInteger;
|
||||
import java.util.concurrent.atomic.AtomicLong;
|
||||
import java.util.regex.Matcher;
|
||||
import java.util.regex.Pattern;
|
||||
|
||||
import com.ankurm.graphrag.EntityType;
|
||||
import com.ankurm.graphrag.RelationType;
|
||||
import org.springframework.ai.chat.messages.AssistantMessage;
|
||||
import org.springframework.ai.chat.model.ChatResponse;
|
||||
import org.springframework.ai.chat.model.Generation;
|
||||
import org.springframework.ai.chat.prompt.Prompt;
|
||||
import tools.jackson.databind.json.JsonMapper;
|
||||
|
||||
/**
|
||||
* A STAND-IN for a language model, not a language model. It reads the document out of the prompt the way
|
||||
* {@code Extractor} writes it, finds sentences built from eight fixed phrases ("depends on", "is owned by",
|
||||
* "is led by", ...) with regular expressions, and replies with the JSON a good model would produce for them.
|
||||
* It exists so that the whole pipeline (prompt, JSON schema, parsing, graph writes) runs offline and the same
|
||||
* way every time. It says nothing about how well a real model extracts entities from real prose, and no
|
||||
* number in this module claims to.
|
||||
*
|
||||
* <p>The {@link Mode} values reproduce replies that real models do produce: JSON wrapped in a markdown
|
||||
* fence, JSON cut off mid-way, a relation type outside the allowed list, a relation whose endpoint was never
|
||||
* listed as an entity.
|
||||
*/
|
||||
public final class RuleBasedExtractorModel implements org.springframework.ai.chat.model.ChatModel {
|
||||
|
||||
public enum Mode { CLEAN, FENCED, PROSE_THEN_FENCE, TRUNCATED, UNKNOWN_RELATION_TYPE, INJECTION_RELATION_TYPE, DANGLING_RELATION }
|
||||
|
||||
private static final String NAME = "[A-Z][\\w-]*(?: [A-Z][\\w-]*)*";
|
||||
private static final Pattern DEPENDS = Pattern.compile("(?:depends on|\\band on) (?:[Tt]he )?(" + NAME + ")");
|
||||
private static final Pattern FIRST_NAME = Pattern.compile("^\\s*(?:[Tt]he )?(" + NAME + ")");
|
||||
private static final JsonMapper JSON = JsonMapper.builder().build();
|
||||
|
||||
private final Mode mode;
|
||||
private final AtomicInteger calls = new AtomicInteger();
|
||||
private final AtomicLong inputChars = new AtomicLong();
|
||||
private volatile String lastUserText = "";
|
||||
|
||||
public RuleBasedExtractorModel() {
|
||||
this(Mode.CLEAN);
|
||||
}
|
||||
|
||||
public RuleBasedExtractorModel(Mode mode) {
|
||||
this.mode = mode;
|
||||
}
|
||||
|
||||
public int calls() {
|
||||
return calls.get();
|
||||
}
|
||||
|
||||
public long inputChars() {
|
||||
return inputChars.get();
|
||||
}
|
||||
|
||||
/** The full user message of the most recent call, including the JSON-schema instructions Spring AI appended. */
|
||||
public String lastUserText() {
|
||||
return lastUserText;
|
||||
}
|
||||
|
||||
@Override
|
||||
public ChatResponse call(Prompt prompt) {
|
||||
calls.incrementAndGet();
|
||||
String user = prompt.getUserMessage().getText();
|
||||
lastUserText = user;
|
||||
inputChars.addAndGet(prompt.getContents().length());
|
||||
int start = user.lastIndexOf("DOCUMENT_START");
|
||||
int end = user.lastIndexOf("DOCUMENT_END");
|
||||
String document = user.substring(start + "DOCUMENT_START".length(), end).strip();
|
||||
String json = render(extract(document));
|
||||
return new ChatResponse(List.of(new Generation(new AssistantMessage(json))));
|
||||
}
|
||||
|
||||
// ---- the "model" ----------------------------------------------------------------------------------
|
||||
|
||||
private record Ent(String name, EntityType type) {
|
||||
}
|
||||
|
||||
private record Rel(String source, String type, String target) {
|
||||
}
|
||||
|
||||
private record Found(Set<Ent> entities, List<Rel> relations) {
|
||||
}
|
||||
|
||||
private Found extract(String document) {
|
||||
String[] sentences = document.split("(?<=\\.)\\s+");
|
||||
String title = sentences[0].replaceAll("\\.$", "");
|
||||
Set<Ent> entities = new LinkedHashSet<>();
|
||||
List<Rel> relations = new ArrayList<>();
|
||||
for (String s : sentences) {
|
||||
s = s.replaceAll("\\.$", "");
|
||||
for (RelationType type : RelationType.values()) {
|
||||
int at = s.indexOf(" " + type.phrase() + " ");
|
||||
if (at < 0) {
|
||||
continue;
|
||||
}
|
||||
String subject = subject(s.substring(0, at), title);
|
||||
if (subject == null) {
|
||||
continue;
|
||||
}
|
||||
List<String> objects = objects(type, s.substring(at));
|
||||
for (String object : objects) {
|
||||
entities.add(new Ent(subject, subjectType(type, subject)));
|
||||
entities.add(new Ent(object, objectType(type)));
|
||||
relations.add(new Rel(subject, type.name(), object));
|
||||
}
|
||||
}
|
||||
}
|
||||
return new Found(entities, relations);
|
||||
}
|
||||
|
||||
private static String subject(String before, String title) {
|
||||
String b = before.strip();
|
||||
if (b.equals("It") || b.equals("She") || b.equals("He")) {
|
||||
return title;
|
||||
}
|
||||
Matcher m = FIRST_NAME.matcher(b);
|
||||
return m.matches() ? m.group(1) : null;
|
||||
}
|
||||
|
||||
private static List<String> objects(RelationType type, String fromVerb) {
|
||||
List<String> out = new ArrayList<>();
|
||||
if (type == RelationType.DEPENDS_ON) {
|
||||
Matcher m = DEPENDS.matcher(fromVerb);
|
||||
while (m.find()) {
|
||||
out.add(m.group(1));
|
||||
}
|
||||
} else {
|
||||
Matcher m = FIRST_NAME.matcher(fromVerb.substring((" " + type.phrase()).length()));
|
||||
if (m.find()) {
|
||||
out.add(m.group(1));
|
||||
}
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
private static EntityType subjectType(RelationType type, String subject) {
|
||||
return switch (type) {
|
||||
case DEPENDS_ON, OWNED_BY, HOSTED_ON -> EntityType.SYSTEM;
|
||||
case LED_BY -> EntityType.TEAM;
|
||||
case AFFECTED, CAUSED_BY -> EntityType.INCIDENT;
|
||||
case REPORTS_TO -> EntityType.PERSON;
|
||||
case LOCATED_IN -> subject.endsWith("Team") ? EntityType.TEAM : EntityType.VENDOR;
|
||||
};
|
||||
}
|
||||
|
||||
private static EntityType objectType(RelationType type) {
|
||||
return switch (type) {
|
||||
case DEPENDS_ON, AFFECTED, CAUSED_BY -> EntityType.SYSTEM;
|
||||
case OWNED_BY -> EntityType.TEAM;
|
||||
case LED_BY, REPORTS_TO -> EntityType.PERSON;
|
||||
case HOSTED_ON -> EntityType.VENDOR;
|
||||
case LOCATED_IN -> EntityType.PLACE;
|
||||
};
|
||||
}
|
||||
|
||||
private String render(Found found) {
|
||||
List<Map<String, String>> ents = new ArrayList<>();
|
||||
for (Ent e : found.entities()) {
|
||||
ents.add(Map.of("name", e.name(), "type", e.type().name()));
|
||||
}
|
||||
List<Map<String, String>> rels = new ArrayList<>();
|
||||
for (Rel r : found.relations()) {
|
||||
Map<String, String> m = new LinkedHashMap<>();
|
||||
m.put("source", r.source());
|
||||
m.put("type", r.type());
|
||||
m.put("target", r.target());
|
||||
rels.add(m);
|
||||
}
|
||||
switch (mode) {
|
||||
case UNKNOWN_RELATION_TYPE -> {
|
||||
if (!rels.isEmpty()) {
|
||||
rels.get(0).put("type", "IS_RESPONSIBLE_FOR");
|
||||
}
|
||||
}
|
||||
case INJECTION_RELATION_TYPE -> {
|
||||
if (!rels.isEmpty()) {
|
||||
rels.get(0).put("type", "OWNED_BY]->(x) DETACH DELETE x //");
|
||||
}
|
||||
}
|
||||
case DANGLING_RELATION -> rels.add(Map.of("source", rels.isEmpty() ? "Nobody" : rels.get(0).get("source"),
|
||||
"type", "OWNED_BY", "target", "Ghost Team"));
|
||||
default -> {
|
||||
}
|
||||
}
|
||||
String json = JSON.writeValueAsString(Map.of("entities", ents, "relations", rels));
|
||||
return switch (mode) {
|
||||
case FENCED -> "```json\n" + json + "\n```";
|
||||
case PROSE_THEN_FENCE -> "Here is the extraction you asked for:\n```json\n" + json + "\n```";
|
||||
case TRUNCATED -> json.substring(0, Math.max(1, json.length() * 2 / 3));
|
||||
default -> json;
|
||||
};
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
Reading the checkout latency dashboard
|
||||
The checkout latency dashboard shows p50, p95 and p99 per region. A red band on the p99 panel usually means the Payments Gateway is slow rather than the Checkout Service itself. Compare against the Inventory Service panel before paging anyone.
|
||||
@@ -0,0 +1,2 @@
|
||||
Orders Database
|
||||
The Orders Database is a PostgreSQL 16 cluster that stores orders and ledger entries. The Orders Database is owned by the Core Data Team. It is hosted on Borealis Hosting. Nightly backups are kept for 35 days.
|
||||
@@ -0,0 +1,2 @@
|
||||
FAQ: why does the stock count look wrong?
|
||||
Stock counts lag the warehouse by up to five minutes. Reserved stock is subtracted immediately at checkout. If a count is negative, file a ticket with the Fulfilment Team queue.
|
||||
@@ -0,0 +1,2 @@
|
||||
Incident process
|
||||
Incidents are numbered INC-NNNN and graded from severity one to four. A write-up is due within five working days. Write-ups are blameless and name systems, never individuals.
|
||||
@@ -0,0 +1,2 @@
|
||||
Month-end reconciliation guide
|
||||
Reconciliation compares settled amounts in the Ledger Service with the bank statements. Differences under one cent are rounded away. Anything larger is raised with the finance operations desk.
|
||||
@@ -0,0 +1,2 @@
|
||||
On-call rotation guide
|
||||
Every team keeps a weekly rotation with a primary and a secondary. Handover happens on Monday mornings. A page that is not acknowledged in ten minutes escalates to the secondary.
|
||||
@@ -0,0 +1,2 @@
|
||||
Warehouse scanner setup
|
||||
Scanners connect over the warehouse wifi and report counts in batches. Replace the battery every eighteen months. The Warehouse Sync job tolerates a scanner being offline for a day.
|
||||
@@ -0,0 +1,2 @@
|
||||
INC-2041
|
||||
INC-2041 was a 47-minute outage on 14 March. INC-2041 affected the Ledger Service. INC-2041 was caused by the Orders Database. A disk-full alert had been silenced two weeks earlier.
|
||||
@@ -0,0 +1,2 @@
|
||||
INC-2057
|
||||
INC-2057 was a stock-count drift that lasted a weekend. INC-2057 affected the Inventory Service. INC-2057 was caused by the Warehouse Sync job. Counts were corrected by hand on Monday.
|
||||
@@ -0,0 +1,2 @@
|
||||
INC-2060
|
||||
INC-2060 was a 12-minute slowdown on card payments. INC-2060 affected the Card Gateway. The Card Gateway is the on-call channel's name for the Payments Gateway. A retry storm after a network blip was the trigger. The write-up recommends capping retries.
|
||||
@@ -0,0 +1,2 @@
|
||||
Warehouse Sync
|
||||
The Warehouse Sync job pulls stock counts from every warehouse scanner every five minutes. The Warehouse Sync job is owned by the Fulfilment Team. It is hosted on Borealis Hosting.
|
||||
@@ -0,0 +1,2 @@
|
||||
Memo: cloud cost review
|
||||
Spend on Arcadia Cloud grew eleven per cent last quarter, mostly in the Notification Service and the Inventory Service. The Platform Team will propose a tagging scheme.
|
||||
@@ -0,0 +1,2 @@
|
||||
Memo: hosting contract renewal
|
||||
The colocation contract with Borealis Hosting renews in the autumn. Finance wants a quote for a second facility. Arcadia Cloud was considered for the databases and rejected for latency reasons.
|
||||
@@ -0,0 +1,2 @@
|
||||
Storefront redesign launch notes
|
||||
The new Checkout Service flow went to ten per cent of customers on Tuesday. The Storefront Team saw a small rise in abandoned carts on narrow phones and fixed the layout the same day.
|
||||
@@ -0,0 +1,2 @@
|
||||
Payments Team offsite notes
|
||||
The Payments Team offsite covered refunds, chargebacks and a proposal to move reporting off the Ledger Service. Nothing was decided. Notes from the whiteboard session are in the team wiki.
|
||||
@@ -0,0 +1,2 @@
|
||||
Engineering onboarding
|
||||
New engineers get a laptop, repository access and a buddy on day one. In the first month they shadow an on-call shift and ship one small change. Questions go to the engineering channel.
|
||||
@@ -0,0 +1,2 @@
|
||||
Elena Petrova
|
||||
Elena Petrova moved from design to engineering management in 2022. Elena Petrova reports to Kwame Mensah. She owns the storefront roadmap.
|
||||
@@ -0,0 +1,2 @@
|
||||
Marcus Webb
|
||||
Marcus Webb has worked on the company's databases since 2018. Marcus Webb reports to Dana Okafor. He chairs the schema review board.
|
||||
@@ -0,0 +1,2 @@
|
||||
Priya Raman
|
||||
Priya Raman joined the company in 2021 as a payments engineer. Priya Raman reports to Dana Okafor. She is the named contact for card network certificate renewals.
|
||||
@@ -0,0 +1,2 @@
|
||||
Tomasz Nowak
|
||||
Tomasz Nowak set up the first warehouse scanners. Tomasz Nowak reports to Kwame Mensah. He runs the fulfilment on-call rotation.
|
||||
@@ -0,0 +1,2 @@
|
||||
Card data handling policy
|
||||
Card numbers must never be written to logs. Only the Payments Gateway may see raw card data; the Checkout Service passes opaque tokens. Annual PCI reviews are scheduled in the autumn.
|
||||
@@ -0,0 +1,2 @@
|
||||
Runbook: restoring the Orders Database from backup
|
||||
Restore the latest nightly backup to a scratch cluster first and check row counts. Replay the write-ahead log to the minute you need. Promote the scratch cluster only after the Ledger Service owners have signed off.
|
||||
@@ -0,0 +1,2 @@
|
||||
Runbook: Payments Gateway timeouts
|
||||
When the Payments Gateway starts timing out, check the card network status pages first. Raise the read timeout only as a stopgap and tell the Payments Team channel. Never retry a charge without an idempotency key.
|
||||
@@ -0,0 +1,2 @@
|
||||
Checkout Service
|
||||
The Checkout Service takes a customer's cart and turns it into an order. It depends on the Payments Gateway to take card payments and on the Inventory Service to reserve stock. The Checkout Service is owned by the Storefront Team. Its latency dashboard lives in the shared Grafana folder.
|
||||
@@ -0,0 +1,2 @@
|
||||
Inventory Service
|
||||
The Inventory Service keeps stock counts per warehouse. It depends on the Warehouse Sync job to receive counts. The Inventory Service is owned by the Fulfilment Team. It is hosted on Arcadia Cloud.
|
||||
@@ -0,0 +1,2 @@
|
||||
Ledger Service
|
||||
The Ledger Service is the system of record for settled money movements. It depends on the Orders Database for storage. The Ledger Service is owned by the Core Data Team. Month-end reconciliation runs against it on the third working day.
|
||||
@@ -0,0 +1,2 @@
|
||||
Notification Service
|
||||
The Notification Service sends order emails and shipping texts. It depends on the Mailer Relay for outbound mail. The Notification Service is owned by the Platform Team. It is hosted on Arcadia Cloud.
|
||||
@@ -0,0 +1,2 @@
|
||||
Payments Gateway
|
||||
The Payments Gateway talks to the card networks on behalf of every other service. The Payments Gateway depends on the Ledger Service to record each authorised charge. The Payments Gateway is owned by the Payments Team. It is hosted on Arcadia Cloud.
|
||||
@@ -0,0 +1,2 @@
|
||||
Core Data Team
|
||||
The Core Data Team looks after the shared databases and the ledger. The Core Data Team is led by Marcus Webb. The Core Data Team is based in Austin. The team keeps a public changelog for schema changes.
|
||||
@@ -0,0 +1,2 @@
|
||||
Fulfilment Team
|
||||
The Fulfilment Team connects the shop to the warehouses. The Fulfilment Team is led by Tomasz Nowak. The Fulfilment Team is based in Gdansk.
|
||||
@@ -0,0 +1,2 @@
|
||||
Payments Team
|
||||
The Payments Team looks after card acceptance and refunds. The Payments Team is led by Priya Raman. The Payments Team is based in Pune. The team holds a weekly review on Thursdays.
|
||||
@@ -0,0 +1,2 @@
|
||||
Platform Team
|
||||
The Platform Team runs the build system and the shared tooling. The Platform Team is led by Aiko Tanaka. The Platform Team is based in Osaka.
|
||||
@@ -0,0 +1,2 @@
|
||||
Storefront Team
|
||||
The Storefront Team builds the web and mobile shopping experience. The Storefront Team is led by Elena Petrova. The Storefront Team is based in Lisbon.
|
||||
@@ -0,0 +1,2 @@
|
||||
Arcadia Cloud
|
||||
Arcadia Cloud is the public cloud used for stateless services. Arcadia Cloud is located in Dublin. Reserved capacity is renegotiated each spring.
|
||||
@@ -0,0 +1,2 @@
|
||||
Borealis Hosting
|
||||
Borealis Hosting runs the colocation facility used for databases and batch jobs. Borealis Hosting is located in Frankfurt. Its support line is staffed around the clock.
|
||||
@@ -0,0 +1,212 @@
|
||||
package com.ankurm.graphrag;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.LinkedHashSet;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.Set;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
import com.ankurm.graphrag.support.GraphFixture;
|
||||
import com.ankurm.graphrag.support.Transcript;
|
||||
import org.junit.jupiter.api.BeforeAll;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.springframework.ai.chat.messages.AssistantMessage;
|
||||
import org.springframework.ai.chat.model.ChatResponse;
|
||||
import org.springframework.ai.chat.model.Generation;
|
||||
|
||||
/** Plain vector RAG against GraphRAG on the same 35 chunks and the same 9 questions. */
|
||||
class ComparisonTest {
|
||||
|
||||
static final int DEPTH = 3;
|
||||
static final int TOP_PATHS = 6;
|
||||
static final double LINK = 0.80;
|
||||
|
||||
static Pipeline p;
|
||||
static List<Question> questions = Question.all();
|
||||
|
||||
@BeforeAll
|
||||
static void build() {
|
||||
p = GraphFixture.pipeline();
|
||||
p.buildGraph(GraphFixture.docs(), GraphFixture.ALIASES, 2.0);
|
||||
}
|
||||
|
||||
private static boolean complete(Question q, Set<String> got) {
|
||||
return got.containsAll(q.evidence());
|
||||
}
|
||||
|
||||
private static String mark(boolean ok) {
|
||||
return ok ? "yes" : "NO ";
|
||||
}
|
||||
|
||||
@Test
|
||||
void vectorAgainstGraph() {
|
||||
GraphRetriever graph = new GraphRetriever(p.graph(), MiniLm.get(), LINK);
|
||||
try (Transcript t = new Transcript("04-vector-vs-graph.txt", "Plain vector RAG against GraphRAG: is every supporting document in the context?")) {
|
||||
t.line("corpus: 35 chunks (one document each); 9 questions; evidence = the documents that together hold every fact needed");
|
||||
t.line("vector: Neo4jVectorStore, top-k chunks. graph: link spans (cosine >= %.2f), walk <= %d hops, keep the %d best-ranked paths, take the chunks that stated their edges", LINK, DEPTH, TOP_PATHS);
|
||||
t.blank();
|
||||
t.line("%-3s %-4s | %-9s %-9s | %-9s %-7s | question", "id", "hops", "vector k4", "vector k8", "graph", "chunks");
|
||||
int v4 = 0;
|
||||
int v8 = 0;
|
||||
int g = 0;
|
||||
int gChunks = 0;
|
||||
int[][] byHop = new int[4][4];
|
||||
int[] hopTotal = new int[4];
|
||||
for (Question q : questions) {
|
||||
boolean ok4 = complete(q, new LinkedHashSet<>(p.vectorSearch(q.text(), 4)));
|
||||
boolean ok8 = complete(q, new LinkedHashSet<>(p.vectorSearch(q.text(), 8)));
|
||||
GraphRetriever.Retrieval r = graph.retrieve(q.text(), DEPTH, TOP_PATHS);
|
||||
boolean okG = complete(q, r.chunkIds());
|
||||
v4 += ok4 ? 1 : 0;
|
||||
v8 += ok8 ? 1 : 0;
|
||||
g += okG ? 1 : 0;
|
||||
gChunks += r.chunkIds().size();
|
||||
hopTotal[q.hops()]++;
|
||||
byHop[q.hops()][0] += ok4 ? 1 : 0;
|
||||
byHop[q.hops()][1] += ok8 ? 1 : 0;
|
||||
byHop[q.hops()][2] += okG ? 1 : 0;
|
||||
t.line("%-3s %-4d | %-9s %-9s | %-9s %-7d | %s", q.id(), q.hops(), mark(ok4), mark(ok8), mark(okG), r.chunkIds().size(), q.text());
|
||||
}
|
||||
t.blank();
|
||||
t.line("complete evidence: vector k=4 %d/9, vector k=8 %d/9, graph %d/9 (graph context averaged %.1f chunks; vector k=8 is always 8)", v4, v8, g, gChunks / 9.0);
|
||||
for (int h = 1; h <= 3; h++) {
|
||||
t.line(" %d-hop questions (%d): vector k=4 %d, vector k=8 %d, graph %d", h, hopTotal[h], byHop[h][0], byHop[h][1], byHop[h][2]);
|
||||
}
|
||||
t.section("what each retriever missed");
|
||||
for (Question q : questions) {
|
||||
Set<String> v = new LinkedHashSet<>(p.vectorSearch(q.text(), 8));
|
||||
Set<String> gr = graph.retrieve(q.text(), DEPTH, TOP_PATHS).chunkIds();
|
||||
Set<String> mv = new LinkedHashSet<>(q.evidence());
|
||||
mv.removeAll(v);
|
||||
Set<String> mg = new LinkedHashSet<>(q.evidence());
|
||||
mg.removeAll(gr);
|
||||
t.line("%s vector k=8 missing %s graph missing %s", q.id(), mv, mg);
|
||||
}
|
||||
assertThat(v4).isEqualTo(2);
|
||||
assertThat(v8).isEqualTo(5);
|
||||
assertThat(g).isEqualTo(7);
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void sweeps() {
|
||||
try (Transcript t = new Transcript("05-graph-sweeps.txt", "How the graph retriever's own settings move the result (complete evidence out of 9)")) {
|
||||
t.section("hops walked x paths kept (link threshold 0.80)");
|
||||
t.line("%-9s | %-12s %-12s %-12s %-12s", "max hops", "6 paths", "10 paths", "20 paths", "40 paths");
|
||||
GraphRetriever graph = new GraphRetriever(p.graph(), MiniLm.get(), LINK);
|
||||
for (int depth = 1; depth <= 3; depth++) {
|
||||
StringBuilder row = new StringBuilder();
|
||||
for (int paths : new int[] {6, 10, 20, 40}) {
|
||||
int ok = 0;
|
||||
int chunks = 0;
|
||||
for (Question q : questions) {
|
||||
GraphRetriever.Retrieval r = graph.retrieve(q.text(), depth, paths);
|
||||
ok += complete(q, r.chunkIds()) ? 1 : 0;
|
||||
chunks += r.chunkIds().size();
|
||||
}
|
||||
if (depth == 3 && paths == 20) {
|
||||
assertThat(ok).isEqualTo(9);
|
||||
}
|
||||
if (depth == 1) {
|
||||
assertThat(ok).isEqualTo(3);
|
||||
}
|
||||
row.append(String.format("%d/9 (%4.1f ch) ", ok, chunks / 9.0));
|
||||
}
|
||||
t.line("%-9d | %s", depth, row);
|
||||
}
|
||||
t.section("entity-link threshold (3 hops, 6 paths)");
|
||||
t.line("%-9s | %-9s %s", "threshold", "complete", "seeds per question (avg)");
|
||||
for (double th : new double[] {0.70, 0.80, 0.85, 0.90, 0.95}) {
|
||||
GraphRetriever gr = new GraphRetriever(p.graph(), MiniLm.get(), th);
|
||||
int ok = 0;
|
||||
int seeds = 0;
|
||||
for (Question q : questions) {
|
||||
GraphRetriever.Retrieval r = gr.retrieve(q.text(), DEPTH, TOP_PATHS);
|
||||
ok += complete(q, r.chunkIds()) ? 1 : 0;
|
||||
seeds += r.seeds().size();
|
||||
}
|
||||
t.line("%-9.2f | %d/9 %.1f", th, ok, seeds / 9.0);
|
||||
}
|
||||
t.section("the seeds found for three questions at 0.80 (span -> entity, cosine)");
|
||||
for (String id : List.of("Q2", "Q5", "Q7")) {
|
||||
Question q = questions.stream().filter(x -> x.id().equals(id)).findFirst().orElseThrow();
|
||||
t.line("%s %s", id, q.text());
|
||||
graph.link(q.text()).forEach(s -> t.line(" '%s' -> %s (%.2f)", s.span(), s.name(), s.score()));
|
||||
}
|
||||
t.section("candidate paths from one seed as the hop limit grows");
|
||||
for (String seed : List.of("Payments Gateway", "Orders Database", "Checkout Service")) {
|
||||
StringBuilder sb = new StringBuilder();
|
||||
for (int h = 1; h <= 4; h++) {
|
||||
sb.append(String.format("%d hop%s: %d ", h, h == 1 ? " " : "s", p.graph().paths(EntityResolver.key(seed), h, 100_000).size()));
|
||||
}
|
||||
t.line("%-17s %s", seed, sb);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void whatTheModelWouldReceive() {
|
||||
List<String> seen = new ArrayList<>();
|
||||
Answerer answerer = new Answerer(new org.springframework.ai.chat.model.ChatModel() {
|
||||
@Override
|
||||
public ChatResponse call(org.springframework.ai.chat.prompt.Prompt prompt) {
|
||||
seen.add(prompt.getContents());
|
||||
return new ChatResponse(List.of(new Generation(new AssistantMessage("(no model attached: this reply is a placeholder)"))));
|
||||
}
|
||||
});
|
||||
GraphRetriever graph = new GraphRetriever(p.graph(), MiniLm.get(), LINK);
|
||||
Question q = questions.stream().filter(x -> x.id().equals("Q4")).findFirst().orElseThrow();
|
||||
GraphRetriever.Retrieval r = graph.retrieve(q.text(), DEPTH, TOP_PATHS);
|
||||
Map<String, String> texts = p.graph().chunkTexts(r.chunkIds());
|
||||
answerer.answer(r, texts);
|
||||
try (Transcript t = new Transcript("06-prompt.txt", "The prompt GraphRAG assembles for Q4 (no model was called; nothing below is a model answer)")) {
|
||||
t.line("%s", Answerer.prompt(r, texts));
|
||||
t.blank();
|
||||
t.line("characters in this prompt: %d; characters in the 8 chunks a k=8 vector search would add instead: %d",
|
||||
Answerer.prompt(r, texts).length(),
|
||||
p.chunks().similaritySearch(org.springframework.ai.vectorstore.SearchRequest.builder().query(q.text()).topK(8).build())
|
||||
.stream().mapToInt(d -> d.getText().length()).sum());
|
||||
}
|
||||
assertThat(seen).hasSize(1);
|
||||
}
|
||||
|
||||
@Test
|
||||
void whenYouKnowTheShapeWriteTheCypher() {
|
||||
try (Transcript t = new Transcript("07-fixed-shape-cypher.txt", "Questions whose shape is known: the Cypher is the retrieval (hand-written, parameterised, run on the real graph)")) {
|
||||
String q7 = """
|
||||
MATCH (db:Entity {name: $name})<-[:DEPENDS_ON*1..5]-(s:Entity)
|
||||
RETURN DISTINCT s.name AS service ORDER BY service""";
|
||||
t.section("Q7 Which services stop working if the Orders Database goes down?");
|
||||
t.line("%s", q7);
|
||||
List<String> rows = p.graph().query(q7, Map.of("name", "Orders Database")).stream().map(r -> r.get("service").asString()).toList();
|
||||
rows.forEach(r -> t.line(" -> %s", r));
|
||||
|
||||
String q6 = """
|
||||
MATCH (c:Entity {name: $name})-[:DEPENDS_ON]->(d:Entity)-[:OWNED_BY]->(team:Entity)-[:LED_BY]->(lead:Entity)
|
||||
RETURN d.name AS dependency, team.name AS team, lead.name AS lead ORDER BY dependency""";
|
||||
t.section("Q6 Who leads the teams that own the services the Checkout Service depends on?");
|
||||
t.line("%s", q6);
|
||||
List<String> leads = new ArrayList<>();
|
||||
p.graph().query(q6, Map.of("name", "Checkout Service")).forEach(r -> {
|
||||
t.line(" -> %s / %s / %s", r.get("dependency").asString(), r.get("team").asString(), r.get("lead").asString());
|
||||
leads.add(r.get("lead").asString());
|
||||
});
|
||||
|
||||
String q8 = """
|
||||
CALL db.index.vector.queryNodes('entity_name_embedding', 1, $vector) YIELD node AS c, score
|
||||
MATCH (c)-[:DEPENDS_ON]->(d:Entity)<-[:AFFECTED]-(i:Entity)
|
||||
RETURN c.name AS seed, score, d.name AS dependency, i.name AS incident ORDER BY incident""";
|
||||
t.section("Q8 entity linking and the traversal in ONE statement: the vector index finds the seed, MATCH walks from it");
|
||||
t.line("%s", q8);
|
||||
t.line(" $vector = embedding of the span \"Checkout Service\"");
|
||||
p.graph().query(q8, Map.of("vector", MiniLm.get().embed("Checkout Service"))).forEach(r ->
|
||||
t.line(" -> seed %s (%.2f): %s affected %s", r.get("seed").asString(), r.get("score").asDouble(), r.get("incident").asString(), r.get("dependency").asString()));
|
||||
|
||||
assertThat(rows).containsExactly("Checkout Service", "Ledger Service", "Payments Gateway");
|
||||
assertThat(leads).containsExactlyInAnyOrder("Priya Raman", "Tomasz Nowak");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,111 @@
|
||||
package com.ankurm.graphrag;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.Assertions.assertThatThrownBy;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
import com.ankurm.graphrag.stub.RuleBasedExtractorModel;
|
||||
import com.ankurm.graphrag.stub.RuleBasedExtractorModel.Mode;
|
||||
import com.ankurm.graphrag.support.Transcript;
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
/** Step 1 with no database: the prompt, the schema, the parsed record, and what bad replies do. */
|
||||
class ExtractionTest {
|
||||
|
||||
private static Corpus.Doc doc(String id) {
|
||||
return Corpus.load().stream().filter(d -> d.id().equals(id)).findFirst().orElseThrow();
|
||||
}
|
||||
|
||||
@Test
|
||||
void whatTheModelSeesAndWhatComesBack() {
|
||||
RuleBasedExtractorModel stub = new RuleBasedExtractorModel();
|
||||
Extractor extractor = new Extractor(stub);
|
||||
try (Transcript t = new Transcript("01-extraction.txt", "Entity and relation extraction through ChatClient.entity() (STUB model, not an LLM)")) {
|
||||
Extractor.Result r = extractor.extract(doc("svc-checkout"));
|
||||
t.section("the user message the model receives for svc-checkout.txt");
|
||||
t.line(stub.lastUserText().strip());
|
||||
t.section("the parsed Extraction record");
|
||||
r.extraction().entities().forEach(e -> t.line("entity %-8s %s", e.type(), e.name()));
|
||||
r.extraction().relations().forEach(x -> t.line("relation %s -[%s]-> %s", x.source(), x.type(), x.target()));
|
||||
assertThat(r.extraction().relations()).hasSize(3);
|
||||
|
||||
RuleBasedExtractorModel all = new RuleBasedExtractorModel();
|
||||
Extractor full = new Extractor(all);
|
||||
int entities = 0;
|
||||
int relations = 0;
|
||||
int empty = 0;
|
||||
for (Corpus.Doc d : Corpus.load()) {
|
||||
Extraction x = full.extract(d).extraction();
|
||||
entities += x.entities().size();
|
||||
relations += x.relations().size();
|
||||
if (x.relations().isEmpty()) {
|
||||
empty++;
|
||||
}
|
||||
}
|
||||
t.section("all 35 chunks");
|
||||
t.line("model calls : %d (one per chunk)", all.calls());
|
||||
t.line("characters sent : %d (about %d tokens at 4 characters per token, an estimate)", all.inputChars(), all.inputChars() / 4);
|
||||
t.line("entity mentions (before merging): %d", entities);
|
||||
t.line("relations : %d", relations);
|
||||
t.line("chunks with no relation at all : %d", empty);
|
||||
assertThat(all.calls()).isEqualTo(35);
|
||||
assertThat(empty).isEqualTo(14);
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void badRepliesAndWhatTheyDo() {
|
||||
try (Transcript t = new Transcript("02-extraction-failures.txt", "What a bad model reply does to extraction (replies produced by the STUB's failure modes)")) {
|
||||
Corpus.Doc d = doc("svc-ledger");
|
||||
|
||||
Extractor.Result fenced = new Extractor(new RuleBasedExtractorModel(Mode.FENCED)).extract(d);
|
||||
t.section("JSON wrapped in a markdown fence (reply starts with the fence)");
|
||||
t.line("parsed OK, %d relations", fenced.extraction().relations().size());
|
||||
|
||||
t.section("the same fence, after one sentence of chat");
|
||||
Throwable prose = catchThrowable(() -> new Extractor(new RuleBasedExtractorModel(Mode.PROSE_THEN_FENCE)).extract(d));
|
||||
t.line("%s", describe(prose));
|
||||
assertThat(prose).isNotNull();
|
||||
|
||||
t.section("JSON cut off two thirds of the way through");
|
||||
Throwable truncated = catchThrowable(() -> new Extractor(new RuleBasedExtractorModel(Mode.TRUNCATED)).extract(d));
|
||||
t.line("%s", describe(truncated));
|
||||
|
||||
t.section("a relation type outside the allowed list (IS_RESPONSIBLE_FOR)");
|
||||
Throwable unknown = catchThrowable(() -> new Extractor(new RuleBasedExtractorModel(Mode.UNKNOWN_RELATION_TYPE)).extract(d));
|
||||
t.line("%s", describe(unknown));
|
||||
|
||||
t.section("a relation type that is a Cypher fragment");
|
||||
Throwable injected = catchThrowable(() -> new Extractor(new RuleBasedExtractorModel(Mode.INJECTION_RELATION_TYPE)).extract(d));
|
||||
t.line("%s", describe(injected));
|
||||
|
||||
t.section("a relation to an entity that was never listed");
|
||||
Extractor.Result dangling = new Extractor(new RuleBasedExtractorModel(Mode.DANGLING_RELATION)).extract(d);
|
||||
t.line("relations kept %d, dropped %d", dangling.extraction().relations().size(), dangling.droppedRelations());
|
||||
|
||||
assertThat(fenced.extraction().relations()).hasSize(2);
|
||||
assertThat(truncated).isNotNull();
|
||||
assertThat(unknown).isNotNull();
|
||||
assertThat(injected).isNotNull();
|
||||
assertThat(dangling.droppedRelations()).isEqualTo(1);
|
||||
}
|
||||
}
|
||||
|
||||
private static Throwable catchThrowable(Runnable r) {
|
||||
try {
|
||||
r.run();
|
||||
return null;
|
||||
} catch (Throwable e) {
|
||||
return e;
|
||||
}
|
||||
}
|
||||
|
||||
private static String describe(Throwable e) {
|
||||
if (e == null) {
|
||||
return "no exception";
|
||||
}
|
||||
String message = e.getMessage() == null ? "" : e.getMessage().lines().findFirst().orElse("");
|
||||
return e.getClass().getName() + ": " + (message.length() > 220 ? message.substring(0, 220) + "..." : message);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,107 @@
|
||||
package com.ankurm.graphrag;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
import com.ankurm.graphrag.support.GraphFixture;
|
||||
import com.ankurm.graphrag.support.Transcript;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.neo4j.driver.Record;
|
||||
|
||||
/** Steps 2 and 3 against a real Neo4j in a container: name merging, the schema, and what ends up in the graph. */
|
||||
class GraphBuildTest {
|
||||
|
||||
private static long count(Pipeline p, String cypher) {
|
||||
return p.graph().query(cypher, Map.of()).get(0).get(0).asLong();
|
||||
}
|
||||
|
||||
private static String neighbours(Pipeline p, String name) {
|
||||
return p.graph().query("""
|
||||
MATCH (e:Entity {name: $n})-[r]-(o:Entity)
|
||||
RETURN CASE WHEN startNode(r) = e THEN '-[' + type(r) + ']-> ' + o.name ELSE '<-[' + type(r) + ']- ' + o.name END AS s
|
||||
ORDER BY s
|
||||
""", Map.of("n", name)).stream().map(r -> r.get("s").asString()).collect(Collectors.joining("; "));
|
||||
}
|
||||
|
||||
@Test
|
||||
void graphAndTheNameProblem() {
|
||||
Pipeline p = GraphFixture.pipeline();
|
||||
List<Corpus.Doc> docs = GraphFixture.docs();
|
||||
try (Transcript t = new Transcript("03-graph-build.txt", "Building the graph in Neo4j; names that are the same thing (STUB extraction)")) {
|
||||
Record version = p.graph().query("CALL dbms.components() YIELD name, versions, edition RETURN name, versions[0] AS v, edition", Map.of()).get(0);
|
||||
t.line("image : %s", GraphFixture.image());
|
||||
t.line("server : %s %s %s", version.get("name").asString(), version.get("v").asString(), version.get("edition").asString());
|
||||
t.section("vector indexes (one from Neo4jVectorStore for chunks, one of ours for entity names)");
|
||||
p.graph().query("SHOW INDEXES YIELD name, type, labelsOrTypes, properties, state WHERE type = 'VECTOR' RETURN name, labelsOrTypes[0] AS label, properties[0] AS property, state ORDER BY name", Map.of())
|
||||
.forEach(r -> t.line("%-22s on (:%s).%s %s", r.get("name").asString(), r.get("label").asString(), r.get("property").asString(), r.get("state").asString()));
|
||||
|
||||
// A: nothing merged
|
||||
Pipeline.Built a = p.buildGraph(docs, Map.of(), 2.0);
|
||||
long entitiesA = count(p, "MATCH (e:Entity) RETURN count(e)");
|
||||
t.section("A. exact-name matching only (no alias table, similarity merge off)");
|
||||
t.line("entities: %d", entitiesA);
|
||||
t.line("Payments Gateway : %s", neighbours(p, "Payments Gateway"));
|
||||
t.line("Card Gateway : %s", neighbours(p, "Card Gateway"));
|
||||
|
||||
t.section("same-type name pairs by cosine of their embeddings, top 8 of " + pairs(p, a).size());
|
||||
List<EntityResolver.Pair> pairs = pairs(p, a);
|
||||
for (int i = 0; i < 8; i++) {
|
||||
EntityResolver.Pair x = pairs.get(i);
|
||||
t.line("%.3f %-7s %s <-> %s", x.cosine(), x.type(), x.a(), x.b());
|
||||
}
|
||||
int rank = 0;
|
||||
for (int i = 0; i < pairs.size(); i++) {
|
||||
EntityResolver.Pair x = pairs.get(i);
|
||||
if (x.a().equals("Card Gateway") || x.b().equals("Card Gateway")) {
|
||||
if (x.a().equals("Payments Gateway") || x.b().equals("Payments Gateway")) {
|
||||
rank = i + 1;
|
||||
t.line("-- Card Gateway <-> Payments Gateway: cosine %.3f, rank %d", x.cosine(), rank);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
t.section("what a similarity-only merge would do at each threshold");
|
||||
for (double th : new double[] {0.90, 0.80, 0.70, 0.60}) {
|
||||
EntityResolver.Resolution r = new EntityResolver(MiniLmHolder.get()).resolve(a.allMentions(), Map.of(), th);
|
||||
t.line("threshold %.2f -> %d merges: %s", th, r.merged().size(),
|
||||
r.merged().stream().map(x -> x.a() + " + " + x.b()).collect(Collectors.joining("; ")));
|
||||
}
|
||||
|
||||
// B: alias table
|
||||
Pipeline.Built b = p.buildGraph(docs, GraphFixture.ALIASES, 2.0);
|
||||
long entitiesB = count(p, "MATCH (e:Entity) RETURN count(e)");
|
||||
t.section("B. plus a one-line alias table (Card Gateway = Payments Gateway)");
|
||||
t.line("entities: %d", entitiesB);
|
||||
t.line("Payments Gateway : %s", neighbours(p, "Payments Gateway"));
|
||||
t.line("chunks: %d, MENTIONS edges: %d, dropped relations: %d", count(p, "MATCH (c:Chunk) RETURN count(c)"),
|
||||
count(p, "MATCH (:Chunk)-[r:MENTIONS]->(:Entity) RETURN count(r)"), b.droppedRelations());
|
||||
t.section("relationships by type");
|
||||
p.graph().query("MATCH (:Entity)-[r]->(:Entity) RETURN type(r) AS t, count(r) AS n ORDER BY n DESC, t", Map.of())
|
||||
.forEach(r -> t.line("%-12s %d", r.get("t").asString(), r.get("n").asLong()));
|
||||
|
||||
assertThat(entitiesA).isEqualTo(entitiesB + 1);
|
||||
assertThat(neighbours(p, "Payments Gateway")).contains("INC-2060");
|
||||
}
|
||||
}
|
||||
|
||||
private static List<EntityResolver.Pair> pairs(Pipeline p, Pipeline.Built built) {
|
||||
EntityResolver er = new EntityResolver(MiniLmHolder.get());
|
||||
Map<String, EntityType> types = new java.util.LinkedHashMap<>();
|
||||
Map<String, String> spelling = new java.util.LinkedHashMap<>();
|
||||
built.allMentions().forEach(e -> {
|
||||
types.putIfAbsent(EntityResolver.key(e.name()), e.type());
|
||||
spelling.putIfAbsent(EntityResolver.key(e.name()), e.name());
|
||||
});
|
||||
return er.similarPairs(List.copyOf(types.keySet()), types, spelling);
|
||||
}
|
||||
|
||||
/** Tiny indirection so the test reads the same embedding model the pipeline uses. */
|
||||
static final class MiniLmHolder {
|
||||
static org.springframework.ai.embedding.EmbeddingModel get() {
|
||||
return MiniLm.get();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
package com.ankurm.graphrag.support;
|
||||
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
import com.ankurm.graphrag.Corpus;
|
||||
import com.ankurm.graphrag.MiniLm;
|
||||
import com.ankurm.graphrag.Pipeline;
|
||||
import com.ankurm.graphrag.stub.RuleBasedExtractorModel;
|
||||
import org.neo4j.driver.AuthTokens;
|
||||
import org.neo4j.driver.Driver;
|
||||
import org.neo4j.driver.GraphDatabase;
|
||||
import org.testcontainers.neo4j.Neo4jContainer;
|
||||
|
||||
/**
|
||||
* One Neo4j container for the whole test run (Testcontainers' reaper removes it when the JVM exits), the
|
||||
* 35 chunks loaded once, and the extraction stub shared so its call counts add up. Set NEO4J_IMAGE to run the
|
||||
* same tests against another Neo4j release.
|
||||
*/
|
||||
public final class GraphFixture {
|
||||
|
||||
public static final String IMAGE = System.getenv().getOrDefault("NEO4J_IMAGE", "neo4j:2026.10.1-community");
|
||||
|
||||
/** The one merge a person has to supply: INC-2060's own text says the on-call channel calls the Payments Gateway the Card Gateway. */
|
||||
public static final Map<String, String> ALIASES = Map.of("Card Gateway", "Payments Gateway");
|
||||
|
||||
private static Neo4jContainer container;
|
||||
private static Driver driver;
|
||||
private static Pipeline pipeline;
|
||||
private static final RuleBasedExtractorModel STUB = new RuleBasedExtractorModel();
|
||||
private static final List<Corpus.Doc> DOCS = Corpus.load();
|
||||
|
||||
private GraphFixture() {
|
||||
}
|
||||
|
||||
public static synchronized Pipeline pipeline() {
|
||||
if (pipeline == null) {
|
||||
container = new Neo4jContainer(IMAGE).withAdminPassword("graphrag-demo");
|
||||
container.start();
|
||||
driver = GraphDatabase.driver(container.getBoltUrl(), AuthTokens.basic("neo4j", container.getAdminPassword()));
|
||||
pipeline = new Pipeline(driver, MiniLm.get(), STUB);
|
||||
pipeline.loadChunks(DOCS);
|
||||
}
|
||||
return pipeline;
|
||||
}
|
||||
|
||||
public static Driver driver() {
|
||||
pipeline();
|
||||
return driver;
|
||||
}
|
||||
|
||||
public static String image() {
|
||||
return IMAGE;
|
||||
}
|
||||
|
||||
public static List<Corpus.Doc> docs() {
|
||||
return DOCS;
|
||||
}
|
||||
|
||||
public static RuleBasedExtractorModel stub() {
|
||||
return STUB;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
package com.ankurm.graphrag.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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user