import re from collections import defaultdict from collections.abc import Sequence from itertools import batched from typing import Any, Self from neo4j import GraphDatabase from techdocs_graphrag.config import Settings from techdocs_graphrag.models import Chunk, GraphExtraction, GraphFact, RetrievedChunk VECTOR_INDEX = "chunk_embeddings" FULLTEXT_INDEX = "entity_names" WRITE_BATCH_SIZE = 200 MIN_SEARCH_TERM_LENGTH = 3 INDEX_TIMEOUT_SECONDS = 300 CYPHER_IDENTIFIER = re.compile(r"^[A-Za-z][A-Za-z0-9_]*$") def cypher_identifier(value: str) -> str: if not CYPHER_IDENTIFIER.fullmatch(value): raise ValueError(f"Unsafe Cypher identifier: {value!r}") return value class Neo4jGraphStore: def __init__(self, settings: Settings) -> None: self._driver = GraphDatabase.driver( settings.neo4j_uri, auth=(settings.neo4j_username, settings.neo4j_password.get_secret_value()), ) self._database = settings.neo4j_database self._embedding_dimensions = settings.embedding_dimensions def __enter__(self) -> Self: return self def __exit__(self, *exc_info: object) -> None: self.close() def close(self) -> None: self._driver.close() def verify_connectivity(self) -> None: self._driver.verify_connectivity() def ensure_schema(self) -> None: statements = ( "CREATE CONSTRAINT document_name IF NOT EXISTS " "FOR (d:Document) REQUIRE d.name IS UNIQUE", "CREATE CONSTRAINT chunk_id IF NOT EXISTS FOR (c:Chunk) REQUIRE c.id IS UNIQUE", "CREATE CONSTRAINT entity_key IF NOT EXISTS FOR (e:Entity) REQUIRE e.key IS UNIQUE", f"CREATE FULLTEXT INDEX {FULLTEXT_INDEX} IF NOT EXISTS " "FOR (e:Entity) ON EACH [e.name]", f""" CREATE VECTOR INDEX {VECTOR_INDEX} IF NOT EXISTS FOR (c:Chunk) ON c.embedding OPTIONS {{indexConfig: {{ `vector.dimensions`: {int(self._embedding_dimensions)}, `vector.similarity_function`: 'cosine' }}}} """, ) for statement in statements: self._run(statement) def await_indexes(self) -> None: self._run("CALL db.awaitIndexes($timeout)", timeout=INDEX_TIMEOUT_SECONDS) def upsert_chunks(self, chunks: Sequence[Chunk], embeddings: Sequence[list[float]]) -> None: rows = [ {**chunk.model_dump(), "embedding": embedding} for chunk, embedding in zip(chunks, embeddings, strict=True) ] for batch in batched(rows, WRITE_BATCH_SIZE): self._run( """ UNWIND $rows AS row MERGE (d:Document {name: row.document}) MERGE (c:Chunk {id: row.id}) SET c.document = row.document, c.position = row.position, c.text = row.text, c.embedding = row.embedding MERGE (d)-[:HAS_CHUNK]->(c) """, rows=list(batch), ) def upsert_extraction(self, chunk: Chunk, extraction: GraphExtraction) -> None: self._upsert_entities(chunk, extraction) self._upsert_relations(chunk, extraction) def _upsert_entities(self, chunk: Chunk, extraction: GraphExtraction) -> None: rows_by_type: dict[str, list[dict[str, str]]] = defaultdict(list) for entity in extraction.entities: rows_by_type[entity.type.value].append({"key": entity.key, "name": entity.name}) for entity_type, rows in rows_by_type.items(): self._run( f""" MATCH (c:Chunk {{id: $chunk_id}}) UNWIND $rows AS row MERGE (e:Entity {{key: row.key}}) ON CREATE SET e.name = row.name, e.type = $entity_type SET e:{cypher_identifier(entity_type)} MERGE (c)-[:MENTIONS]->(e) """, chunk_id=chunk.id, rows=rows, entity_type=entity_type, ) def _upsert_relations(self, chunk: Chunk, extraction: GraphExtraction) -> None: rows_by_type: dict[str, list[dict[str, str]]] = defaultdict(list) for relation in extraction.relations: rows_by_type[relation.relation].append( { "source": relation.source_key, "target": relation.target_key, "evidence": relation.evidence, "document": chunk.document, } ) for relation_type, rows in rows_by_type.items(): self._run( f""" UNWIND $rows AS row MATCH (s:Entity {{key: row.source}}) MATCH (t:Entity {{key: row.target}}) MERGE (s)-[r:{cypher_identifier(relation_type)}]->(t) ON CREATE SET r.evidence = [], r.documents = [] SET r.evidence = CASE WHEN row.evidence = '' OR row.evidence IN r.evidence THEN r.evidence ELSE r.evidence + row.evidence END, r.documents = CASE WHEN row.document IN r.documents THEN r.documents ELSE r.documents + row.document END """, rows=rows, ) def vector_search(self, embedding: list[float], top_k: int) -> list[RetrievedChunk]: rows = self._run( """ CALL db.index.vector.queryNodes($index, $top_k, $embedding) YIELD node, score OPTIONAL MATCH (node)-[:MENTIONS]->(entity:Entity) RETURN node.id AS id, node.document AS document, node.position AS position, node.text AS text, score, collect(entity.key) AS entity_keys ORDER BY score DESC """, index=VECTOR_INDEX, top_k=top_k, embedding=embedding, ) return [RetrievedChunk.model_validate(row) for row in rows] def search_entities(self, text: str, limit: int) -> list[str]: terms = sorted( {term.lower() for term in re.findall(r"\w+", text) if len(term) >= MIN_SEARCH_TERM_LENGTH} ) if not terms: return [] rows = self._run( """ CALL db.index.fulltext.queryNodes($index, $query) YIELD node, score RETURN node.key AS key ORDER BY score DESC LIMIT $limit """, index=FULLTEXT_INDEX, query=" OR ".join(terms), limit=limit, ) return [row["key"] for row in rows] def expand(self, seed_keys: Sequence[str], max_hops: int, limit: int) -> list[GraphFact]: if not seed_keys: return [] rows = self._run( f""" MATCH (seed:Entity) WHERE seed.key IN $keys MATCH path = (seed)-[*1..{int(max_hops)}]-(:Entity) WHERE all(node IN nodes(path) WHERE node:Entity) UNWIND relationships(path) AS rel WITH rel, min(length(path)) AS distance RETURN startNode(rel).name AS source, type(rel) AS relation, endNode(rel).name AS target, coalesce(rel.evidence, []) AS evidence, coalesce(rel.documents, []) AS documents, distance ORDER BY distance, source, relation, target LIMIT $limit """, keys=list(seed_keys), limit=limit, ) return [GraphFact.model_validate(row) for row in rows] def stats(self) -> dict[str, int]: rows = self._run( """ RETURN COUNT { (:Document) } AS documents, COUNT { (:Chunk) } AS chunks, COUNT { (:Entity) } AS entities, COUNT { (:Entity)-[]->(:Entity) } AS relations """ ) return rows[0] def reset(self) -> None: self._run("MATCH (n) DETACH DELETE n") def _run(self, query: str, **parameters: Any) -> list[dict[str, Any]]: records, _, _ = self._driver.execute_query( query, parameters_=parameters, database_=self._database, ) return [record.data() for record in records]