Merge pull request #109 from neo4j-labs/update_model_relationship

Update model, Create relation between chunks of each document as NEXT…

Author: karanchellani

backend/score.py

@@ -144,12 +144,15 @@ async def get_source_list(uri:str,
     return result
 
 @app.post("/update_similarity_graph")
-async def update_similarity_graph():
+async def update_similarity_graph(uri=Form(None),
+    userName=Form(None),
+    password=Form(None),
+    database=Form(None)):
     """
     Calls 'update_graph' which post the query to update the similiar nodes in the graph
     """
 
-    result = await asyncio.to_thread(update_graph)
+    result = await asyncio.to_thread(update_graph,uri,userName,password,database)
     return result
 
 def decode_password(pwd):

backend/src/diffbot_transformer.py

@@ -5,6 +5,7 @@
 from typing import List
 import os
 import logging
+import uuid
 
 logging.basicConfig(format='%(asctime)s - %(message)s',level='INFO')
 def extract_graph_from_diffbot(graph: Neo4jGraph, 
@@ -31,16 +32,21 @@ def extract_graph_from_diffbot(graph: Neo4jGraph,
     graph_document_list = []
 
     logging.info(f"create relationship between source,chunk and entity nodes created from Diffbot")
+    current_chunk_id = ''
+    relationship_cypher_list = []
     for i,chunk in enumerate(chunks):
+        previous_chunk_id = current_chunk_id
+        current_chunk_id = str(uuid.uuid1())
         if i == 0:
             firstChunk = True
         else:
             firstChunk = False
         graph_document = diffbot_nlp.convert_to_graph_documents([chunk])
         graph.add_graph_documents(graph_document)
-        create_source_chunk_entity_relationship(file_name,graph,graph_document,chunk,uri,userName,password,firstChunk)
-        graph_document_list.append(graph_document[0]) 
+        lst_cypher_queries_chunk_relationship = create_source_chunk_entity_relationship(file_name,graph,graph_document,chunk,uri,userName,password,firstChunk,current_chunk_id,
+        previous_chunk_id)
+        graph_document_list.append(graph_document[0])
+        relationship_cypher_list.extend(lst_cypher_queries_chunk_relationship)
 
     graph.refresh_schema()
-    return graph_document_list
-    
+    return graph_document_list, relationship_cypher_list

backend/src/main.py

@@ -362,18 +362,19 @@ def extract_graph_from_file(uri, userName, password, model, db_name=None, file=N
 
     logging.info("Get graph document list from models")
     if model == 'Diffbot' :
-      graph_documents = extract_graph_from_diffbot(graph,chunks,file_name,uri,userName,password)
+      graph_documents, cypher_list = extract_graph_from_diffbot(graph,chunks,file_name,uri,userName,password)
 
     elif model == 'OpenAI GPT 3.5':
       model_version = 'gpt-3.5-turbo-16k'
-      graph_documents = extract_graph_from_OpenAI(model_version,graph,chunks,file_name,uri,userName,password)
+      graph_documents, cypher_list = extract_graph_from_OpenAI(model_version,graph,chunks,file_name,uri,userName,password)
 
     elif model == 'OpenAI GPT 4':
       model_version = 'gpt-4-0125-preview' 
-      graph_documents = extract_graph_from_OpenAI(model_version,graph,chunks,file_name,uri,userName,password)
+      graph_documents, cypher_list = extract_graph_from_OpenAI(model_version,graph,chunks,file_name,uri,userName,password)
 
-    #update_similarity_graph for the KNN Graph
-    update_graph(graph)
+    #create relation between chunks (FIRST_CHUNK and NEXT_CHUNK)
+    for query in cypher_list:
+       graph.query(query)
 
     distinct_nodes = set()
     relations = []
@@ -498,15 +499,15 @@ def get_source_list_from_graph(uri,userName,password,db_name=None):
     logging.exception(f'Exception:{error_message}')
     return create_api_response(job_status,message=message,error=error_message)
 
-def update_graph(graph):
+def update_graph(uri,userName,password,db_name):
   """
   Update the graph node with SIMILAR relationship where embedding scrore match
   """
   try:   
     knn_min_score = os.environ.get('KNN_MIN_SCORE')
 
     query = "WHERE node <> c and score >= {} MERGE (c)-[rel:SIMILAR]-(node) SET rel.score = score"
-    # graph = Neo4jGraph()
+    graph = Neo4jGraph(url=uri, database=db_name, username=userName, password=password)
     result = graph.query("""MATCH (c:Chunk)
                 WHERE c.embedding IS NOT NULL AND count { (c)-[:SIMILAR]-() } < 5
                 CALL db.index.vector.queryNodes('vector', 6, c.embedding) yield node, score """+ query.format(knn_min_score))

backend/src/make_relationships.py

@@ -15,7 +15,9 @@ def create_source_chunk_entity_relationship(source_file_name :str,
                                             uri : str,
                                             userName : str,
                                             password : str,
-                                            isFirstChunk : bool):
+                                            isFirstChunk : bool,
+                                            current_chunk_id:uuid,
+                                            previous_chunk_id:uuid)-> list:
     """ Create relationship between source, chunk and entity nodes
     Args:
         source_file_name (str): file name of input source
@@ -26,14 +28,16 @@ def create_source_chunk_entity_relationship(source_file_name :str,
         userName: Username to use for graph creation ( if None will use username from config file )
         password: Password to use for graph creation ( if None will use password from config file )
         isFirstChunk : It's bool value to create FIRST_CHUNK AND NEXT_CHUNK relationship between chunk and document node.
+        current_chunk_id : Unique id of chunk
+        previous_chunk_id : Unique id of previous chunk
     """
     source_node = 'fileName: "{}"'
+    lst_cypher_queries_chunk_relationship = []
     # logging.info(f'Graph Document print{graph_document}')
     # openai_api_key = os.environ.get('OPENAI_API_KEY')
     embedding_model = os.environ.get('EMBEDDING_MODEL')
     isEmbedding = os.environ.get('IS_EMBEDDING')
 
-    chunk_uuid = str(uuid.uuid1())
     chunk_node_id_set = 'id:"{}"'
     update_chunk_node_prop = ' SET c.text = "{}"'
     if isEmbedding:
@@ -43,27 +47,29 @@ def create_source_chunk_entity_relationship(source_file_name :str,
             url=uri,
             username=userName,
             password=password,
-            ids=[chunk_uuid]
+            ids=[current_chunk_id]
         )
     else:
-        graph.query('CREATE(c:Chunk {id:"'+ chunk_uuid+'"})' + update_chunk_node_prop.format(chunk.page_content))
+        graph.query('MERGE(c:Chunk {id:"'+ current_chunk_id+'"})' + update_chunk_node_prop.format(chunk.page_content))
 
     logging.info("make PART_OF relationship between chunk node and document node")
-    graph.query('MATCH(d:Document {'+source_node.format(source_file_name)+'}) ,(c:Chunk {'+chunk_node_id_set.format(chunk_uuid)+'}) CREATE (c)-[:PART_OF]->(d)')
+    graph.query('MATCH(d:Document {'+source_node.format(source_file_name)+'}) ,(c:Chunk {'+chunk_node_id_set.format(current_chunk_id)+'}) MERGE (c)-[:PART_OF]->(d)')
 
-    logging.info("make FIRST_CHUNK, NEXT_CHUNK relationship between chunk node and document node")
+    # logging.info("make FIRST_CHUNK, NEXT_CHUNK relationship between chunk node and document node")
     if isFirstChunk:
-        graph.query('MATCH(d:Document {'+source_node.format(source_file_name)+'}) ,(c:Chunk {'+chunk_node_id_set.format(chunk_uuid)+'}) CREATE (d)-[:FIRST_CHUNK]->(c)')
+        lst_cypher_queries_chunk_relationship.append('MATCH(d:Document {'+source_node.format(source_file_name)+'}) ,(c:Chunk {'+chunk_node_id_set.format(current_chunk_id)+'}) MERGE (d)-[:FIRST_CHUNK]->(c)')
+        # graph.query('MATCH(d:Document {'+source_node.format(source_file_name)+'}) ,(c:Chunk {'+chunk_node_id_set.format(current_chunk_id)+'}) CREATE (d)-[:FIRST_CHUNK]->(c)')
     else:
-        graph.query('MATCH(d:Document {'+source_node.format(source_file_name)+'}) ,(c:Chunk {'+chunk_node_id_set.format(chunk_uuid)+'}) CREATE (d)-[:NEXT_CHUNK]->(c)')
+        lst_cypher_queries_chunk_relationship.append('MATCH(pc:Chunk {'+chunk_node_id_set.format(previous_chunk_id)+'}) ,(cc:Chunk {'+chunk_node_id_set.format(current_chunk_id)+'}) MERGE (pc)-[:NEXT_CHUNK]->(cc)')
+        # graph.query('MATCH(pc:Chunk {'+chunk_node_id_set.format(previous_chunk_id)+'}) ,(cc:Chunk {'+chunk_node_id_set.format(current_chunk_id)+'}) CREATE (pc)-[:NEXT_CHUNK]->(cc)')
     # dict = {}
     # nodes_list = []
     for node in graph_document[0].nodes:
         node_id = node.id
-        result = graph.query('MATCH(c:Chunk {'+chunk_node_id_set.format(chunk_uuid)+'}), (n:'+ node.type +'{ id: "'+node_id+'"}) CREATE (c)-[:HAS_ENTITY]->(n)')
+        result = graph.query('MATCH(c:Chunk {'+chunk_node_id_set.format(current_chunk_id)+'}), (n:'+ node.type +'{ id: "'+node_id+'"}) MERGE (c)-[:HAS_ENTITY]->(n)')
     #     json_obj = {'node_id': node_id, 'node_type' : node.type, 'uuid' : chunk_uuid}
     #     nodes_list.append(json_obj)
-
+    return lst_cypher_queries_chunk_relationship
     # dict['chunk_doc'] = chunk.page_content
     # dict['rel_chunk_entity_node'] = nodes_list
     # dict['nodes_created_in_chunk'] = len(graph_document[0].nodes)

backend/src/openAI_llm.py

@@ -25,6 +25,7 @@
 import concurrent.futures
 from concurrent.futures import ThreadPoolExecutor
 import threading
+import uuid
 
 load_dotenv()
 logging.basicConfig(format='%(asctime)s - %(message)s',level='INFO')
@@ -190,6 +191,13 @@ def extract_and_store_graph(
     model_version,
     graph: Neo4jGraph,
     document: Document,
+    file_name: str,
+    uri: str,
+    userName:str,
+    password:str,
+    firstChunk:bool,
+    current_chunk_id:uuid,
+    previous_chunk_id:uuid,
     nodes:Optional[List[str]] = None,
     rels:Optional[List[str]]=None) -> None:
 
@@ -198,10 +206,18 @@ def extract_and_store_graph(
      store the result into a Neo4jGraph.
      
      Args:
-     	 graph: Neo4j graph to store the data into
-     	 document: Langchain document to extract data from
-     	 nodes: List of nodes to extract ( default : None )
-     	 rels: List of relationships to extract ( default : None )
+        model_version: LLM model version
+        graph: Neo4j graph to store the data into
+        document: Langchain document to extract data from
+        file_name (str): file name of input source
+        uri: URI of the graph to extract
+        userName: Username to use for graph creation ( if None will use username from config file )
+        password: Password to use for graph creation ( if None will use password from config file )
+        firstChunk : It's bool value to create FIRST_CHUNK AND NEXT_CHUNK relationship between chunk and document node.
+        current_chunk_id : Unique id of chunk
+        previous_chunk_id : Unique id of previous chunk
+        nodes: List of nodes to extract ( default : None )
+        rels: List of relationships to extract ( default : None )
      
      Returns: 
      	 The GraphDocument that was extracted and stored into the Neo4jgraph
@@ -221,7 +237,9 @@ def extract_and_store_graph(
     )]   
 
     graph.add_graph_documents(graph_document)
-    return graph_document   
+    lst_cypher_queries_chunk_relationship = create_source_chunk_entity_relationship(file_name,graph,graph_document,document,uri,userName,password,firstChunk,current_chunk_id,
+    previous_chunk_id)
+    return graph_document, lst_cypher_queries_chunk_relationship
 
 
 def extract_graph_from_OpenAI(model_version,
@@ -248,19 +266,28 @@ def extract_graph_from_OpenAI(model_version,
     """
     openai_api_key = os.environ.get('OPENAI_API_KEY')
     graph_document_list = []
+    relationship_cypher_list = []
     futures=[]
     logging.info(f"create relationship between source,chunk and entity nodes created from {model_version}")
 
     with ThreadPoolExecutor(max_workers=10) as executor:
+        current_chunk_id= ''
         for i, chunk_document in tqdm(enumerate(chunks), total=len(chunks)):
+            previous_chunk_id = current_chunk_id
+            current_chunk_id = str(uuid.uuid1())
+            position = i+1
             if i == 0:
                 firstChunk = True
             else:
                 firstChunk = False
-            futures.append(executor.submit(extract_and_store_graph,model_version,graph,chunk_document))   
+            metadata = {"position": position,"length": len(chunk_document.page_content)}
+            chunk_document = Document(page_content=chunk_document.page_content,metadata = metadata)
+            
+            futures.append(executor.submit(extract_and_store_graph,model_version,graph,chunk_document,file_name,uri,userName,password,firstChunk,current_chunk_id,previous_chunk_id))   
         for future in concurrent.futures.as_completed(futures):
-            graph_document = future.result()
-            create_source_chunk_entity_relationship(file_name,graph,graph_document,chunk_document,uri,userName,password,firstChunk)
+            graph_document,lst_cypher_queries_chunk_relationship = future.result()
+            
             graph_document_list.append(graph_document[0])
+            relationship_cypher_list.extend(lst_cypher_queries_chunk_relationship)
 
-    return graph_document_list
+    return graph_document_list, relationship_cypher_list