Merge branch 'dev' of https://github.com/neo4j-labs/llm-graph-builder into backend_connection_config

Author: prakriti-solankey

backend/example.env

@@ -49,3 +49,5 @@ BEDROCK_EMBEDDING_MODEL="model_name,aws_access_key,aws_secret_key,region_name"
 LLM_MODEL_CONFIG_bedrock_nova_micro_v1="model_name,aws_access_key,aws_secret_key,region_name"        #model_name="amazon.nova-micro-v1:0"
 LLM_MODEL_CONFIG_bedrock_nova_lite_v1="model_name,aws_access_key,aws_secret_key,region_name"         #model_name="amazon.nova-lite-v1:0"
 LLM_MODEL_CONFIG_bedrock_nova_pro_v1="model_name,aws_access_key,aws_secret_key,region_name"          #model_name="amazon.nova-pro-v1:0"
+LLM_MODEL_CONFIG_fireworks_deepseek_r1="model_name,fireworks_api_key"      #model_name="accounts/fireworks/models/deepseek-r1"
+LLM_MODEL_CONFIG_fireworks_deepseek_v3="model_name,fireworks_api_key"      #model_name="accounts/fireworks/models/deepseek-v3"

backend/requirements.txt

@@ -1,33 +1,33 @@
 asyncio==3.4.3
-boto3==1.35.90
-botocore==1.35.90
+boto3==1.36.2
+botocore==1.36.2
 certifi==2024.8.30
 fastapi==0.115.6
 fastapi-health==0.4.0
-google-api-core==2.23.0
-google-auth==2.36.0
+google-api-core==2.24.0
+google-auth==2.37.0
 google_auth_oauthlib==1.2.1
 google-cloud-core==2.4.1
 json-repair==0.30.2
 pip-install==1.3.5
-langchain==0.3.13
-langchain-aws==0.2.10
-langchain-anthropic==0.3.0
-langchain-fireworks==0.2.5
-langchain-community==0.3.13
-langchain-core==0.3.28
+langchain==0.3.15
+langchain-aws==0.2.11
+langchain-anthropic==0.3.3
+langchain-fireworks==0.2.6
+langchain-community==0.3.15
+langchain-core==0.3.31
 langchain-experimental==0.3.4
-langchain-google-vertexai==2.0.7
-langchain-groq==0.2.1
-langchain-openai==0.2.14
-langchain-text-splitters==0.3.4
+langchain-google-vertexai==2.0.11
+langchain-groq==0.2.3
+langchain-openai==0.3.1
+langchain-text-splitters==0.3.5
 langchain-huggingface==0.1.2
 langdetect==1.0.9
-langsmith==0.2.4
-langserve==0.3.0
+langsmith==0.2.11
+langserve==0.3.1
 neo4j-rust-ext
 nltk==3.9.1
-openai==1.58.1
+openai==1.59.9
 opencv-python==4.10.0.84
 psutil==6.1.0
 pydantic==2.9.2
@@ -56,7 +56,6 @@ google-cloud-logging==3.11.3
 pypandoc==1.13
 graphdatascience==1.12
 Secweb==1.11.0
-ragas==0.2.6
+ragas==0.2.11
 rouge_score==0.1.2
-langchain-neo4j==0.2.0
-
+langchain-neo4j==0.3.0

backend/score.py

@@ -11,7 +11,7 @@
 from langchain_google_vertexai import ChatVertexAI
 from src.api_response import create_api_response
 from src.graphDB_dataAccess import graphDBdataAccess
-from src.graph_query import get_graph_results,get_chunktext_results
+from src.graph_query import get_graph_results,get_chunktext_results,visualize_schema
 from src.chunkid_entities import get_entities_from_chunkids
 from src.post_processing import create_vector_fulltext_indexes, create_entity_embedding, graph_schema_consolidation
 from sse_starlette.sse import EventSourceResponse
@@ -82,7 +82,7 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send):
 app = FastAPI()
 app.add_middleware(XContentTypeOptions)
 app.add_middleware(XFrame, Option={'X-Frame-Options': 'DENY'})
-app.add_middleware(CustomGZipMiddleware, minimum_size=1000, compresslevel=5,paths=["/sources_list","/url/scan","/extract","/chat_bot","/chunk_entities","/get_neighbours","/graph_query","/schema","/populate_graph_schema","/get_unconnected_nodes_list","/get_duplicate_nodes","/fetch_chunktext"])
+app.add_middleware(CustomGZipMiddleware, minimum_size=1000, compresslevel=5,paths=["/sources_list","/url/scan","/extract","/chat_bot","/chunk_entities","/get_neighbours","/graph_query","/schema","/populate_graph_schema","/get_unconnected_nodes_list","/get_duplicate_nodes","/fetch_chunktext","/schema_visualization"])
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -1033,6 +1033,32 @@ async def backend_connection_configuration():
         return create_api_response(job_status, message=message, error=error_message.rstrip('.') + ', or fill from the login dialog.', data=graph_connection)
     finally:
         gc.collect()
-        
+    
+@app.post("/schema_visualization")
+async def get_schema_visualization(uri=Form(None), userName=Form(None), password=Form(None), database=Form(None)):
+    try:
+        start = time.time()
+        result = await asyncio.to_thread(visualize_schema,
+           uri=uri,
+           userName=userName,
+           password=password,
+           database=database)
+        if result:
+            logging.info("Graph schema visualization query successful")
+        end = time.time()
+        elapsed_time = end - start
+        logging.info(f'Schema result from DB: {result}')
+        json_obj = {'api_name':'schema_visualization','db_url':uri, 'userName':userName, 'database':database, '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:
+        message="Unable to get schema visualization from neo4j database"
+        error_message = str(e)
+        logging.info(message)
+        logging.exception(f'Exception:{error_message}')
+        return create_api_response("Failed", message=message, error=error_message)
+    finally:
+        gc.collect()
+
 if __name__ == "__main__":
     uvicorn.run(app)

backend/src/diffbot_transformer.py

@@ -8,4 +8,4 @@ def get_graph_from_diffbot(graph,chunkId_chunkDoc_list:List):
     combined_chunk_document_list = get_combined_chunks(chunkId_chunkDoc_list)
     llm,model_name = get_llm('diffbot')
     graph_documents = llm.convert_to_graph_documents(combined_chunk_document_list)
-    return graph_documents
+    return graph_documents

backend/src/document_sources/local_file.py

@@ -62,4 +62,4 @@ def get_pages_with_page_numbers(unstructured_pages):
                             'filetype':page.metadata['filetype']}
             if page == unstructured_pages[-1]:
                     pages.append(Document(page_content = page_content, metadata=metadata_with_custom_page_number))
-    return pages                
+    return pages                

backend/src/document_sources/web_pages.py

@@ -13,4 +13,4 @@ def get_documents_from_web_page(source_url:str):
       file_name = last_url_segment(source_url)
     return file_name, pages
   except Exception as e:
-    raise LLMGraphBuilderException(str(e))
+    raise LLMGraphBuilderException(str(e))

backend/src/graph_query.py

@@ -3,7 +3,8 @@
 from neo4j import GraphDatabase
 import os
 import json
-from src.shared.constants import GRAPH_CHUNK_LIMIT,GRAPH_QUERY,CHUNK_TEXT_QUERY,COUNT_CHUNKS_QUERY
+
+from src.shared.constants import GRAPH_CHUNK_LIMIT,GRAPH_QUERY,CHUNK_TEXT_QUERY,COUNT_CHUNKS_QUERY,SCHEMA_VISUALIZATION_QUERY
 
 def get_graphDB_driver(uri, username, password,database="neo4j"):
     """
@@ -259,4 +260,23 @@ def get_chunktext_results(uri, username, password, database, document_name, page
        raise Exception("An error occurred in get_chunktext_results. Please check the logs for more details.") from e
    finally:
        if driver:
-           driver.close()
+           driver.close()
+
+
+def visualize_schema(uri, userName, password, database):
+   """Retrieves graph schema"""
+   driver = None
+   try:
+       logging.info("Starting visualizing graph schema")
+       driver = GraphDatabase.driver(uri, auth=(userName, password),database=database)
+       records, summary, keys = driver.execute_query(SCHEMA_VISUALIZATION_QUERY)
+       nodes = records[0].get("nodes", [])
+       relationships = records[0].get("relationships", [])
+       result = {"nodes": nodes, "relationships": relationships}
+       return result
+   except Exception as e:
+       logging.error(f"An error occurred schema retrieval. Error: {str(e)}")
+       raise Exception(f"An error occurred schema retrieval. Error: {str(e)}")
+   finally:
+       if driver:
+           driver.close()

backend/src/make_relationships.py

@@ -161,7 +161,8 @@ def create_chunk_vector_index(graph):
                                     graph=graph,
                                     node_label="Chunk", 
                                     embedding_node_property="embedding",
-                                    index_name="vector"
+                                    index_name="vector",
+                                    embedding_dimension=EMBEDDING_DIMENSION
                                     )
             vector_store.create_new_index()
             logging.info(f"Index created successfully. Time taken: {time.time() - start_time:.2f} seconds")

backend/src/post_processing.py

@@ -233,6 +233,3 @@ def graph_schema_consolidation(graph):
         graph.query(query)
 
     return None
-                      
-    
-    

backend/src/shared/constants.py

@@ -1,4 +1,3 @@
-
 OPENAI_MODELS = ["openai-gpt-3.5", "openai-gpt-4o", "openai-gpt-4o-mini"]
 GEMINI_MODELS = ["gemini-1.0-pro", "gemini-1.5-pro", "gemini-1.5-flash"]
 GROQ_MODELS = ["groq-llama3"]
@@ -893,3 +892,19 @@
 types such as dates, numbers, revenues, and other non-entity information are not extracted as separate nodes.
 Instead, treat these as properties associated with the relevant entities."""
 
+SCHEMA_VISUALIZATION_QUERY = """
+CALL db.schema.visualization() YIELD nodes, relationships
+RETURN
+  [n IN nodes | {
+      element_id: elementId(n),
+      labels: labels(n),
+      properties: apoc.any.properties(n)
+  }] AS nodes,
+  [r IN relationships | {
+      type: type(r),
+      properties: apoc.any.properties(r),
+      element_id: elementId(r),
+      start_node_element_id: elementId(startNode(r)),
+      end_node_element_id: elementId(endNode(r))
+  }] AS relationships;
+"""