Adding Links to get neighboring nodes (#796)

* addition of link

* added neighbours query

* implemented with driver

* updated the query

* communitiesInfo name change

* communities.tsx removed

* api integration

* modified response

* entities change

* chunk and communities

* chunk space removal

* added element id to chunks

* loading on click

* format changes

* added file name for Dcoumrnt node

* chat token cut off model name update

* icon change

* duplicate sources removal

* Entity change

---------

Co-authored-by: vasanthasaikalluri <[email protected]>

Author: 2 people

backend/score.py

@@ -17,6 +17,7 @@
 from src.post_processing import create_vector_fulltext_indexes, create_entity_embedding
 from sse_starlette.sse import EventSourceResponse
 from src.communities import create_communities
+from src.neighbours import get_neighbour_nodes
 import json
 from typing import List, Mapping
 from starlette.middleware.sessions import SessionMiddleware
@@ -369,6 +370,25 @@ async def chunk_entities(uri=Form(),userName=Form(), password=Form(), database=F
     finally:
         gc.collect()
 
+@app.post("/get_neighbours")
+async def get_neighbours(uri=Form(),userName=Form(), password=Form(), database=Form(), elementId=Form(None)):
+    try:
+        start = time.time()
+        result = await asyncio.to_thread(get_neighbour_nodes,uri=uri, username=userName, password=password,database=database, element_id=elementId)
+        end = time.time()
+        elapsed_time = end - start
+        json_obj = {'api_name':'get_neighbours','db_url':uri, 'logging_time': formatted_time(datetime.now(timezone.utc)), 'elapsed_api_time':f'{elapsed_time:.2f}'}
+        logger.log_struct(json_obj, "INFO")
+        return create_api_response('Success',data=result,message=f"Total elapsed API time {elapsed_time:.2f}")
+    except Exception as e:
+        job_status = "Failed"
+        message="Unable to extract neighbour nodes for given element ID"
+        error_message = str(e)
+        logging.exception(f'Exception in get neighbours :{error_message}')
+        return create_api_response(job_status, message=message, error=error_message)
+    finally:
+        gc.collect()
+
 @app.post("/graph_query")
 async def graph_query(
     uri: str = Form(),

backend/src/chunkid_entities.py

@@ -1,6 +1,7 @@
 import logging
 from src.graph_query import *
 from src.shared.constants import * 
+import re
 
 def process_records(records):
     """
@@ -191,7 +192,11 @@ def get_entities_from_chunkids(uri, username, password, database ,nodedetails,en
             if "entitydetails" in nodedetails and nodedetails["entitydetails"]:
                 entity_ids = [item["id"] for item in nodedetails["entitydetails"]]
                 logging.info(f"chunkid_entities module: Starting for entity ids: {entity_ids}")
-                return process_entityids(driver, entity_ids)
+                result = process_entityids(driver, entity_ids)
+                if "chunk_data" in result.keys():
+                    for chunk in result["chunk_data"]:
+                        chunk["text"] = re.sub(r'\s+', ' ', chunk["text"])
+                return result
             else:
                 logging.info("chunkid_entities module: No entity ids are passed")
                 return default_response  
@@ -201,7 +206,11 @@ def get_entities_from_chunkids(uri, username, password, database ,nodedetails,en
             if "chunkdetails" in nodedetails and nodedetails["chunkdetails"]:
                 chunk_ids = [item["id"] for item in nodedetails["chunkdetails"]]
                 logging.info(f"chunkid_entities module: Starting for chunk ids: {chunk_ids}")
-                return process_chunkids(driver, chunk_ids, entities)
+                result = process_chunkids(driver, chunk_ids, entities)
+                if "chunk_data" in result.keys():
+                    for chunk in result["chunk_data"]:
+                        chunk["text"] = re.sub(r'\s+', ' ', chunk["text"])
+                return result
             else:
                 logging.info("chunkid_entities module: No chunk ids are passed")
                 return default_response

backend/src/neighbours.py

@@ -0,0 +1,63 @@
+import logging
+from src.graph_query import *
+
+NEIGHBOURS_FROM_ELEMENT_ID_QUERY = """
+MATCH (n) 
+WHERE elementId(n) = $element_id
+
+MATCH (n)<-[rels]->(m)  
+WITH n, 
+     ([n] + COLLECT(DISTINCT m)) AS allNodes, 
+     COLLECT(DISTINCT rels) AS allRels
+
+RETURN 
+    [node IN allNodes | 
+        node {
+            .*,
+            embedding: null,
+            text: null,
+            summary: null,
+            labels: [coalesce(apoc.coll.removeAll(labels(node), ['__Entity__'])[0], "*")],
+            element_id: elementId(node),
+            properties: { 
+                id: CASE WHEN node.id IS NOT NULL THEN node.id ELSE node.fileName END
+            }
+        }
+    ] AS nodes,
+    
+    [r IN allRels | 
+        {
+            start_node_element_id: elementId(startNode(r)),
+            end_node_element_id: elementId(endNode(r)),
+            type: type(r),
+            element_id: elementId(r)
+        }
+    ] AS relationships
+"""
+
+
+def get_neighbour_nodes(uri, username, password, database, element_id, query=NEIGHBOURS_FROM_ELEMENT_ID_QUERY):
+    driver = None
+
+    try:
+        logging.info(f"Querying neighbours for element_id: {element_id}")
+        driver = get_graphDB_driver(uri, username, password, database)
+        driver.verify_connectivity()
+        logging.info("Database connectivity verified.")
+
+        records, summary, keys = driver.execute_query(query,element_id=element_id)
+        nodes = records[0].get("nodes", [])
+        relationships = records[0].get("relationships", [])
+        result = {"nodes": nodes, "relationships": relationships}
+        
+        logging.info(f"Successfully retrieved neighbours for element_id: {element_id}")
+        return result
+    
+    except Exception as e:
+        logging.error(f"Error retrieving neighbours for element_id: {element_id}: {e}")
+        return {"nodes": [], "relationships": []}
+    
+    finally:
+        if driver is not None:
+            driver.close()
+            logging.info("Database driver closed.")

backend/src/shared/constants.py

@@ -121,7 +121,7 @@
 RETURN 
     d AS doc, 
     [chunk IN chunks | 
-        chunk {.*, embedding: null}
+        chunk {.*, embedding: null, element_id: elementId(chunk)}
     ] AS chunks,
     [
         node IN nodes | 
@@ -168,10 +168,10 @@
 CHAT_EMBEDDING_FILTER_SCORE_THRESHOLD = 0.10
 
 CHAT_TOKEN_CUT_OFF = {
-     ("openai-gpt-3.5",'azure_ai_gpt_35',"gemini-1.0-pro","gemini-1.5-pro", "gemini-1.5-flash","groq-llama3",'groq_llama3_70b','anthropic_claude_3_5_sonnet','fireworks_llama_v3_70b','bedrock_claude_3_5_sonnet', ) : 4, 
-     ("openai-gpt-4","diffbot" ,'azure_ai_gpt_4o',"openai-gpt-4o", "openai-gpt-4o-mini") : 28,
+     ('openai_gpt_3.5','azure_ai_gpt_35',"gemini_1.0_pro","gemini_1.5_pro", "gemini_1.5_flash","groq-llama3",'groq_llama3_70b','anthropic_claude_3_5_sonnet','fireworks_llama_v3_70b','bedrock_claude_3_5_sonnet', ) : 4, 
+     ("openai-gpt-4","diffbot" ,'azure_ai_gpt_4o',"openai_gpt_4o", "openai_gpt_4o_mini") : 28,
      ("ollama_llama3") : 2  
-} 
+}  
 
 ### CHAT TEMPLATES 
 CHAT_SYSTEM_TEMPLATE = """
@@ -473,14 +473,16 @@
             .*,
             embedding: null,
             fileName: d.fileName,
-            fileSource: d.fileSource
+            fileSource: d.fileSource, 
+            element_id: elementId(c)
         }
     ] AS chunks,
     [
         community IN communities WHERE community IS NOT NULL | 
         community {
             .*,
-            embedding: null
+            embedding: null,
+            element_id:elementId(community)
         }
     ] AS communities,
     [
@@ -551,7 +553,7 @@
 WHERE elementId(community) IN $communityids
 WITH collect(distinct community) AS communities
 RETURN [community IN communities | 
-        community {.*, embedding: null, elementid: elementId(community)}] AS communities
+        community {.*, embedding: null, element_id: elementId(community)}] AS communities
 """
 
 ## CHAT MODES 

frontend/src/components/ChatBot/ChatInfoModal.tsx

@@ -23,7 +23,7 @@ import { tokens } from '@neo4j-ndl/base';
 import ChunkInfo from './ChunkInfo';
 import EntitiesInfo from './EntitiesInfo';
 import SourcesInfo from './SourcesInfo';
-import CommunitiesInfo from './Communities';
+import CommunitiesInfo from './CommunitiesInfo';
 import { chatModeLables, supportedLLmsForRagas } from '../../utils/Constants';
 import { Relationship } from '@neo4j-nvl/base';
 import { getChatMetrics } from '../../services/GetRagasMetric';
@@ -74,6 +74,8 @@ const ChatInfoModal: React.FC<chatInfoMessage> = ({
   const [copiedText, setcopiedText] = useState<boolean>(false);
   const [showMetricsTable, setShowMetricsTable] = useState<boolean>(Boolean(metricDetails));
 
+  console.log('node', nodeDetails);
+
   const actions: CypherCodeBlockProps['actions'] = useMemo(
     () => [
       {
@@ -348,9 +350,14 @@ const ChatInfoModal: React.FC<chatInfoMessage> = ({
           <></>
         )}
       </Flex>
-      {activeTab == 4 && nodes?.length && relationships?.length ? (
+      {activeTab == 4 && nodes?.length && relationships?.length && mode !== chatModeLables.graph ? (
         <Box className='button-container flex mt-2 justify-center'>
-          <GraphViewButton nodeValues={nodes} relationshipValues={relationships} />
+          <GraphViewButton
+            nodeValues={nodes}
+            relationshipValues={relationships}
+            label='Graph Entities used for Answer Generation'
+            viewType='chatInfoView'
+          />
         </Box>
       ): <></>}
     </Box>

frontend/src/components/ChatBot/ChatModesSwitch.tsx

@@ -24,6 +24,7 @@ export default function ChatModesSwitch({
         size='small'
         clean
         onClick={() => switchToOtherMode(currentModeIndex - 1)}
+        aria-label='left'
       >
         <ChevronLeftIconSolid />
       </IconButton>
@@ -39,6 +40,7 @@ export default function ChatModesSwitch({
         size='small'
         clean
         onClick={() => switchToOtherMode(currentModeIndex + 1)}
+        aria-label='right'
       >
         <ChevronRightIconSolid />
       </IconButton>