Text2Cypher & Graph Retrievers

import os

import neo4j
from neo4j_graphrag.embeddings import OllamaEmbeddings
from neo4j_graphrag.generation import GraphRAG
from neo4j_graphrag.llm import OllamaLLM
from neo4j_graphrag.retrievers import VectorCypherRetriever

driver = neo4j.GraphDatabase.driver(os.environ["NEO4J_URI"],
                                    auth=(os.environ["NEO4J_USER"], os.environ["NEO4J_PASSWORD"]))

RETRIEVAL_QUERY = """
MATCH (node)-[:FROM_DOCUMENT]->(d:Document)
WHERE node.deletedAt IS NULL AND d.acl IN $readable_acls
OPTIONAL MATCH (e:__Entity__)-[:FROM_CHUNK]->(node)
OPTIONAL MATCH (e)-[r:DEPENDS_ON|EXPLAINS|IMPLEMENTS|USES]-(n:__Entity__)
WITH node, d, score, collect(DISTINCT e.name + ' ' + type(r) + ' ' + n.name)[..10] AS facts
RETURN node.text AS text, d.source AS source, node.id AS chunk_id, facts, score
ORDER BY score DESC
"""

retriever = VectorCypherRetriever(
    driver,
    index_name="chunk_embeddings",
    retrieval_query=RETRIEVAL_QUERY,
    embedder=OllamaEmbeddings(model=os.environ["EMBEDDING_MODEL"]),
)


def search(question: str, readable_acls: list[str], top_k: int = 8):
    # query_params are passed to the Cypher; the ACL list comes from the verified caller, never from the question
    return retriever.search(query_text=question, top_k=top_k, effective_search_ratio=4,
                            query_params={"readable_acls": readable_acls})


rag = GraphRAG(retriever=retriever, llm=OllamaLLM(model_name=os.environ["CHAT_MODEL"]))

MATCH (c:Component)-[:DEPENDS_ON]->(:Component {name: 'Irurueta Sorting'})
RETURN c.name AS component ORDER BY component

import json
import os
import re

import neo4j
from neo4j_graphrag.exceptions import Text2CypherRetrievalError
from neo4j_graphrag.llm import OllamaLLM
from neo4j_graphrag.retrievers import Text2CypherRetriever
from neo4j_graphrag.retrievers.text2cypher import extract_cypher
from neo4j_graphrag.types import RetrieverResultItem

# a separate driver as the read-only graph user, not the ingestion (writer) driver of retrievers.py
readonly_driver = neo4j.GraphDatabase.driver(
    os.environ["NEO4J_URI"], auth=(os.environ["NEO4J_READONLY_USER"], os.environ["NEO4J_READONLY_PASSWORD"]))

ALLOWED_PROCEDURES = {"db.index.fulltext.querynodes"}            # entity-name lookup; lower case
CALL = re.compile(r"\bCALL\b\s*([{(]|[\w.`]+)?", re.IGNORECASE)


def check_clauses(cypher: str) -> None:
    """EXPLAIN reports LOAD CSV and read procedures such as apoc.load.* as read-only, so refuse them by name."""
    if re.search(r"\bLOAD\b", cypher, re.IGNORECASE):
        raise Text2CypherRetrievalError(f"Refusing LOAD: {cypher}")
    for m in CALL.finditer(cypher):
        target = m.group(1)
        if target in ("{", "("):
            continue                                                # a CALL { ... } subquery
        if target is None or target.replace("`", "").lower() not in ALLOWED_PROCEDURES:
            raise Text2CypherRetrievalError(f"Refusing procedure call: {cypher}")


class ClauseCheckedLLM(OllamaLLM):
    """Checks the generated Cypher before the retriever runs EXPLAIN and the query."""

    def invoke(self, input, *args, **kwargs):
        response = super().invoke(input, *args, **kwargs)
        check_clauses(extract_cypher(response.content))
        return response


SCHEMA = """
Node properties:
Component {name: STRING, aliases: LIST<STRING>}
Concept {name: STRING}
API {name: STRING, kind: STRING}
Document {source: STRING, acl: STRING, updatedAt: DATE_TIME}
Relationships:
(:Component)-[:DEPENDS_ON]->(:Component)
(:Component)-[:EXPLAINS]->(:Concept)
(:API)-[:IMPLEMENTS]->(:Concept)
(:Component)-[:USES]->(:Tool)
"""

EXAMPLES = [
    "USER INPUT: 'Which components depend on sorting?' QUERY: MATCH (c:Component)-[:DEPENDS_ON]->(s:Component) "
    "WHERE s.name = 'Irurueta Sorting' OR 'sorting' IN s.aliases RETURN c.name",
    "USER INPUT: 'How many concepts does the vector-rag section explain?' QUERY: MATCH (:Component {name: "
    "'Vector Databases & RAG'})-[:EXPLAINS]->(k:Concept) RETURN count(k) AS concepts",
]

text2cypher = Text2CypherRetriever(
    driver=readonly_driver,
    llm=ClauseCheckedLLM(model_name=os.environ["CYPHER_MODEL"], model_params={"options": {"temperature": 0}}),
    neo4j_schema=SCHEMA,
    examples=EXAMPLES,
    neo4j_database="docsgraph",
    # keep each row as JSON instead of the default str(record)
    result_formatter=lambda record: RetrieverResultItem(content=json.dumps(dict(record), default=str)),
)


def graph_answer_rows(question: str) -> tuple[list[dict], str]:
    result = text2cypher.search(query_text=question)
    rows = [json.loads(item.content) for item in result.items]
    return rows, result.metadata.get("cypher", "")          # the executed query, for citations and audit

import os

from langchain.chat_models import init_chat_model
from langchain_neo4j import GraphCypherQAChain, Neo4jGraph

graph = Neo4jGraph(url=os.environ["NEO4J_URI"], username=os.environ["NEO4J_READONLY_USER"],
                   password=os.environ["NEO4J_READONLY_PASSWORD"], database="docsgraph",
                   timeout=5.0, enhanced_schema=False)

chain = GraphCypherQAChain.from_llm(
    cypher_llm=init_chat_model(os.environ["CYPHER_MODEL"], temperature=0),
    qa_llm=init_chat_model(os.environ["CHAT_MODEL"], temperature=0),
    graph=graph,
    include_types=["Component", "Concept", "API", "DEPENDS_ON", "EXPLAINS", "IMPLEMENTS"],
    validate_cypher=True,
    top_k=25,
    return_intermediate_steps=True,          # keep the generated Cypher for citations and audit
    allow_dangerous_requests=True,           # required: the model's query runs against the database
)

result = chain.invoke({"query": "Which components depend on the sorting library?"})
cypher = result["intermediate_steps"][0]["query"]

package com.example.docsassistant.l4j.graph;

import java.util.List;

import dev.langchain4j.community.rag.content.retriever.neo4j.Neo4jGraph;
import dev.langchain4j.community.rag.content.retriever.neo4j.Neo4jText2CypherRetriever;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import org.neo4j.driver.Driver;

public final class GraphRetrievers {

    private GraphRetrievers() {
    }

    /** The driver must authenticate as the read-only graph user. */
    public static ContentRetriever text2Cypher(Driver readOnlyDriver, ChatModel cypherModel) {
        Neo4jGraph graph = Neo4jGraph.builder().driver(readOnlyDriver).build();
        return Neo4jText2CypherRetriever.builder()
                .graph(graph)
                .chatModel(cypherModel)
                .examples(List.of(
                        "Which components depend on sorting? -> MATCH (c:Component)-[:DEPENDS_ON]->(s:Component) "
                                + "WHERE s.name = 'Irurueta Sorting' RETURN c.name",
                        "How many concepts does vector-rag explain? -> MATCH (:Component {name: 'Vector Databases & RAG'})"
                                + "-[:EXPLAINS]->(k:Concept) RETURN count(k)"))
                .maxRetries(2)
                .build();
    }
}

package com.example.docsassistant.rag.graph;

import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Set;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

import org.neo4j.driver.Driver;
import org.neo4j.driver.SessionConfig;
import org.neo4j.driver.TransactionConfig;
import org.neo4j.driver.summary.QueryType;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.document.Document;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.retrieval.search.DocumentRetriever;

public class Text2CypherRetriever implements DocumentRetriever {

    record CypherQuery(String cypher) {
    }

    private static final Pattern LOAD = Pattern.compile("\\bLOAD\\b", Pattern.CASE_INSENSITIVE);
    private static final Pattern CALL = Pattern.compile("\\bCALL\\b\\s*([{(]|[\\w.`]+)?", Pattern.CASE_INSENSITIVE);
    private static final Set<String> ALLOWED_PROCEDURES = Set.of("db.index.fulltext.querynodes");   // lower case

    private final ChatClient chatClient;
    private final Driver readOnlyDriver;
    private final String schemaAndExamples;

    public Text2CypherRetriever(ChatClient.Builder builder, Driver readOnlyDriver, String schemaAndExamples) {
        this.chatClient = builder.build();
        this.readOnlyDriver = readOnlyDriver;
        this.schemaAndExamples = schemaAndExamples;
    }

    @Override
    public List<Document> retrieve(Query query) {
        CypherQuery generated = chatClient.prompt()
                .system("Write one read-only Cypher query that answers the question, using only this schema. "
                        + "Return at most 50 rows.\n" + schemaAndExamples)
                .user(query.text())
                .call()
                .entity(CypherQuery.class);
        if (generated == null || generated.cypher() == null) {
            return List.of();
        }
        String cypher = generated.cypher();
        checkClauses(cypher);
        try (var session = readOnlyDriver.session(SessionConfig.forDatabase("docsgraph"))) {
            QueryType type = session.executeRead(tx -> tx.run("EXPLAIN " + cypher).consume().queryType());
            if (type != QueryType.READ_ONLY) {
                throw new IllegalArgumentException("Refusing non-read-only Cypher: " + cypher);
            }
            List<Map<String, Object>> rows = session.executeRead(
                    tx -> tx.run(cypher).list(r -> r.asMap()),
                    TransactionConfig.builder().withTimeout(java.time.Duration.ofSeconds(5)).build());
            String table = rows.stream().map(Map::toString).collect(Collectors.joining("\n"));
            return List.of(Document.builder()
                    .text("Graph query result:\n" + table)
                    .metadata(Map.of("source", "knowledge-graph", "cypher", cypher))
                    .build());
        }
    }

    /** EXPLAIN reports LOAD CSV and read procedures such as apoc.load.* as READ_ONLY, so refuse them by name. */
    static void checkClauses(String cypher) {
        if (LOAD.matcher(cypher).find()) {
            throw new IllegalArgumentException("Refusing LOAD: " + cypher);
        }
        Matcher call = CALL.matcher(cypher);
        while (call.find()) {
            String target = call.group(1);
            if ("{".equals(target) || "(".equals(target)) {
                continue;                                               // a CALL { ... } subquery
            }
            if (target == null
                    || !ALLOWED_PROCEDURES.contains(target.replace("`", "").toLowerCase(Locale.ROOT))) {
                throw new IllegalArgumentException("Refusing procedure call: " + cypher);
            }
        }
    }
}