build_knowledge_graph.py

"""
Build Knowledge Graph - Full Integration Script
Processes documents, extracts entities and relationships, and builds Neo4j graph
"""

import sys
from pathlib import Path
import json
from typing import List, Dict
import time

# Add project root to path
project_root = Path(__file__).parent
sys.path.insert(0, str(project_root))

from backend.core.document_processor import DocumentProcessor
from backend.core.nlp_processor import NLPProcessor
from backend.core.relationship_extractor import RelationshipExtractor
from backend.database.neo4j_manager import Neo4jManager
from loguru import logger


def build_knowledge_graph(
    sample_dir: str = "data/sample/txt",
    clear_database: bool = True
):
    """
    Build complete knowledge graph from sample documents
    
    Args:
        sample_dir: Directory containing sample documents
        clear_database: Whether to clear existing database
    """
    
    print("\n" + "="*80)
    print("BUILDING KNOWLEDGE GRAPH")
    print("="*80)
    
    # Initialize processors
    print("\n[1/6] Initializing processors...")
    doc_processor = DocumentProcessor()
    nlp_processor = NLPProcessor()
    rel_extractor = RelationshipExtractor(nlp=nlp_processor.nlp)
    neo4j_manager = Neo4jManager()
    
    print("[OK] All processors initialized")
    
    # Clear database if requested
    if clear_database:
        print("\n[2/6] Clearing existing database...")
        neo4j_manager.clear_database()
        print("[OK] Database cleared")
    else:
        print("\n[2/6] Keeping existing database content")
    
    # Get all sample files
    sample_path = Path(sample_dir)
    if not sample_path.exists():
        print(f"[ERROR] Sample directory not found: {sample_dir}")
        return
    
    sample_files = list(sample_path.glob("*.txt"))
    if not sample_files:
        print(f"[ERROR] No .txt files found in {sample_dir}")
        return
    
    print(f"\n[3/6] Processing {len(sample_files)} documents...")
    
    all_entities = []
    all_relationships = []
    
    # Process each document
    for i, file_path in enumerate(sample_files, 1):
        print(f"\n  [{i}/{len(sample_files)}] {file_path.name}")
        
        # Extract text
        doc_result = doc_processor.process_document(str(file_path))
        if not doc_result.get('success'):
            print(f"    [ERROR] Failed to extract text")
            continue
        
        text = doc_result['text']
        print(f"    - Extracted {len(text)} characters")
        
        # Extract entities
        doc_entities = nlp_processor.process_document(
            text=text,
            document_id=f"doc_{i}",
            filename=file_path.name
        )
        
        entities = doc_entities.entities
        print(f"    - Extracted {len(entities)} entities")
        
        # Extract relationships
        relationships = rel_extractor.extract_relationships(text, entities)
        print(f"    - Extracted {len(relationships)} relationships")
        
        all_entities.extend(entities)
        all_relationships.extend(relationships)
    
    print(f"\n[OK] Processed all documents")
    print(f"  Total entities: {len(all_entities)}")
    print(f"  Total relationships: {len(all_relationships)}")
    
    # Create entity nodes in Neo4j
    print(f"\n[4/6] Creating entity nodes in Neo4j...")
    start_time = time.time()
    created_entities = neo4j_manager.batch_create_entities(all_entities)
    elapsed = time.time() - start_time
    print(f"[OK] Created {created_entities} entity nodes in {elapsed:.2f}s")
    
    # Create relationships in Neo4j
    print(f"\n[5/6] Creating relationships in Neo4j...")
    start_time = time.time()
    created_rels = neo4j_manager.batch_create_relationships(all_relationships)
    elapsed = time.time() - start_time
    print(f"[OK] Created {created_rels} relationships in {elapsed:.2f}s")
    
    # Get final statistics
    print(f"\n[6/6] Knowledge Graph Statistics:")
    stats = neo4j_manager.get_graph_stats()
    
    print(f"\n  Nodes:")
    print(f"    Total: {stats['total_nodes']}")
    for node_type, count in stats['nodes_by_type'].items():
        print(f"    - {node_type}: {count}")
    
    print(f"\n  Relationships:")
    print(f"    Total: {stats['total_relationships']}")
    for rel_type, count in stats['relationships_by_type'].items():
        print(f"    - {rel_type}: {count}")
    
    # Save summary to JSON
    output_dir = Path("data/processed")
    output_dir.mkdir(parents=True, exist_ok=True)
    
    summary_file = output_dir / "knowledge_graph_summary.json"
    
    with open(summary_file, 'w', encoding='utf-8') as f:
        json.dump({
            "stats": stats,
            "documents_processed": len(sample_files),
            "entities_created": created_entities,
            "relationships_created": created_rels,
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
        }, f, indent=2)
    
    print(f"\n[OK] Summary saved to: {summary_file}")
    
    # Close Neo4j connection
    neo4j_manager.close()
    
    print("\n" + "="*80)
    print("[OK] KNOWLEDGE GRAPH BUILD COMPLETE!")
    print("="*80)
    
    print("\nNext steps:")
    print("1. Open Neo4j Browser: http://localhost:7474")
    print("2. Run query: MATCH (n) RETURN n LIMIT 100")
    print("3. Explore the graph visually!")
    print("4. Try queries like:")
    print("   - MATCH (p:Person)-[r:WORKS_FOR]->(o:Organization) RETURN p, r, o")
    print("   - MATCH (o:Organization)-[r:LOCATED_IN]->(l:Location) RETURN o, r, l")
    print("\n")


def query_example_paths():
    """
    Example: Query some interesting paths in the graph
    """
    print("\n" + "="*80)
    print("EXAMPLE GRAPH QUERIES")
    print("="*80)
    
    neo4j_manager = Neo4jManager()
    
    # Example queries
    examples = [
        ("Maria Weber", "Munich"),
        ("Helsing", "NATO"),
        ("ARX Robotics", "Munich"),
    ]
    
    print("\nFinding shortest paths:")
    for source, target in examples:
        print(f"\n  Path from '{source}' to '{target}':")
        path = neo4j_manager.find_shortest_path(source, target)
        
        if path:
            print(f"    Length: {path['length']} hops")
            print(f"    Nodes: {' -> '.join(n['name'] for n in path['nodes'])}")
        else:
            print(f"    No path found")
    
    neo4j_manager.close()
    
    print("\n" + "="*80)


if __name__ == "__main__":
    # Build the knowledge graph
    build_knowledge_graph(
        sample_dir="data/sample/txt",
        clear_database=True
    )
    
    # Show example queries
    # query_example_paths()