{ "cells": [ { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "%pip install --upgrade langchain langchain-experimental langchain-openai python-dotenv pyvis" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "from dotenv import load_dotenv\n", "import os\n", "\n", "# Load the .env file\n", "load_dotenv()\n", "# Get API key from environment variable \n", "# (make sure the key is present in .env file in the project directory)\n", "api_key = os.getenv(\"OPENAI_API_KEY\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### LLM Graph Transformer\n", "Using GPT-4o in all examples." ] }, { "cell_type": "code", "execution_count": 8, "metadata": {}, "outputs": [], "source": [ "from langchain_experimental.graph_transformers import LLMGraphTransformer\n", "from langchain_core.documents import Document\n", "from langchain_openai import ChatOpenAI\n", "\n", "llm = ChatOpenAI(temperature=0, model_name=\"gpt-4o\")\n", "\n", "graph_transformer = LLMGraphTransformer(llm=llm)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Extract graph data" ] }, { "cell_type": "code", "execution_count": 9, "metadata": {}, "outputs": [], "source": [ "text = \"\"\"\n", "Albert Einstein[a] (14 March 1879 – 18 April 1955) was a German-born theoretical physicist who is best known for developing the theory of relativity. Einstein also made important contributions to quantum mechanics.[1][5] His mass–energy equivalence formula E = mc2, which arises from special relativity, has been called \"the world's most famous equation\".[6] He received the 1921 Nobel Prize in Physics for his services to theoretical physics, and especially for his discovery of the law of the photoelectric effect.[7]\n", "\n", "Born in the German Empire, Einstein moved to Switzerland in 1895, forsaking his German citizenship (as a subject of the Kingdom of Württemberg)[note 1] the following year. In 1897, at the age of seventeen, he enrolled in the mathematics and physics teaching diploma program at the Swiss federal polytechnic school in Zurich, graduating in 1900. He acquired Swiss citizenship a year later, which he kept for the rest of his life, and afterwards secured a permanent position at the Swiss Patent Office in Bern. In 1905, he submitted a successful PhD dissertation to the University of Zurich. In 1914, he moved to Berlin to join the Prussian Academy of Sciences and the Humboldt University of Berlin, becoming director of the Kaiser Wilhelm Institute for Physics in 1917; he also became a German citizen again, this time as a subject of the Kingdom of Prussia.[note 1] In 1933, while Einstein was visiting the United States, Adolf Hitler came to power in Germany. Horrified by the Nazi persecution of his fellow Jews,[8] he decided to remain in the US, and was granted American citizenship in 1940.[9] On the eve of World War II, he endorsed a letter to President Franklin D. Roosevelt alerting him to the potential German nuclear weapons program and recommending that the US begin similar research.\n", "\n", "In 1905, sometimes described as his annus mirabilis (miracle year), he published four groundbreaking papers.[10] In them, he outlined a theory of the photoelectric effect, explained Brownian motion, introduced his special theory of relativity, and demonstrated that if the special theory is correct, mass and energy are equivalent to each other. In 1915, he proposed a general theory of relativity that extended his system of mechanics to incorporate gravitation. A cosmological paper that he published the following year laid out the implications of general relativity for the modeling of the structure and evolution of the universe as a whole.[11][12] In 1917, Einstein wrote a paper which introduced the concepts of spontaneous emission and stimulated emission, the latter of which is the core mechanism behind the laser and maser, and which contained a trove of information that would be beneficial to developments in physics later on, such as quantum electrodynamics and quantum optics.[13]\n", "\n", "In the middle part of his career, Einstein made important contributions to statistical mechanics and quantum theory. Especially notable was his work on the quantum physics of radiation, in which light consists of particles, subsequently called photons. With physicist Satyendra Nath Bose, he laid the groundwork for Bose–Einstein statistics. For much of the last phase of his academic life, Einstein worked on two endeavors that ultimately proved unsuccessful. First, he advocated against quantum theory's introduction of fundamental randomness into science's picture of the world, objecting that God does not play dice.[14] Second, he attempted to devise a unified field theory by generalizing his geometric theory of gravitation to include electromagnetism. As a result, he became increasingly isolated from mainstream modern physics.\n", "\"\"\"" ] }, { "cell_type": "code", "execution_count": 21, "metadata": {}, "outputs": [], "source": [ "documents = [Document(page_content=text)]\n", "graph_documents = await graph_transformer.aconvert_to_graph_documents(documents)" ] }, { "cell_type": "code", "execution_count": 22, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Nodes:[Node(id='Albert Einstein', type='Person', properties={}), Node(id='Theory Of Relativity', type='Concept', properties={}), Node(id='Quantum Mechanics', type='Concept', properties={}), Node(id='Mass–Energy Equivalence Formula', type='Concept', properties={}), Node(id='E = Mc2', type='Concept', properties={}), Node(id='1921 Nobel Prize In Physics', type='Award', properties={}), Node(id='Law Of The Photoelectric Effect', type='Concept', properties={}), Node(id='German Empire', type='Place', properties={}), Node(id='Switzerland', type='Place', properties={}), Node(id='Swiss Federal Polytechnic School In Zurich', type='Organization', properties={}), Node(id='Swiss Patent Office In Bern', type='Organization', properties={}), Node(id='University Of Zurich', type='Organization', properties={}), Node(id='Prussian Academy Of Sciences', type='Organization', properties={}), Node(id='Humboldt University Of Berlin', type='Organization', properties={}), Node(id='Kaiser Wilhelm Institute For Physics', type='Organization', properties={}), Node(id='United States', type='Place', properties={}), Node(id='Adolf Hitler', type='Person', properties={}), Node(id='Franklin D. Roosevelt', type='Person', properties={}), Node(id='World War Ii', type='Event', properties={}), Node(id='Annus Mirabilis', type='Event', properties={}), Node(id='Brownian Motion', type='Concept', properties={}), Node(id='Special Theory Of Relativity', type='Concept', properties={}), Node(id='General Theory Of Relativity', type='Concept', properties={}), Node(id='Spontaneous Emission', type='Concept', properties={}), Node(id='Stimulated Emission', type='Concept', properties={}), Node(id='Laser', type='Concept', properties={}), Node(id='Maser', type='Concept', properties={}), Node(id='Statistical Mechanics', type='Concept', properties={}), Node(id='Quantum Theory', type='Concept', properties={}), Node(id='Quantum Physics Of Radiation', type='Concept', properties={}), Node(id='Photons', type='Concept', properties={}), Node(id='Satyendra Nath Bose', type='Person', properties={}), Node(id='Bose–Einstein Statistics', type='Concept', properties={}), Node(id='Unified Field Theory', type='Concept', properties={}), Node(id='Electromagnetism', type='Concept', properties={})]\n", "Relationships:[Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Theory Of Relativity', type='Concept', properties={}), type='DEVELOPED', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Quantum Mechanics', type='Concept', properties={}), type='CONTRIBUTED_TO', properties={}), Relationship(source=Node(id='Mass–Energy Equivalence Formula', type='Concept', properties={}), target=Node(id='E = Mc2', type='Concept', properties={}), type='REPRESENTED_BY', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='1921 Nobel Prize In Physics', type='Award', properties={}), type='RECEIVED', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Law Of The Photoelectric Effect', type='Concept', properties={}), type='DISCOVERED', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='German Empire', type='Place', properties={}), type='BORN_IN', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Switzerland', type='Place', properties={}), type='MOVED_TO', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Swiss Federal Polytechnic School In Zurich', type='Organization', properties={}), type='ENROLLED_IN', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Swiss Patent Office In Bern', type='Organization', properties={}), type='WORKED_AT', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='University Of Zurich', type='Organization', properties={}), type='SUBMITTED_PHD_TO', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Prussian Academy Of Sciences', type='Organization', properties={}), type='JOINED', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Humboldt University Of Berlin', type='Organization', properties={}), type='JOINED', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Kaiser Wilhelm Institute For Physics', type='Organization', properties={}), type='DIRECTOR_OF', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='United States', type='Place', properties={}), type='MOVED_TO', properties={}), Relationship(source=Node(id='Adolf Hitler', type='Person', properties={}), target=Node(id='Germany', type='Place', properties={}), type='CAME_TO_POWER_IN', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Franklin D. Roosevelt', type='Person', properties={}), type='ENDORSED_LETTER_TO', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='World War Ii', type='Event', properties={}), type='RELATED_TO', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Annus Mirabilis', type='Event', properties={}), type='ASSOCIATED_WITH', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Brownian Motion', type='Concept', properties={}), type='EXPLAINED', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Special Theory Of Relativity', type='Concept', properties={}), type='INTRODUCED', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='General Theory Of Relativity', type='Concept', properties={}), type='PROPOSED', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Spontaneous Emission', type='Concept', properties={}), type='INTRODUCED', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Stimulated Emission', type='Concept', properties={}), type='INTRODUCED', properties={}), Relationship(source=Node(id='Stimulated Emission', type='Concept', properties={}), target=Node(id='Laser', type='Concept', properties={}), type='CORE_MECHANISM_OF', properties={}), Relationship(source=Node(id='Stimulated Emission', type='Concept', properties={}), target=Node(id='Maser', type='Concept', properties={}), type='CORE_MECHANISM_OF', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Statistical Mechanics', type='Concept', properties={}), type='CONTRIBUTED_TO', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Quantum Theory', type='Concept', properties={}), type='CONTRIBUTED_TO', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Quantum Physics Of Radiation', type='Concept', properties={}), type='WORKED_ON', properties={}), Relationship(source=Node(id='Quantum Physics Of Radiation', type='Concept', properties={}), target=Node(id='Photons', type='Concept', properties={}), type='CONSISTS_OF', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Satyendra Nath Bose', type='Person', properties={}), type='COLLABORATED_WITH', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Bose–Einstein Statistics', type='Concept', properties={}), type='LAID_GROUNDWORK_FOR', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Unified Field Theory', type='Concept', properties={}), type='ATTEMPTED_TO_DEVISE', properties={}), Relationship(source=Node(id='Unified Field Theory', type='Concept', properties={}), target=Node(id='Electromagnetism', type='Concept', properties={}), type='INCLUDES', properties={})]\n" ] } ], "source": [ "print(f\"Nodes:{graph_documents[0].nodes}\")\n", "print(f\"Relationships:{graph_documents[0].relationships}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "#### Visualize graph" ] }, { "cell_type": "code", "execution_count": 39, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Graph saved to /Users/thuvu/Documents/vlogging/Research/knowledge_graph_app/knowledge_graph.html\n" ] } ], "source": [ "from pyvis.network import Network\n", "\n", "def visualize_graph(graph_documents):\n", "\n", " # Create network\n", " net = Network(height=\"1200px\", width=\"100%\", directed=True,\n", " notebook=False, bgcolor=\"#222222\", font_color=\"white\")\n", " \n", " nodes = graph_documents[0].nodes\n", " relationships = graph_documents[0].relationships\n", "\n", " # Build lookup for valid nodes\n", " node_dict = {node.id: node for node in nodes}\n", " \n", " # Filter out invalid edges and collect valid node IDs\n", " valid_edges = []\n", " valid_node_ids = set()\n", " for rel in relationships:\n", " if rel.source.id in node_dict and rel.target.id in node_dict:\n", " valid_edges.append(rel)\n", " valid_node_ids.update([rel.source.id, rel.target.id])\n", "\n", "\n", " # Track which nodes are part of any relationship\n", " connected_node_ids = set()\n", " for rel in relationships:\n", " connected_node_ids.add(rel.source.id)\n", " connected_node_ids.add(rel.target.id)\n", "\n", " # Add valid nodes\n", " for node_id in valid_node_ids:\n", " node = node_dict[node_id]\n", " try:\n", " net.add_node(node.id, label=node.id, title=node.type, group=node.type)\n", " except:\n", " continue # skip if error\n", "\n", " # Add valid edges\n", " for rel in valid_edges:\n", " try:\n", " net.add_edge(rel.source.id, rel.target.id, label=rel.type.lower())\n", " except:\n", " continue # skip if error\n", "\n", " # Configure physics\n", " net.set_options(\"\"\"\n", " {\n", " \"physics\": {\n", " \"forceAtlas2Based\": {\n", " \"gravitationalConstant\": -100,\n", " \"centralGravity\": 0.01,\n", " \"springLength\": 200,\n", " \"springConstant\": 0.08\n", " },\n", " \"minVelocity\": 0.75,\n", " \"solver\": \"forceAtlas2Based\"\n", " }\n", " }\n", " \"\"\")\n", " \n", " output_file = \"knowledge_graph.html\"\n", " net.save_graph(output_file)\n", " print(f\"Graph saved to {os.path.abspath(output_file)}\")\n", "\n", " # Try to open in browser\n", " try:\n", " import webbrowser\n", " webbrowser.open(f\"file://{os.path.abspath(output_file)}\")\n", " except:\n", " print(\"Could not open browser automatically\")\n", " \n", "# Run the function\n", "visualize_graph(graph_documents)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Extract specific types of nodes" ] }, { "cell_type": "code", "execution_count": 40, "metadata": {}, "outputs": [], "source": [ "allowed_nodes = [\"Person\", \"Organization\", \"Location\", \"Award\", \"ResearchField\"]\n", "graph_transformer_nodes_defined = LLMGraphTransformer(llm=llm, allowed_nodes=allowed_nodes)\n", "graph_documents_nodes_defined = await graph_transformer_nodes_defined.aconvert_to_graph_documents(documents)" ] }, { "cell_type": "code", "execution_count": 41, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Nodes:[Node(id='Albert Einstein', type='Person', properties={}), Node(id='German Empire', type='Location', properties={}), Node(id='Switzerland', type='Location', properties={}), Node(id='University Of Zurich', type='Organization', properties={}), Node(id='Prussian Academy Of Sciences', type='Organization', properties={}), Node(id='Humboldt University Of Berlin', type='Organization', properties={}), Node(id='Kaiser Wilhelm Institute For Physics', type='Organization', properties={}), Node(id='United States', type='Location', properties={}), Node(id='Franklin D. Roosevelt', type='Person', properties={}), Node(id='Nobel Prize In Physics', type='Award', properties={}), Node(id='Theory Of Relativity', type='Researchfield', properties={}), Node(id='Quantum Mechanics', type='Researchfield', properties={}), Node(id='Photoelectric Effect', type='Researchfield', properties={}), Node(id='Special Relativity', type='Researchfield', properties={}), Node(id='General Relativity', type='Researchfield', properties={}), Node(id='Quantum Electrodynamics', type='Researchfield', properties={}), Node(id='Quantum Optics', type='Researchfield', properties={}), Node(id='Statistical Mechanics', type='Researchfield', properties={}), Node(id='Quantum Theory', type='Researchfield', properties={}), Node(id='Unified Field Theory', type='Researchfield', properties={})]\n", "Relationships:[Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='German Empire', type='Location', properties={}), type='BORN_IN', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Switzerland', type='Location', properties={}), type='MOVED_TO', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='University Of Zurich', type='Organization', properties={}), type='PHD_FROM', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Prussian Academy Of Sciences', type='Organization', properties={}), type='MEMBER_OF', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Humboldt University Of Berlin', type='Organization', properties={}), type='MEMBER_OF', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Kaiser Wilhelm Institute For Physics', type='Organization', properties={}), type='DIRECTOR_OF', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='United States', type='Location', properties={}), type='MOVED_TO', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Franklin D. Roosevelt', type='Person', properties={}), type='ENDORSED_LETTER_TO', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Nobel Prize In Physics', type='Award', properties={}), type='AWARDED', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Theory Of Relativity', type='Researchfield', properties={}), type='CONTRIBUTED_TO', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Quantum Mechanics', type='Researchfield', properties={}), type='CONTRIBUTED_TO', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Photoelectric Effect', type='Researchfield', properties={}), type='DISCOVERED', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Special Relativity', type='Researchfield', properties={}), type='CONTRIBUTED_TO', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='General Relativity', type='Researchfield', properties={}), type='PROPOSED', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Quantum Electrodynamics', type='Researchfield', properties={}), type='CONTRIBUTED_TO', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Quantum Optics', type='Researchfield', properties={}), type='CONTRIBUTED_TO', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Statistical Mechanics', type='Researchfield', properties={}), type='CONTRIBUTED_TO', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Quantum Theory', type='Researchfield', properties={}), type='CONTRIBUTED_TO', properties={}), Relationship(source=Node(id='Albert Einstein', type='Person', properties={}), target=Node(id='Unified Field Theory', type='Researchfield', properties={}), type='ATTEMPTED', properties={})]\n" ] } ], "source": [ "print(f\"Nodes:{graph_documents_nodes_defined[0].nodes}\")\n", "print(f\"Relationships:{graph_documents_nodes_defined[0].relationships}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Extract specific types of relationships" ] }, { "cell_type": "code", "execution_count": 42, "metadata": {}, "outputs": [], "source": [ "allowed_nodes = [\"Person\", \"Organization\", \"Location\", \"Award\", \"ResearchField\"]\n", "allowed_relationships = [\n", " (\"Person\", \"WORKS_AT\", \"Organization\"),\n", " (\"Person\", \"SPOUSE\", \"Person\"),\n", " (\"Person\", \"AWARD\", \"Award\"),\n", " (\"Organization\", \"IN_LOCATION\", \"Location\"),\n", " (\"Person\", \"FIELD_OF_RESEARCH\", \"ResearchField\")\n", "]\n", "graph_transformer_rel_defined = LLMGraphTransformer(\n", " llm=llm,\n", " allowed_nodes=allowed_nodes,\n", " allowed_relationships=allowed_relationships\n", ")\n", "graph_documents_rel_defined = await graph_transformer_rel_defined.aconvert_to_graph_documents(documents)" ] }, { "cell_type": "code", "execution_count": 43, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Graph saved to /Users/thuvu/Documents/vlogging/Research/knowledge_graph_app/knowledge_graph.html\n" ] } ], "source": [ "# Visualize graph\n", "visualize_graph(graph_documents_rel_defined)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": ".venv", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.12.6" } }, "nbformat": 4, "nbformat_minor": 2 }