final.py

from flask import Flask, jsonify
import os
import glob
from langchain_neo4j import Neo4jGraph
from langchain_groq import ChatGroq
from langchain_core.documents import Document
from langchain_experimental.graph_transformers import LLMGraphTransformer
import textwrap
import uuid
import re

app = Flask(__name__)

# Neo4j and Groq configuration
NEO4J_URI = ""
NEO4J_USERNAME = "neo4j"
NEO4J_PASSWORD = ""
GROQ_API_KEY = ""

# Set environment variables
os.environ["NEO4J_URI"] = NEO4J_URI
os.environ["NEO4J_USERNAME"] = NEO4J_USERNAME
os.environ["NEO4J_PASSWORD"] = NEO4J_PASSWORD

# Initialize Neo4j graph
try:
    graph = Neo4jGraph(
        url=NEO4J_URI,
        username=NEO4J_USERNAME,
        password=NEO4J_PASSWORD
    )
    print("Neo4j graph initialized successfully")
except Exception as e:
    print(f"Failed to initialize Neo4j graph: {str(e)}")
    graph = None

# Initialize LLM
llm = ChatGroq(groq_api_key=GROQ_API_KEY, model_name="Gemma2-9b-It")
llm_transformer = LLMGraphTransformer(llm=llm)

# Function to read text files from dataset folder
def read_text_files_from_dataset(dataset_folder="dataset"):
    text = ""
    text_files = glob.glob(os.path.join(dataset_folder, "*.txt"))
    
    if not text_files:
        raise FileNotFoundError(f"No .txt files found in the {dataset_folder} folder")
    
    for file_path in text_files:
        with open(file_path, 'r', encoding='utf-8') as file:
            text += file.read() + "\n\n"
    
    return text.strip()

# Split text into chunks
def split_text_into_chunks(text, max_chunk_size=1000):
    paragraphs = text.split('\n\n')
    chunks = []
    current_chunk = ""
    
    for paragraph in paragraphs:
        if len(current_chunk) + len(paragraph) + 1 <= max_chunk_size:
            current_chunk += paragraph + "\n\n"
        else:
            if current_chunk:
                chunks.append(current_chunk.strip())
            current_chunk = paragraph + "\n\n"
    
    if current_chunk:
        chunks.append(current_chunk.strip())
    
    return chunks

# Load dataset and split into chunks
try:
    text = read_text_files_from_dataset()
    chunks = split_text_into_chunks(text)
    print(f"Loaded {len(chunks)} chunks from dataset")
except Exception as e:
    chunks = []
    print(f"Error reading dataset: {str(e)}")

# Store current chunk index
current_chunk_index = 0

# Track changes for undo functionality
change_log = []

# Test route to verify app is running
@app.route('/')
def home():
    return "Flask app is running!"

@app.route('/process_chunk', methods=['GET'])
def process_chunk():
    global current_chunk_index
    
    if not graph:
        return jsonify({"status": "error", "message": "Neo4j graph not initialized"})
    
    if not chunks:
        return jsonify({"status": "error", "message": "No text files found in dataset folder"})
    
    if current_chunk_index >= len(chunks):
        return jsonify({"status": "error", "message": "No more chunks to process"})
    
    chunk_text = chunks[current_chunk_index]
    document = [Document(page_content=chunk_text)]
    
    try:
        graph_doc = llm_transformer.convert_to_graph_documents(document)
        
        # Track nodes and relationships added in this chunk
        added_nodes = []
        added_relationships = []
        
        for node in graph_doc[0].nodes:
            merge_query = f"""
            MERGE (n:`{node.type}` {{id: $id}})
            SET n += $properties
            """
            properties = node.properties.copy()
            properties['id'] = node.id
            graph.query(merge_query, params={'id': node.id, 'properties': properties})
            added_nodes.append(node.id)
        
        for rel in graph_doc[0].relationships:
            if rel.source.id == rel.target.id:
                continue
            rel_query = f"""
            MATCH (source:`{rel.source.type}` {{id: $source_id}}),
                  (target:`{rel.target.type}` {{id: $target_id}})
            CREATE (source)-[r:`{rel.type}`]->(target)
            SET r += $properties
            """
            properties = rel.properties
            graph.query(rel_query, params={'source_id': rel.source.id, 'target_id': rel.target.id, 'properties': properties})
            added_relationships.append((rel.source.id, rel.type, rel.target.id))
        
        # Log the changes for undo
        change_log.append({
            'chunk_index': current_chunk_index,
            'added_nodes': added_nodes,
            'added_relationships': added_relationships
        })
        
        current_chunk_index += 1
        
        return jsonify({
            "status": "success",
            "chunk_index": current_chunk_index - 1,
            "chunk_content": chunk_text[:100] + "..." if len(chunk_text) > 100 else chunk_text,
            "nodes_added": len(added_nodes),
            "relationships_added": len(added_relationships)
        })
    
    except Exception as e:
        return jsonify({"status": "error", "message": str(e)})

@app.route('/undo', methods=['GET'])
def undo():
    if not graph:
        return jsonify({"status": "error", "message": "Neo4j graph not initialized"})
    
    if not change_log:
        return jsonify({"status": "error", "message": "No changes to undo"})
    
    try:
        last_change = change_log.pop()
        
        # Delete relationships added in the last chunk
        for rel in last_change['added_relationships']:
            source_id, rel_type, target_id = rel
            delete_rel_query = f"""
            MATCH (source {{id: $source_id}})-[r:`{rel_type}`]->(target {{id: $target_id}})
            DELETE r
            """
            graph.query(delete_rel_query, params={'source_id': source_id, 'target_id': target_id})
        
        # Delete nodes added in the last chunk (if they have no relationships)
        for node_id in last_change['added_nodes']:
            delete_node_query = f"""
            MATCH (n {{id: $node_id}})
            WHERE NOT EXISTS ((n)--())
            DELETE n
            """
            graph.query(delete_node_query, params={'node_id': node_id})
        
        return jsonify({
            "status": "success",
            "undone_chunk_index": last_change['chunk_index']
        })
    
    except Exception as e:
        return jsonify({"status": "error", "message": str(e)})

if __name__ == '__main__':
    app.run(debug=True, host='localhost', port=5000)