Failed file copy to another bucket, user agent configurable and added some extra logs

Author: praveshkumar1988

backend/example.env

@@ -21,4 +21,5 @@ LANGCHAIN_PROJECT = ""
 LANGCHAIN_TRACING_V2 = ""
 LANGCHAIN_ENDPOINT = ""
 GCS_FILE_CACHE = "" #save the file into GCS or local, SHould be True or False
-NEO4J_USER_AGENT=""
+NEO4J_USER_AGENT = ""
+ENABLE_USER_AGENT = ""

backend/score.py

@@ -114,7 +114,7 @@ async def create_source_knowledge_graph_url(
             return create_api_response('Failed',message='source_type is other than accepted source')
 
         message = f"Source Node created successfully for source type: {source_type} and source: {source}"
-        josn_obj = {'api_name':'url_scan','db_url':uri,'url_scanned_file':lst_file_name}
+        josn_obj = {'api_name':'url_scan','db_url':uri,'url_scanned_file':lst_file_name, 'source_url':source_url, 'wiki_query':wiki_query}
         logger.log_struct(josn_obj)
         return create_api_response("Success",message=message,success_count=success_count,failed_count=failed_count,file_name=lst_file_name)    
     except Exception as e:
@@ -197,6 +197,9 @@ async def extract_knowledge_graph_from_file(
         if result is not None:
             result['db_url'] = uri
             result['api_name'] = 'extract'
+            result['source_url'] = source_url
+            result['wiki_query'] = wiki_query
+            result['source_type'] = source_type
         logger.log_struct(result)
         return create_api_response('Success', data=result, file_source= source_type)
     except Exception as e:
@@ -207,11 +210,13 @@ async def extract_knowledge_graph_from_file(
         if source_type == 'local file':
             if gcs_file_cache == 'True':
                 folder_name = create_gcs_bucket_folder_name_hashed(uri,file_name)
+                copy_failed_file(BUCKET_UPLOAD, BUCKET_FAILED_FILE, folder_name, file_name)
+                time.sleep(5)
                 delete_file_from_gcs(BUCKET_UPLOAD,folder_name,file_name)
             else:
                 logging.info(f'Deleted File Path: {merged_file_path} and Deleted File Name : {file_name}')
                 delete_uploaded_local_file(merged_file_path,file_name)
-        josn_obj = {'message':message,'error_message':error_message, 'file_name': file_name,'status':'Failed','db_url':uri,'failed_count':1, 'source_type': source_type}
+        josn_obj = {'message':message,'error_message':error_message, 'file_name': file_name,'status':'Failed','db_url':uri,'failed_count':1, 'source_type': source_type, 'source_url':source_url, 'wiki_query':wiki_query}
         logger.log_struct(josn_obj)
         logging.exception(f'File Failed in extraction: {josn_obj}')
         return create_api_response('Failed', message=message + error_message[:100], error=error_message, file_name = file_name)
@@ -284,7 +289,7 @@ async def chat_bot(uri=Form(None),model=Form(None),userName=Form(None), password
         logging.info(f"Total Response time is  {total_call_time:.2f} seconds")
         result["info"]["response_time"] = round(total_call_time, 2)
 
-        josn_obj = {'api_name':'chat_bot','db_url':uri}
+        josn_obj = {'api_name':'chat_bot','db_url':uri,'session_id':session_id}
         logger.log_struct(josn_obj)
         return create_api_response('Success',data=result)
     except Exception as e:
@@ -331,7 +336,7 @@ async def graph_query(
             query_type=query_type,
             document_names=document_names
         )
-        josn_obj = {'api_name':'graph_query','db_url':uri}
+        josn_obj = {'api_name':'graph_query','db_url':uri,'document_names':document_names}
         logger.log_struct(josn_obj)
         return create_api_response('Success', data=result)
     except Exception as e:

backend/src/document_sources/gcs_bucket.py

@@ -139,3 +139,14 @@ def delete_file_from_gcs(bucket_name,folder_name, file_name):
     logging.info('File deleted from GCS successfully')
   except Exception as e:
     raise Exception(e)
+  
+def copy_failed_file(source_bucket_name,dest_bucket_name,folder_name, file_name):
+  try:
+    storage_client = storage.Client()
+    bucket = storage_client.bucket(source_bucket_name)
+    folder_file_name = folder_name +'/'+file_name
+    source_blob = bucket.blob(folder_file_name)
+    bucket.copy_blob(source_blob,dest_bucket_name,file_name)
+    logging.info(f'Failed file {file_name} copied to {dest_bucket_name} from {source_bucket_name} in GCS successfully')
+  except Exception as e:
+    raise Exception(e)  

backend/src/graph_query.py

@@ -55,7 +55,11 @@ def get_graphDB_driver(uri, username, password):
     """
     try:
         logging.info(f"Attempting to connect to the Neo4j database at {uri}")
-        driver = GraphDatabase.driver(uri, auth=(username, password), user_agent=os.environ.get('NEO4J_USER_AGENT'))
+        enable_user_agent = os.environ.get("ENABLE_USER_AGENT", "False").lower() in ("true", "1", "yes")
+        if enable_user_agent:
+            driver = GraphDatabase.driver(uri, auth=(username, password), user_agent=os.environ.get('NEO4J_USER_AGENT'))
+        else:
+            driver = GraphDatabase.driver(uri, auth=(username, password))
         logging.info("Connection successful")
         return driver
     except Exception as e:

backend/src/shared/common_fn.py

@@ -77,7 +77,11 @@ def get_chunk_and_graphDocument(graph_document_list, chunkId_chunkDoc_list):
   return lst_chunk_chunkId_document  
 
 def create_graph_database_connection(uri, userName, password, database):
-  graph = Neo4jGraph(url=uri, database=database, username=userName, password=password, refresh_schema=False, sanitize=True,driver_config={'user_agent':os.environ.get('NEO4J_USER_AGENT')})  
+  enable_user_agent = os.environ.get("ENABLE_USER_AGENT", "False").lower() in ("true", "1", "yes")
+  if enable_user_agent:
+    graph = Neo4jGraph(url=uri, database=database, username=userName, password=password, refresh_schema=False, sanitize=True,driver_config={'user_agent':os.environ.get('NEO4J_USER_AGENT')})  
+  else:
+    graph = Neo4jGraph(url=uri, database=database, username=userName, password=password, refresh_schema=False, sanitize=True)    
   return graph
 
 

backend/src/shared/constants.py

@@ -11,6 +11,7 @@
 GEMINI_MODELS = ["gemini-1.0-pro", "gemini-1.5-pro"]
 GROQ_MODELS = ["groq-llama3"]
 BUCKET_UPLOAD = 'llm-graph-builder-upload'
+BUCKET_FAILED_FILE = 'llm-graph-builder-failed'
 PROJECT_ID = 'llm-experiments-387609'