import os import logging import threading from concurrent.futures import ThreadPoolExecutor, as_completed from typing import List from ..models import Document from .neo4j_client import Neo4jClient from .entity_extractor import EntityExtractor from .relationship_extractor import RelationshipExtractor from .entity_resolver import EntityResolver from .vector_retriever import VectorRetriever logger = logging.getLogger(__name__) class GraphBuilder: def __init__(self): logger.info("Initializing GraphBuilder orchestrator service.") self.neo4j_client = Neo4jClient() self.entity_extractor = EntityExtractor() self.relationship_extractor = RelationshipExtractor() self.entity_resolver = EntityResolver() self.vector_retriever = VectorRetriever() def _update_progress(self, doc, step: str, progress: int): """Update document processing progress for frontend polling.""" doc.processing_step = step doc.processing_progress = progress doc.save(update_fields=['processing_step', 'processing_progress']) def _batch_create_entities(self, entities: List[dict], user_id: str, document_id: str): """Batch create entities using UNWIND to reduce N+1 queries.""" if not entities: return batch_size = 50 for i in range(0, len(entities), batch_size): batch = entities[i:i + batch_size] query = ( "UNWIND $entities AS ent " "MERGE (e:Entity {name: ent.name, user_id: $user_id}) " "ON CREATE SET e.type = ent.type, e.description = ent.description, " " e.source_doc = ent.source_doc, e.source_doc_id = $document_id, e.page = ent.page, e.created_at = timestamp() " "ON MATCH SET e.description = coalesce(e.description, ent.description), e.source_doc_id = $document_id, e.source_doc = ent.source_doc" ) params = { "entities": [ { "name": e["name"].strip(), "type": e["type"].strip(), "description": e["description"].strip(), "source_doc": e.get("source_doc", ""), "page": e.get("page", 0) } for e in batch ], "user_id": str(user_id), "document_id": str(document_id) } self.neo4j_client.execute_query(query, params) def _batch_create_relationships(self, relationships: List[dict], user_id: str, document_id: str): """Batch create relationships grouped by type using UNWIND.""" if not relationships: return from .neo4j_client import VALID_RELATIONSHIP_TYPES # Group relationships by type grouped: dict[str, list] = {} for rel in relationships: rel_type = rel["relationship_type"].upper().strip() if rel_type not in VALID_RELATIONSHIP_TYPES: rel_type = "RELATED_TO" grouped.setdefault(rel_type, []).append(rel) for rel_type, rels in grouped.items(): batch_size = 50 for i in range(0, len(rels), batch_size): batch = rels[i:i + batch_size] query = ( "UNWIND $rels AS rel " "MATCH (source:Entity {name: rel.source, user_id: $user_id}) " "MATCH (target:Entity {name: rel.target, user_id: $user_id}) " f"MERGE (source)-[r:{rel_type}]->(target) " "ON CREATE SET r.description = rel.description, r.confidence = rel.confidence, " " r.source_doc = rel.source_doc, r.source_doc_id = $document_id, r.page = rel.page, r.created_at = timestamp() " "ON MATCH SET r.source_doc_id = $document_id, r.source_doc = rel.source_doc " "RETURN r LIMIT 1" ) params = { "rels": [ { "source": rel["source_entity"].strip(), "target": rel["target_entity"].strip(), "description": rel["description"].strip(), "confidence": float(rel["confidence"]), "source_doc": rel.get("source_doc", ""), "page": rel.get("page", 0) } for rel in batch ], "user_id": str(user_id), "document_id": str(document_id) } self.neo4j_client.execute_query(query, params) def process_document(self, document_id, user_id): """ Orchestrates the entire GraphRAG ingestion pipeline. Reads file, extracts entities & relationships, resolves duplicates, and writes to Neo4j. Updates processing_progress and processing_step at each stage for frontend polling. """ try: doc = Document.objects.get(id=document_id) except Document.DoesNotExist: logger.error("Document with ID %s does not exist. Ingestion aborted.", document_id) return logger.info("Beginning background graph building for Document: %s (User ID: %s)", doc.name, user_id) # 1. Update status to PROCESSING doc.status = Document.Status.PROCESSING doc.save() try: filepath = doc.file.path if not os.path.exists(filepath): raise FileNotFoundError(f"File not found on disk: {filepath}") # 2. Parse file into sections/pages self._update_progress(doc, "Parsing document...", 5) sections = self._parse_file_to_sections(filepath) logger.info("Parsed document into %d sections for analysis.", len(sections)) # 2b. Index document text in ChromaDB vector store self._update_progress(doc, "Indexing vectors in ChromaDB...", 15) full_text = "\n\n".join([sec["text"] for sec in sections]) logger.info("Indexing document text in ChromaDB (Doc: %s, User: %s)...", doc.name, user_id) self.vector_retriever.index_document( text_content=full_text, doc_name=doc.name, user_id=user_id ) all_entities = [] all_relationships = [] all_entities_lock = threading.Lock() all_relationships_lock = threading.Lock() total_sections = len(sections) def process_section(sec: dict) -> tuple[List[dict], List[dict]]: """Process a single section: extract entities and relationships.""" import time text = sec["text"] page = sec["page"] time.sleep(2) ents = self.entity_extractor.extract_entities(text) for e in ents: e["page"] = page e["source_doc"] = doc.name time.sleep(2) rels = self.relationship_extractor.extract_relationships(text) for r in rels: r["page"] = page r["source_doc"] = doc.name return ents, rels # 3. Perform Entity and Relationship Extraction per section (parallel) completed = 0 failed_count = 0 last_error = None with ThreadPoolExecutor(max_workers=1) as executor: futures = {executor.submit(process_section, sec): sec for sec in sections} for future in as_completed(futures): try: ents, rels = future.result() with all_entities_lock: all_entities.extend(ents) with all_relationships_lock: all_relationships.extend(rels) except Exception as e: logger.error("Section processing failed: %s", str(e)) failed_count += 1 last_error = e completed += 1 extraction_progress = 20 + int(55 * (completed / total_sections)) if total_sections > 0 else 20 self._update_progress(doc, f"Extracting entities... ({completed}/{total_sections})", extraction_progress) if failed_count == total_sections and total_sections > 0: raise RuntimeError(f"All sections failed to process. Last error: {last_error}") # 4. Run entity resolution (deduplicate entities and rewrite relationships) self._update_progress(doc, "Resolving duplicates...", 80) resolved_ents, rewritten_rels = self.entity_resolver.resolve_entities( all_entities, all_relationships ) # 5. Batch store resolved nodes inside Neo4j self._update_progress(doc, "Building knowledge graph...", 85) logger.info("Writing %d resolved entities to Neo4j...", len(resolved_ents)) self._batch_create_entities(resolved_ents, user_id, document_id) # 6. Batch store rewritten edges inside Neo4j self._update_progress(doc, "Writing relationships...", 92) logger.info("Writing %d rewritten relationships to Neo4j...", len(rewritten_rels)) self._batch_create_relationships(rewritten_rels, user_id, document_id) # 7. Update status to COMPLETED and record counts doc.entity_count = len(resolved_ents) doc.relationship_count = len(rewritten_rels) doc.status = Document.Status.COMPLETED doc.error_message = None doc.processing_progress = 100 doc.processing_step = "Complete" doc.save() logger.info("Successfully finished building knowledge graph for Document: %s", doc.name) except Exception as e: logger.error("Failed to process document: %s. Error: %s", doc.name, str(e), exc_info=True) doc.status = Document.Status.FAILED doc.error_message = str(e) doc.processing_step = f"Failed: {str(e)[:100]}" doc.save() def delete_document_data(self, document_id, user_id): """ Cleans up and deletes associated Neo4j node/edge elements for a deleted document. Attempts both graph and vector cleanup independently to avoid orphaned data. """ try: doc = Document.objects.get(id=document_id) logger.info("Triggering graph wipe for Document: %s (User ID: %s)", doc.name, user_id) # Attempt both cleanups independently neo4j_ok = True vector_ok = True try: self.neo4j_client.delete_document_nodes(doc.id, user_id) except Exception as e: logger.error("Neo4j cleanup failed for Document: %s. Error: %s", doc.name, str(e)) neo4j_ok = False try: self.vector_retriever.delete_document_vectors(doc.name, user_id) except Exception as e: logger.error("ChromaDB cleanup failed for Document: %s. Error: %s", doc.name, str(e)) vector_ok = False if neo4j_ok and vector_ok: logger.info("Finished Graph cleanup for Document: %s", doc.name) else: logger.warning("Partial cleanup for Document: %s (Neo4j: %s, Vector: %s)", doc.name, "OK" if neo4j_ok else "FAIL", "OK" if vector_ok else "FAIL") except Document.DoesNotExist: logger.error("Document with ID %s does not exist. Cleanup aborted.", document_id) except Exception as e: logger.error("Failed to clean up graph data for Document ID: %s. Error: %s", document_id, str(e), exc_info=True) def _parse_file_to_sections(self, filepath: str) -> List[dict]: """ Loads document file and splits content into page/paragraph sections. """ ext = filepath.split(".")[-1].lower() sections = [] if ext == "pdf": import pypdf reader = pypdf.PdfReader(filepath) for idx, page in enumerate(reader.pages): text = page.extract_text() if text and text.strip(): sections.append({ "text": text.strip(), "page": idx + 1 }) elif ext in ["docx", "doc"]: import docx doc = docx.Document(filepath) current_chunk = [] current_len = 0 section_idx = 1 for p in doc.paragraphs: txt = p.text.strip() if p.text else "" if txt: current_chunk.append(txt) current_len += len(txt) if current_len >= 1500: sections.append({ "text": "\n".join(current_chunk), "page": section_idx }) current_chunk = [] current_len = 0 section_idx += 1 if current_chunk: sections.append({ "text": "\n".join(current_chunk), "page": section_idx }) else: # Default fallback for TXT, Markdown, etc. with open(filepath, "r", encoding="utf-8", errors="ignore") as f: content = f.read() # Split by double newlines paragraphs = [p.strip() for p in content.split("\n\n") if p.strip()] current_chunk = [] current_len = 0 section_idx = 1 for p in paragraphs: current_chunk.append(p) current_len += len(p) if current_len >= 1500: sections.append({ "text": "\n\n".join(current_chunk), "page": section_idx }) current_chunk = [] current_len = 0 section_idx += 1 if current_chunk: sections.append({ "text": "\n\n".join(current_chunk), "page": section_idx }) return sections