├── .gitignore ├── README.md ├── examples └── wikipedia-abstracts-v0_0_1.ndjson ├── img └── kg.png ├── requirements.txt ├── run.py └── src ├── __init__.py └── neo4j.py /.gitignore: -------------------------------------------------------------------------------- 1 | # System files 2 | .DS_Store 3 | 4 | # Byte-compiled / optimized / DLL files 5 | __pycache__/ 6 | *.py[cod] 7 | *$py.class 8 | 9 | # C extensions 10 | *.so 11 | 12 | # Distribution / packaging 13 | .Python 14 | venv/ 15 | develop-eggs/ 16 | dist/ 17 | downloads/ 18 | eggs/ 19 | .eggs/ 20 | lib/ 21 | lib64/ 22 | parts/ 23 | sdist/ 24 | var/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | 29 | # PyInstaller 30 | # Usually these files are written by a python script from a template 31 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 32 | *.manifest 33 | *.spec 34 | 35 | # Installer logs 36 | pip-$$LOG_FILE 37 | pip-delete-this-directory.txt 38 | 39 | # Unit test / coverage reports 40 | htmlcov/ 41 | .tox/ 42 | .coverage 43 | .coverage.* 44 | .cache 45 | nosetests.xml 46 | coverage.xml 47 | *,cover 48 | .hypothesis/ 49 | 50 | # Translations 51 | *.mo 52 | *.pot 53 | 54 | # Django stuff: 55 | *.log 56 | 57 | # Sphinx documentation 58 | docs/_build/ 59 | 60 | # PyBuilder 61 | target/ 62 | 63 | # VS Code 64 | .vscode 65 | 66 | # Jupyter Notebook 67 | .ipynb_checkpoints 68 | 69 | # pyenv python configuration file 70 | .python-version 71 | .venv* 72 | venv* 73 | 74 | # Profiling tools 75 | profile* 76 | fil* 77 | summary_* 78 | 79 | # Environment variables 80 | .env* 81 | env.sh 82 | 83 | # PyCharm 84 | .idea/ 85 | 86 | # Secrets 87 | *.pem 88 | 89 | # AWS SAM/CloudFormation 90 | .aws-sam 91 | packaged.yaml 92 | samconfig.toml 93 | 94 | # Terraform 95 | .terraform/ 96 | .terraform.lock.hcl 97 | 98 | # Project Files 99 | NOTES.md -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # dspy-neo4j-knowledge-graph 2 | LLM-driven automated knowledge graph construction from text using DSPy and Neo4j. 3 | 4 | ![Knowledge Graph](img/kg.png) 5 | 6 | ## Project Structure 7 | ```sh 8 | dspy-neo4j-knowledge-graph/ 9 | ├── README.md 10 | ├── examples 11 | ├── requirements.txt 12 | ├── run.py 13 | └── src 14 | ``` 15 | 16 | ## Description 17 | 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. 18 | 19 | ### Optimized Schema Context 20 | 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. 21 | 22 | ## Quick Start 23 | 1. Clone the repository. 24 | 2. Create a [Python virtual environment](#python-virtual-environment) and install the required packages. 25 | 3. Create a `.env` file and add the required [environment variables](#environment-variables). 26 | 4. [Run Neo4j using Docker](#usage). 27 | 5. Run `python3 run.py` and paste your text in the prompt. 28 | 6. Navigate to `http://localhost:7474/browser/` to view the Knowledge Graph in Neo4j Browser. 29 | 30 | ## Installation 31 | 32 | ### Prerequisites 33 | * Python 3.12 34 | * OpenAI API Key 35 | * Docker 36 | 37 | ### Environment Variables 38 | Before you begin, make sure to create a `.env` file and add your OpenAI API key. 39 | ```sh 40 | NEO4J_URI=bolt://localhost:7687 41 | OPENAI_API_KEY= 42 | ``` 43 | 44 | ### Python Virtual Environment 45 | Create a Python virtual environment and install the required packages. 46 | ```sh 47 | python3 -m venv .venv 48 | source .venv/bin/activate 49 | pip install --upgrade pip 50 | pip install -r requirements.txt 51 | ``` 52 | 53 | ## Usage 54 | Run Neo4j using Docker. 55 | ```sh 56 | docker run \ 57 | --name dspy-kg \ 58 | --publish=7474:7474 \ 59 | --publish=7687:7687 \ 60 | --env "NEO4J_AUTH=none" \ 61 | neo4j:5.15 62 | ``` 63 | 64 | ## Clean Up 65 | Stop and remove the Neo4j container. 66 | ```sh 67 | docker stop dspy-kg 68 | docker rm dspy-kg 69 | ``` 70 | 71 | Deactivate the Python virtual environment. 72 | ```sh 73 | deactivate 74 | rm -rf .venv 75 | ``` 76 | 77 | ## License 78 | This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details. 79 | 80 | ## References 81 | - [DSPy docs](https://dspy-docs.vercel.app/docs/intro) 82 | - [Neo4j docs](https://neo4j.com/docs/) 83 | 84 | ## Contact 85 | **Primary Contact:** [@chrisammon3000](https://github.com/chrisammon3000) 86 | -------------------------------------------------------------------------------- /examples/wikipedia-abstracts-v0_0_1.ndjson: -------------------------------------------------------------------------------- 1 | {"text": "John Singer Sargent (/ˈsɑːrdʒənt/; January 12, 1856 – April 14, 1925)[1] was an American expatriate artist, considered the \"leading portrait painter of his generation\" for his evocations of Edwardian-era luxury.[2][3] He created roughly 900 oil paintings and more than 2,000 watercolors, as well as countless sketches and charcoal drawings. His oeuvre documents worldwide travel, from Venice to the Tyrol, Corfu, Spain, the Middle East, Montana, Maine, and Florida.","statement": "```\nMERGE (p:Person {name: 'John Singer Sargent'})\nMERGE (c1:City {name: 'Venice'})\nMERGE (c2:City {name: 'Tyrol'})\nMERGE (c3:City {name: 'Corfu'})\nMERGE (c4:City {name: 'Spain'})\nMERGE (c5:City {name: 'Middle East'})\nMERGE (c6:City {name: 'Montana'})\nMERGE (c7:City {name: 'Maine'})\nMERGE (c8:City {name: 'Florida'})\nMERGE (a1:Artwork {type: 'oil painting'})\nMERGE (a2:Artwork {type: 'watercolor'})\nMERGE (a3:Artwork {type: 'sketch'})\nMERGE (a4:Artwork {type: 'charcoal drawing'})\nMERGE (p)-[:TRAVELED_TO]->(c1)\nMERGE (p)-[:TRAVELED_TO]->(c2)\nMERGE (p)-[:TRAVELED_TO]->(c3)\nMERGE (p)-[:TRAVELED_TO]->(c4)\nMERGE (p)-[:TRAVELED_TO]->(c5)\nMERGE (p)-[:TRAVELED_TO]->(c6)\nMERGE (p)-[:TRAVELED_TO]->(c7)\nMERGE (p)-[:TRAVELED_TO]->(c8)\nMERGE (p)-[:CREATED]->(a1)\nMERGE (p)-[:CREATED]->(a2)\nMERGE (p)-[:CREATED]->(a3)\nMERGE (p)-[:CREATED]->(a4)\n```"} 2 | {"text": "Lyman Frank Baum (/bɔːm/;[1] May 15, 1856 – May 6, 1919) was an American author best known for his children's fantasy books, particularly The Wonderful Wizard of Oz, part of a series. In addition to the 14 Oz books, Baum penned 41 other novels (not including four lost, unpublished novels), 83 short stories, over 200 poems, and at least 42 scripts. He made numerous attempts to bring his works to the stage and screen; the 1939 adaptation of the first Oz book became a landmark of 20th-century cinema.", "statement": "```\nMERGE (p:Person {name: 'Lyman Frank Baum', dob: 'May 15, 1856', dod: 'May 6, 1919', nationality: 'American'})\nMERGE (b:Book {name: 'The Wonderful Wizard of Oz'})\nMERGE (s:Series {name: 'Oz'})\nMERGE (c:Cinema {name: '1939 adaptation of the first Oz book', year: 1939})\nMERGE (p)-[:AUTHORED]->(b)\nMERGE (b)-[:PART_OF]->(s)\nMERGE (p)-[:ATTEMPTED]->(:Adaptation {name: 'Stage and Screen Adaptation'})\nMERGE (c)-[:ADAPTATION_OF]->(b)\n```"} 3 | {"text": "John Davison Rockefeller III (March 21, 1906 – July 10, 1978) was an American philanthropist. Rockefeller was the eldest son and second child of John D. Rockefeller Jr. and Abby Aldrich Rockefeller as well as a grandson of Standard Oil co-founder John D. Rockefeller. He was engaged in a wide range of philanthropic projects, many of which his family had launched, as well as supporting organizations related to East Asian affairs. Rockefeller was also a major supporter of the Population Council, and the committee that created the Lincoln Center in Manhattan.", "statement": "```\nMERGE (jdr3:Person {name: 'John Davison Rockefeller III'})\nMERGE (jdrjr:Person {name: 'John D. Rockefeller Jr.'})\nMERGE (aar:Person {name: 'Abby Aldrich Rockefeller'})\nMERGE (jdr:Person {name: 'John D. Rockefeller'})\nMERGE (so:Company {name: 'Standard Oil'})\nMERGE (pc:Council {name: 'Population Council'})\nMERGE (lc:Center {name: 'Lincoln Center'})\nMERGE (ea:Affairs {name: 'East Asian affairs'})\nMERGE (jdr3)-[:SON_OF]->(jdrjr)\nMERGE (jdr3)-[:SON_OF]->(aar)\nMERGE (jdr3)-[:GRANDSON_OF]->(jdr)\nMERGE (jdr3)-[:SUPPORTER_OF]->(pc)\nMERGE (jdr3)-[:SUPPORTER_OF]->(lc)\nMERGE (jdr3)-[:RELATED_TO]->(ea)\nMERGE (jdr)-[:CO_FOUNDER_OF]->(so)\n```"} 4 | {"text": "Samuel Langhorne Clemens (November 30, 1835 – April 21, 1910),[1] known by the pen name Mark Twain, was an American writer, humorist and essayist. He was praised as the \"greatest humorist the United States has produced,\"[2] with William Faulkner calling him \"the father of American literature.\"[3] His novels include The Adventures of Tom Sawyer (1876) and its sequel, Adventures of Huckleberry Finn (1884),[4] with the latter often called the \"Great American Novel.\" Twain also wrote A Connecticut Yankee in King Arthur's Court (1889) and Pudd'nhead Wilson (1894), and co-wrote The Gilded Age: A Tale of Today (1873) with Charles Dudley Warner.", "statement": "```\nMERGE (p:Person {name: 'Samuel Langhorne Clemens', dob: '1835-11-30', dod: '1910-04-21', penName: 'Mark Twain'})\nMERGE (c:Country {name: 'United States'})\nMERGE (p)-[:IS_FROM]->(c)\nMERGE (w:Writer {name: 'William Faulkner'})\nMERGE (w)-[:CALLED {description: 'the father of American literature'}]->(p)\nMERGE (b1:Book {name: 'The Adventures of Tom Sawyer', year: 1876})\nMERGE (b2:Book {name: 'Adventures of Huckleberry Finn', year: 1884})\nMERGE (b3:Book {name: 'A Connecticut Yankee in King Arthur\\'s Court', year: 1889})\nMERGE (b4:Book {name: 'Pudd\\'nhead Wilson', year: 1894})\nMERGE (b5:Book {name: 'The Gilded Age: A Tale of Today', year: 1873})\nMERGE (p)-[:WROTE]->(b1)\nMERGE (p)-[:WROTE]->(b2)\nMERGE (p)-[:WROTE]->(b3)\nMERGE (p)-[:WROTE]->(b4)\nMERGE (p)-[:CO_WROTE {with: 'Charles Dudley Warner'}]->(b5)\n```"} 5 | {"text": "Henry Huttleston Rogers (January 29, 1840 – May 19, 1909) was an American industrialist and financier. He made his fortune in the oil refining business, becoming a leader at Standard Oil. He also played a major role in numerous corporations and business enterprises in the gas industry, copper, and railroads. He became a close friend of Mark Twain.", "statement": "```\nMERGE (p1:Person {name: 'Henry Huttleston Rogers', dob: 'January 29, 1840', dod: 'May 19, 1909', occupation: 'Industrialist and Financier'})\nMERGE (p2:Person {name: 'Mark Twain'})\nMERGE (c:Company {name: 'Standard Oil'})\nMERGE (i:Industry {name: 'Oil Refining'})\nMERGE (i2:Industry {name: 'Gas'})\nMERGE (i3:Industry {name: 'Copper'})\nMERGE (i4:Industry {name: 'Railroads'})\nMERGE (p1)-[:MADE_FORTUNE_IN]->(i)\nMERGE (p1)-[:BECAME_LEADER_AT]->(c)\nMERGE (p1)-[:PLAYED_MAJOR_ROLE_IN]->(i2)\nMERGE (p1)-[:PLAYED_MAJOR_ROLE_IN]->(i3)\nMERGE (p1)-[:PLAYED_MAJOR_ROLE_IN]->(i4)\nMERGE (p1)-[:FRIEND_OF]->(p2)\n```"} 6 | {"text": "John Pierpont Morgan (April 17, 1837 – March 31, 1913)[1] was an American financier and investment banker who dominated corporate finance on Wall Street throughout the Gilded Age and Progressive Era. As the head of the banking firm that ultimately became known as J.P. Morgan and Co., he was a driving personal force behind the wave of industrial consolidations in the United States at the turn of the twentieth century.", "statement": "```\nMERGE (p:Person {name: 'John Pierpont Morgan', birthDate: 'April 17, 1837', deathDate: 'March 31, 1913', nationality: 'American', profession: 'financier and investment banker'})\nMERGE (c:Company {name: 'J.P. Morgan and Co.'})\nMERGE (e:Event {name: 'Industrial consolidations in the United States at the turn of the twentieth century'})\nMERGE (p)-[:HEAD_OF]->(c)\nMERGE (p)-[:INFLUENCED]->(e)\n```"} 7 | {"text": "John Charles Frémont or Fremont (January 21, 1813 – July 13, 1890) was an American explorer, military officer, and politician. He was a United States senator from California and was the first Republican nominee for president of the U.S. in 1856 and founder of the California Republican Party when he was nominated. He lost the election to Democrat James Buchanan when the vote was split by Know Nothings.", "statement": "```\nMERGE (p:Person {name: 'John Charles Frémont', birthdate: '1813-01-21', deathdate: '1890-07-13'})\nMERGE (c:Country {name: 'United States'})\nMERGE (s:State {name: 'California'})\nMERGE (pp:PoliticalParty {name: 'Republican Party'})\nMERGE (e:Election {year: 1856})\nMERGE (op:Person {name: 'James Buchanan'})\nMERGE (p)-[:ROLE]->(:Role {name: 'Explorer'})\nMERGE (p)-[:ROLE]->(:Role {name: 'Military Officer'})\nMERGE (p)-[:ROLE]->(:Role {name: 'Politician'})\nMERGE (p)-[:SENATOR_FROM]->(s)\nMERGE (s)-[:PART_OF]->(c)\nMERGE (p)-[:NOMINEE_FOR {party: 'Republican', year: 1856}]->(pp)\nMERGE (p)-[:FOUNDER_OF]->(:PoliticalParty {name: 'California Republican Party'})\nMERGE (p)-[:LOST_TO]->(op)\nMERGE (op)-[:WON]->(e)\n```"} 8 | {"text": "Geronimo (Mescalero-Chiricahua: Goyaałé, Athapascan pronunciation: [kòjàːɬɛ́], lit. 'the one who yawns'; June 16, 1829 – February 17, 1909) was a military leader and medicine man from the Bedonkohe band of the Ndendahe Apache people. From 1850 to 1886, Geronimo joined with members of three other Central Apache bands – the Tchihende, the Tsokanende (called Chiricahua by Americans) and the Nednhi – to carry out numerous raids, as well as fight against Mexican and U.S. military campaigns in the northern Mexico states of Chihuahua and Sonora and in the southwestern American territories of New Mexico and Arizona.", "statement": "```\nMERGE (p:Person {name: 'Geronimo', birthdate: '1829-06-16', deathdate: '1909-02-17', nationality: 'Apache'})\nMERGE (c1:Country {name: 'Mexico'})\nMERGE (s1:State {name: 'Chihuahua'})\nMERGE (s2:State {name: 'Sonora'})\nMERGE (c2:Country {name: 'United States'})\nMERGE (s3:State {name: 'New Mexico'})\nMERGE (s4:State {name: 'Arizona'})\nMERGE (b1:Band {name: 'Bedonkohe'})\nMERGE (b2:Band {name: 'Tchihende'})\nMERGE (b3:Band {name: 'Tsokanende'})\nMERGE (b4:Band {name: 'Nednhi'})\nMERGE (p)-[:IS_FROM]->(b1)\nMERGE (p)-[:RELATED_TO]->(b2)\nMERGE (p)-[:RELATED_TO]->(b3)\nMERGE (p)-[:RELATED_TO]->(b4)\nMERGE (p)-[:TRAVELED_TO]->(s1)\nMERGE (p)-[:TRAVELED_TO]->(s2)\nMERGE (p)-[:TRAVELED_TO]->(s3)\nMERGE (p)-[:TRAVELED_TO]->(s4)\n```"} -------------------------------------------------------------------------------- /img/kg.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/chrisammon3000/dspy-neo4j-knowledge-graph/2856b2dfe81fc801601c2a9c4b439429a4562a98/img/kg.png -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | python-dotenv~=1.0.1 2 | dspy-ai~=2.4.0 3 | neo4j~=5.18.0 -------------------------------------------------------------------------------- /run.py: -------------------------------------------------------------------------------- 1 | """ 2 | Text -> Knowledge Graph 3 | 1. text -> cypher 4 | 5 | Constraints: 6 | - Use the existing schema before creating new nodes and relationships. 7 | """ 8 | import os 9 | from dotenv import find_dotenv, load_dotenv 10 | load_dotenv(find_dotenv()); 11 | import dspy 12 | from src.neo4j import Neo4j 13 | 14 | # set up Neo4j using NEO4J_URI 15 | neo4j = Neo4j(uri=os.getenv("NEO4J_URI"), user=os.getenv("NEO4J_USER"), password=os.getenv("NEO4J_PASSWORD")) 16 | 17 | lm = dspy.OpenAI( 18 | model="gpt-4", 19 | max_tokens=1024, 20 | ) 21 | dspy.configure(lm=lm) 22 | 23 | class CypherFromText(dspy.Signature): 24 | """Instructions: 25 | Create a Cypher MERGE statement to model all entities and relationships found in the text following these guidelines: 26 | - Refer to the provided schema and use existing or similar nodes, properties or relationships before creating new ones. 27 | - Use generic categories for node and relationship labels.""" 28 | 29 | text = dspy.InputField(desc="Text to model using nodes, properties and relationships.") 30 | neo4j_schema = dspy.InputField(desc="Current graph schema in Neo4j as a list of NODES and RELATIONSHIPS.") 31 | statement = dspy.OutputField(desc="Cypher statement to merge nodes and relationships found in the text.") 32 | 33 | generate_cypher = dspy.ChainOfThought(CypherFromText) 34 | 35 | if __name__ == "__main__": 36 | from pathlib import Path 37 | # import json 38 | 39 | # examples_path = Path(__file__).parent / "examples" / "wikipedia-abstracts-v0_0_1.ndjson" 40 | # with open(examples_path, "r") as f: 41 | # # process line by line 42 | # for line in f: 43 | # data = json.loads(line) 44 | # text = data["text"] 45 | # print(text[:50]) 46 | # cypher = generate_cypher(text=text, neo4j_schema=neo4j.fmt_schema()) 47 | # neo4j.query(cypher.statement.replace('```', '')) 48 | 49 | while True: 50 | try: 51 | text = input("\nEnter text: ") 52 | cypher = generate_cypher(text=text.replace("\n", " "), neo4j_schema=neo4j.fmt_schema()) 53 | neo4j.query(cypher.statement.replace('```', '')) 54 | 55 | except Exception as e: 56 | print(e) 57 | print("Please input one paragraph at a time.") 58 | continue 59 | 60 | except KeyboardInterrupt: 61 | break -------------------------------------------------------------------------------- /src/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/chrisammon3000/dspy-neo4j-knowledge-graph/2856b2dfe81fc801601c2a9c4b439429a4562a98/src/__init__.py -------------------------------------------------------------------------------- /src/neo4j.py: -------------------------------------------------------------------------------- 1 | import json 2 | import neo4j 3 | 4 | def parse_relationships(schema: dict) -> str: 5 | # Parse the JSON string into a Python object if it's not already 6 | if isinstance(schema, str): 7 | data = json.loads(schema) 8 | else: 9 | data = schema 10 | 11 | data = data[0]['relationships'] 12 | 13 | # Initialize a list to hold the formatted relationship strings 14 | relationships = [] 15 | 16 | # Iterate through each relationship in the data 17 | for relationship in data: 18 | entity1, relation, entity2 = relationship 19 | # Extract the names of the entities and the relationship 20 | entity1_name = entity1['name'] 21 | entity2_name = entity2['name'] 22 | # Format the string as specified and add it to the list 23 | formatted_relationship = f"{entity1_name}-{relation}->{entity2_name}" 24 | relationships.append(formatted_relationship) 25 | 26 | # Join all formatted strings with a newline character 27 | result = "\n".join(relationships) 28 | return result 29 | 30 | def parse_nodes(schema): 31 | schema = schema 32 | nodes = [node['name'] for node in schema[0]['nodes']] 33 | return "\n".join(nodes) 34 | 35 | def parse_node_properties(node_properties): 36 | # Initialize a dictionary to accumulate node details 37 | node_details = {} 38 | 39 | # Iterate through each item in the input JSON 40 | for item in node_properties: 41 | node_label = item["nodeLabels"][0] # Assuming there's always one label 42 | prop_name = item["propertyName"] 43 | mandatory = "required" if item["mandatory"] else "optional" 44 | 45 | # Prepare the property string 46 | property_str = f"{prop_name} ({mandatory})" if item["mandatory"] else prop_name 47 | 48 | # If the node label exists, append the property; otherwise, create a new entry 49 | if node_label in node_details: 50 | node_details[node_label].append(property_str) 51 | else: 52 | node_details[node_label] = [property_str] 53 | 54 | # Format the output 55 | output_lines = [] 56 | for node, properties in node_details.items(): 57 | output_lines.append(f"{node}") 58 | for prop in properties: 59 | prop_line = f" - {prop}" if "required" in prop else f" - {prop}" 60 | output_lines.append(prop_line) 61 | 62 | return "\n".join(output_lines) 63 | 64 | 65 | def parse_rel_properties(rel_properties): 66 | # Initialize a dictionary to accumulate relationship details 67 | rel_details = {} 68 | 69 | # Iterate through each item in the input JSON 70 | for item in rel_properties: 71 | # Extract relationship type name, removing :` and ` 72 | rel_type = item["relType"][2:].strip("`") 73 | prop_name = item["propertyName"] 74 | mandatory = "required" if item["mandatory"] else "optional" 75 | 76 | # If propertyName is not None, prepare the property string 77 | if prop_name is not None: 78 | property_str = f"{prop_name} ({mandatory})" 79 | # If the relationship type exists, append the property; otherwise, create a new entry 80 | if rel_type in rel_details: 81 | rel_details[rel_type].append(property_str) 82 | else: 83 | rel_details[rel_type] = [property_str] 84 | else: 85 | # For relationships without properties, ensure the relationship is listed 86 | rel_details.setdefault(rel_type, []) 87 | 88 | # Format the output 89 | output_lines = [] 90 | for rel_type, properties in rel_details.items(): 91 | output_lines.append(f"{rel_type}") 92 | for prop in properties: 93 | output_lines.append(f" - {prop}") 94 | 95 | return "\n".join(output_lines) 96 | 97 | 98 | class Neo4j: 99 | def __init__(self, uri, user: str = None, password: str = None): 100 | self._uri = uri 101 | self._user = user 102 | self._password = password 103 | self._auth = None if (self._user is None and self._password is None) else (self._user, self._password) 104 | self._driver = neo4j.GraphDatabase.driver(self._uri, auth=(self._user, self._password)) 105 | 106 | self._verify_connection() 107 | 108 | def close(self): 109 | self._driver.close() 110 | 111 | def _verify_connection(self): 112 | with self._driver as driver: 113 | driver.verify_connectivity() 114 | 115 | def query(self, query, parameters=None, db=None): 116 | assert db is None, "The Neo4j implementation does not support multiple databases." 117 | with self._driver.session(database=db) as session: 118 | result = session.run(query, parameters) 119 | return result.data() 120 | 121 | def schema(self, parsed=False): 122 | query = """ 123 | CALL db.schema.visualization() 124 | """ 125 | schema = self.query(query) 126 | 127 | if parsed: 128 | return parse_nodes(schema), parse_relationships(schema) 129 | 130 | return schema 131 | 132 | def schema_properties(self, parsed=False): 133 | props = self._schema_node_properties(), self._schema_relationship_properties() 134 | if parsed: 135 | return parse_node_properties(props[0]), parse_rel_properties(props[1]) 136 | 137 | return props 138 | 139 | def _schema_node_properties(self): 140 | query = """ 141 | CALL db.schema.nodeTypeProperties() 142 | """ 143 | return self.query(query) 144 | 145 | def _schema_relationship_properties(self): 146 | query = """ 147 | CALL db.schema.relTypeProperties() 148 | """ 149 | return self.query(query) 150 | 151 | def fmt_schema(self): 152 | parsed_schema = self.schema(parsed=True) 153 | parsed_props = self.schema_properties(parsed=True) 154 | parsed = (*parsed_props, parsed_schema[1]) 155 | return "\n".join([f"{element}:\n{parsed[idx]}\n" for idx, element in enumerate(["NODE LABELS & PROPERTIES", "RELATIONSHIP LABELS & PROPERTIES", "RELATIONSHIPS"])]) --------------------------------------------------------------------------------