Staging to Production

backend/score.py

@@ -4,6 +4,7 @@
 from src.main import *
 from src.QA_integration import *
 from src.shared.common_fn import *
+from src.shared.llm_graph_builder_exception import LLMGraphBuilderException
 import uvicorn
 import asyncio
 import base64
@@ -307,7 +308,8 @@ async def extract_knowledge_graph_from_file(
         node_detail = graphDb_data_Access.get_current_status_document_node(file_name)
         # Set the status "Completed" in logging becuase we are treating these error already handled by application as like custom errors.
         json_obj = {'api_name':'extract','message':error_message,'file_created_at':formatted_time(node_detail[0]['created_time']),'error_message':error_message, 'file_name': file_name,'status':'Completed',
-                    'db_url':uri, 'userName':userName, 'database':database,'success_count':1, 'source_type': source_type, 'source_url':source_url, 'wiki_query':wiki_query, 'logging_time': formatted_time(datetime.now(timezone.utc)),'email':email}
+                    'db_url':uri, 'userName':userName, 'database':database,'success_count':1, 'source_type': source_type, 'source_url':source_url, 'wiki_query':wiki_query, 'logging_time': formatted_time(datetime.now(timezone.utc)),'email':email,
+                    'allowedNodes': allowedNodes, 'allowedRelationship': allowedRelationship}
         logger.log_struct(json_obj, "INFO")
         logging.exception(f'File Failed in extraction: {e}')
         return create_api_response("Failed", message = error_message, error=error_message, file_name=file_name)
@@ -322,7 +324,8 @@ async def extract_knowledge_graph_from_file(
         node_detail = graphDb_data_Access.get_current_status_document_node(file_name)
 
         json_obj = {'api_name':'extract','message':message,'file_created_at':formatted_time(node_detail[0]['created_time']),'error_message':error_message, 'file_name': file_name,'status':'Failed',
-                    'db_url':uri, 'userName':userName, 'database':database,'failed_count':1, 'source_type': source_type, 'source_url':source_url, 'wiki_query':wiki_query, 'logging_time': formatted_time(datetime.now(timezone.utc)),'email':email}
+                    'db_url':uri, 'userName':userName, 'database':database,'failed_count':1, 'source_type': source_type, 'source_url':source_url, 'wiki_query':wiki_query, 'logging_time': formatted_time(datetime.now(timezone.utc)),'email':email,
+                    'allowedNodes': allowedNodes, 'allowedRelationship': allowedRelationship}
         logger.log_struct(json_obj, "ERROR")
         logging.exception(f'File Failed in extraction: {e}')
         return create_api_response('Failed', message=message + error_message[:100], error=error_message, file_name = file_name)
@@ -341,14 +344,6 @@ async def get_source_list(
     """
     try:
         start = time.time()
-        # if password is not None and password != "null":
-        #     decoded_password = decode_password(password)
-        # else:
-        #     decoded_password = None
-        #     userName = None
-        #     database = None
-        # if " " in uri:
-        #     uri = uri.replace(" ","+")
         result = await asyncio.to_thread(get_source_list_from_graph,uri,userName,password,database)
         end = time.time()
         elapsed_time = end - start

backend/src/graphDB_dataAccess.py

@@ -29,10 +29,10 @@ def update_exception_db(self, file_name, exp_msg, retry_condition=None):
             if retry_condition is not None: 
                 retry_condition = None
                 self.graph.query("""MERGE(d:Document {fileName :$fName}) SET d.status = $status, d.errorMessage = $error_msg, d.retry_condition = $retry_condition""",
-                            {"fName":file_name, "status":job_status, "error_msg":exp_msg, "retry_condition":retry_condition})
+                            {"fName":file_name, "status":job_status, "error_msg":exp_msg, "retry_condition":retry_condition},session_params={"database":self.graph._database})
             else :    
                 self.graph.query("""MERGE(d:Document {fileName :$fName}) SET d.status = $status, d.errorMessage = $error_msg""",
-                            {"fName":file_name, "status":job_status, "error_msg":exp_msg})
+                            {"fName":file_name, "status":job_status, "error_msg":exp_msg},session_params={"database":self.graph._database})
         except Exception as e:
             error_message = str(e)
             logging.error(f"Error in updating document node status as failed: {error_message}")
@@ -66,7 +66,7 @@ def create_source_node(self, obj_source_node:sourceNode):
                             "entityEntityRelCount":obj_source_node.entityEntityRelCount,
                             "communityNodeCount":obj_source_node.communityNodeCount,
                             "communityRelCount":obj_source_node.communityRelCount
-                            })
+                            },session_params={"database":self.graph._database})
         except Exception as e:
             error_message = str(e)
             logging.info(f"error_message = {error_message}")
@@ -118,7 +118,7 @@ def update_source_node(self, obj_source_node:sourceNode):
             logging.info(f'Base Param value 1 : {param}')
             query = "MERGE(d:Document {fileName :$props.fileName}) SET d += $props"
             logging.info("Update source node properties")
-            self.graph.query(query,param)
+            self.graph.query(query,param,session_params={"database":self.graph._database})
         except Exception as e:
             error_message = str(e)
             self.update_exception_db(self,self.file_name,error_message)
@@ -139,15 +139,15 @@ def get_source_list(self):
         """
         logging.info("Get existing files list from graph")
         query = "MATCH(d:Document) WHERE d.fileName IS NOT NULL RETURN d ORDER BY d.updatedAt DESC"
-        result = self.graph.query(query)
+        result = self.graph.query(query,session_params={"database":self.graph._database})
         list_of_json_objects = [entry['d'] for entry in result]
         return list_of_json_objects
 
     def update_KNN_graph(self):
         """
         Update the graph node with SIMILAR relationship where embedding scrore match
         """
-        index = self.graph.query("""show indexes yield * where type = 'VECTOR' and name = 'vector'""")
+        index = self.graph.query("""show indexes yield * where type = 'VECTOR' and name = 'vector'""",session_params={"database":self.graph._database})
         # logging.info(f'show index vector: {index}')
         knn_min_score = os.environ.get('KNN_MIN_SCORE')
         if len(index) > 0:
@@ -158,14 +158,14 @@ def update_KNN_graph(self):
                                     WHERE node <> c and score >= $score MERGE (c)-[rel:SIMILAR]-(node) SET rel.score = score
                                 """,
                                 {"score":float(knn_min_score)}
-                                )
+                                ,session_params={"database":self.graph._database})
         else:
             logging.info("Vector index does not exist, So KNN graph not update")
 
     def check_account_access(self, database):
         try:
             query_dbms_componenet = "call dbms.components() yield edition"
-            result_dbms_componenet = self.graph.query(query_dbms_componenet)
+            result_dbms_componenet = self.graph.query(query_dbms_componenet,session_params={"database":self.graph._database})
 
             if  result_dbms_componenet[0]["edition"] == "enterprise":
                 query = """
@@ -177,7 +177,7 @@ def check_account_access(self, database):
 
                 logging.info(f"Checking access for database: {database}")
 
-                result = self.graph.query(query, params={"database": database})
+                result = self.graph.query(query, params={"database": database},session_params={"database":self.graph._database})
                 read_access_count = result[0]["readAccessCount"] if result else 0
 
                 logging.info(f"Read access count: {read_access_count}")
@@ -202,7 +202,7 @@ def check_gds_version(self):
             gds_procedure_count = """
             SHOW FUNCTIONS YIELD name WHERE name STARTS WITH 'gds.version' RETURN COUNT(*) AS totalGdsProcedures
             """
-            result = self.graph.query(gds_procedure_count)
+            result = self.graph.query(gds_procedure_count,session_params={"database":self.graph._database})
             total_gds_procedures = result[0]['totalGdsProcedures'] if result else 0
 
             if total_gds_procedures > 0:
@@ -231,11 +231,11 @@ def connection_check_and_get_vector_dimensions(self,database):
         db_vector_dimension = self.graph.query("""SHOW INDEXES YIELD *
                                     WHERE type = 'VECTOR' AND name = 'vector'
                                     RETURN options.indexConfig['vector.dimensions'] AS vector_dimensions
-                                """)
+                                """,session_params={"database":self.graph._database})
 
         result_chunks = self.graph.query("""match (c:Chunk) return size(c.embedding) as embeddingSize, count(*) as chunks, 
                                                     count(c.embedding) as hasEmbedding
-                                """)
+                                """,session_params={"database":self.graph._database})
 
         embedding_model = os.getenv('EMBEDDING_MODEL')
         embeddings, application_dimension = load_embedding_model(embedding_model)
@@ -260,7 +260,7 @@ def execute_query(self, query, param=None,max_retries=3, delay=2):
         retries = 0
         while retries < max_retries:
             try:
-                return self.graph.query(query, param)
+                return self.graph.query(query, param,session_params={"database":self.graph._database})
             except TransientError as e:
                 if "DeadlockDetected" in str(e):
                     retries += 1
@@ -473,8 +473,8 @@ def drop_create_vector_index(self, isVectorIndexExist):
         embeddings, dimension = load_embedding_model(embedding_model)
 
         if isVectorIndexExist == 'true':
-            self.graph.query("""drop index vector""")
-        # self.graph.query("""drop index vector""")
+            self.graph.query("""drop index vector""",session_params={"database":self.graph._database})
+
         self.graph.query("""CREATE VECTOR INDEX `vector` if not exists for (c:Chunk) on (c.embedding)
                             OPTIONS {indexConfig: {
                             `vector.dimensions`: $dimensions,
@@ -483,7 +483,7 @@ def drop_create_vector_index(self, isVectorIndexExist):
                         """,
                         {
                             "dimensions" : dimension
-                        }
+                        },session_params={"database":self.graph._database}
                         )
         return "Drop and Re-Create vector index succesfully"
 

backend/src/llm.py

@@ -14,8 +14,9 @@
 import boto3
 import google.auth
 from src.shared.constants import ADDITIONAL_INSTRUCTIONS
+from src.shared.llm_graph_builder_exception import LLMGraphBuilderException
 import re
-import json
+from typing import List
 
 def get_llm(model: str):
     """Retrieve the specified language model based on the model name."""
@@ -127,6 +128,14 @@ def get_llm(model: str):
     logging.info(f"Model created - Model Version: {model}")
     return llm, model_name
 
+def get_llm_model_name(llm):
+    """Extract name of llm model from llm object"""
+    for attr in ["model_name", "model", "model_id"]:
+        model_name = getattr(llm, attr, None)
+        if model_name:
+            return model_name.lower()
+    print("Could not determine model name; defaulting to empty string")
+    return ""
 
 def get_combined_chunks(chunkId_chunkDoc_list, chunks_to_combine):
     combined_chunk_document_list = []
@@ -181,8 +190,9 @@ async def get_graph_document_list(
             node_properties = ["description"]
             relationship_properties = ["description"]
         TOOL_SUPPORTED_MODELS = {"qwen3", "deepseek"} 
-        model_name = llm.model_name.lower() 
+        model_name = get_llm_model_name(llm)
         ignore_tool_usage = not any(pattern in model_name for pattern in TOOL_SUPPORTED_MODELS)
+        logging.info(f"Keeping ignore tool usage parameter as {ignore_tool_usage}")
         llm_transformer = LLMGraphTransformer(
             llm=llm,
             node_properties=node_properties,
@@ -200,21 +210,45 @@ async def get_graph_document_list(
     return graph_document_list
 
 async def get_graph_from_llm(model, chunkId_chunkDoc_list, allowedNodes, allowedRelationship, chunks_to_combine, additional_instructions=None):
+   try:
+       llm, model_name = get_llm(model)
+       logging.info(f"Using model: {model_name}")
 
-    llm, model_name = get_llm(model)
-    combined_chunk_document_list = get_combined_chunks(chunkId_chunkDoc_list, chunks_to_combine)
+       combined_chunk_document_list = get_combined_chunks(chunkId_chunkDoc_list, chunks_to_combine)
+       logging.info(f"Combined {len(combined_chunk_document_list)} chunks")
 
-    allowedNodes = allowedNodes.split(',') if allowedNodes else []
+       allowed_nodes = [node.strip() for node in allowedNodes.split(',') if node.strip()]
+       logging.info(f"Allowed nodes: {allowed_nodes}")
+
+       allowed_relationships = []
+       if allowedRelationship:
+           items = [item.strip() for item in allowedRelationship.split(',') if item.strip()]
+           if len(items) % 3 != 0:
+               raise LLMGraphBuilderException("allowedRelationship must be a multiple of 3 (source, relationship, target)")
+           for i in range(0, len(items), 3):
+               source, relation, target = items[i:i + 3]
+               if source not in allowed_nodes or target not in allowed_nodes:
+                   raise LLMGraphBuilderException(
+                       f"Invalid relationship ({source}, {relation}, {target}): "
+                       f"source or target not in allowedNodes"
+                   )
+               allowed_relationships.append((source, relation, target))
+           logging.info(f"Allowed relationships: {allowed_relationships}")
+       else:
+           logging.info("No allowed relationships provided")
 
-    if not allowedRelationship:
-        allowedRelationship = []
-    else:
-        items = allowedRelationship.split(',')
-        allowedRelationship = [tuple(items[i:i+3]) for i in range(0, len(items), 3)]
-    graph_document_list = await get_graph_document_list(
-        llm, combined_chunk_document_list, allowedNodes, allowedRelationship, additional_instructions
-    )
-    return graph_document_list
+       graph_document_list = await get_graph_document_list(
+           llm,
+           combined_chunk_document_list,
+           allowed_nodes,
+           allowed_relationships,
+           additional_instructions
+       )
+       logging.info(f"Generated {len(graph_document_list)} graph documents")
+       return graph_document_list
+   except Exception as e:
+       logging.error(f"Error in get_graph_from_llm: {e}", exc_info=True)
+       raise LLMGraphBuilderException(f"Error in getting graph from llm: {e}")
 
 def sanitize_additional_instruction(instruction: str) -> str:
    """

frontend/src/components/ChatBot/ChunkInfo.tsx

@@ -15,7 +15,6 @@ import { handleGraphNodeClick } from './chatInfo';
 import { IconButtonWithToolTip } from '../UI/IconButtonToolTip';
 import remarkGfm from 'remark-gfm';
 import rehypeRaw from 'rehype-raw';
-
 const ChunkInfo: FC<ChunkProps> = ({ loading, chunks, mode }) => {
   const themeUtils = useContext(ThemeWrapperContext);
   const [neoNodes, setNeoNodes] = useState<any[]>([]);

frontend/src/components/Content.tsx

@@ -929,7 +929,12 @@ const Content: React.FC<ContentProps> = ({
           <div className='connectionstatus__container'>
             <span className='h6 px-1'>Neo4j connection {isReadOnlyUser ? '(Read only Mode)' : ''}</span>
             <Typography variant='body-medium'>
-              <DatabaseStatusIcon isConnected={connectionStatus} isGdsActive={isGdsActive} uri={userCredentials?.uri} />
+              <DatabaseStatusIcon
+                isConnected={connectionStatus}
+                isGdsActive={isGdsActive}
+                uri={userCredentials?.uri}
+                database={userCredentials?.database}
+              />
               <div className='pt-1 flex! gap-1 items-center'>
                 <div>{!hasSelections ? <StatusIndicator type='danger' /> : <StatusIndicator type='success' />}</div>
                 <div>

frontend/src/components/Layout/PageLayout.tsx

@@ -22,7 +22,8 @@ import LoadDBSchemaDialog from '../Popups/GraphEnhancementDialog/EnitityExtracti
 import PredefinedSchemaDialog from '../Popups/GraphEnhancementDialog/EnitityExtraction/PredefinedSchemaDialog';
 import { SKIP_AUTH } from '../../utils/Constants';
 import { useNavigate } from 'react-router';
-import { deduplicateByRelationshipTypeOnly, deduplicateNodeByValue } from '../../utils/Utils';
+import { deduplicateByFullPattern, deduplicateNodeByValue } from '../../utils/Utils';
+
 
 const GCSModal = lazy(() => import('../DataSources/GCS/GCSModal'));
 const S3Modal = lazy(() => import('../DataSources/AWS/S3Modal'));
@@ -378,7 +379,7 @@ const PageLayout: React.FC = () => {
       setSchemaValRels(rels);
       setCombinedRelsVal((prevRels: OptionType[]) => {
         const combined = [...rels, ...prevRels];
-        return deduplicateByRelationshipTypeOnly(combined);
+        return deduplicateByFullPattern(combined);
       });
       setSchemaView('text');
       localStorage.setItem(LOCAL_KEYS.source, JSON.stringify(updatedSource));
@@ -418,7 +419,7 @@ const PageLayout: React.FC = () => {
       setDbRels(rels);
       setCombinedRelsVal((prevRels: OptionType[]) => {
         const combined = [...rels, ...prevRels];
-        return deduplicateByRelationshipTypeOnly(combined);
+        return deduplicateByFullPattern(combined);
       });
       localStorage.setItem(LOCAL_KEYS.source, JSON.stringify(updatedSource));
       localStorage.setItem(LOCAL_KEYS.type, JSON.stringify(updatedType));
@@ -456,7 +457,7 @@ const PageLayout: React.FC = () => {
       setPreDefinedRels(rels);
       setCombinedRelsVal((prevRels: OptionType[]) => {
         const combined = [...rels, ...prevRels];
-        return deduplicateByRelationshipTypeOnly(combined);
+        return deduplicateByFullPattern(combined);
       });
       localStorage.setItem(LOCAL_KEYS.source, JSON.stringify(updatedSource));
       localStorage.setItem(LOCAL_KEYS.type, JSON.stringify(updatedType));

frontend/src/components/Popups/GraphEnhancementDialog/EnitityExtraction/NewEntityExtractionSetting.tsx

@@ -13,7 +13,7 @@ import {
   updateLocalStorage,
   extractOptions,
   parseRelationshipString,
-  deduplicateByRelationshipTypeOnly,
+  deduplicateByFullPattern,
   deduplicateNodeByValue,
 } from '../../../../utils/Utils';
 import TooltipWrapper from '../../../UI/TipWrapper';
@@ -175,15 +175,15 @@ export default function NewEntityExtractionSetting({
         });
         setUserDefinedRels((prev: OptionType[]) => {
           const combined = [...prev, ...relationshipTypeOptions];
-          return deduplicateByRelationshipTypeOnly(combined);
+          return deduplicateByFullPattern(combined);
         });
         setCombinedNodes((prev: OptionType[]) => {
           const combined = [...prev, ...nodeLabelOptions];
           return deduplicateNodeByValue(combined);
         });
         setCombinedRels((prev: OptionType[]) => {
           const combined = [...prev, ...relationshipTypeOptions];
-          return deduplicateByRelationshipTypeOnly(combined);
+          return deduplicateByFullPattern(combined);
         });
         setTupleOptions((prev) => [...updatedTuples, ...prev]);
       } else {