Spaces:
Sleeping
Sleeping
| 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] | |