Compare commits
8 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 501c224540 | |||
| 64c45ee66c | |||
| b221ae4b42 | |||
| 31e1186573 | |||
| 1decc02218 | |||
| 5cdedc3403 | |||
| 194597adbc | |||
| 38e9965d2e |
162
README.md
162
README.md
@@ -1,76 +1,159 @@
|
|||||||
# dspy-neo4j-knowledge-graph
|
# text-to-cypher
|
||||||
LLM-driven automated knowledge graph construction from text using DSPy and Neo4j.
|
LLM-driven automated knowledge graph construction from text using DSPy and Neo4j.
|
||||||
|
|
||||||

|
|
||||||
|
|
||||||
## Project Structure
|
## Project Structure
|
||||||
```sh
|
```sh
|
||||||
dspy-neo4j-knowledge-graph/
|
text-to-cypher/
|
||||||
├── README.md
|
├── README.md
|
||||||
├── examples
|
├── main.py
|
||||||
├── requirements.txt
|
├── pyproject.toml
|
||||||
├── run.py
|
├── uv.lock
|
||||||
└── src
|
└── src/
|
||||||
|
├── __init__.py
|
||||||
|
└── neo4j.py
|
||||||
```
|
```
|
||||||
|
|
||||||
## Description
|
## Description
|
||||||
Model entities and relationships and build a Knowledge Graph using DSPy, Neo4j, and OpenAI's GPT-4. When given a paragraph or block of text, the app uses the DSPy library and OpenAI's GPT-4 to extract entities and relationships and generate a Cypher statement which is run in Neo4j to create the Knowledge Graph.
|
Build knowledge graphs automatically from text using DSPy, Modaic, and Neo4j. This implementation uses OpenAI's GPT-4o to extract entities and relationships from Wikipedia abstracts, generating Cypher statements that create structured knowledge graphs in Neo4j.
|
||||||
|
|
||||||
|
### Key Features
|
||||||
|
- **DSPy-Powered**: Uses DSPy's Chain of Thought for structured entity and relationship extraction
|
||||||
|
- **Modaic Integration**: Leverages Modaic's PrecompiledProgram for reusable, shareable DSPy programs
|
||||||
|
- **Schema-Aware**: Passes the current Neo4j graph schema to the model, enabling it to reuse existing nodes and relationships
|
||||||
|
- **Batch Processing**: Processes multiple text samples from NDJSON files
|
||||||
|
- **Hugging Face Hub**: Push trained programs to the Hub for sharing and versioning
|
||||||
|
|
||||||
### Optimized Schema Context
|
### Optimized Schema Context
|
||||||
The current graph schema is passed to the model as a list of nodes, relationships and properties in the context of the prompt. This allows the model to use elements from the existing schema and make connections between existing entities and relationships.
|
The current graph schema is passed to the model as a list of nodes, relationships and properties in the context of the prompt. This allows the model to use elements from the existing schema and make connections between existing entities and relationships.
|
||||||
|
|
||||||
## Quick Start
|
## Quick Start
|
||||||
1. Clone the repository.
|
1. Clone the repository
|
||||||
2. Create a [Python virtual environment](#python-virtual-environment) and install the required packages.
|
2. Install dependencies using [uv](#installation-with-uv)
|
||||||
3. Create a `.env` file and add the required [environment variables](#environment-variables).
|
3. Create a `.env` file and add the required [environment variables](#environment-variables)
|
||||||
4. [Run Neo4j using Docker](#usage).
|
4. Set up [Neo4j](#neo4j-setup) (local Docker or cloud-hosted)
|
||||||
5. Run `python3 run.py` and paste your text in the prompt.
|
5. Run `uv run main.py` to process example Wikipedia abstracts
|
||||||
6. Navigate to `http://localhost:7474/browser/` to view the Knowledge Graph in Neo4j Browser.
|
6. View your Knowledge Graph in Neo4j Browser
|
||||||
|
|
||||||
## Installation
|
## Installation
|
||||||
|
|
||||||
### Prerequisites
|
### Prerequisites
|
||||||
* Python 3.12
|
* Python 3.13+
|
||||||
* OpenAI API Key
|
* OpenAI API Key
|
||||||
* Docker
|
* [uv](https://docs.astral.sh/uv/) (Python package manager)
|
||||||
|
* Neo4j instance (local Docker or cloud-hosted)
|
||||||
|
|
||||||
|
### Installation with uv
|
||||||
|
Install dependencies using uv:
|
||||||
|
```sh
|
||||||
|
# Install uv if you haven't already
|
||||||
|
curl -LsSf https://astral.sh/uv/install.sh | sh
|
||||||
|
|
||||||
|
# Install project dependencies
|
||||||
|
uv sync
|
||||||
|
```
|
||||||
|
|
||||||
### Environment Variables
|
### Environment Variables
|
||||||
Before you begin, make sure to create a `.env` file and add your OpenAI API key.
|
Create a `.env` file in the project root with the following variables:
|
||||||
|
|
||||||
```sh
|
```sh
|
||||||
NEO4J_URI=bolt://localhost:7687
|
# OpenAI API Key
|
||||||
OPENAI_API_KEY=<your-api-key>
|
OPENAI_API_KEY=<your-openai-api-key>
|
||||||
|
|
||||||
|
# Neo4j Configuration
|
||||||
|
NEO4J_URI=bolt://localhost:7687 # or neo4j+s://xxx.databases.neo4j.io for cloud
|
||||||
|
NEO4J_USER=neo4j # optional for local Docker with NEO4J_AUTH=none
|
||||||
|
NEO4J_PASSWORD=<your-password> # optional for local Docker with NEO4J_AUTH=none
|
||||||
|
|
||||||
|
# Modaic Token (optional, for pushing to Hub)
|
||||||
|
MODAIC_TOKEN=<your-modaic-token>
|
||||||
```
|
```
|
||||||
|
|
||||||
### Python Virtual Environment
|
## Neo4j Setup
|
||||||
Create a Python virtual environment and install the required packages.
|
|
||||||
```sh
|
|
||||||
python3 -m venv .venv
|
|
||||||
source .venv/bin/activate
|
|
||||||
pip install --upgrade pip
|
|
||||||
pip install -r requirements.txt
|
|
||||||
```
|
|
||||||
|
|
||||||
## Usage
|
### Option 1: Local Docker (Development)
|
||||||
Run Neo4j using Docker.
|
Run Neo4j locally using Docker:
|
||||||
```sh
|
```sh
|
||||||
docker run \
|
docker run \
|
||||||
--name dspy-kg \
|
--name text-to-cypher \
|
||||||
--publish=7474:7474 \
|
--publish=7474:7474 \
|
||||||
--publish=7687:7687 \
|
--publish=7687:7687 \
|
||||||
--env "NEO4J_AUTH=none" \
|
--env "NEO4J_AUTH=none" \
|
||||||
neo4j:5.15
|
neo4j:5.15
|
||||||
```
|
```
|
||||||
|
|
||||||
## Clean Up
|
Access Neo4j Browser at `http://localhost:7474`
|
||||||
Stop and remove the Neo4j container.
|
|
||||||
|
### Option 2: Neo4j Aura (Cloud)
|
||||||
|
1. Create a free instance at [neo4j.com/cloud/aura](https://neo4j.com/cloud/aura)
|
||||||
|
2. Get your connection URI (e.g., `neo4j+s://xxx.databases.neo4j.io`)
|
||||||
|
3. Add credentials to your `.env` file
|
||||||
|
|
||||||
|
## Usage
|
||||||
|
|
||||||
|
### Process Wikipedia Abstracts
|
||||||
|
Run the main script to process example Wikipedia abstracts and build a knowledge graph:
|
||||||
```sh
|
```sh
|
||||||
docker stop dspy-kg
|
uv run main.py
|
||||||
docker rm dspy-kg
|
|
||||||
```
|
```
|
||||||
|
|
||||||
Deactivate the Python virtual environment.
|
This will:
|
||||||
|
1. Load Wikipedia abstracts from `examples/wikipedia-abstracts-v0_0_1.ndjson`
|
||||||
|
2. For each abstract, generate a Cypher statement using GPT-4o
|
||||||
|
3. Execute the Cypher statement in Neo4j
|
||||||
|
4. Build a connected knowledge graph
|
||||||
|
|
||||||
|
### View Your Knowledge Graph
|
||||||
|
Navigate to Neo4j Browser:
|
||||||
|
- Local: `http://localhost:7474/browser/`
|
||||||
|
- Cloud: Your Neo4j Aura console URL
|
||||||
|
|
||||||
|
Run Cypher queries to explore your graph:
|
||||||
|
```cypher
|
||||||
|
MATCH (n) RETURN n LIMIT 25
|
||||||
|
MATCH (p:Person)-[r]->(n) RETURN p, r, n LIMIT 50
|
||||||
|
```
|
||||||
|
|
||||||
|
## Development
|
||||||
|
|
||||||
|
### Push to Hugging Face Hub
|
||||||
|
To share your trained DSPy program on Hugging Face Hub:
|
||||||
|
|
||||||
|
```python
|
||||||
|
# In main.py, uncomment the push_to_hub section
|
||||||
|
generate_cypher.push_to_hub(
|
||||||
|
"your-username/text-to-cypher",
|
||||||
|
with_code=True,
|
||||||
|
tag="v0.0.1",
|
||||||
|
commit_message="Initial release"
|
||||||
|
)
|
||||||
|
```
|
||||||
|
|
||||||
|
### Customize the Model
|
||||||
|
Modify the `GenerateCypherConfig` in `main.py` to customize:
|
||||||
|
```python
|
||||||
|
class GenerateCypherConfig(PrecompiledConfig):
|
||||||
|
model: str = "openai/gpt-4o" # Change model
|
||||||
|
max_tokens: int = 1024 # Adjust token limit
|
||||||
|
```
|
||||||
|
|
||||||
|
### Process Custom Text
|
||||||
|
Modify `main.py` to process your own text:
|
||||||
|
```python
|
||||||
|
text = "Your custom text here..."
|
||||||
|
cypher = generate_cypher(text=text, neo4j_schema=neo4j.fmt_schema())
|
||||||
|
neo4j.query(cypher.statement.replace('```', ''))
|
||||||
|
```
|
||||||
|
|
||||||
|
## Clean Up
|
||||||
|
|
||||||
|
### Stop Neo4j Docker Container
|
||||||
|
```sh
|
||||||
|
docker stop text-to-cypher
|
||||||
|
docker rm text-to-cypher
|
||||||
|
```
|
||||||
|
|
||||||
|
### Remove Virtual Environment
|
||||||
```sh
|
```sh
|
||||||
deactivate
|
|
||||||
rm -rf .venv
|
rm -rf .venv
|
||||||
```
|
```
|
||||||
|
|
||||||
@@ -79,7 +162,6 @@ This project is licensed under the MIT License - see the [LICENSE](LICENSE) file
|
|||||||
|
|
||||||
## References
|
## References
|
||||||
- [DSPy docs](https://dspy-docs.vercel.app/docs/intro)
|
- [DSPy docs](https://dspy-docs.vercel.app/docs/intro)
|
||||||
|
- [Modaic docs](https://docs.modaic.com/)
|
||||||
- [Neo4j docs](https://neo4j.com/docs/)
|
- [Neo4j docs](https://neo4j.com/docs/)
|
||||||
|
- [uv docs](https://docs.astral.sh/uv/)
|
||||||
## Contact
|
|
||||||
**Primary Contact:** [@chrisammon3000](https://github.com/chrisammon3000)
|
|
||||||
|
|||||||
@@ -1,5 +1,4 @@
|
|||||||
{
|
{
|
||||||
"model": "gpt-4",
|
"model": "openai/gpt-4o",
|
||||||
"neo4j_schema": [],
|
|
||||||
"max_tokens": 1024
|
"max_tokens": 1024
|
||||||
}
|
}
|
||||||
80
main.py
80
main.py
@@ -1,47 +1,81 @@
|
|||||||
from dotenv import load_dotenv
|
import os
|
||||||
import dspy
|
import dspy
|
||||||
|
from dotenv import load_dotenv
|
||||||
from modaic import PrecompiledProgram, PrecompiledConfig
|
from modaic import PrecompiledProgram, PrecompiledConfig
|
||||||
|
from src.neo4j import Neo4j
|
||||||
|
|
||||||
load_dotenv()
|
load_dotenv()
|
||||||
|
|
||||||
class CypherFromText(dspy.Signature):
|
# set up Neo4j using NEO4J_URI
|
||||||
"""Instructions:
|
neo4j = Neo4j(
|
||||||
Create a Cypher MERGE statement to model all entities and relationships found in the text following these guidelines:
|
uri=os.getenv("NEO4J_URI"),
|
||||||
- Refer to the provided schema and use existing or similar nodes, properties or relationships before creating new ones.
|
user=os.getenv("NEO4J_USER"),
|
||||||
- Use generic categories for node and relationship labels."""
|
password=os.getenv("NEO4J_PASSWORD"),
|
||||||
text = dspy.InputField(
|
)
|
||||||
desc="Text to model using nodes, properties and relationships."
|
|
||||||
|
|
||||||
|
class CypherFromQuestion(dspy.Signature):
|
||||||
|
"""Task: Generate Cypher statement to query a graph database.
|
||||||
|
Instructions: Use only the provided relationship types and properties in the schema.
|
||||||
|
Do not use any other relationship types or properties that are not provided in the schema.
|
||||||
|
Do not include any explanations or apologies in your responses.
|
||||||
|
Do not respond to any questions that might ask anything else than for you to construct a Cypher statement.
|
||||||
|
Do not include any text except the generated Cypher statement.
|
||||||
|
"""
|
||||||
|
|
||||||
|
question = dspy.InputField(
|
||||||
|
desc="Question to model using a cypher statement. Use only the provided relationship types and properties in the schema."
|
||||||
)
|
)
|
||||||
neo4j_schema = dspy.InputField(
|
neo4j_schema = dspy.InputField(
|
||||||
desc="Current graph schema in Neo4j as a list of NODES and RELATIONSHIPS."
|
desc="Current graph schema in Neo4j as a list of NODES and RELATIONSHIPS."
|
||||||
)
|
)
|
||||||
statement = dspy.OutputField(
|
statement = dspy.OutputField(desc="Cypher statement to query the graph database.")
|
||||||
desc="Cypher statement to merge nodes and relationships found in the text."
|
|
||||||
)
|
|
||||||
|
|
||||||
class GenerateCypherConfig(PrecompiledConfig):
|
class GenerateCypherConfig(PrecompiledConfig):
|
||||||
neo4j_schema: list[str] = []
|
model: str = "openai/gpt-4o"
|
||||||
model: str = "gpt-4"
|
|
||||||
max_tokens: int = 1024
|
max_tokens: int = 1024
|
||||||
|
|
||||||
|
|
||||||
class GenerateCypher(PrecompiledProgram):
|
class GenerateCypher(PrecompiledProgram):
|
||||||
config: GenerateCypherConfig
|
config: GenerateCypherConfig
|
||||||
|
|
||||||
def _init_(self, config: GenerateCypherConfig, **kwargs):
|
def __init__(self, config: GenerateCypherConfig, **kwargs):
|
||||||
super()._init_(**kwargs)
|
super().__init__(config=config, **kwargs)
|
||||||
self.lm = dspy.LM(
|
self.lm = dspy.LM(
|
||||||
model=config.model,
|
model=config.model,
|
||||||
max_tokens=config.max_tokens,
|
max_tokens=config.max_tokens,
|
||||||
)
|
)
|
||||||
self.generate_cypher = dspy.ChainOfThought(CypherFromText)
|
self.generate_cypher = dspy.ChainOfThought(CypherFromQuestion)
|
||||||
|
self.generate_cypher.set_lm(self.lm)
|
||||||
def forward(self, text: str, neo4j_schema: list[str]):
|
|
||||||
return self.generate_cypher(text=text, neo4j_schema=neo4j_schema)
|
|
||||||
|
|
||||||
|
def forward(self, question: str, neo4j_schema: list[str]):
|
||||||
|
return self.generate_cypher(question=question, neo4j_schema=neo4j_schema)
|
||||||
|
|
||||||
|
|
||||||
|
generate_cypher = GenerateCypher(GenerateCypherConfig())
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
generate_cypher = GenerateCypher(GenerateCypherConfig())
|
"""
|
||||||
generate_cypher.push_to_hub("farouk1/text-to-cypher", with_code=True, tag="v0.0.1", commit_message="init")
|
from pathlib import Path
|
||||||
|
import json
|
||||||
|
|
||||||
|
examples_path = Path(__file__).parent / "examples" / "wikipedia-abstracts-v0_0_1.ndjson"
|
||||||
|
with open(examples_path, "r") as f:
|
||||||
|
for line in f:
|
||||||
|
data = json.loads(line)
|
||||||
|
text = data["text"]
|
||||||
|
print("TEXT TO PROCESS:\n", text[:50])
|
||||||
|
cypher = generate_cypher(text=text, neo4j_schema=neo4j.fmt_schema())
|
||||||
|
neo4j.query(cypher.statement.replace('```', ''))
|
||||||
|
print("CYPHER STATEMENT:\n", cypher.statement)
|
||||||
|
|
||||||
|
schema = neo4j.fmt_schema()
|
||||||
|
print("SCHEMA:\n", schema)
|
||||||
|
"""
|
||||||
|
generate_cypher.push_to_hub(
|
||||||
|
"farouk1/text-to-cypher",
|
||||||
|
with_code=True,
|
||||||
|
tag="v0.0.9",
|
||||||
|
commit_message="Update README.md",
|
||||||
|
)
|
||||||
|
|||||||
37
program.json
37
program.json
@@ -1,4 +1,41 @@
|
|||||||
{
|
{
|
||||||
|
"generate_cypher.predict": {
|
||||||
|
"traces": [],
|
||||||
|
"train": [],
|
||||||
|
"demos": [],
|
||||||
|
"signature": {
|
||||||
|
"instructions": "Task: Generate Cypher statement to query a graph database.\nInstructions: Use only the provided relationship types and properties in the schema.\nDo not use any other relationship types or properties that are not provided in the schema.\nDo not include any explanations or apologies in your responses.\nDo not respond to any questions that might ask anything else than for you to construct a Cypher statement.\nDo not include any text except the generated Cypher statement.",
|
||||||
|
"fields": [
|
||||||
|
{
|
||||||
|
"prefix": "Question:",
|
||||||
|
"description": "Question to model using a cypher statement. Use only the provided relationship types and properties in the schema."
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"prefix": "Neo 4 J Schema:",
|
||||||
|
"description": "Current graph schema in Neo4j as a list of NODES and RELATIONSHIPS."
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"prefix": "Reasoning: Let's think step by step in order to",
|
||||||
|
"description": "${reasoning}"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"prefix": "Statement:",
|
||||||
|
"description": "Cypher statement to query the graph database."
|
||||||
|
}
|
||||||
|
]
|
||||||
|
},
|
||||||
|
"lm": {
|
||||||
|
"model": "openai/gpt-4o",
|
||||||
|
"model_type": "chat",
|
||||||
|
"cache": true,
|
||||||
|
"num_retries": 3,
|
||||||
|
"finetuning_model": null,
|
||||||
|
"launch_kwargs": {},
|
||||||
|
"train_kwargs": {},
|
||||||
|
"temperature": null,
|
||||||
|
"max_tokens": 1024
|
||||||
|
}
|
||||||
|
},
|
||||||
"metadata": {
|
"metadata": {
|
||||||
"dependency_versions": {
|
"dependency_versions": {
|
||||||
"python": "3.13",
|
"python": "3.13",
|
||||||
|
|||||||
@@ -4,4 +4,4 @@ version = "0.1.0"
|
|||||||
description = "Add your description here"
|
description = "Add your description here"
|
||||||
readme = "README.md"
|
readme = "README.md"
|
||||||
requires-python = ">=3.13"
|
requires-python = ">=3.13"
|
||||||
dependencies = ["dspy>=3.0.4", "modaic>=0.8.2", "neo4j~=5.18.0", "python-dotenv~=1.0.1"]
|
dependencies = ["datasets>=4.4.2", "dspy>=3.0.4", "modaic>=0.8.2", "neo4j~=5.18.0", "python-dotenv~=1.0.1", "sacrebleu>=2.5.1"]
|
||||||
|
|||||||
0
src/__init__.py
Normal file
0
src/__init__.py
Normal file
180
src/neo4j.py
Normal file
180
src/neo4j.py
Normal file
@@ -0,0 +1,180 @@
|
|||||||
|
import json
|
||||||
|
import neo4j
|
||||||
|
|
||||||
|
|
||||||
|
def parse_relationships(schema: dict) -> str:
|
||||||
|
# Parse the JSON string into a Python object if it's not already
|
||||||
|
if isinstance(schema, str):
|
||||||
|
data = json.loads(schema)
|
||||||
|
else:
|
||||||
|
data = schema
|
||||||
|
|
||||||
|
data = data[0]["relationships"]
|
||||||
|
|
||||||
|
# Initialize a list to hold the formatted relationship strings
|
||||||
|
relationships = []
|
||||||
|
|
||||||
|
# Iterate through each relationship in the data
|
||||||
|
for relationship in data:
|
||||||
|
entity1, relation, entity2 = relationship
|
||||||
|
# Extract the names of the entities and the relationship
|
||||||
|
entity1_name = entity1["name"]
|
||||||
|
entity2_name = entity2["name"]
|
||||||
|
# Format the string as specified and add it to the list
|
||||||
|
formatted_relationship = f"{entity1_name}-{relation}->{entity2_name}"
|
||||||
|
relationships.append(formatted_relationship)
|
||||||
|
|
||||||
|
# Join all formatted strings with a newline character
|
||||||
|
result = "\n".join(relationships)
|
||||||
|
return result
|
||||||
|
|
||||||
|
|
||||||
|
def parse_nodes(schema):
|
||||||
|
schema = schema
|
||||||
|
nodes = [node["name"] for node in schema[0]["nodes"]]
|
||||||
|
return "\n".join(nodes)
|
||||||
|
|
||||||
|
|
||||||
|
def parse_node_properties(node_properties):
|
||||||
|
# Initialize a dictionary to accumulate node details
|
||||||
|
node_details = {}
|
||||||
|
|
||||||
|
# Iterate through each item in the input JSON
|
||||||
|
for item in node_properties:
|
||||||
|
node_label = item["nodeLabels"][0] # Assuming there's always one label
|
||||||
|
prop_name = item["propertyName"]
|
||||||
|
mandatory = "required" if item["mandatory"] else "optional"
|
||||||
|
|
||||||
|
# Prepare the property string
|
||||||
|
property_str = f"{prop_name} ({mandatory})" if item["mandatory"] else prop_name
|
||||||
|
|
||||||
|
# If the node label exists, append the property; otherwise, create a new entry
|
||||||
|
if node_label in node_details:
|
||||||
|
node_details[node_label].append(property_str)
|
||||||
|
else:
|
||||||
|
node_details[node_label] = [property_str]
|
||||||
|
|
||||||
|
# Format the output
|
||||||
|
output_lines = []
|
||||||
|
for node, properties in node_details.items():
|
||||||
|
output_lines.append(f"{node}")
|
||||||
|
for prop in properties:
|
||||||
|
prop_line = f" - {prop}" if "required" in prop else f" - {prop}"
|
||||||
|
output_lines.append(prop_line)
|
||||||
|
|
||||||
|
return "\n".join(output_lines)
|
||||||
|
|
||||||
|
|
||||||
|
def parse_rel_properties(rel_properties):
|
||||||
|
# Initialize a dictionary to accumulate relationship details
|
||||||
|
rel_details = {}
|
||||||
|
|
||||||
|
# Iterate through each item in the input JSON
|
||||||
|
for item in rel_properties:
|
||||||
|
# Extract relationship type name, removing :` and `
|
||||||
|
rel_type = item["relType"][2:].strip("`")
|
||||||
|
prop_name = item["propertyName"]
|
||||||
|
mandatory = "required" if item["mandatory"] else "optional"
|
||||||
|
|
||||||
|
# If propertyName is not None, prepare the property string
|
||||||
|
if prop_name is not None:
|
||||||
|
property_str = f"{prop_name} ({mandatory})"
|
||||||
|
# If the relationship type exists, append the property; otherwise, create a new entry
|
||||||
|
if rel_type in rel_details:
|
||||||
|
rel_details[rel_type].append(property_str)
|
||||||
|
else:
|
||||||
|
rel_details[rel_type] = [property_str]
|
||||||
|
else:
|
||||||
|
# For relationships without properties, ensure the relationship is listed
|
||||||
|
rel_details.setdefault(rel_type, [])
|
||||||
|
|
||||||
|
# Format the output
|
||||||
|
output_lines = []
|
||||||
|
for rel_type, properties in rel_details.items():
|
||||||
|
output_lines.append(f"{rel_type}")
|
||||||
|
for prop in properties:
|
||||||
|
output_lines.append(f" - {prop}")
|
||||||
|
|
||||||
|
return "\n".join(output_lines)
|
||||||
|
|
||||||
|
|
||||||
|
class Neo4j:
|
||||||
|
def __init__(self, uri, user: str = None, password: str = None):
|
||||||
|
self._uri = uri
|
||||||
|
self._user = user
|
||||||
|
self._password = password
|
||||||
|
self._auth = (
|
||||||
|
None
|
||||||
|
if (self._user is None and self._password is None)
|
||||||
|
else (self._user, self._password)
|
||||||
|
)
|
||||||
|
self._driver = neo4j.GraphDatabase.driver(
|
||||||
|
self._uri, auth=(self._user, self._password)
|
||||||
|
)
|
||||||
|
|
||||||
|
self._verify_connection()
|
||||||
|
print("CONNECTION ESTABLISHED")
|
||||||
|
|
||||||
|
def close(self):
|
||||||
|
self._driver.close()
|
||||||
|
print("CONNECTION CLOSED")
|
||||||
|
|
||||||
|
def _verify_connection(self):
|
||||||
|
with self._driver as driver:
|
||||||
|
driver.verify_connectivity()
|
||||||
|
print("CONNECTION VERIFIED")
|
||||||
|
|
||||||
|
def query(self, query, parameters=None, db=None):
|
||||||
|
assert db is None, (
|
||||||
|
"The Neo4j implementation does not support multiple databases."
|
||||||
|
)
|
||||||
|
with self._driver.session(database=db) as session:
|
||||||
|
result = session.run(query, parameters)
|
||||||
|
return result.data()
|
||||||
|
|
||||||
|
def schema(self, parsed=False):
|
||||||
|
query = """
|
||||||
|
CALL db.schema.visualization()
|
||||||
|
"""
|
||||||
|
schema = self.query(query)
|
||||||
|
|
||||||
|
if parsed:
|
||||||
|
return parse_nodes(schema), parse_relationships(schema)
|
||||||
|
|
||||||
|
return schema
|
||||||
|
|
||||||
|
def schema_properties(self, parsed=False):
|
||||||
|
props = self._schema_node_properties(), self._schema_relationship_properties()
|
||||||
|
if parsed:
|
||||||
|
return parse_node_properties(props[0]), parse_rel_properties(props[1])
|
||||||
|
|
||||||
|
return props
|
||||||
|
|
||||||
|
def _schema_node_properties(self):
|
||||||
|
query = """
|
||||||
|
CALL db.schema.nodeTypeProperties()
|
||||||
|
"""
|
||||||
|
return self.query(query)
|
||||||
|
|
||||||
|
def _schema_relationship_properties(self):
|
||||||
|
query = """
|
||||||
|
CALL db.schema.relTypeProperties()
|
||||||
|
"""
|
||||||
|
return self.query(query)
|
||||||
|
|
||||||
|
def fmt_schema(self):
|
||||||
|
parsed_schema = self.schema(parsed=True)
|
||||||
|
parsed_props = self.schema_properties(parsed=True)
|
||||||
|
parsed = (*parsed_props, parsed_schema[1])
|
||||||
|
return "\n".join(
|
||||||
|
[
|
||||||
|
f"{element}:\n{parsed[idx]}\n"
|
||||||
|
for idx, element in enumerate(
|
||||||
|
[
|
||||||
|
"NODE LABELS & PROPERTIES",
|
||||||
|
"RELATIONSHIP LABELS & PROPERTIES",
|
||||||
|
"RELATIONSHIPS",
|
||||||
|
]
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
Reference in New Issue
Block a user