soumya-ai/Knowledge-Graph-Ingest
0
1"""2graph_builder.py3 4Extracts entities and relationships from raw text using OpenAI,5then writes the knowledge graph into Neo4j Aura (LangChain LCEL).6 7Uses: OpenAI Chat — entity/relation extraction8 OpenAI Embeddings — node & chunk vectors9 Neo4j — graph storage + vector indexes10"""11 12import json13import re14from typing import Any15import unicodedata16 17from langchain_core.documents import Document18from langchain_core.embeddings import Embeddings19from langchain_core.prompts import ChatPromptTemplate20from langchain_openai import ChatOpenAI21from langchain_neo4j import Neo4jGraph22from langchain_text_splitters import RecursiveCharacterTextSplitter23import hashlib24 25from src.config import require_neo4j_config, require_openai_key, settings26from src.embeddings import get_embedder27from src.neo4j_client import get_graph28 29def canonical_id(name: str) -> str:30 """Deterministic id from display name — mirrors the prompt rule."""31 name = unicodedata.normalize("NFKD", name)32 name = name.lower().strip()33 name = re.sub(r"[''`]", "", name) # strip apostrophes34 name = re.sub(r"[^a-z0-9]+", "_", name) # everything else → underscore35 name = name.strip("_")36 return name37 38 39def normalize_entities(entities: list[dict]) -> list[dict]:40 """41 1. Override whatever id the LLM invented with the canonical one.42 2. Merge duplicate entities within the same chunk (same canonical id).43 """44 seen: dict[str, dict] = {}45 for e in entities:46 cid = canonical_id(e["name"])47 e["id"] = cid48 if cid not in seen:49 seen[cid] = e50 else:51 # Keep the longer description52 if len(e.get("description", "")) > len(seen[cid].get("description", "")):53 seen[cid]["description"] = e["description"]54 return list(seen.values())55 56 57def normalize_relationships(relationships: list[dict], entities: list[dict]) -> list[dict]:58 """Remap source/target ids to match the normalized entity ids."""59 valid_ids = {e["id"] for e in entities}60 out = []61 for rel in relationships:62 src = canonical_id_from_old(rel["source"], entities)63 tgt = canonical_id_from_old(rel["target"], entities)64 if src in valid_ids and tgt in valid_ids:65 rel["source"] = src66 rel["target"] = tgt67 out.append(rel)68 return out69 70 71def canonical_id_from_old(old_id: str, entities: list[dict]) -> str:72 """Map an LLM-invented id back to the canonical one via entity lookup."""73 for e in entities:74 if e["id"] == old_id or canonical_id(e.get("name","")) == canonical_id(old_id):75 return e["id"]76 return canonical_id(old_id) # fallback: normalise the id string itself77 78# ── Prompts ───────────────────────────────────────────────────────────────────79 80EXTRACTION_PROMPT = ChatPromptTemplate.from_messages(81 [82 (83 "system",84 """You are an expert knowledge graph builder.85Extract ALL entities and relationships from the text below.86 87Return ONLY valid JSON — no markdown, no explanation, no code fences.88 89Format:90{{91 "entities": [92 {{"id": "unique_snake_case_id", "type": "Person|Organization|Concept|Technology|Location|Event|Other", "name": "display name", "description": "one sentence"}}93 ],94 "relationships": [95 {{"source": "entity_id", "target": "entity_id", "type": "RELATION_TYPE_CAPS", "description": "short description"}}96 ]97}}98 99Rules:100- entity id = name.lower(), spaces→underscores, remove all punctuation101 e.g. "OpenAI" → "openai", "GPT-4" → "gpt4", "Sam Altman" → "sam_altman"102- NEVER append _company, _org, _inc, _person — use the bare canonical name only103- If two entities have the same name, they ARE the same entity — use the same id104- relationship type must be SCREAMING_SNAKE_CASE (e.g. DEVELOPED_BY, USES, PART_OF)105- only include entities explicitly mentioned in the text106- only include relationships you are confident about""",107 ),108 ("human", "Text:\n{text}"),109 ]110)111 112 113# ── Neo4j helpers ─────────────────────────────────────────────────────────────114 115 116def deduplicate_entities(graph: Neo4jGraph) -> int:117 """118 Merge Entity nodes with identical names (case-insensitive).119 Returns the number of merges performed.120 Requires APOC (already in docker-compose).121 """122 result = graph.query("""123 MATCH (a:Entity), (b:Entity)124 WHERE id(a) < id(b)125 AND toLower(trim(a.name)) = toLower(trim(b.name))126 WITH a, b127 CALL apoc.refactor.mergeNodes([a, b], {properties: 'discard', mergeRels: true})128 YIELD node129 RETURN count(node) AS merges130 """)131 132 count = result[0]["merges"] if result else 0133 print(f"[dedup] Merged {count} duplicate entity node(s).")134 return count135 136# ── Schema setup ──────────────────────────────────────────────────────────────137 138def _schema_queries() -> list[str]:139 dim = settings.OPENAI_EMBED_DIMENSIONS140 return [141 "CREATE CONSTRAINT entity_id IF NOT EXISTS FOR (e:Entity) REQUIRE e.id IS UNIQUE",142 "CREATE CONSTRAINT chunk_id IF NOT EXISTS FOR (c:Chunk) REQUIRE c.id IS UNIQUE",143 f"""144 CREATE VECTOR INDEX entity_embedding IF NOT EXISTS145 FOR (e:Entity) ON (e.embedding)146 OPTIONS {{indexConfig: {{`vector.dimensions`: {dim}, `vector.similarity_function`: 'cosine'}}}}147 """,148 f"""149 CREATE VECTOR INDEX chunk_embedding IF NOT EXISTS150 FOR (c:Chunk) ON (c.embedding)151 OPTIONS {{indexConfig: {{`vector.dimensions`: {dim}, `vector.similarity_function`: 'cosine'}}}}152 """,153 ]154 155 156def setup_schema(graph: Neo4jGraph) -> None:157 """Create constraints and vector indexes if they don't exist."""158 for q in _schema_queries():159 try:160 graph.query(q)161 except Exception as e:162 # Index/constraint may already exist — safe to ignore163 if "already exists" not in str(e).lower():164 print(f"[schema] warning: {e}")165 print("[schema] Neo4j schema ready.")166 167 168# ── Extraction chain ──────────────────────────────────────────────────────────169 170 171def build_extraction_chain():172 """LCEL chain: text → OpenAI → parsed entities/relations dict."""173 llm = ChatOpenAI(174 model=settings.OPENAI_EXTRACT_MODEL,175 api_key=settings.OPENAI_API_KEY,176 temperature=0,177 model_kwargs={"response_format": {"type": "json_object"}},178 )179 180 def safe_parse(ai_message) -> dict:181 """Parse JSON from LLM output, with fallback for fenced code."""182 if isinstance(ai_message, dict):183 return ai_message184 text = ai_message.content if hasattr(ai_message, "content") else str(ai_message)185 # Strip markdown fences if model ignores format="json"186 text = re.sub(r"```(?:json)?", "", text).strip().rstrip("`").strip()187 try:188 return json.loads(text)189 except json.JSONDecodeError:190 # Try to extract the first {...} block191 match = re.search(r"\{.*\}", text, re.DOTALL)192 if match:193 return json.loads(match.group())194 return {"entities": [], "relationships": []}195 196 return EXTRACTION_PROMPT | llm | safe_parse197 198 199# ── Graph writer ──────────────────────────────────────────────────────────────200 201 202def write_entities(203 graph: Neo4jGraph, embedder: Embeddings, entities: list[dict]204) -> None:205 """Upsert entity nodes with embeddings."""206 for entity in entities:207 text_for_embed = f"{entity['name']}: {entity.get('description', '')}"208 try:209 embedding = embedder.embed_query(text_for_embed)210 except Exception:211 embedding = []212 213 graph.query(214 """215 MERGE (e:Entity {id: $id})216 SET e.name = $name,217 e.type = $type,218 e.description = $description,219 e.embedding = $embedding220 WITH e221 CALL apoc.create.addLabels(e, [$type]) YIELD node222 RETURN node223 """,224 {225 "id": entity["id"],226 "name": entity["name"],227 "type": entity.get("type", "Other"),228 "description": entity.get("description", ""),229 "embedding": embedding,230 },231 )232 233 234def write_relationships(graph: Neo4jGraph, relationships: list[dict]) -> None:235 """Create typed relationships between entity nodes."""236 for rel in relationships:237 rel_type = re.sub(r"[^A-Z0-9_]", "_", rel["type"].upper())238 graph.query(239 f"""240 MATCH (a:Entity {{id: $source}})241 MATCH (b:Entity {{id: $target}})242 MERGE (a)-[r:{rel_type}]->(b)243 SET r.description = $description244 """,245 {246 "source": rel["source"],247 "target": rel["target"],248 "description": rel.get("description", ""),249 },250 )251 252 253def write_chunk(254 graph: Neo4jGraph,255 embedder: Embeddings,256 chunk_id: str,257 text: str,258 entity_ids: list[str],259) -> None:260 """Store the raw text chunk and link it to its extracted entities."""261 try:262 embedding = embedder.embed_query(text)263 except Exception:264 embedding = []265 266 graph.query(267 """268 MERGE (c:Chunk {id: $id})269 SET c.text = $text,270 c.embedding = $embedding271 """,272 {"id": chunk_id, "text": text, "embedding": embedding},273 )274 for eid in entity_ids:275 graph.query(276 """277 MATCH (c:Chunk {id: $chunk_id})278 MATCH (e:Entity {id: $entity_id})279 MERGE (c)-[:MENTIONS]->(e)280 """,281 {"chunk_id": chunk_id, "entity_id": eid},282 )283 284 285# ── Public API ────────────────────────────────────────────────────────────────286 287 288def build_graph_from_documents(documents: list[Document]) -> None:289 """290 Main entry point.291 292 1. Splits documents into chunks293 2. Runs OpenAI extraction chain on each chunk294 3. Writes entities, relationships, and chunks to Neo4j Aura295 """296 require_openai_key()297 require_neo4j_config()298 graph = get_graph()299 embedder = get_embedder()300 setup_schema(graph)301 302 splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=100)303 chain = build_extraction_chain()304 305 chunks = splitter.split_documents(documents)306 print(307 f"[builder] Processing {len(chunks)} chunks with {settings.OPENAI_EXTRACT_MODEL}..."308 )309 310 for i, chunk in enumerate(chunks):311 print(f"[builder] Chunk {i+1}/{len(chunks)} — extracting entities...")312 try:313 result: dict[str, Any] = chain.invoke({"text": chunk.page_content})314 except Exception as e:315 print(f"[builder] Chunk {i+1} extraction failed: {e}")316 continue317 318 entities: list[dict] = normalize_entities(result.get("entities", []))319 relationships: list[dict] = normalize_relationships(result.get("relationships", []), entities)320 321 if entities:322 write_entities(graph, embedder, entities)323 if relationships:324 write_relationships(graph, relationships)325 326 chunk_id = "chunk_" + hashlib.md5(chunk.page_content.encode()).hexdigest()[:12]327 entity_ids = [e["id"] for e in entities]328 write_chunk(graph, embedder, chunk_id, chunk.page_content, entity_ids)329 330 print(331 f"[builder] → {len(entities)} entities, {len(relationships)} relationships"332 )333 334 print("[builder] Knowledge graph build complete.")335 deduplicate_entities(graph)