Spaces:
Sleeping
Sleeping
File size: 8,561 Bytes
da21448 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 | 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]
|