├── .dockerignore ├── .env.example ├── .gitignore ├── .prettierrc ├── LICENSE ├── README.md ├── api ├── .gcloudignore ├── Dockerfile ├── __init__.py ├── requirements.txt └── src │ ├── __init__.py │ ├── components │ ├── __init__.py │ ├── base_component.py │ ├── company_report.py │ ├── data_disambiguation.py │ ├── data_to_csv.py │ ├── graph_data_update.py │ ├── graph_extraction.py │ ├── question_proposal_generator.py │ ├── summarize_cypher_result.py │ ├── text2cypher.py │ ├── unstructured_data_extractor.py │ └── vector_search.py │ ├── domains │ ├── __init__.py │ └── graph_domain.py │ ├── driver │ ├── __init__.py │ └── neo4j.py │ ├── embedding │ ├── __init__.py │ ├── base_embedding.py │ └── openai.py │ ├── fewshot_examples.py │ ├── llm │ ├── __init__.py │ ├── basellm.py │ └── openai.py │ ├── main.py │ └── utils │ └── unstructured_data_utils.py ├── docker-compose.yml ├── graph_schema.png └── ui ├── .env.example ├── .eslintrc.cts ├── .gcloudignore ├── .gitignore ├── Dockerfile ├── index.html ├── package-lock.json ├── package.json ├── postcss.config.js ├── src ├── chat-with-kg │ ├── App.tsx │ ├── ChatContainer.tsx │ ├── ChatInput.tsx │ ├── ChatMessage.tsx │ ├── images │ │ ├── chatbot.png │ │ └── user.png │ ├── index.css │ ├── main.tsx │ └── types │ │ └── websocketTypes.ts ├── components │ ├── graph-data-modal.tsx │ ├── keymodal.tsx │ ├── neo-graph-2d.css │ ├── neo-graph-2d.tsx │ ├── neo-graph.ts │ └── switch.tsx ├── unstructured-import │ ├── App.css │ ├── App.tsx │ ├── README.md │ ├── index.css │ ├── main.tsx │ ├── types │ │ └── respons-types.ts │ └── utils │ │ ├── cypher-utils.ts │ │ ├── fetch-utils.ts │ │ ├── file-utils.ts │ │ └── graph-schema-utils.ts └── use-case-selector │ ├── index.css │ └── main.tsx ├── tailwind.config.js ├── tsconfig.json ├── use-cases ├── chat-with-kg │ └── index.html └── unstructured-import │ └── index.html └── vite.config.js /.dockerignore: -------------------------------------------------------------------------------- 1 | node_modules -------------------------------------------------------------------------------- /.env.example: -------------------------------------------------------------------------------- 1 | OPENAI_API_KEY=OPENAI 2 | NEO4J_URL=neo4j+s://demo.neo4jlabs.com 3 | NEO4J_USER=companies 4 | NEO4J_PASS=companies 5 | NEO4J_DATABASE=companies -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | node_modules 2 | coverage 3 | .DS_Store 4 | dist 5 | 6 | dist/ 7 | params/ 8 | *.bak 9 | # Byte-compiled / optimized / DLL files 10 | __pycache__/ 11 | *.py[cod] 12 | *$py.class 13 | 14 | *.S 15 | # C extensions 16 | *.so 17 | 18 | 19 | *.ll 20 | .npm 21 | # Distribution / packaging 22 | .Python 23 | env/ 24 | build/ 25 | build-*/ 26 | develop-eggs/ 27 | dist/ 28 | downloads/ 29 | eggs/ 30 | .eggs/ 31 | lib/ 32 | lib64/ 33 | parts/ 34 | sdist/ 35 | var/ 36 | wheels/ 37 | pip-wheel-metadata/ 38 | share/python-wheels/ 39 | *.egg-info/ 40 | .installed.cfg 41 | *.egg 42 | MANIFEST 43 | .vscode 44 | .env 45 | **/neo4j/data/ 46 | 47 | # IDE - VSCode 48 | .idea/ 49 | .venv/ -------------------------------------------------------------------------------- /.prettierrc: -------------------------------------------------------------------------------- 1 | { 2 | "trailingComma": "es5", 3 | "tabWidth": 2, 4 | "semi": true, 5 | "singleQuote": false 6 | } 7 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright 2023 - present Neo4j Inc, 2 | 3 | Apache License 4 | Version 2.0, January 2004 5 | http://www.apache.org/licenses/ 6 | 7 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 8 | 9 | 1. Definitions. 10 | 11 | "License" shall mean the terms and conditions for use, reproduction, 12 | and distribution as defined by Sections 1 through 9 of this document. 13 | 14 | "Licensor" shall mean the copyright owner or entity authorized by 15 | the copyright owner that is granting the License. 16 | 17 | "Legal Entity" shall mean the union of the acting entity and all 18 | other entities that control, are controlled by, or are under common 19 | control with that entity. For the purposes of this definition, 20 | "control" means (i) the power, direct or indirect, to cause the 21 | direction or management of such entity, whether by contract or 22 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 23 | outstanding shares, or (iii) beneficial ownership of such entity. 24 | 25 | "You" (or "Your") shall mean an individual or Legal Entity 26 | exercising permissions granted by this License. 27 | 28 | "Source" form shall mean the preferred form for making modifications, 29 | including but not limited to software source code, documentation 30 | source, and configuration files. 31 | 32 | "Object" form shall mean any form resulting from mechanical 33 | transformation or translation of a Source form, including but 34 | not limited to compiled object code, generated documentation, 35 | and conversions to other media types. 36 | 37 | "Work" shall mean the work of authorship, whether in Source or 38 | Object form, made available under the License, as indicated by a 39 | copyright notice that is included in or attached to the work 40 | (an example is provided in the Appendix below). 41 | 42 | "Derivative Works" shall mean any work, whether in Source or Object 43 | form, that is based on (or derived from) the Work and for which the 44 | editorial revisions, annotations, elaborations, or other modifications 45 | represent, as a whole, an original work of authorship. For the purposes 46 | of this License, Derivative Works shall not include works that remain 47 | separable from, or merely link (or bind by name) to the interfaces of, 48 | the Work and Derivative Works thereof. 49 | 50 | "Contribution" shall mean any work of authorship, including 51 | the original version of the Work and any modifications or additions 52 | to that Work or Derivative Works thereof, that is intentionally 53 | submitted to Licensor for inclusion in the Work by the copyright owner 54 | or by an individual or Legal Entity authorized to submit on behalf of 55 | the copyright owner. For the purposes of this definition, "submitted" 56 | means any form of electronic, verbal, or written communication sent 57 | to the Licensor or its representatives, including but not limited to 58 | communication on electronic mailing lists, source code control systems, 59 | and issue tracking systems that are managed by, or on behalf of, the 60 | Licensor for the purpose of discussing and improving the Work, but 61 | excluding communication that is conspicuously marked or otherwise 62 | designated in writing by the copyright owner as "Not a Contribution." 63 | 64 | "Contributor" shall mean Licensor and any individual or Legal Entity 65 | on behalf of whom a Contribution has been received by Licensor and 66 | subsequently incorporated within the Work. 67 | 68 | 2. Grant of Copyright License. Subject to the terms and conditions of 69 | this License, each Contributor hereby grants to You a perpetual, 70 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 71 | copyright license to reproduce, prepare Derivative Works of, 72 | publicly display, publicly perform, sublicense, and distribute the 73 | Work and such Derivative Works in Source or Object form. 74 | 75 | 3. Grant of Patent License. Subject to the terms and conditions of 76 | this License, each Contributor hereby grants to You a perpetual, 77 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 78 | (except as stated in this section) patent license to make, have made, 79 | use, offer to sell, sell, import, and otherwise transfer the Work, 80 | where such license applies only to those patent claims licensable 81 | by such Contributor that are necessarily infringed by their 82 | Contribution(s) alone or by combination of their Contribution(s) 83 | with the Work to which such Contribution(s) was submitted. If You 84 | institute patent litigation against any entity (including a 85 | cross-claim or counterclaim in a lawsuit) alleging that the Work 86 | or a Contribution incorporated within the Work constitutes direct 87 | or contributory patent infringement, then any patent licenses 88 | granted to You under this License for that Work shall terminate 89 | as of the date such litigation is filed. 90 | 91 | 4. Redistribution. You may reproduce and distribute copies of the 92 | Work or Derivative Works thereof in any medium, with or without 93 | modifications, and in Source or Object form, provided that You 94 | meet the following conditions: 95 | 96 | (a) You must give any other recipients of the Work or 97 | Derivative Works a copy of this License; and 98 | 99 | (b) You must cause any modified files to carry prominent notices 100 | stating that You changed the files; and 101 | 102 | (c) You must retain, in the Source form of any Derivative Works 103 | that You distribute, all copyright, patent, trademark, and 104 | attribution notices from the Source form of the Work, 105 | excluding those notices that do not pertain to any part of 106 | the Derivative Works; and 107 | 108 | (d) If the Work includes a "NOTICE" text file as part of its 109 | distribution, then any Derivative Works that You distribute must 110 | include a readable copy of the attribution notices contained 111 | within such NOTICE file, excluding those notices that do not 112 | pertain to any part of the Derivative Works, in at least one 113 | of the following places: within a NOTICE text file distributed 114 | as part of the Derivative Works; within the Source form or 115 | documentation, if provided along with the Derivative Works; or, 116 | within a display generated by the Derivative Works, if and 117 | wherever such third-party notices normally appear. The contents 118 | of the NOTICE file are for informational purposes only and 119 | do not modify the License. You may add Your own attribution 120 | notices within Derivative Works that You distribute, alongside 121 | or as an addendum to the NOTICE text from the Work, provided 122 | that such additional attribution notices cannot be construed 123 | as modifying the License. 124 | 125 | You may add Your own copyright statement to Your modifications and 126 | may provide additional or different license terms and conditions 127 | for use, reproduction, or distribution of Your modifications, or 128 | for any such Derivative Works as a whole, provided Your use, 129 | reproduction, and distribution of the Work otherwise complies with 130 | the conditions stated in this License. 131 | 132 | 5. Submission of Contributions. Unless You explicitly state otherwise, 133 | any Contribution intentionally submitted for inclusion in the Work 134 | by You to the Licensor shall be under the terms and conditions of 135 | this License, without any additional terms or conditions. 136 | Notwithstanding the above, nothing herein shall supersede or modify 137 | the terms of any separate license agreement you may have executed 138 | with Licensor regarding such Contributions. 139 | 140 | 6. Trademarks. This License does not grant permission to use the trade 141 | names, trademarks, service marks, or product names of the Licensor, 142 | except as required for reasonable and customary use in describing the 143 | origin of the Work and reproducing the content of the NOTICE file. 144 | 145 | 7. Disclaimer of Warranty. Unless required by applicable law or 146 | agreed to in writing, Licensor provides the Work (and each 147 | Contributor provides its Contributions) on an "AS IS" BASIS, 148 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 149 | implied, including, without limitation, any warranties or conditions 150 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 151 | PARTICULAR PURPOSE. You are solely responsible for determining the 152 | appropriateness of using or redistributing the Work and assume any 153 | risks associated with Your exercise of permissions under this License. 154 | 155 | 8. Limitation of Liability. In no event and under no legal theory, 156 | whether in tort (including negligence), contract, or otherwise, 157 | unless required by applicable law (such as deliberate and grossly 158 | negligent acts) or agreed to in writing, shall any Contributor be 159 | liable to You for damages, including any direct, indirect, special, 160 | incidental, or consequential damages of any character arising as a 161 | result of this License or out of the use or inability to use the 162 | Work (including but not limited to damages for loss of goodwill, 163 | work stoppage, computer failure or malfunction, or any and all 164 | other commercial damages or losses), even if such Contributor 165 | has been advised of the possibility of such damages. 166 | 167 | 9. Accepting Warranty or Additional Liability. While redistributing 168 | the Work or Derivative Works thereof, You may choose to offer, 169 | and charge a fee for, acceptance of support, warranty, indemnity, 170 | or other liability obligations and/or rights consistent with this 171 | License. However, in accepting such obligations, You may act only 172 | on Your own behalf and on Your sole responsibility, not on behalf 173 | of any other Contributor, and only if You agree to indemnify, 174 | defend, and hold each Contributor harmless for any liability 175 | incurred by, or claims asserted against, such Contributor by reason 176 | of your accepting any such warranty or additional liability. 177 | 178 | END OF TERMS AND CONDITIONS 179 | 180 | APPENDIX: How to apply the Apache License to your work. 181 | 182 | To apply the Apache License to your work, attach the following 183 | boilerplate notice, with the fields enclosed by brackets "{}" 184 | replaced with your own identifying information. (Don't include 185 | the brackets!) The text should be enclosed in the appropriate 186 | comment syntax for the file format. We also recommend that a 187 | file or class name and description of purpose be included on the 188 | same "printed page" as the copyright notice for easier 189 | identification within third-party archives. 190 | 191 | Copyright 2023 - present Neo4j Inc, 192 | 193 | Licensed under the Apache License, Version 2.0 (the "License"); 194 | you may not use this file except in compliance with the License. 195 | You may obtain a copy of the License at 196 | 197 | http://www.apache.org/licenses/LICENSE-2.0 198 | 199 | Unless required by applicable law or agreed to in writing, software 200 | distributed under the License is distributed on an "AS IS" BASIS, 201 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 202 | See the License for the specific language governing permissions and 203 | limitations under the License. 204 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Project Knowledge Graph and LLM 2 | 3 | Welcome to the Knowledge Graph + LLM project, where we explore the intersection of large language models and knowledge graphs. 4 | In the project, we explore how to use large language models to augment knowledge graphs and how to use natural language to query knowledge graphs. 5 | 6 | ## Repository Structure 7 | 8 | Our repository is designed with an efficient and logical structure for ease of navigation: 9 | 10 | - **Backend Code**: The backend code is found in the api folder in the main.py file you can find all endpoints and their corresponding functions. All LLM functionality is split into different components which have thier own purpose. 11 | 12 | - **Frontend Code**: The frontend code is organized into two folders - one for each use case these can be found in ui/src. Each folder contains separate React applications that are independent from each other. 13 | 14 | ## Running the Demos 15 | 16 | To simplify the process of running the demos, we have incorporated scripts that generate Docker images. To use these, you'll need to: 17 | 18 | 1. Navigate into the root directory. 19 | 2. Create an env file. You can use the env.example file as a template. (The open API key is optional and can be provided from the UI instead) 20 | 3. run `docker-compose up` to build the images. 21 | 22 | This will start the backend and frontend servers, and you can access the demos at the following URLs: 23 | 24 | - user interface: http://localhost:4173/ 25 | 26 | - backend: localhost:7860 27 | 28 | Please note that you'll need Docker installed on your machine to build and run these images. If you haven't already, you can download Docker from [here](https://www.docker.com/products/docker-desktop). 29 | 30 | ## Demo database 31 | Please go to https://sandbox.neo4j.com/ and create a new neo4j database. 32 | ``` 33 | URI: neo4j+s://demo.neo4jlabs.com 34 | username: companies 35 | password: companies 36 | database: companies 37 | ``` 38 | 39 | ## Running the backend and frontend separately 40 | ```shell 41 | cd $PROJECT_ROOT/api 42 | python main.py 43 | 44 | cd $PROJECT_ROOT/ui 45 | npm run dev 46 | 47 | access url: http://localhost:8080/ 48 | ``` 49 | 50 | -------------------------------------------------------------------------------- /api/.gcloudignore: -------------------------------------------------------------------------------- 1 | # This file specifies files that are *not* uploaded to Google Cloud 2 | # using gcloud. It follows the same syntax as .gitignore, with the addition of 3 | # "#!include" directives (which insert the entries of the given .gitignore-style 4 | # file at that point). 5 | # 6 | # For more information, run: 7 | # $ gcloud topic gcloudignore 8 | # 9 | .gcloudignore 10 | # If you would like to upload your .git directory, .gitignore file or files 11 | # from your .gitignore file, remove the corresponding line 12 | # below: 13 | .git 14 | .gitignore 15 | 16 | # Python pycache: 17 | __pycache__/ 18 | # Ignored by the build system 19 | /setup.cfg -------------------------------------------------------------------------------- /api/Dockerfile: -------------------------------------------------------------------------------- 1 | # Use base shared Python image 2 | FROM python:3.11 3 | 4 | ENV PYTHONUNBUFFERED=1 5 | 6 | 7 | # Relevant folder 8 | ARG FOLDER=/api 9 | 10 | # Create a folder 11 | RUN mkdir -p $FOLDER 12 | 13 | # Install packages 14 | COPY ./requirements.txt $FOLDER/requirements.txt 15 | RUN pip install -r $FOLDER/requirements.txt 16 | 17 | # Copy the project files into the container 18 | COPY ./src $FOLDER/src 19 | 20 | # Expose any necessary ports 21 | EXPOSE 7860 22 | 23 | # Set the working directory 24 | WORKDIR $FOLDER/src 25 | 26 | # Start the application 27 | CMD ["uvicorn", "--host", "0.0.0.0", "--port", "7860", "--reload", "--reload-dir", "/api", "main:app"] -------------------------------------------------------------------------------- /api/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ryang420/Knowledge-Graph-LLM/4083ba6278c24597aca18a98b4f8ce248f8bfb65/api/__init__.py -------------------------------------------------------------------------------- /api/requirements.txt: -------------------------------------------------------------------------------- 1 | fastapi==0.95.1 2 | logger==1.4 3 | neo4j==5.8.0 4 | uvicorn==0.22.0 5 | openai==0.27.6 6 | retry==0.9.2 7 | tiktoken==0.4.0 8 | python-dotenv==1.0.0 9 | websockets===11.0.3 10 | gunicorn===20.1.0 -------------------------------------------------------------------------------- /api/src/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ryang420/Knowledge-Graph-LLM/4083ba6278c24597aca18a98b4f8ce248f8bfb65/api/src/__init__.py -------------------------------------------------------------------------------- /api/src/components/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ryang420/Knowledge-Graph-LLM/4083ba6278c24597aca18a98b4f8ce248f8bfb65/api/src/components/__init__.py -------------------------------------------------------------------------------- /api/src/components/base_component.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | from typing import List, Union 3 | 4 | 5 | class BaseComponent(ABC): 6 | """""" 7 | 8 | @abstractmethod 9 | def run( 10 | self, 11 | input: Union[str, List[float]], 12 | ) -> str: 13 | """Comment""" 14 | 15 | def run_async( 16 | self, 17 | input: Union[str, List[float]], 18 | ) -> str: 19 | """Comment""" 20 | -------------------------------------------------------------------------------- /api/src/components/company_report.py: -------------------------------------------------------------------------------- 1 | from components.base_component import BaseComponent 2 | from components.summarize_cypher_result import SummarizeCypherResult 3 | from driver.neo4j import Neo4jDatabase 4 | from llm.basellm import BaseLLM 5 | 6 | HARD_LIMIT_CONTEXT_RECORDS = 10 7 | 8 | 9 | class CompanyReport(BaseComponent): 10 | def __init__( 11 | self, 12 | database: Neo4jDatabase, 13 | company: str, 14 | llm: BaseLLM, 15 | ) -> None: 16 | self.database = database 17 | self.company = company 18 | self.llm = llm 19 | 20 | def run(self): 21 | summarize_results = SummarizeCypherResult( 22 | llm=self.llm, 23 | ) 24 | print("CompanyReport") 25 | company_data = self.database.query( 26 | "MATCH (n {name:$companyName}) return n.summary, n.isDissolved, n.nbrEmployees, n.name, n.motto, n.isPublic, n.revenue", 27 | {"companyName": self.company}, 28 | ) 29 | print(company_data) 30 | relation_data = self.database.query( 31 | "MATCH (n {name:$companyName})-[r]->(m) WHERE NOT m:Article OPTIONAL MATCH (m)-[:IN_COUNTRY]->(c:Country) WITH r,m,c return r,m,c", 32 | {"companyName": self.company}, 33 | ) 34 | print(relation_data) 35 | company_data_output = { 36 | "name": company_data[0]["n.name"], 37 | "motto": company_data[0]["n.motto"], 38 | "summary": company_data[0]["n.summary"], 39 | "isDissolved": company_data[0]["n.isDissolved"], 40 | "nbrEmployees": company_data[0]["n.nbrEmployees"], 41 | "isPublic": company_data[0]["n.isPublic"], 42 | "revenue": company_data[0].get("n.revenue", None), 43 | } 44 | print(company_data_output) 45 | print("all data fetched") 46 | offices = [] 47 | suppliers = [] 48 | subsidiaries = [] 49 | for relation in relation_data: 50 | print(relation) 51 | relation_type = relation["r"][1] 52 | if relation_type == "IN_CITY": 53 | offices.append( 54 | { 55 | "city": relation["m"].get("name", None), 56 | "country": relation.get("c") 57 | and relation["c"].get("name", None), 58 | } 59 | ) 60 | elif relation_type == "HAS_CATEGORY": 61 | company_data_output["industry"] = relation["m"]["name"] 62 | elif relation_type == "HAS_SUPPLIER": 63 | category_result = self.database.query( 64 | "MATCH (n {name:$companyName})-[HAS_CATEGORY]-(c:IndustryCategory) return c.name LIMIT 1", 65 | {"companyName": relation["m"].get("name", None)}, 66 | ) 67 | category = None 68 | if len(category_result) > 0: 69 | category = category_result[0]["c.name"] 70 | 71 | suppliers.append( 72 | { 73 | "summary": relation["m"].get("summary", None), 74 | "revenue": relation["m"].get("revenue", None), 75 | "isDissolved": relation["m"].get("isDissolved", None), 76 | "name": relation["m"].get("name", None), 77 | "isPublic": relation["m"].get("isPublic", None), 78 | "category": category, 79 | } 80 | ) 81 | elif relation_type == "HAS_SUBSIDIARY": 82 | category_result = self.database.query( 83 | "MATCH (n {name:$companyName})-[HAS_CATEGORY]-(c:IndustryCategory) return c.name LIMIT 1", 84 | {"companyName": relation["m"].get("name", None)}, 85 | ) 86 | category = None 87 | if len(category_result) > 0: 88 | category = category_result[0]["c.name"] 89 | article_data = self.database.query( 90 | "MATCH p=(n {name:$companyName})<-[:MENTIONS]-(a:Article)-[:HAS_CHUNK]->(c:Chunk) return c.text, a.title, a.siteName", 91 | {"companyName": relation["m"].get("name", None)}, 92 | ) 93 | print("Article data: " + str(article_data)) 94 | 95 | output = "There is not articles about this company." 96 | if len(article_data) > 0: 97 | output = summarize_results.run( 98 | "Can you summarize the following articles in 50 words about " 99 | + relation["m"].get("name", None) 100 | + " ?", 101 | article_data[:HARD_LIMIT_CONTEXT_RECORDS], 102 | ) 103 | subsidiaries.append( 104 | { 105 | "summary": relation["m"].get("summary", None), 106 | "revenue": relation["m"].get("revenue", None), 107 | "isDissolved": relation["m"].get("isDissolved", None), 108 | "name": relation["m"].get("name", None), 109 | "isPublic": relation["m"].get("isPublic", None), 110 | "category": category, 111 | "articleSummary": output, 112 | } 113 | ) 114 | elif relation_type == "HAS_CEO": 115 | company_data_output["ceo"] = relation["m"]["name"] 116 | company_data_output["offices"] = offices 117 | article_data = self.database.query( 118 | "MATCH p=(n {name:$companyName})<-[:MENTIONS]-(a:Article)-[:HAS_CHUNK]->(c:Chunk) return c.text, a.title, a.siteName", 119 | {"companyName": self.company}, 120 | ) 121 | 122 | output = summarize_results.run( 123 | "Can you summarize the following articles about " + self.company + " ?", 124 | article_data[:HARD_LIMIT_CONTEXT_RECORDS], 125 | ) 126 | print("output: " + output) 127 | return { 128 | "company": company_data_output, 129 | "subsidiaries": subsidiaries, 130 | "suppliers": suppliers, 131 | "articleSummary": output, 132 | } 133 | -------------------------------------------------------------------------------- /api/src/components/data_disambiguation.py: -------------------------------------------------------------------------------- 1 | import json 2 | import re 3 | from itertools import groupby 4 | 5 | from components.base_component import BaseComponent 6 | from utils.unstructured_data_utils import ( 7 | nodesTextToListOfDict, 8 | relationshipTextToListOfDict, 9 | ) 10 | 11 | 12 | def generate_system_message_for_nodes() -> str: 13 | return """Your task is to identify if there are duplicated nodes and if so merge them into one nod. Only merge the nodes that refer to the same entity. 14 | You will be given different datasets of nodes and some of these nodes may be duplicated or refer to the same entity. 15 | The datasets contains nodes in the form [ENTITY_ID, TYPE, PROPERTIES]. When you have completed your task please give me the 16 | resulting nodes in the same format. Only return the nodes and relationships no other text. If there is no duplicated nodes return the original nodes. 17 | 18 | Here is an example of the input you will be given: 19 | ["alice", "Person", {"age": 25, "occupation": "lawyer", "name":"Alice"}], ["bob", "Person", {"occupation": "journalist", "name": "Bob"}], ["alice.com", "Webpage", {"url": "www.alice.com"}], ["bob.com", "Webpage", {"url": "www.bob.com"}] 20 | """ 21 | 22 | 23 | def generate_system_message_for_relationships() -> str: 24 | return """ 25 | Your task is to identify if a set of relationships make sense. 26 | If they do not make sense please remove them from the dataset. 27 | Some relationships may be duplicated or refer to the same entity. 28 | Please merge relationships that refer to the same entity. 29 | The datasets contains relationships in the form [ENTITY_ID_1, RELATIONSHIP, ENTITY_ID_2, PROPERTIES]. 30 | You will also be given a set of ENTITY_IDs that are valid. 31 | Some relationships may use ENTITY_IDs that are not in the valid set but refer to a entity in the valid set. 32 | If a relationships refer to a ENTITY_ID in the valid set please change the ID so it matches the valid ID. 33 | When you have completed your task please give me the valid relationships in the same format. Only return the relationships no other text. 34 | 35 | Here is an example of the input you will be given: 36 | ["alice", "roommate", "bob", {"start": 2021}], ["alice", "owns", "alice.com", {}], ["bob", "owns", "bob.com", {}] 37 | """ 38 | 39 | 40 | def generate_prompt(data) -> str: 41 | return f""" Here is the data: 42 | {data} 43 | """ 44 | 45 | 46 | internalRegex = "\[(.*?)\]" 47 | 48 | 49 | class DataDisambiguation(BaseComponent): 50 | def __init__(self, llm) -> None: 51 | self.llm = llm 52 | 53 | def run(self, data: dict) -> str: 54 | nodes = sorted(data["nodes"], key=lambda x: x["label"]) 55 | relationships = data["relationships"] 56 | new_nodes = [] 57 | new_relationships = [] 58 | 59 | node_groups = groupby(nodes, lambda x: x["label"]) 60 | for group in node_groups: 61 | disString = "" 62 | nodes_in_group = list(group[1]) 63 | if len(nodes_in_group) == 1: 64 | new_nodes.extend(nodes_in_group) 65 | continue 66 | 67 | for node in nodes_in_group: 68 | disString += ( 69 | '["' 70 | + node["name"] 71 | + '", "' 72 | + node["label"] 73 | + '", ' 74 | + json.dumps(node["properties"]) 75 | + "]\n" 76 | ) 77 | 78 | messages = [ 79 | {"role": "system", "content": generate_system_message_for_nodes()}, 80 | {"role": "user", "content": generate_prompt(disString)}, 81 | ] 82 | rawNodes = self.llm.generate(messages) 83 | 84 | n = re.findall(internalRegex, rawNodes) 85 | 86 | new_nodes.extend(nodesTextToListOfDict(n)) 87 | 88 | relationship_data = "Relationships:\n" 89 | for relation in relationships: 90 | relationship_data += ( 91 | '["' 92 | + relation["start"] 93 | + '", "' 94 | + relation["type"] 95 | + '", "' 96 | + relation["end"] 97 | + '", ' 98 | + json.dumps(relation["properties"]) 99 | + "]\n" 100 | ) 101 | 102 | node_labels = [node["name"] for node in new_nodes] 103 | relationship_data += "Valid Nodes:\n" + "\n".join(node_labels) 104 | 105 | messages = [ 106 | { 107 | "role": "system", 108 | "content": generate_system_message_for_relationships(), 109 | }, 110 | {"role": "user", "content": generate_prompt(relationship_data)}, 111 | ] 112 | rawRelationships = self.llm.generate(messages) 113 | rels = re.findall(internalRegex, rawRelationships) 114 | new_relationships.extend(relationshipTextToListOfDict(rels)) 115 | return {"nodes": new_nodes, "relationships": new_relationships} 116 | -------------------------------------------------------------------------------- /api/src/components/data_to_csv.py: -------------------------------------------------------------------------------- 1 | from typing import List 2 | 3 | from components.base_component import BaseComponent 4 | 5 | 6 | def generate_system_message() -> str: 7 | return f""" 8 | You will be given a dataset of nodes and relationships. Your task is to covert this data into a CSV format. 9 | Return only the data in the CSV format and nothing else. Return a CSV file for every type of node and relationship. 10 | The data you will be given is in the form [ENTITY, TYPE, PROPERTIES] and a set of relationships in the form [ENTITY1, RELATIONSHIP, ENTITY2, PROPERTIES]. 11 | Important: If you don't get any data or data that does not follow the previously mentioned format return "No data" and nothing else. This is very important. If you don't follow this instruction you will get a 0. 12 | """ 13 | 14 | 15 | def generate_prompt(data) -> str: 16 | return f""" Here is the data: 17 | {data} 18 | """ 19 | 20 | 21 | class DataToCSV(BaseComponent): 22 | def __init__(self, llm) -> None: 23 | self.llm = llm 24 | 25 | def run(self, data: List[str]) -> str: 26 | messages = [ 27 | {"role": "system", "content": generate_system_message()}, 28 | {"role": "user", "content": generate_prompt(data)}, 29 | ] 30 | output = self.llm.generate(messages) 31 | return output 32 | -------------------------------------------------------------------------------- /api/src/components/graph_data_update.py: -------------------------------------------------------------------------------- 1 | from langchain.prompts import ChatPromptTemplate 2 | from langchain.chat_models import ChatOpenAI 3 | from langchain.chains.openai_functions import ( 4 | create_structured_output_chain, 5 | ) 6 | 7 | from langchain_core.output_parsers import StrOutputParser 8 | 9 | 10 | def graph_data_augmentation(llm: ChatOpenAI, 11 | graph_data: str, 12 | user_input: str): 13 | """Augment graph data with user input.""" 14 | prompt = ChatPromptTemplate.from_messages([ 15 | ( 16 | "system", 17 | """ 18 | Your task is to make changes on the graph data by user's input and return a json string. 19 | The graph data is a json format string. It contains two keys: nodes and links. 20 | The user's input is to add, update or delete the nodes and links. 21 | If the deleted node has any relationships or links on it, delete them as well. 22 | Return me the json format string of the updated graph data only. No additional information included! 23 | 24 | """ 25 | ), 26 | ("human", "Use the given context to update {graph_data} by {user_input}") 27 | ]) 28 | 29 | output_parser = StrOutputParser() 30 | chain = prompt | llm | output_parser 31 | 32 | return chain.invoke({"graph_data": graph_data, "user_input": user_input}) 33 | -------------------------------------------------------------------------------- /api/src/components/graph_extraction.py: -------------------------------------------------------------------------------- 1 | from langchain.chains.openai_functions import ( 2 | create_structured_output_chain, 3 | ) 4 | from langchain.chat_models import ChatOpenAI 5 | from langchain.graphs.graph_document import GraphDocument 6 | from langchain.prompts import ChatPromptTemplate 7 | from typing import List, Dict, Any, Optional 8 | from domains.graph_domain import KnowledgeGraph, map_to_base_relationship, map_to_base_node 9 | from langchain.schema import Document 10 | 11 | 12 | def extract_graph( 13 | llm: ChatOpenAI, 14 | document: Document, 15 | nodes: Optional[List[str]] = None, 16 | rels: Optional[List[str]] = None) -> GraphDocument: 17 | # Extract graph data using OpenAI functions 18 | extract_chain = get_extraction_chain(llm, nodes, rels) 19 | data = extract_chain.run(document.page_content) 20 | 21 | # set page_content to empty string to avoid sending it to the client 22 | document.page_content = "" 23 | 24 | # Construct a graph document 25 | graph_document = GraphDocument( 26 | nodes=[map_to_base_node(node) for node in data.nodes], 27 | relationships=[map_to_base_relationship(rel) for rel in data.rels], 28 | source=document 29 | ) 30 | 31 | return graph_document 32 | 33 | 34 | def get_extraction_chain( 35 | llm: ChatOpenAI, 36 | allowed_nodes: Optional[List[str]] = None, 37 | allowed_rels: Optional[List[str]] = None 38 | ): 39 | prompt = ChatPromptTemplate.from_messages( 40 | [( 41 | "system", 42 | f"""# Knowledge Graph Instructions for GPT-4 43 | ## 1. Overview 44 | You are a top-tier algorithm designed for extracting information in structured formats to build a knowledge graph. 45 | - **Nodes** represent entities and concepts. They're akin to Wikipedia nodes. 46 | - The aim is to achieve simplicity and clarity in the knowledge graph, making it accessible for a vast audience. 47 | ## 2. Labeling Nodes 48 | - **Consistency**: Ensure you use basic or elementary types for node labels. 49 | - For example, when you identify an entity representing a person, always label it as **"person"**. Avoid using more specific terms like "mathematician" or "scientist". 50 | - **Node IDs**: Never utilize integers as node IDs. Node IDs should be names or human-readable identifiers found in the text. 51 | {'- **Allowed Node Labels:**' + ", ".join(allowed_nodes) if allowed_nodes else ""} 52 | {'- **Allowed Relationship Types**:' + ", ".join(allowed_rels) if allowed_rels else ""} 53 | ## 3. Handling Numerical Data and Dates 54 | - Numerical data, like age or other related information, should be incorporated as attributes or properties of the respective nodes. 55 | - **No Separate Nodes for Dates/Numbers**: Do not create separate nodes for dates or numerical values. Always attach them as attributes or properties of nodes. 56 | - **Property Format**: Properties must be in a key-value format. 57 | - **Quotation Marks**: Never use escaped single or double quotes within property values. 58 | - **Naming Convention**: Use camelCase for property keys, e.g., `birthDate`. Use UPPER_CASE for relationship type, e.g., `FOUNDED_BY`. 59 | ## 4. Co-reference Resolution 60 | - **Maintain Entity Consistency**: When extracting entities, it's vital to ensure consistency. 61 | If an entity, such as "John Doe", is mentioned multiple times in the text but is referred to by different names or pronouns (e.g., "Joe", "he"), 62 | always use the most complete identifier for that entity throughout the knowledge graph. In this example, use "John Doe" as the entity ID. 63 | Remember, the knowledge graph should be coherent and easily understandable, so maintaining consistency in entity references is crucial. 64 | ## 5. Strict Compliance 65 | Adhere to the rules strictly. Non-compliance will result in termination. 66 | """), 67 | ("human", "Use the given format to extract information from the following input: {input}"), 68 | ("human", "Tip: Make sure to answer in the correct format"), 69 | ]) 70 | return create_structured_output_chain(KnowledgeGraph, llm, prompt, verbose=False) 71 | -------------------------------------------------------------------------------- /api/src/components/question_proposal_generator.py: -------------------------------------------------------------------------------- 1 | from typing import Any, Dict, List, Union 2 | 3 | from components.base_component import BaseComponent 4 | from driver.neo4j import Neo4jDatabase 5 | from llm.basellm import BaseLLM 6 | import re 7 | 8 | 9 | class QuestionProposalGenerator(BaseComponent): 10 | def __init__( 11 | self, 12 | llm: BaseLLM, 13 | database: Neo4jDatabase, 14 | ) -> None: 15 | self.llm = llm 16 | self.database = database 17 | 18 | def get_system_message(self) -> str: 19 | system = f""" 20 | Your task is to come up with questions someone might as about the content of a Neo4j database. Try to make the questions as different as possible. 21 | The questions should be separated by a new line and each line should only contain one question. 22 | To do this, you need to understand the schema of the database. Therefore it's very important that you read the schema carefully. You can find the schema below. 23 | Schema: 24 | {self.database.schema} 25 | """ 26 | 27 | return system 28 | 29 | def get_database_sample(self) -> str: 30 | return self.database.query( 31 | """MATCH (n) 32 | WITH n 33 | WHERE rand() < 0.3 34 | RETURN apoc.map.removeKey(n, 'embedding') AS properties, LABELS(n) as labels 35 | LIMIT 5""" 36 | ) 37 | 38 | def run(self) -> Dict[str, Union[str, List[Dict[str, Any]]]]: 39 | messages = [{"role": "system", "content": self.get_system_message()}] 40 | sample = self.get_database_sample() 41 | messages.append( 42 | { 43 | "role": "user", 44 | "content": f"""Please generate 5 questions about the content of the database. Here is a sample of the database you can use when generating questions: {sample}""", 45 | } 46 | ) 47 | print(messages) 48 | questionsString = self.llm.generate(messages) 49 | questions = [ 50 | # remove number and dot from the beginning of the question 51 | re.sub(r"\A\d\.?\s*", "", question) 52 | for question in questionsString.split("\n") 53 | ] 54 | return { 55 | "output": questions, 56 | } 57 | -------------------------------------------------------------------------------- /api/src/components/summarize_cypher_result.py: -------------------------------------------------------------------------------- 1 | from typing import Any, Awaitable, Callable, Dict, List 2 | 3 | from components.base_component import BaseComponent 4 | from llm.basellm import BaseLLM 5 | 6 | system = f""" 7 | You are an assistant that helps to generate text to form nice and human understandable answers based. 8 | The latest prompt contains the information, and you need to generate a human readable response based on the given information. 9 | Make the answer sound as a response to the question. Do not mention that you based the result on the given information. 10 | Do not add any additional information that is not explicitly provided in the latest prompt. 11 | I repeat, do not add any information that is not explicitly given. 12 | Make the answer as concise as possible and do not use more than 50 words. 13 | """ 14 | 15 | 16 | def remove_large_lists(d: Dict[str, Any]) -> Dict[str, Any]: 17 | """ 18 | The idea is to remove all properties that have large lists (embeddings) or text as values 19 | """ 20 | LIST_CUTOFF = 56 21 | CHARACTER_CUTOFF = 5000 22 | # iterate over all key-value pairs in the dictionary 23 | for key, value in d.items(): 24 | # if the value is a list and has more than list cutoff elements 25 | if isinstance(value, list) and len(value) > LIST_CUTOFF: 26 | d[key] = None 27 | # if the value is a string and has more than list cutoff elements 28 | if isinstance(value, str) and len(value) > CHARACTER_CUTOFF: 29 | d[key] = d[key][:CHARACTER_CUTOFF] 30 | # if the value is a dictionary 31 | elif isinstance(value, dict): 32 | # recurse into the nested dictionary 33 | remove_large_lists(d[key]) 34 | return d 35 | 36 | 37 | class SummarizeCypherResult(BaseComponent): 38 | llm: BaseLLM 39 | exclude_embeddings: bool 40 | 41 | def __init__(self, llm: BaseLLM, exclude_embeddings: bool = True) -> None: 42 | self.llm = llm 43 | self.exclude_embeddings = exclude_embeddings 44 | 45 | def generate_user_prompt(self, question: str, results: List[Dict[str, str]]) -> str: 46 | return f""" 47 | The question was {question} 48 | Answer the question by using the following results: 49 | {[remove_large_lists(el) for el in results] if self.exclude_embeddings else results} 50 | """ 51 | 52 | def run( 53 | self, 54 | question: str, 55 | results: List[Dict[str, Any]], 56 | ) -> Dict[str, str]: 57 | messages = [ 58 | {"role": "system", "content": system}, 59 | {"role": "user", "content": self.generate_user_prompt(question, results)}, 60 | ] 61 | 62 | output = self.llm.generate(messages) 63 | return output 64 | 65 | async def run_async( 66 | self, 67 | question: str, 68 | results: List[Dict[str, Any]], 69 | callback: Callable[[str], Awaitable[Any]] = None, 70 | ) -> Dict[str, str]: 71 | messages = [ 72 | {"role": "system", "content": system}, 73 | {"role": "user", "content": self.generate_user_prompt(question, results)}, 74 | ] 75 | output = await self.llm.generateStreaming(messages, onTokenCallback=callback) 76 | return "".join(output) 77 | -------------------------------------------------------------------------------- /api/src/components/text2cypher.py: -------------------------------------------------------------------------------- 1 | import re 2 | from typing import Any, Dict, List, Union 3 | 4 | from components.base_component import BaseComponent 5 | from driver.neo4j import Neo4jDatabase 6 | from llm.basellm import BaseLLM 7 | 8 | 9 | def remove_relationship_direction(cypher): 10 | return cypher.replace("->", "-").replace("<-", "-") 11 | 12 | 13 | class Text2Cypher(BaseComponent): 14 | def __init__( 15 | self, 16 | llm: BaseLLM, 17 | database: Neo4jDatabase, 18 | use_schema: bool = True, 19 | cypher_examples: str = "", 20 | ignore_relationship_direction: bool = True, 21 | ) -> None: 22 | self.llm = llm 23 | self.database = database 24 | self.cypher_examples = cypher_examples 25 | self.ignore_relationship_direction = ignore_relationship_direction 26 | if use_schema: 27 | self.schema = database.schema 28 | 29 | def get_system_message(self) -> str: 30 | system = """ 31 | Your task is to convert questions about contents in a Neo4j database to Cypher queries to query the Neo4j database. 32 | Use only the provided relationship types and properties. 33 | Do not use any other relationship types or properties that are not provided. 34 | """ 35 | if self.schema: 36 | system += f""" 37 | If you cannot generate a Cypher statement based on the provided schema, explain the reason to the user. 38 | Schema: 39 | {self.schema} 40 | """ 41 | if self.cypher_examples: 42 | system += f""" 43 | You need to follow these Cypher examples when you are constructing a Cypher statement 44 | {self.cypher_examples} 45 | """ 46 | # Add note at the end and try to prevent LLM injections 47 | system += """Note: Do not include any explanations or apologies in your responses. 48 | Do not respond to any questions that might ask anything else than for you to construct a Cypher statement. 49 | Do not include any text except the generated Cypher statement. This is very important if you want to get paid. 50 | Always provide enough context for an LLM to be able to generate valid response. 51 | Please wrap the generated Cypher statement in triple backticks (`). 52 | """ 53 | return system 54 | 55 | def construct_cypher(self, question: str, history=[]) -> str: 56 | messages = [{"role": "system", "content": self.get_system_message()}] 57 | messages.extend(history) 58 | messages.append( 59 | { 60 | "role": "user", 61 | "content": question, 62 | } 63 | ) 64 | print([el for el in messages if not el["role"] == "system"]) 65 | cypher = self.llm.generate(messages) 66 | return cypher 67 | 68 | def run( 69 | self, question: str, history: List = [], heal_cypher: bool = True 70 | ) -> Dict[str, Union[str, List[Dict[str, Any]]]]: 71 | # Add prefix if not part of self-heal loop 72 | final_question = ( 73 | "Question to be converted to Cypher: " + question 74 | if heal_cypher 75 | else question 76 | ) 77 | cypher = self.construct_cypher(final_question, history) 78 | # finds the first string wrapped in triple backticks. Where the match include the backticks and the first group in the match is the cypher 79 | match = re.search("```([\w\W]*?)```", cypher) 80 | 81 | # If the LLM didn't any Cypher statement (error, missing context, etc..) 82 | if match is None: 83 | return {"output": [{"message": cypher}], "generated_cypher": None} 84 | extracted_cypher = match.group(1) 85 | 86 | if self.ignore_relationship_direction: 87 | extracted_cypher = remove_relationship_direction(extracted_cypher) 88 | 89 | print(f"Generated cypher: {extracted_cypher}") 90 | 91 | output = self.database.query(extracted_cypher) 92 | # Catch Cypher syntax error 93 | if heal_cypher and output and output[0].get("code") == "invalid_cypher": 94 | syntax_messages = [{"role": "system", "content": self.get_system_message()}] 95 | syntax_messages.extend( 96 | [ 97 | {"role": "user", "content": question}, 98 | {"role": "assistant", "content": cypher}, 99 | ] 100 | ) 101 | # Try to heal Cypher syntax only once 102 | return self.run( 103 | output[0].get("message"), syntax_messages, heal_cypher=False 104 | ) 105 | 106 | return { 107 | "output": output, 108 | "generated_cypher": extracted_cypher, 109 | } 110 | -------------------------------------------------------------------------------- /api/src/components/unstructured_data_extractor.py: -------------------------------------------------------------------------------- 1 | import re 2 | import os 3 | from typing import List 4 | 5 | from components.base_component import BaseComponent 6 | from llm.basellm import BaseLLM 7 | from utils.unstructured_data_utils import ( 8 | nodesTextToListOfDict, 9 | relationshipTextToListOfDict, 10 | ) 11 | 12 | 13 | def generate_system_message_with_schema() -> str: 14 | return """ 15 | You are a data scientist working for a company that is building a graph database. Your task is to extract information from data and convert it into a graph database. 16 | Provide a set of Nodes in the form [ENTITY, TYPE, PROPERTIES] and a set of relationships in the form [ENTITY1, RELATIONSHIP, ENTITY2, PROPERTIES]. 17 | Pay attention to the type of the properties, if you can't find data for a property set it to null. Don't make anything up and don't add any extra data. If you can't find any data for a node or relationship don't add it. 18 | Only add nodes and relationships that are part of the schema. If you don't get any relationships in the schema only add nodes. 19 | 20 | Example: 21 | Schema: Nodes: [Person {age: integer, name: string}] Relationships: [Person, roommate, Person] 22 | Alice is 25 years old and Bob is her roommate since 1999. 23 | Nodes: 24 | [["Alice", "Person", {"age": 25, "name": "Alice", "gender": "Female"}], 25 | ["Bob", "Person", {"name": "Bob", "gender": "Male"}]] 26 | Relationships: [["Alice", "roommate", "Bob", {"start": 1999}]] 27 | """ 28 | 29 | 30 | def generate_system_message() -> str: 31 | return """ 32 | You are a data scientist working for a company that is building a graph database. Your task is to extract information from data and convert it into a graph database. 33 | Provide a set of Nodes in the form [ENTITY_ID, TYPE, PROPERTIES] and a set of relationships in the form [ENTITY_ID_1, RELATIONSHIP, ENTITY_ID_2, PROPERTIES]. 34 | It is important that the ENTITY_ID_1 and ENTITY_ID_2 exists as nodes with a matching ENTITY_ID. If you can't pair a relationship with a pair of nodes don't add it. 35 | When you find a node or relationship you want to add try to create a generic TYPE for it that describes the entity you can also think of it as a label. 36 | 37 | Example: 38 | Data: Alice lawyer and is 25 years old and Bob is her roommate since 2001. Bob works as a journalist. Alice owns a the webpage www.alice.com and Bob owns the webpage www.bob.com. 39 | Nodes: ["alice", "Person", {"age": 25, "occupation": "lawyer", "name":"Alice"}], ["bob", "Person", {"occupation": "journalist", "name": "Bob"}], ["alice.com", "Webpage", {"url": "www.alice.com"}], ["bob.com", "Webpage", {"url": "www.bob.com"}] 40 | Relationships: ["alice", "roommate", "bob", {"start": 2021}], ["alice", "owns", "alice.com", {}], ["bob", "owns", "bob.com", {}] 41 | """ 42 | 43 | 44 | def generate_system_message_with_labels() -> str: 45 | return """ 46 | You are a data scientist working for a company that is building a graph database. Your task is to extract information from data and convert it into a graph database. 47 | Provide a set of Nodes in the form [ENTITY_ID, TYPE, PROPERTIES] and a set of relationships in the form [ENTITY_ID_1, RELATIONSHIP, ENTITY_ID_2, PROPERTIES]. 48 | It is important that the ENTITY_ID_1 and ENTITY_ID_2 exists as nodes with a matching ENTITY_ID. If you can't pair a relationship with a pair of nodes don't add it. 49 | When you find a node or relationship you want to add try to create a generic TYPE for it that describes the entity you can also think of it as a label. 50 | You will be given a list of types that you should try to use when creating the TYPE for a node. If you can't find a type that fits the node you can create a new one. 51 | 52 | Example: 53 | Data: Alice lawyer and is 25 years old and Bob is her roommate since 2001. Bob works as a journalist. Alice owns a the webpage www.alice.com and Bob owns the webpage www.bob.com. 54 | Types: ["Person", "Webpage"] 55 | Nodes: ["alice", "Person", {"age": 25, "occupation": "lawyer", "name":"Alice"}], ["bob", "Person", {"occupation": "journalist", "name": "Bob"}], ["alice.com", "Webpage", {"url": "www.alice.com"}], ["bob.com", "Webpage", {"url": "www.bob.com"}] 56 | Relationships: ["alice", "roommate", "bob", {"start": 2021}], ["alice", "owns", "alice.com", {}], ["bob", "owns", "bob.com", {}] 57 | """ 58 | 59 | 60 | def generate_prompt(data) -> str: 61 | return f""" 62 | Data: {data}""" 63 | 64 | 65 | def generate_prompt_with_schema(data, schema) -> str: 66 | return f""" 67 | Schema: {schema} 68 | Data: {data}""" 69 | 70 | 71 | def generate_prompt_with_labels(data, labels) -> str: 72 | return f""" 73 | Data: {data} 74 | Types: {labels}""" 75 | 76 | 77 | def splitString(string, max_length) -> List[str]: 78 | return [string[i : i + max_length] for i in range(0, len(string), max_length)] 79 | 80 | 81 | def splitStringToFitTokenSpace( 82 | llm: BaseLLM, string: str, token_use_per_string: int 83 | ) -> List[str]: 84 | allowed_tokens = llm.max_allowed_token_length() - token_use_per_string 85 | chunked_data = splitString(string, 500) 86 | combined_chunks = [] 87 | current_chunk = "" 88 | for chunk in chunked_data: 89 | if ( 90 | llm.num_tokens_from_string(current_chunk) 91 | + llm.num_tokens_from_string(chunk) 92 | < allowed_tokens 93 | ): 94 | current_chunk += chunk 95 | else: 96 | combined_chunks.append(current_chunk) 97 | current_chunk = chunk 98 | combined_chunks.append(current_chunk) 99 | 100 | return combined_chunks 101 | 102 | 103 | def getNodesAndRelationshipsFromResult(result): 104 | regex = "Nodes:\s+(.*?)\s?\s?Relationships:\s?\s?(.*)" 105 | internalRegex = "\[(.*?)\]" 106 | nodes = [] 107 | relationships = [] 108 | for row in result: 109 | parsing = re.match(regex, row, flags=re.S) 110 | if parsing == None: 111 | continue 112 | rawNodes = str(parsing.group(1)) 113 | rawRelationships = parsing.group(2) 114 | nodes.extend(re.findall(internalRegex, rawNodes)) 115 | relationships.extend(re.findall(internalRegex, rawRelationships)) 116 | 117 | result = dict() 118 | result["nodes"] = [] 119 | result["relationships"] = [] 120 | result["nodes"].extend(nodesTextToListOfDict(nodes)) 121 | result["relationships"].extend(relationshipTextToListOfDict(relationships)) 122 | return result 123 | 124 | 125 | class DataExtractor(BaseComponent): 126 | llm: BaseLLM 127 | 128 | def __init__(self, llm: BaseLLM) -> None: 129 | self.llm = llm 130 | 131 | def process(self, chunk): 132 | messages = [ 133 | {"role": "system", "content": generate_system_message()}, 134 | {"role": "user", "content": generate_prompt(chunk)}, 135 | ] 136 | print(messages) 137 | output = self.llm.generate(messages) 138 | return output 139 | 140 | def process_with_labels(self, chunk, labels): 141 | messages = [ 142 | {"role": "system", "content": generate_system_message_with_schema()}, 143 | {"role": "user", "content": generate_prompt_with_labels(chunk, labels)}, 144 | ] 145 | print(messages) 146 | output = self.llm.generate(messages) 147 | return output 148 | 149 | def run(self, data: str) -> List[str]: 150 | system_message = generate_system_message() 151 | prompt_string = generate_prompt("") 152 | token_usage_per_prompt = self.llm.num_tokens_from_string( 153 | system_message + prompt_string 154 | ) 155 | chunked_data = splitStringToFitTokenSpace( 156 | llm=self.llm, string=data, token_use_per_string=token_usage_per_prompt 157 | ) 158 | 159 | results = [] 160 | labels = set() 161 | print("Starting chunked processing") 162 | for chunk in chunked_data: 163 | proceededChunk = self.process_with_labels(chunk, list(labels)) 164 | print("proceededChunk", proceededChunk) 165 | chunkResult = getNodesAndRelationshipsFromResult([proceededChunk]) 166 | print("chunkResult", chunkResult) 167 | newLabels = [node["label"] for node in chunkResult["nodes"]] 168 | print("newLabels", newLabels) 169 | results.append(proceededChunk) 170 | labels.update(newLabels) 171 | 172 | return getNodesAndRelationshipsFromResult(results) 173 | 174 | 175 | class DataExtractorWithSchema(BaseComponent): 176 | llm: BaseLLM 177 | 178 | def __init__(self, llm) -> None: 179 | self.llm = llm 180 | 181 | def run(self, data: str, schema: str) -> List[str]: 182 | system_message = generate_system_message_with_schema() 183 | prompt_string = ( 184 | generate_system_message_with_schema() 185 | + generate_prompt_with_schema(schema=schema, data="") 186 | ) 187 | token_usage_per_prompt = self.llm.num_tokens_from_string( 188 | system_message + prompt_string 189 | ) 190 | 191 | chunked_data = splitStringToFitTokenSpace( 192 | llm=self.llm, string=data, token_use_per_string=token_usage_per_prompt 193 | ) 194 | result = [] 195 | print("Starting chunked processing") 196 | 197 | for chunk in chunked_data: 198 | print("prompt", generate_prompt_with_schema(chunk, schema)) 199 | messages = [ 200 | { 201 | "role": "system", 202 | "content": system_message, 203 | }, 204 | {"role": "user", "content": generate_prompt_with_schema(chunk, schema)}, 205 | ] 206 | output = self.llm.generate(messages) 207 | result.append(output) 208 | return getNodesAndRelationshipsFromResult(result) 209 | -------------------------------------------------------------------------------- /api/src/components/vector_search.py: -------------------------------------------------------------------------------- 1 | from typing import Dict, List, Union 2 | 3 | from components.base_component import BaseComponent 4 | from driver.neo4j import Neo4jDatabase 5 | 6 | 7 | def construct_cypher(label, property, k) -> str: 8 | return f""" 9 | MATCH (n:`{label}`) 10 | WHERE n.`{property}` IS NOT NULL 11 | WITH n, gds.similarity.cosine($input_vector, n.`{property}`) AS similarity 12 | ORDER BY similarity DESC 13 | LIMIT {k} 14 | RETURN apoc.map.removeKey(properties(n), "{property}") AS output 15 | """ 16 | 17 | 18 | class VectorSearch(BaseComponent): 19 | def __init__( 20 | self, database: Neo4jDatabase, label: str, property: str, k: int 21 | ) -> None: 22 | self.database = database 23 | self.generated_cypher = construct_cypher(label, property, k) 24 | 25 | def run(self, input: str) -> Dict[str, Union[str, List[Dict[str, str]]]]: 26 | try: 27 | return { 28 | "output": [ 29 | str(el["output"]) 30 | for el in self.database.query( 31 | self.generated_cypher, {"input_vector": input} 32 | ) 33 | ], 34 | "generated_cypher": self.generated_cypher, 35 | } 36 | except Exception as e: 37 | return e 38 | -------------------------------------------------------------------------------- /api/src/domains/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ryang420/Knowledge-Graph-LLM/4083ba6278c24597aca18a98b4f8ce248f8bfb65/api/src/domains/__init__.py -------------------------------------------------------------------------------- /api/src/domains/graph_domain.py: -------------------------------------------------------------------------------- 1 | from langchain.graphs.graph_document import ( 2 | Node as BaseNode, 3 | Relationship as BaseRelationship, 4 | GraphDocument, 5 | ) 6 | from langchain.schema import Document 7 | from typing import List, Dict, Any, Optional 8 | from langchain.pydantic_v1 import Field, BaseModel 9 | 10 | 11 | class Property(BaseModel): 12 | """A single property consisting of key and value""" 13 | key: str = Field(..., description="key") 14 | value: str = Field(..., description="value") 15 | 16 | 17 | class Node(BaseNode): 18 | properties: Optional[List[Property]] = Field( 19 | None, description="List of node properties") 20 | 21 | 22 | class Relationship(BaseRelationship): 23 | properties: Optional[List[Property]] = Field( 24 | None, description="List of relationship properties" 25 | ) 26 | 27 | 28 | class KnowledgeGraph(BaseModel): 29 | """Generate a knowledge graph with entities and relationships.""" 30 | nodes: List[Node] = Field( 31 | [], description="List of nodes in the knowledge graph") 32 | rels: List[Relationship] = Field( 33 | [], description="List of relationships in the knowledge graph" 34 | ) 35 | 36 | 37 | class GraphData(BaseModel): 38 | """Generate a knowledge graph with entities and relationships.""" 39 | nodes: List[Node] = Field( 40 | [], description="List of nodes in the knowledge graph") 41 | links: List[Relationship] = Field( 42 | [], description="List of relationships in the knowledge graph" 43 | ) 44 | 45 | 46 | def format_property_key(s: str) -> str: 47 | words = s.split() 48 | if not words: 49 | return s 50 | first_word = words[0].lower() 51 | capitalized_words = [word.capitalize() for word in words[1:]] 52 | return "".join([first_word] + capitalized_words) 53 | 54 | 55 | def props_to_dict(props) -> dict: 56 | """Convert properties to a dictionary.""" 57 | properties = {} 58 | if not props: 59 | return properties 60 | for p in props: 61 | properties[format_property_key(p.key)] = p.value 62 | return properties 63 | 64 | 65 | def map_to_base_node(node: Node) -> BaseNode: 66 | """Map the KnowledgeGraph Node to the base Node.""" 67 | properties = props_to_dict(node.properties) if node.properties else {} 68 | # Add name property for better Cypher statement generation 69 | properties["name"] = node.id.title() 70 | return BaseNode( 71 | id=node.id.title(), type=node.type.capitalize(), properties=properties 72 | ) 73 | 74 | 75 | def map_to_base_relationship(rel: Relationship) -> BaseRelationship: 76 | """Map the KnowledgeGraph Relationship to the base Relationship.""" 77 | source = map_to_base_node(rel.source) 78 | target = map_to_base_node(rel.target) 79 | properties = props_to_dict(rel.properties) if rel.properties else {} 80 | return BaseRelationship( 81 | source=source, target=target, type=rel.type, properties=properties 82 | ) 83 | -------------------------------------------------------------------------------- /api/src/driver/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ryang420/Knowledge-Graph-LLM/4083ba6278c24597aca18a98b4f8ce248f8bfb65/api/src/driver/__init__.py -------------------------------------------------------------------------------- /api/src/driver/neo4j.py: -------------------------------------------------------------------------------- 1 | from typing import Any, Dict, List, Optional 2 | 3 | from neo4j import GraphDatabase, exceptions 4 | 5 | node_properties_query = """ 6 | CALL apoc.meta.data() 7 | YIELD label, other, elementType, type, property 8 | WHERE NOT type = "RELATIONSHIP" AND elementType = "node" 9 | WITH label AS nodeLabels, collect({property:property, type:type}) AS properties 10 | RETURN {labels: nodeLabels, properties: properties} AS output 11 | 12 | """ 13 | 14 | rel_properties_query = """ 15 | CALL apoc.meta.data() 16 | YIELD label, other, elementType, type, property 17 | WHERE NOT type = "RELATIONSHIP" AND elementType = "relationship" 18 | WITH label AS nodeLabels, collect({property:property, type:type}) AS properties 19 | RETURN {type: nodeLabels, properties: properties} AS output 20 | """ 21 | 22 | rel_query = """ 23 | CALL apoc.meta.data() 24 | YIELD label, other, elementType, type, property 25 | WHERE type = "RELATIONSHIP" AND elementType = "node" 26 | RETURN "(:" + label + ")-[:" + property + "]->(:" + toString(other[0]) + ")" AS output 27 | """ 28 | 29 | 30 | def schema_text(node_props, rel_props, rels) -> str: 31 | return f""" 32 | This is the schema representation of the Neo4j database. 33 | Node properties are the following: 34 | {node_props} 35 | Relationship properties are the following: 36 | {rel_props} 37 | The relationships are the following 38 | {rels} 39 | """ 40 | 41 | 42 | class Neo4jDatabase: 43 | def __init__( 44 | self, 45 | host: str = "neo4j://localhost:7687", 46 | user: str = "neo4j", 47 | password: str = "pleaseletmein", 48 | database: str = "neo4j", 49 | read_only: bool = True, 50 | ) -> None: 51 | """Initialize a neo4j database""" 52 | self._driver = GraphDatabase.driver(host, auth=(user, password)) 53 | self._database = database 54 | self._read_only = read_only 55 | self.schema = "" 56 | # Verify connection 57 | try: 58 | self._driver.verify_connectivity() 59 | except exceptions.ServiceUnavailable: 60 | raise ValueError( 61 | "Could not connect to Neo4j database. " 62 | "Please ensure that the url is correct" 63 | ) 64 | except exceptions.AuthError: 65 | raise ValueError( 66 | "Could not connect to Neo4j database. " 67 | "Please ensure that the username and password are correct" 68 | ) 69 | try: 70 | self.refresh_schema() 71 | except: 72 | raise ValueError("Missing APOC Core plugin") 73 | 74 | @staticmethod 75 | def _execute_read_only_query(tx, cypher_query: str, params: Optional[Dict] = {}): 76 | result = tx.run(cypher_query, params) 77 | return [r.data() for r in result] 78 | 79 | def query( 80 | self, cypher_query: str, params: Optional[Dict] = {} 81 | ) -> List[Dict[str, Any]]: 82 | with self._driver.session(database=self._database) as session: 83 | try: 84 | if self._read_only: 85 | result = session.read_transaction( 86 | self._execute_read_only_query, cypher_query, params 87 | ) 88 | return result 89 | else: 90 | result = session.run(cypher_query, params) 91 | # Limit to at most 10 results 92 | return [r.data() for r in result] 93 | 94 | # Catch Cypher syntax errors 95 | except exceptions.CypherSyntaxError as e: 96 | return [ 97 | { 98 | "code": "invalid_cypher", 99 | "message": f"Invalid Cypher statement due to an error: {e}", 100 | } 101 | ] 102 | 103 | except exceptions.ClientError as e: 104 | # Catch access mode errors 105 | if e.code == "Neo.ClientError.Statement.AccessMode": 106 | return [ 107 | { 108 | "code": "error", 109 | "message": "Couldn't execute the query due to the read only access to Neo4j", 110 | } 111 | ] 112 | else: 113 | return [{"code": "error", "message": e}] 114 | 115 | def refresh_schema(self) -> None: 116 | node_props = [el["output"] for el in self.query(node_properties_query)] 117 | rel_props = [el["output"] for el in self.query(rel_properties_query)] 118 | rels = [el["output"] for el in self.query(rel_query)] 119 | schema = schema_text(node_props, rel_props, rels) 120 | self.schema = schema 121 | print(schema) 122 | 123 | def check_if_empty(self) -> bool: 124 | data = self.query( 125 | """ 126 | MATCH (n) 127 | WITH count(n) as c 128 | RETURN CASE WHEN c > 0 THEN true ELSE false END AS output 129 | """ 130 | ) 131 | return data[0]["output"] 132 | -------------------------------------------------------------------------------- /api/src/embedding/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ryang420/Knowledge-Graph-LLM/4083ba6278c24597aca18a98b4f8ce248f8bfb65/api/src/embedding/__init__.py -------------------------------------------------------------------------------- /api/src/embedding/base_embedding.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | 3 | 4 | class BaseEmbedding(ABC): 5 | """""" 6 | 7 | @abstractmethod 8 | async def generate( 9 | self, 10 | input: str, 11 | ) -> str: 12 | """Comment""" 13 | -------------------------------------------------------------------------------- /api/src/embedding/openai.py: -------------------------------------------------------------------------------- 1 | import openai 2 | from embedding.base_embedding import BaseEmbedding 3 | 4 | 5 | class OpenAIEmbedding(BaseEmbedding): 6 | """Wrapper around OpenAI embedding models.""" 7 | 8 | def __init__( 9 | self, openai_api_key: str, model_name: str = "text-embedding-ada-002" 10 | ) -> None: 11 | openai.api_key = openai_api_key 12 | self.model = model_name 13 | 14 | def generate( 15 | self, 16 | input: str, 17 | ) -> str: 18 | embedding = openai.Embedding.create(input=input, model=self.model) 19 | return embedding["data"][0]["embedding"] 20 | -------------------------------------------------------------------------------- /api/src/fewshot_examples.py: -------------------------------------------------------------------------------- 1 | def get_fewshot_examples(openai_api_key): 2 | return f""" 3 | #How is Emil Eifrem connected to Michael Hunger? 4 | MATCH (p1:Person {{name:"Emil Eifrem"}}), (p2:Person {{name:"Michael Hunger"}}) 5 | MATCH p=shortestPath((p1)-[*]-(p2)) 6 | RETURN p 7 | #What are the latest news regarding Google? 8 | MATCH (o:Organization {{name:"Google"}})<-[:MENTIONS]-(a:Article)-[:HAS_CHUNK]->(c) 9 | RETURN a.title AS title, c.text AS text, c.date AS date 10 | ORDER BY date DESC LIMIT 3 11 | #Are there any news regarding return to office policies? 12 | CALL apoc.ml.openai.embedding(["Are there any news regarding return to office policies?"], 13 | "{openai_api_key}") YIELD embedding 14 | MATCH (c:Chunk) 15 | WITH c, gds.similarity.cosine(c.embedding, embedding) AS score 16 | ORDER BY score DESC LIMIT 3 17 | RETURN c.text, score 18 | #What is Microsoft policy regarding to the return to office? 19 | CALL apoc.ml.openai.embedding(["What is Microsoft policy regarding to the return to office?"], "{openai_api_key}") YIELD embedding 20 | MATCH (o:Organization {{name:"Microsoft"}})<-[:MENTIONS]-()-[:HAS_CHUNK]->(c) 21 | WITH distinct c, embedding 22 | WITH c, gds.similarity.cosine(c.embedding, embedding) AS score 23 | ORDER BY score DESC LIMIT 3 24 | RETURN c.text, score 25 | 26 | When searching for specific information in the text chunks, never use the CONTAINS clause, but always use the apoc.ml.openai.embedding 27 | and gds.similarity.cosine functions as shown in the examples. 28 | When returning text chunks, always return exactly three chunks, no more, no less. 29 | Remember, instead of using CONTAINS to find information within text chunks use the apoc.ml.openai.embedding and gds.similarity.cosine functions. 30 | """ 31 | -------------------------------------------------------------------------------- /api/src/llm/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ryang420/Knowledge-Graph-LLM/4083ba6278c24597aca18a98b4f8ce248f8bfb65/api/src/llm/__init__.py -------------------------------------------------------------------------------- /api/src/llm/basellm.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | from typing import ( 3 | Any, 4 | List, 5 | ) 6 | 7 | 8 | def raise_(ex): 9 | raise ex 10 | 11 | 12 | class BaseLLM(ABC): 13 | """LLM wrapper should take in a prompt and return a string.""" 14 | 15 | @abstractmethod 16 | def generate(self, messages: List[str]) -> str: 17 | """Comment""" 18 | 19 | @abstractmethod 20 | async def generateStreaming( 21 | self, messages: List[str], onTokenCallback 22 | ) -> List[Any]: 23 | """Comment""" 24 | 25 | @abstractmethod 26 | async def num_tokens_from_string( 27 | self, 28 | string: str, 29 | ) -> str: 30 | """Given a string returns the number of tokens the given string consists of""" 31 | 32 | @abstractmethod 33 | async def max_allowed_token_length( 34 | self, 35 | ) -> int: 36 | """Returns the maximum number of tokens the LLM can handle""" 37 | -------------------------------------------------------------------------------- /api/src/llm/openai.py: -------------------------------------------------------------------------------- 1 | from typing import ( 2 | Callable, 3 | List, 4 | ) 5 | 6 | import openai 7 | import tiktoken 8 | from llm.basellm import BaseLLM 9 | from retry import retry 10 | 11 | 12 | class OpenAIChat(BaseLLM): 13 | """Wrapper around OpenAI Chat large language models.""" 14 | 15 | def __init__( 16 | self, 17 | openai_api_key: str, 18 | model_name: str = "gpt-3.5-turbo", 19 | max_tokens: int = 1000, 20 | temperature: float = 0.0, 21 | ) -> None: 22 | openai.api_key = openai_api_key 23 | self.model = model_name 24 | self.max_tokens = max_tokens 25 | self.temperature = temperature 26 | 27 | @retry(tries=3, delay=1) 28 | def generate( 29 | self, 30 | messages: List[str], 31 | ) -> str: 32 | try: 33 | completions = openai.ChatCompletion.create( 34 | model=self.model, 35 | temperature=self.temperature, 36 | max_tokens=self.max_tokens, 37 | messages=messages, 38 | ) 39 | return completions.choices[0].message.content 40 | # catch context length / do not retry 41 | except openai.error.InvalidRequestError as e: 42 | return str(f"Error: {e}") 43 | # catch authorization errors / do not retry 44 | except openai.error.AuthenticationError as e: 45 | return "Error: The provided OpenAI API key is invalid" 46 | except Exception as e: 47 | print(f"Retrying LLM call {e}") 48 | raise Exception() 49 | 50 | async def generateStreaming( 51 | self, 52 | messages: List[str], 53 | onTokenCallback=Callable[[str], None], 54 | ) -> str: 55 | result = [] 56 | completions = openai.ChatCompletion.create( 57 | model=self.model, 58 | temperature=self.temperature, 59 | max_tokens=self.max_tokens, 60 | messages=messages, 61 | stream=True, 62 | ) 63 | result = [] 64 | for message in completions: 65 | # Process the streamed messages or perform any other desired action 66 | delta = message["choices"][0]["delta"] 67 | if "content" in delta: 68 | result.append(delta["content"]) 69 | await onTokenCallback(message) 70 | return result 71 | 72 | def num_tokens_from_string(self, string: str) -> int: 73 | encoding = tiktoken.encoding_for_model(self.model) 74 | num_tokens = len(encoding.encode(string)) 75 | return num_tokens 76 | 77 | def max_allowed_token_length(self) -> int: 78 | # TODO: list all models and their max tokens from api 79 | return 2049 80 | -------------------------------------------------------------------------------- /api/src/main.py: -------------------------------------------------------------------------------- 1 | import json 2 | import os 3 | from typing import Optional 4 | from langchain.graphs.graph_document import GraphDocument 5 | from langchain.text_splitter import TokenTextSplitter 6 | from langchain.chat_models import ChatOpenAI 7 | from components.company_report import CompanyReport 8 | 9 | from components.data_disambiguation import DataDisambiguation 10 | from components.graph_data_update import graph_data_augmentation 11 | from components.question_proposal_generator import ( 12 | QuestionProposalGenerator, 13 | ) 14 | from components.summarize_cypher_result import SummarizeCypherResult 15 | from components.text2cypher import Text2Cypher 16 | from components.unstructured_data_extractor import ( 17 | DataExtractor, 18 | DataExtractorWithSchema, 19 | ) 20 | from driver.neo4j import Neo4jDatabase 21 | from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect 22 | from fastapi.middleware.cors import CORSMiddleware 23 | from fastapi.responses import JSONResponse 24 | from fewshot_examples import get_fewshot_examples 25 | from llm.openai import OpenAIChat 26 | from pydantic import BaseModel 27 | from components.graph_extraction import extract_graph 28 | from langchain.schema import Document 29 | 30 | # load environment variables from .env 31 | from dotenv import load_dotenv 32 | 33 | # .env is in the root directory 34 | load_dotenv(dotenv_path="../../.env") 35 | 36 | 37 | class Payload(BaseModel): 38 | question: str 39 | api_key: Optional[str] 40 | model_name: Optional[str] 41 | 42 | 43 | class ImportPayload(BaseModel): 44 | input: str 45 | neo4j_schema: Optional[str] 46 | api_key: Optional[str] 47 | 48 | 49 | class UpdateGraphPayload(BaseModel): 50 | graph_data: str 51 | user_input: str 52 | api_key: Optional[str] 53 | 54 | 55 | class questionProposalPayload(BaseModel): 56 | api_key: Optional[str] 57 | 58 | 59 | # Maximum number of records used in the context 60 | HARD_LIMIT_CONTEXT_RECORDS = 10 61 | 62 | HOST = os.getenv("NEO4J_URL") 63 | USER = os.getenv("NEO4J_USER") 64 | PASSWORD = os.getenv("NEO4J_PASS") 65 | DATABASE = os.getenv("NEO4J_DATABASE") 66 | 67 | neo4j_connection = Neo4jDatabase( 68 | host=HOST, 69 | user=USER, 70 | password=PASSWORD, 71 | database=DATABASE, 72 | ) 73 | from langchain.graphs import Neo4jGraph 74 | 75 | graph = Neo4jGraph(url=HOST, username=USER, password=PASSWORD, database=DATABASE) 76 | 77 | # Initialize LLM modules 78 | openai_api_key = os.environ.get("OPENAI_API_KEY", None) 79 | 80 | # Define FastAPI endpoint 81 | app = FastAPI() 82 | 83 | origins = [ 84 | "*", 85 | ] 86 | 87 | app.add_middleware( 88 | CORSMiddleware, 89 | allow_origins=origins, 90 | allow_credentials=True, 91 | allow_methods=["*"], 92 | allow_headers=["*"], 93 | ) 94 | 95 | 96 | @app.post("/questionProposalsForCurrentDb") 97 | async def questionProposalsForCurrentDb(payload: questionProposalPayload): 98 | if not openai_api_key and not payload.api_key: 99 | raise HTTPException( 100 | status_code=422, 101 | detail="Please set OPENAI_API_KEY environment variable or send it as api_key in the request body", 102 | ) 103 | api_key = openai_api_key if openai_api_key else payload.api_key 104 | 105 | questionProposalGenerator = QuestionProposalGenerator( 106 | database=neo4j_connection, 107 | llm=OpenAIChat( 108 | openai_api_key=api_key, 109 | model_name="gpt-3.5-turbo-0613", 110 | max_tokens=512, 111 | temperature=0.8, 112 | ), 113 | ) 114 | 115 | return questionProposalGenerator.run() 116 | 117 | 118 | @app.get("/hasapikey") 119 | async def hasApiKey(): 120 | return JSONResponse(content={"output": openai_api_key is not None}) 121 | 122 | 123 | @app.websocket("/text2text") 124 | async def websocket_endpoint(websocket: WebSocket): 125 | async def sendDebugMessage(message): 126 | await websocket.send_json({"type": "debug", "detail": message}) 127 | 128 | async def sendErrorMessage(message): 129 | await websocket.send_json({"type": "error", "detail": message}) 130 | 131 | async def onToken(token): 132 | delta = token["choices"][0]["delta"] 133 | if "content" not in delta: 134 | return 135 | content = delta["content"] 136 | if token["choices"][0]["finish_reason"] == "stop": 137 | await websocket.send_json({"type": "end", "output": content}) 138 | else: 139 | await websocket.send_json({"type": "stream", "output": content}) 140 | 141 | # await websocket.send_json({"token": token}) 142 | 143 | await websocket.accept() 144 | await sendDebugMessage("connected") 145 | chatHistory = [] 146 | try: 147 | while True: 148 | data = await websocket.receive_json() 149 | if not openai_api_key and not data.get("api_key"): 150 | raise HTTPException( 151 | status_code=422, 152 | detail="Please set OPENAI_API_KEY environment variable or send it as api_key in the request body", 153 | ) 154 | api_key = openai_api_key if openai_api_key else data.get("api_key") 155 | 156 | default_llm = OpenAIChat( 157 | openai_api_key=api_key, 158 | model_name=data.get("model_name", "gpt-3.5-turbo-0613"), 159 | ) 160 | summarize_results = SummarizeCypherResult( 161 | llm=OpenAIChat( 162 | openai_api_key=api_key, 163 | model_name="gpt-3.5-turbo-0613", 164 | max_tokens=128, 165 | ) 166 | ) 167 | 168 | text2cypher = Text2Cypher( 169 | database=neo4j_connection, 170 | llm=default_llm, 171 | cypher_examples=get_fewshot_examples(api_key), 172 | ) 173 | 174 | if "type" not in data: 175 | await websocket.send_json({"error": "missing type"}) 176 | continue 177 | if data["type"] == "question": 178 | try: 179 | question = data["question"] 180 | chatHistory.append({"role": "user", "content": question}) 181 | await sendDebugMessage("received question: " + question) 182 | results = None 183 | try: 184 | results = text2cypher.run(question, chatHistory) 185 | print("results", results) 186 | except Exception as e: 187 | await sendErrorMessage(str(e)) 188 | continue 189 | if results == None: 190 | await sendErrorMessage("Could not generate Cypher statement") 191 | continue 192 | 193 | await websocket.send_json( 194 | { 195 | "type": "start", 196 | } 197 | ) 198 | output = await summarize_results.run_async( 199 | question, 200 | results["output"][:HARD_LIMIT_CONTEXT_RECORDS], 201 | callback=onToken, 202 | ) 203 | chatHistory.append({"role": "system", "content": output}) 204 | await websocket.send_json( 205 | { 206 | "type": "end", 207 | "output": output, 208 | "generated_cypher": results["generated_cypher"], 209 | } 210 | ) 211 | except Exception as e: 212 | await sendErrorMessage(str(e)) 213 | await sendDebugMessage("output done") 214 | except WebSocketDisconnect: 215 | print("disconnected") 216 | 217 | 218 | @app.post("/update_graph") 219 | async def update_graph(payload: UpdateGraphPayload): 220 | if not openai_api_key and not payload.api_key: 221 | raise HTTPException( 222 | status_code=422, 223 | detail="Please set OPENAI_API_KEY environment variable or send it as api_key in the request body", 224 | ) 225 | api_key = openai_api_key if openai_api_key else payload.api_key 226 | try: 227 | llm = ChatOpenAI(openai_api_key=api_key, model="gpt-3.5-turbo-16k", temperature=0) 228 | updated_graph = graph_data_augmentation(llm, payload.graph_data, payload.user_input) 229 | 230 | return {"status": "ok", "data": json.loads(updated_graph)} 231 | except json.JSONDecodeError as e: 232 | print(e) 233 | return {"status": "failed"} 234 | 235 | 236 | @app.post("/save_graph") 237 | async def save_graph(graph_doc: GraphDocument): 238 | graph.add_graph_documents([graph_doc]) 239 | return JSONResponse(content={"output": "ok"}) 240 | 241 | 242 | @app.post("/text2graph") 243 | async def text_to_graph(payload: ImportPayload): 244 | """ 245 | Takes an input and return a graph 246 | """ 247 | if not openai_api_key and not payload.api_key: 248 | raise HTTPException( 249 | status_code=422, 250 | detail="Please set OPENAI_API_KEY environment variable or send it as api_key in the request body", 251 | ) 252 | api_key = openai_api_key if openai_api_key else payload.api_key 253 | 254 | try: 255 | llm = ChatOpenAI(openai_api_key=api_key, model="gpt-3.5-turbo-16k", temperature=0) 256 | 257 | data = payload.input 258 | # transform data to a document object 259 | document = Document(page_content=data) 260 | 261 | # Define chunking strategy 262 | text_splitter = TokenTextSplitter(chunk_size=2048, chunk_overlap=24) 263 | 264 | # split the document by text_splitter 265 | documents = text_splitter.split_documents([document]) 266 | 267 | for d in documents: 268 | graph_data = extract_graph(llm, d) 269 | 270 | return {"data": graph_data} 271 | 272 | except Exception as e: 273 | print(e) 274 | return f"Error: {e}" 275 | 276 | 277 | @app.post("/data2cypher") 278 | async def root(payload: ImportPayload): 279 | """ 280 | Takes an input and created a Cypher query 281 | """ 282 | if not openai_api_key and not payload.api_key: 283 | raise HTTPException( 284 | status_code=422, 285 | detail="Please set OPENAI_API_KEY environment variable or send it as api_key in the request body", 286 | ) 287 | api_key = openai_api_key if openai_api_key else payload.api_key 288 | 289 | try: 290 | result = "" 291 | 292 | llm = OpenAIChat( 293 | openai_api_key=api_key, model_name="gpt-3.5-turbo-16k", max_tokens=4000 294 | ) 295 | 296 | if not payload.neo4j_schema: 297 | extractor = DataExtractor(llm=llm) 298 | result = extractor.run(data=payload.input) 299 | else: 300 | extractor = DataExtractorWithSchema(llm=llm) 301 | result = extractor.run(schema=payload.neo4j_schema, data=payload.input) 302 | 303 | print("Extracted result: " + str(result)) 304 | 305 | disambiguation = DataDisambiguation(llm=llm) 306 | disambiguation_result = disambiguation.run(result) 307 | 308 | print("Disambiguation result " + str(disambiguation_result)) 309 | 310 | return {"data": disambiguation_result} 311 | 312 | except Exception as e: 313 | print(e) 314 | return f"Error: {e}" 315 | 316 | 317 | class companyReportPayload(BaseModel): 318 | company: str 319 | api_key: Optional[str] 320 | 321 | 322 | # This endpoint is database specific and only works with the Demo database. 323 | @app.post("/companyReport") 324 | async def companyInformation(payload: companyReportPayload): 325 | api_key = openai_api_key if openai_api_key else payload.api_key 326 | if not openai_api_key and not payload.api_key: 327 | raise HTTPException( 328 | status_code=422, 329 | detail="Please set OPENAI_API_KEY environment variable or send it as api_key in the request body", 330 | ) 331 | api_key = openai_api_key if openai_api_key else payload.api_key 332 | 333 | llm = OpenAIChat( 334 | openai_api_key=api_key, 335 | model_name="gpt-3.5-turbo-16k-0613", 336 | max_tokens=512, 337 | ) 338 | print("Running company report for " + payload.company) 339 | company_report = CompanyReport(neo4j_connection, payload.company, llm) 340 | result = company_report.run() 341 | 342 | return JSONResponse(content={"output": result}) 343 | 344 | 345 | @app.post("/companyReport/list") 346 | async def companyReportList(): 347 | company_data = neo4j_connection.query( 348 | "MATCH (n:Organization) WITH n WHERE rand() < 0.01 return n.name LIMIT 5", 349 | ) 350 | 351 | return JSONResponse(content={"output": [x["n.name"] for x in company_data]}) 352 | 353 | 354 | @app.get("/health") 355 | async def health(): 356 | return {"status": "ok"} 357 | 358 | 359 | @app.get("/ready") 360 | async def readiness_check(): 361 | return {"status": "ok"} 362 | 363 | 364 | if __name__ == "__main__": 365 | import uvicorn 366 | from pathlib import Path 367 | 368 | uvicorn.run(f"{Path(__file__).stem}:app", port=int(os.environ.get("PORT", 7860)), host="0.0.0.0", reload=True) 369 | -------------------------------------------------------------------------------- /api/src/utils/unstructured_data_utils.py: -------------------------------------------------------------------------------- 1 | import json 2 | import re 3 | 4 | regex = "Nodes:\s+(.*?)\s?\s?Relationships:\s+(.*)" 5 | internalRegex = "\[(.*?)\]" 6 | jsonRegex = "\{.*\}" 7 | 8 | 9 | def nodesTextToListOfDict(nodes): 10 | result = [] 11 | for node in nodes: 12 | nodeList = node.split(",") 13 | if len(nodeList) < 2: 14 | continue 15 | 16 | name = nodeList[0].strip().replace('"', "") 17 | label = nodeList[1].strip().replace('"', "") 18 | properties = re.search(jsonRegex, node) 19 | if properties == None: 20 | properties = "{}" 21 | else: 22 | properties = properties.group(0) 23 | properties = properties.replace("True", "true") 24 | try: 25 | properties = json.loads(properties) 26 | except: 27 | properties = {} 28 | result.append({"name": name, "label": label, "properties": properties}) 29 | return result 30 | 31 | 32 | def relationshipTextToListOfDict(relationships): 33 | result = [] 34 | for relation in relationships: 35 | relationList = relation.split(",") 36 | if len(relation) < 3: 37 | continue 38 | start = relationList[0].strip().replace('"', "") 39 | end = relationList[2].strip().replace('"', "") 40 | type = relationList[1].strip().replace('"', "") 41 | 42 | properties = re.search(jsonRegex, relation) 43 | if properties == None: 44 | properties = "{}" 45 | else: 46 | properties = properties.group(0) 47 | properties = properties.replace("True", "true") 48 | try: 49 | properties = json.loads(properties) 50 | except: 51 | properties = {} 52 | result.append( 53 | {"start": start, "end": end, "type": type, "properties": properties} 54 | ) 55 | return result 56 | -------------------------------------------------------------------------------- /docker-compose.yml: -------------------------------------------------------------------------------- 1 | version: "3.7" 2 | services: 3 | backend: 4 | build: 5 | context: ./api 6 | hostname: api 7 | restart: always 8 | container_name: api 9 | ports: 10 | - 7860:7860 11 | env_file: 12 | - .env 13 | volumes: 14 | - ./api:/api 15 | frontend: 16 | build: 17 | context: ./ui 18 | hostname: ui 19 | restart: always 20 | container_name: ui 21 | ports: 22 | - 4173:4173 23 | -------------------------------------------------------------------------------- /graph_schema.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ryang420/Knowledge-Graph-LLM/4083ba6278c24597aca18a98b4f8ce248f8bfb65/graph_schema.png -------------------------------------------------------------------------------- /ui/.env.example: -------------------------------------------------------------------------------- 1 | VITE_UNSTRUCTURED_IMPORT_BACKEND_ENDPOINT=http://localhost:7860/ 2 | VITE_KG_CHAT_BACKEND_ENDPOINT=ws://localhost:7860/text2text 3 | VITE_HAS_API_KEY_ENDPOINT=http://localhost:7860/hasapikey 4 | VITE_KG_CHAT_SAMPLE_QUESTIONS_ENDPOINT=http://localhost:7860/questionProposalsForCurrentDb 5 | VITE_REPORT_DATA_ENDPOINT=http://localhost:7860 -------------------------------------------------------------------------------- /ui/.eslintrc.cts: -------------------------------------------------------------------------------- 1 | export default { 2 | env: { browser: true, es2020: true }, 3 | extends: [ 4 | "eslint:recommended", 5 | "plugin:react/recommended", 6 | "plugin:react/jsx-runtime", 7 | "plugin:react-hooks/recommended", 8 | ], 9 | parserOptions: { ecmaVersion: "latest", sourceType: "module" }, 10 | settings: { react: { version: "17.0" } }, 11 | plugins: [], 12 | rules: {}, 13 | }; 14 | -------------------------------------------------------------------------------- /ui/.gcloudignore: -------------------------------------------------------------------------------- 1 | # This file specifies files that are _not_ uploaded to Google Cloud 2 | 3 | # using gcloud. It follows the same syntax as .gitignore, with the addition of 4 | 5 | # "#!include" directives (which insert the entries of the given .gitignore-style 6 | 7 | # file at that point). 8 | 9 | # 10 | 11 | # For more information, run: 12 | 13 | # $ gcloud topic gcloudignore 14 | 15 | # 16 | 17 | .gcloudignore 18 | 19 | # If you would like to upload your .git directory, .gitignore file or files 20 | 21 | # from your .gitignore file, remove the corresponding line 22 | 23 | # below: 24 | 25 | .git 26 | .gitignore 27 | 28 | # Python pycache: 29 | 30 | **pycache**/ 31 | 32 | # Ignored by the build system 33 | 34 | /setup.cfg 35 | node_modules/ 36 | dist/ 37 | -------------------------------------------------------------------------------- /ui/.gitignore: -------------------------------------------------------------------------------- 1 | # Logs 2 | logs 3 | *.log 4 | npm-debug.log* 5 | yarn-debug.log* 6 | yarn-error.log* 7 | pnpm-debug.log* 8 | lerna-debug.log* 9 | 10 | node_modules 11 | dist 12 | dist-ssr 13 | *.local 14 | 15 | # Editor directories and files 16 | .vscode/* 17 | !.vscode/extensions.json 18 | .idea 19 | .DS_Store 20 | *.suo 21 | *.ntvs* 22 | *.njsproj 23 | *.sln 24 | *.sw? 25 | -------------------------------------------------------------------------------- /ui/Dockerfile: -------------------------------------------------------------------------------- 1 | FROM node:lts-alpine 2 | 3 | ARG FOLDER=/ui 4 | 5 | RUN mkdir -p $FOLDER 6 | 7 | COPY . /ui/. 8 | 9 | #CMD ["cd", "ui"] 10 | 11 | WORKDIR $FOLDER 12 | 13 | RUN npm install 14 | 15 | EXPOSE 4173 16 | 17 | RUN npm run build 18 | 19 | CMD ["npm", "run", "preview"] 20 | -------------------------------------------------------------------------------- /ui/index.html: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | Freyr KG + LLM Demo 7 | 8 | 9 |
10 | 11 | 12 | 13 | -------------------------------------------------------------------------------- /ui/package.json: -------------------------------------------------------------------------------- 1 | { 2 | "name": "kgllm-ui", 3 | "version": "1.0.0", 4 | "description": "UI template for the KGLLM project", 5 | "scripts": { 6 | "dev": "vite", 7 | "start": "serve dist", 8 | "build": "tsc && vite build", 9 | "lint": "eslint src --ext js,jsx --report-unused-disable-directives --max-warnings 0", 10 | "preview": "vite preview" 11 | }, 12 | "main": "dist/index.js", 13 | "repository": { 14 | "type": "git", 15 | "url": "git+https://github.com/ryang420/Freyr-KG-LLM.git" 16 | }, 17 | "keywords": [ 18 | "Neo4j", 19 | "LLM", 20 | "KGLLM" 21 | ], 22 | "author": "Neo4j", 23 | "license": "MIT", 24 | "bugs": { 25 | "url": "https://github.com/ryang420/Freyr-KG-LLM/issues" 26 | }, 27 | "homepage": "https://github.com/ryang420/Freyr-KG-LLM#readme", 28 | "dependencies": { 29 | "@neo4j-ndl/base": "1.4.0", 30 | "@neo4j/graph-schema-utils": "^1.0.0-next.9", 31 | "@types/react-modal": "^3.16.0", 32 | "file-saver": "^2.0.5", 33 | "jszip": "^3.10.1", 34 | "nanoid": "^4.0.2", 35 | "react": "^18.2.0", 36 | "react-dom": "^18.2.0", 37 | "react-force-graph-2d": "^1.25.3", 38 | "react-minimal-pie-chart": "^8.4.0", 39 | "react-modal": "^3.16.1", 40 | "react-use-websocket": "^4.3.1", 41 | "serve": "^14.2.0", 42 | "typescript": "^5.0.4" 43 | }, 44 | "devDependencies": { 45 | "@tailwindcss/typography": "^0.5.9", 46 | "@types/file-saver": "^2.0.5", 47 | "@types/react": "^18.2.7", 48 | "@types/react-dom": "^18.2.4", 49 | "@types/sockjs-client": "^1.5.1", 50 | "@vitejs/plugin-react": "^4.0.0", 51 | "autoprefixer": "^10.4.14", 52 | "daisyui": "^3.1.7", 53 | "eslint": "^8.41.0", 54 | "eslint-plugin-react": "^7.32.2", 55 | "eslint-plugin-react-hooks": "^4.6.0", 56 | "eslint-plugin-react-refresh": "^0.4.1", 57 | "postcss": "^8.4.23", 58 | "tailwindcss": "^3.3.2", 59 | "vite": "^4.3.8" 60 | }, 61 | "engines": { 62 | "node": "^18.0.0" 63 | } 64 | } 65 | -------------------------------------------------------------------------------- /ui/postcss.config.js: -------------------------------------------------------------------------------- 1 | module.exports = { 2 | plugins: { 3 | tailwindcss: {}, 4 | autoprefixer: {}, 5 | }, 6 | }; 7 | -------------------------------------------------------------------------------- /ui/src/chat-with-kg/App.tsx: -------------------------------------------------------------------------------- 1 | import { useCallback, useEffect, useState, ChangeEvent } from "react"; 2 | import ChatContainer from "./ChatContainer"; 3 | import type { ChatMessageObject } from "./ChatMessage"; 4 | import ChatInput from "./ChatInput"; 5 | import useWebSocket, { ReadyState } from "react-use-websocket"; 6 | import KeyModal from "../components/keymodal"; 7 | import type { 8 | ConversationState, 9 | WebSocketRequest, 10 | WebSocketResponse, 11 | } from "./types/websocketTypes"; 12 | 13 | const SEND_REQUESTS = true; 14 | 15 | const chatMessageObjects: ChatMessageObject[] = SEND_REQUESTS 16 | ? [] 17 | : [ 18 | { 19 | id: 0, 20 | type: "input", 21 | sender: "self", 22 | message: 23 | "This is the first message which has decently long text and would denote something typed by the user", 24 | complete: true, 25 | }, 26 | { 27 | id: 1, 28 | type: "text", 29 | sender: "bot", 30 | message: 31 | "And here is another message which would denote a response from the server, which for now will only be text", 32 | complete: true, 33 | }, 34 | ]; 35 | 36 | const URI = 37 | import.meta.env.VITE_KG_CHAT_BACKEND_ENDPOINT ?? 38 | "ws://localhost:7860/text2text"; 39 | 40 | const HAS_API_KEY_URI = 41 | import.meta.env.VITE_HAS_API_KEY_ENDPOINT ?? 42 | "http://localhost:7860/hasapikey"; 43 | 44 | const QUESTIONS_URI = 45 | import.meta.env.VITE_KG_CHAT_SAMPLE_QUESTIONS_ENDPOINT ?? 46 | "http://localhost:7860/questionProposalsForCurrentDb"; 47 | 48 | function loadKeyFromStorage() { 49 | return localStorage.getItem("api_key"); 50 | } 51 | 52 | const QUESTION_PREFIX_REGEXP = /^[0-9]{1,2}[\w]*[\.\)\-]*[\w]*/; 53 | 54 | function stripQuestionPrefix(question: string): string { 55 | if (question.match(QUESTION_PREFIX_REGEXP)) { 56 | return question.replace(QUESTION_PREFIX_REGEXP, ""); 57 | } 58 | return question; 59 | } 60 | 61 | function App() { 62 | const initialChatMessage: ChatMessageObject = { 63 | id: 0, 64 | type: 'text', 65 | sender: 'bot', 66 | message: 'Hello! How can I assist you today?', 67 | complete: true, 68 | }; 69 | 70 | const [serverAvailable, setServerAvailable] = useState(true); 71 | const [needsApiKeyLoading, setNeedsApiKeyLoading] = useState(true); 72 | const [needsApiKey, setNeedsApiKey] = useState(true); 73 | const [chatMessages, setChatMessages] = useState([initialChatMessage]); 74 | const [conversationState, setConversationState] = 75 | useState("ready"); 76 | const { sendJsonMessage, lastMessage, readyState } = useWebSocket(URI, { 77 | shouldReconnect: () => true, 78 | reconnectInterval: 5000, 79 | }); 80 | const [errorMessage, setErrorMessage] = useState(null); 81 | const [modalIsOpen, setModalIsOpen] = useState(false); 82 | const [apiKey, setApiKey] = useState(loadKeyFromStorage() || ""); 83 | const [sampleQuestions, setSampleQuestions] = useState([]); 84 | const [text2cypherModel, setText2cypherModel] = useState("gpt-3.5-turbo-0613"); 85 | 86 | const showContent = serverAvailable && !needsApiKeyLoading; 87 | 88 | function loadSampleQuestions() { 89 | const body = { 90 | api_key: apiKey, 91 | }; 92 | const options = { 93 | method: "POST", 94 | headers: { 95 | "Content-Type": "application/json", 96 | }, 97 | body: JSON.stringify(body), 98 | }; 99 | fetch(QUESTIONS_URI, options).then( 100 | (response) => { 101 | response.json().then( 102 | (result) => { 103 | if (result.output && result.output.length > 0) { 104 | setSampleQuestions(result.output.map(stripQuestionPrefix)); 105 | } else { 106 | setSampleQuestions([]); 107 | } 108 | }, 109 | (error) => { 110 | setSampleQuestions([]); 111 | } 112 | ); 113 | }, 114 | (error) => { 115 | setSampleQuestions([]); 116 | } 117 | ); 118 | } 119 | 120 | useEffect(() => { 121 | fetch(HAS_API_KEY_URI).then( 122 | (response) => { 123 | response.json().then( 124 | (result) => { 125 | // const needsKey = result.output; 126 | const needsKey = !result.output; 127 | setNeedsApiKey(needsKey); 128 | setNeedsApiKeyLoading(false); 129 | if (needsKey) { 130 | const api_key = loadKeyFromStorage(); 131 | if (api_key) { 132 | setApiKey(api_key); 133 | loadSampleQuestions(); 134 | } else { 135 | setModalIsOpen(true); 136 | } 137 | } else { 138 | loadSampleQuestions(); 139 | } 140 | }, 141 | (error) => { 142 | setNeedsApiKeyLoading(false); 143 | setServerAvailable(false); 144 | } 145 | ); 146 | }, 147 | (error) => { 148 | setNeedsApiKeyLoading(false); 149 | setServerAvailable(false); 150 | } 151 | ); 152 | }, []); 153 | 154 | useEffect(() => { 155 | if (!lastMessage || !serverAvailable) { 156 | return; 157 | } 158 | 159 | const websocketResponse = JSON.parse(lastMessage.data) as WebSocketResponse; 160 | 161 | if (websocketResponse.type === "debug") { 162 | console.log(websocketResponse.detail); 163 | } else if (websocketResponse.type === "error") { 164 | setConversationState("error"); 165 | setErrorMessage(websocketResponse.detail); 166 | console.error(websocketResponse.detail); 167 | } else if (websocketResponse.type === "start") { 168 | setConversationState("streaming"); 169 | 170 | setChatMessages((chatMessages) => [ 171 | ...chatMessages, 172 | { 173 | id: chatMessages.length, 174 | type: "text", 175 | sender: "bot", 176 | message: "", 177 | complete: false, 178 | }, 179 | ]); 180 | } else if (websocketResponse.type === "stream") { 181 | setChatMessages((chatMessages) => { 182 | const lastChatMessage = chatMessages[chatMessages.length - 1]; 183 | const rest = chatMessages.slice(0, -1); 184 | 185 | return [ 186 | ...rest, 187 | { 188 | ...lastChatMessage, 189 | message: lastChatMessage.message + websocketResponse.output, 190 | }, 191 | ]; 192 | }); 193 | } else if (websocketResponse.type === "end") { 194 | setChatMessages((chatMessages) => { 195 | const lastChatMessage = chatMessages[chatMessages.length - 1]; 196 | const rest = chatMessages.slice(0, -1); 197 | return [ 198 | ...rest, 199 | { 200 | ...lastChatMessage, 201 | complete: true, 202 | cypher: websocketResponse.generated_cypher, 203 | }, 204 | ]; 205 | }); 206 | setConversationState("ready"); 207 | } 208 | }, [lastMessage]); 209 | 210 | useEffect(() => { 211 | if (conversationState === "error") { 212 | const timeout = setTimeout(() => { 213 | setConversationState("ready"); 214 | }, 1000); 215 | return () => clearTimeout(timeout); 216 | } 217 | }, [conversationState]); 218 | 219 | const sendQuestion = (question: string) => { 220 | const webSocketRequest: WebSocketRequest = { 221 | type: "question", 222 | question: question, 223 | }; 224 | if (serverAvailable && !needsApiKeyLoading && needsApiKey && apiKey) { 225 | webSocketRequest.api_key = apiKey; 226 | } 227 | webSocketRequest.model_name = text2cypherModel; 228 | sendJsonMessage(webSocketRequest); 229 | }; 230 | 231 | const onChatInput = (message: string) => { 232 | if (conversationState === "ready") { 233 | setChatMessages((chatMessages) => 234 | chatMessages.concat([ 235 | { 236 | id: chatMessages.length, 237 | type: "input", 238 | sender: "self", 239 | message: message, 240 | complete: true, 241 | }, 242 | ]) 243 | ); 244 | if (SEND_REQUESTS) { 245 | setConversationState("waiting"); 246 | sendQuestion(message); 247 | } 248 | setErrorMessage(null); 249 | } 250 | }; 251 | 252 | const openModal = () => { 253 | setModalIsOpen(true); 254 | }; 255 | 256 | const onCloseModal = () => { 257 | setModalIsOpen(false); 258 | if (apiKey && sampleQuestions.length === 0) { 259 | loadSampleQuestions(); 260 | } 261 | }; 262 | 263 | const onApiKeyChange = (newApiKey: string) => { 264 | setApiKey(newApiKey); 265 | localStorage.setItem("api_key", newApiKey); 266 | }; 267 | 268 | const handleModelChange = (e: ChangeEvent) => { 269 | setText2cypherModel(e.target.value) 270 | } 271 | 272 | return ( 273 |
274 | {needsApiKey && ( 275 |
276 | 277 |
278 | )} 279 |
280 | 284 |
285 |
286 | {!serverAvailable && ( 287 |
Server is unavailable, please reload the page to try again.
288 | )} 289 | {serverAvailable && needsApiKeyLoading &&
Initializing...
} 290 | 296 | {showContent && readyState === ReadyState.OPEN && ( 297 | <> 298 | 302 | 307 | {errorMessage} 308 | 309 | )}{" "} 310 | {showContent && readyState === ReadyState.CONNECTING && ( 311 |
Connecting...
312 | )} 313 | {showContent && readyState === ReadyState.CLOSED && ( 314 |
315 |
Could not connect to server, reconnecting...
316 |
317 | )} 318 |
319 |
320 | ); 321 | } 322 | 323 | export default App; 324 | -------------------------------------------------------------------------------- /ui/src/chat-with-kg/ChatContainer.tsx: -------------------------------------------------------------------------------- 1 | import ChatMessage, { ChatMessageObject } from "./ChatMessage"; 2 | 3 | export type ChatContainerProps = { 4 | chatMessages?: ChatMessageObject[]; 5 | loading?: boolean; 6 | }; 7 | 8 | const loadingMessage: ChatMessageObject = { 9 | id: -1, 10 | type: "text", 11 | message: "Loading...", 12 | sender: "bot", 13 | complete: true, 14 | }; 15 | 16 | function ChatContainer(props: ChatContainerProps) { 17 | const { chatMessages = [], loading } = props; 18 | return ( 19 |
20 |
21 |
22 | {chatMessages.map((chatMessage) => ( 23 | 24 | ))} 25 | {loading && } 26 |
27 |
28 |
29 | ); 30 | } 31 | 32 | export default ChatContainer; 33 | -------------------------------------------------------------------------------- /ui/src/chat-with-kg/ChatInput.tsx: -------------------------------------------------------------------------------- 1 | import { useState, useEffect, KeyboardEvent } from "react"; 2 | export type ChatInputProps = { 3 | onChatInput?: (chatText: string) => void; 4 | loading?: boolean; 5 | sampleQuestions: string[] 6 | }; 7 | 8 | //Needed since the types for react don't include enterKeyHint 9 | declare module "react" { 10 | interface TextareaHTMLAttributes extends HTMLAttributes { 11 | enterKeyHint?: 12 | | "enter" 13 | | "done" 14 | | "go" 15 | | "next" 16 | | "previous" 17 | | "search" 18 | | "send"; 19 | } 20 | } 21 | 22 | function ChatInput(props: ChatInputProps) { 23 | const { onChatInput, loading, sampleQuestions } = props; 24 | const [inputText, setInputText] = useState(""); 25 | 26 | const [sampleQuestionIndex, setSampleQuestionIndex] = useState(0); 27 | 28 | useEffect(() => { 29 | setSampleQuestionIndex(0); 30 | }, [sampleQuestions]) 31 | 32 | const onInputKeyPress = (event: KeyboardEvent) => { 33 | if (!loading && event.key === "Enter") { 34 | handleSend(); 35 | } 36 | }; 37 | 38 | const handleSend = () => { 39 | if (!loading && inputText !== "" && onChatInput) { 40 | onChatInput(inputText); 41 | setInputText(""); 42 | } 43 | }; 44 | 45 | const sampleQuestionLeft = () => { 46 | if (sampleQuestionIndex > 0) { 47 | setSampleQuestionIndex(sampleQuestionIndex - 1); 48 | } 49 | }; 50 | 51 | const sampleQuestionRight = () => { 52 | if (sampleQuestionIndex < sampleQuestions.length - 1) { 53 | setSampleQuestionIndex(sampleQuestionIndex + 1); 54 | } 55 | }; 56 | 57 | const onSampleQuestionClick = () => { 58 | const sampleQuestion = sampleQuestions[sampleQuestionIndex]; 59 | if (onChatInput && sampleQuestion !== undefined) { 60 | onChatInput(sampleQuestion); 61 | } 62 | }; 63 | 64 | return ( 65 |
66 |
67 | {/* @ts-ignore */} 68 | 78 | 100 |
101 | {sampleQuestions && sampleQuestions.length > 0 && ( 102 | <> 103 |
Sample Questions
104 |
105 |
106 | 117 |
118 |
{sampleQuestions[sampleQuestionIndex]}
119 |
120 | 131 |
132 |
133 | 134 | )} 135 |
136 | ); 137 | } 138 | 139 | export default ChatInput; 140 | -------------------------------------------------------------------------------- /ui/src/chat-with-kg/ChatMessage.tsx: -------------------------------------------------------------------------------- 1 | import chatImg from "./images/chatbot.png"; 2 | import userImg from "./images/user.png" 3 | 4 | export type ChatMessageObject = { 5 | id: number; 6 | type: "input" | "text" | "error"; 7 | message: string; 8 | cypher?: string | null; 9 | sender: "bot" | "self"; 10 | complete: boolean; 11 | }; 12 | 13 | export type ChatMessageProps = { 14 | chatMessage: ChatMessageObject; 15 | }; 16 | 17 | function ChatMessage(props: ChatMessageProps) { 18 | const { chatMessage } = props; 19 | const { type, message, sender, cypher } = chatMessage; 20 | const chatClass = `flex flex-row relative max-w-full ${ 21 | sender === "bot" ? "self-start mr-10" : "ml-10 self-end" 22 | }`; 23 | 24 | return ( 25 |
26 | {sender === "bot" && Head} 27 |
30 | {message} 31 | {sender === "bot" && cypher && } 32 |
33 | {sender === "self" && Head} 34 |
35 | ); 36 | } 37 | 38 | function ChatMessageTail({ side }: { side: "left" | "right" }) { 39 | const chatTailStyle: React.CSSProperties = { 40 | width: "0.75rem", 41 | height: "0.75rem", 42 | WebkitMaskImage: `url(data:image/svg+xml;base64,PHN2ZyB3aWR0aD0nMycgaGVpZ2h0PSczJyB4bWxucz0naHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmcnPjxwYXRoIGZpbGw9J2JsYWNrJyBkPSdtIDAgMyBMIDMgMyBMIDMgMCBDIDMgMSAxIDMgMCAzJy8+PC9zdmc+)`, 43 | maskImage: `url(data:image/svg+xml;base64,PHN2ZyB3aWR0aD0nMycgaGVpZ2h0PSczJyB4bWxucz0naHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmcnPjxwYXRoIGZpbGw9J2JsYWNrJyBkPSdtIDAgMyBMIDMgMyBMIDMgMCBDIDMgMSAxIDMgMCAzJy8+PC9zdmc+)`, 44 | WebkitMaskPosition: "center", 45 | maskPosition: "center", 46 | maskSize: "contain", 47 | WebkitMaskSize: "contain", 48 | }; 49 | 50 | if (side === "left") { 51 | chatTailStyle["left"] = "-0.75rem"; 52 | } else { 53 | chatTailStyle["right"] = "-0.75rem"; 54 | chatTailStyle["WebkitTransform"] = "scaleX(-1)"; 55 | } 56 | 57 | return ( 58 |
62 | ); 63 | } 64 | 65 | function ChatCypherDetail({ cypher }: { cypher: string }) { 66 | return ( 67 |
68 | Cypher 69 |
70 | {cypher} 71 |
72 |
73 | ); 74 | } 75 | 76 | export default ChatMessage; 77 | -------------------------------------------------------------------------------- /ui/src/chat-with-kg/images/chatbot.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ryang420/Knowledge-Graph-LLM/4083ba6278c24597aca18a98b4f8ce248f8bfb65/ui/src/chat-with-kg/images/chatbot.png -------------------------------------------------------------------------------- /ui/src/chat-with-kg/images/user.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ryang420/Knowledge-Graph-LLM/4083ba6278c24597aca18a98b4f8ce248f8bfb65/ui/src/chat-with-kg/images/user.png -------------------------------------------------------------------------------- /ui/src/chat-with-kg/index.css: -------------------------------------------------------------------------------- 1 | @tailwind base; 2 | @tailwind components; 3 | @tailwind utilities; 4 | 5 | :root { 6 | font-family: Inter, system-ui, Avenir, Helvetica, Arial, sans-serif; 7 | line-height: 1.5; 8 | font-weight: 400; 9 | 10 | color-scheme: light dark; 11 | 12 | font-synthesis: none; 13 | text-rendering: optimizeLegibility; 14 | -webkit-font-smoothing: antialiased; 15 | -moz-osx-font-smoothing: grayscale; 16 | -webkit-text-size-adjust: 100%; 17 | } -------------------------------------------------------------------------------- /ui/src/chat-with-kg/main.tsx: -------------------------------------------------------------------------------- 1 | import React from "react"; 2 | import { createRoot } from "react-dom/client"; 3 | import App from "./App.js"; 4 | import Modal from "react-modal"; 5 | 6 | import "@neo4j-ndl/base/lib/neo4j-ds-styles.css"; 7 | import "./index.css"; 8 | 9 | const container = document.getElementById("root"); 10 | const root = createRoot(container!); 11 | 12 | if (container) { 13 | Modal.setAppElement(container); 14 | } 15 | 16 | root.render( 17 | 18 | 19 | 20 | ); 21 | -------------------------------------------------------------------------------- /ui/src/chat-with-kg/types/websocketTypes.ts: -------------------------------------------------------------------------------- 1 | export type WebSocketRequest = { 2 | type: "question"; 3 | question: string; 4 | api_key?: string; 5 | model_name?: string; 6 | }; 7 | 8 | export type WebSocketResponse = 9 | | { type: "start" } 10 | | { 11 | type: "stream"; 12 | output: string; 13 | } 14 | | { 15 | type: "end"; 16 | output: string; 17 | generated_cypher: string | null; 18 | } 19 | | { 20 | type: "error"; 21 | detail: string; 22 | } 23 | | { 24 | type: "debug"; 25 | detail: string; 26 | }; 27 | 28 | export type ConversationState = "waiting" | "streaming" | "ready" | "error"; 29 | -------------------------------------------------------------------------------- /ui/src/components/graph-data-modal.tsx: -------------------------------------------------------------------------------- 1 | import Modal from "react-modal"; 2 | import "./neo-graph-2d.css"; 3 | import { Node, Link } from "./neo-graph"; 4 | import { NodeType } from "../unstructured-import/types/respons-types"; 5 | 6 | interface GraphModalProps { 7 | selectedItem: Node | Link | null; 8 | setSelectedItem: (item: Node | Link | null) => void; 9 | type: "node" | "link" | null; 10 | } 11 | 12 | export const GraphModal: React.FC = ({ selectedItem, setSelectedItem, type }) => { 13 | const closeModal = () => { 14 | setSelectedItem(null); 15 | }; 16 | 17 | const renderProperties = (properties: { [key: string]: any }) => { 18 | return Object.entries(properties).map(([key, value]) => ( 19 |

20 | {key}: {value} 21 |

22 | )); 23 | }; 24 | 25 | return ( 26 | 34 | {selectedItem && ( 35 |
36 |
37 | X 38 |
39 | {type == "node" ? ( 40 |
41 |

Node Details

42 |

Id: {(selectedItem as Node).id}

43 |

Type: {(selectedItem as Node).type}

44 | {renderProperties(selectedItem.properties)} 45 |
46 | ) : type == "link" ? ( 47 |
48 |

Link Details

49 |

Source: {(selectedItem as Link).source.id}

50 |

Target: {(selectedItem as Link).target.id}

51 | {renderProperties(selectedItem.properties)} 52 |
53 | ) : null} 54 |
55 | )} 56 | 57 | ); 58 | }; -------------------------------------------------------------------------------- /ui/src/components/keymodal.tsx: -------------------------------------------------------------------------------- 1 | import Modal from "react-modal"; 2 | 3 | type KeyModalProps = { 4 | isOpen: boolean; 5 | apiKey: string; 6 | onRequestClose?: () => void; 7 | onCloseModal: () => void; 8 | onApiKeyChanged: (key: string) => void; 9 | }; 10 | 11 | const customModalStyles = { 12 | content: { 13 | top: '50%', 14 | left: '50%', 15 | right: 'auto', 16 | bottom: 'auto', 17 | marginRight: '-50%', 18 | transform: 'translate(-50%, -50%)', 19 | }, 20 | }; 21 | const inputClassNames = ""; 22 | 23 | 24 | function KeyModal(props: KeyModalProps) { 25 | const { isOpen, apiKey, onRequestClose, onCloseModal, onApiKeyChanged } = props; 26 | return ( 27 | 34 | 41 |
42 |
43 |

API Key

44 |
Please enter your OpenAI API Key
45 |
46 |
47 |
48 |
49 | {/* 50 | // @ts-ignore */} 51 | { onApiKeyChanged(e.target.value); }} className={inputClassNames} value={apiKey} /> 52 |
53 |
54 | 55 | ); 56 | } 57 | 58 | export default KeyModal; 59 | -------------------------------------------------------------------------------- /ui/src/components/neo-graph-2d.css: -------------------------------------------------------------------------------- 1 | .custom-modal { 2 | position: fixed; 3 | top: 50%; 4 | left: 50%; 5 | transform: translate(-50%, -50%); 6 | background-color: white; 7 | padding: 20px; 8 | max-width: 500px; /* Set your desired max-width */ 9 | width: 100%; 10 | border: 1px solid #ccc; 11 | border-radius: 5px; 12 | outline: none; 13 | z-index: 1000; 14 | } 15 | 16 | .custom-modal h2 { 17 | font-weight: bold; 18 | } 19 | 20 | .custom-overlay { 21 | position: fixed; 22 | top: 0; 23 | left: 0; 24 | right: 0; 25 | bottom: 0; 26 | background-color: rgba(0, 0, 0, 0.5); 27 | z-index: 999; /* Ensure the overlay is on top of other elements */ 28 | } 29 | 30 | .close-button { 31 | position: absolute; 32 | top: 10px; 33 | right: 10px; 34 | cursor: pointer; 35 | font-size: 18px; 36 | color: #333; 37 | } 38 | -------------------------------------------------------------------------------- /ui/src/components/neo-graph-2d.tsx: -------------------------------------------------------------------------------- 1 | import ForceGraph2D from "react-force-graph-2d"; 2 | import { useRef, useState } from "react"; 3 | import "./neo-graph-2d.css"; 4 | import { GraphModal } from "./graph-data-modal"; 5 | import { GraphData, Node, Link } from "./neo-graph"; 6 | import { ImportResult } from "../unstructured-import/types/respons-types"; 7 | 8 | // define a function to transform raw data into GraphData format 9 | const transform_raw_graph_data = (raw_data: any) => { 10 | const links = raw_data["relationships"] 11 | ? raw_data["relationships"].map((link: any) => { 12 | return { 13 | source: link["source"]["id"], 14 | target: link["target"]["id"], 15 | type: link["type"], 16 | properties: link["properties"], 17 | }; 18 | }) 19 | : []; 20 | 21 | return { 22 | nodes: raw_data["nodes"], 23 | links: links, 24 | }; 25 | }; 26 | 27 | export const NeoGraph2D = ({ graph_raw_data }: { graph_raw_data: ImportResult }) => { 28 | const fgRef = useRef(); 29 | const [selectedNode, setSelectedNode] = useState(null); 30 | const [selectedLink, setSelectedLink] = useState(null); 31 | const [selectedType, setSelectedType] = useState<"node" | "link" | null>(null); 32 | 33 | const handleNodeClick = (node: Node) => { 34 | setSelectedNode(node); 35 | setSelectedType("node"); 36 | }; 37 | 38 | const handleLinkClick = (link: Link) => { 39 | setSelectedLink(link); 40 | setSelectedType("link"); 41 | }; 42 | 43 | const transformedData = transform_raw_graph_data(graph_raw_data); 44 | console.log("updated graph data: ", transformedData); 45 | 46 | return ( 47 | <> 48 | "after"} 59 | nodeCanvasObject={(node: Node, ctx: any, globalScale) => { 60 | const label = node.id; 61 | const fontSize = 12 / globalScale; 62 | ctx.font = `${fontSize}px Sans-Serif`; 63 | ctx.textAlign = "center"; 64 | ctx.fillStyle = "black"; 65 | ctx.fillText(label, node.x ? node.x : 0, node.y ? node.y + 1 : 0); 66 | }} 67 | linkCanvasObjectMode={() => "after"} 68 | linkCanvasObject={(link: Link, ctx: any, globalScale) => { 69 | const label = link.type; 70 | const fontSize = 10 / globalScale; 71 | 72 | // Calculate the middle point of the link 73 | if (link.source.x !== undefined 74 | && link.target.x !== undefined 75 | && link.source.y !== undefined 76 | && link.target.y !== undefined) { 77 | const middleX = (link.source.x + link.target.x) / 2; 78 | const middleY = (link.source.y + link.target.y) / 2; 79 | 80 | // Draw link label 81 | ctx.font = `${fontSize}px Sans-Serif`; 82 | ctx.fillStyle = "black"; 83 | ctx.fillText(label, middleX, middleY); 84 | } 85 | }} 86 | onNodeDragEnd={(node) => { 87 | node.fx = node.x; 88 | node.fy = node.y; 89 | }} 90 | onNodeClick={handleNodeClick} 91 | onLinkClick={handleLinkClick} 92 | /> 93 | 94 | { 97 | setSelectedNode(null); 98 | setSelectedLink(null); 99 | }} 100 | type={selectedType} 101 | /> 102 | 103 | ); 104 | }; -------------------------------------------------------------------------------- /ui/src/components/neo-graph.ts: -------------------------------------------------------------------------------- 1 | /** 2 | * A node or relationship as selected in the graph. 3 | */ 4 | export interface GraphEntity { 5 | properties: any; 6 | type: string; 7 | } 8 | 9 | export interface Node extends GraphEntity { 10 | id: string; 11 | x?: number; 12 | y?: number; 13 | fx?: number; 14 | fy?: number; 15 | } 16 | 17 | export interface Link extends GraphEntity { 18 | source: Node; 19 | target: Node; 20 | } 21 | 22 | export interface GraphData { 23 | nodes: Node[]; 24 | links: Link[]; 25 | } -------------------------------------------------------------------------------- /ui/src/components/switch.tsx: -------------------------------------------------------------------------------- 1 | type SwitchProps = { 2 | label: string; 3 | checked: boolean; 4 | onChange: () => void; 5 | }; 6 | 7 | export function Switch({ label, checked, onChange }: SwitchProps) { 8 | return ( 9 |
10 |
11 | 22 |
23 |
24 | ); 25 | } 26 | -------------------------------------------------------------------------------- /ui/src/unstructured-import/App.css: -------------------------------------------------------------------------------- 1 | .grey-background { 2 | background-color: grey; 3 | margin-top: 1vw; 4 | border-radius: 3px; 5 | } 6 | 7 | .save-modal-content { 8 | position: fixed; 9 | top: 50%; 10 | left: 50%; 11 | transform: translate(-50%, -50%); 12 | background-color: white; 13 | padding: 20px; 14 | max-width: 500px; /* Set your desired max-width */ 15 | width: 100%; 16 | border: 1px solid #ccc; 17 | border-radius: 5px; 18 | outline: none; 19 | z-index: 1000; 20 | } 21 | 22 | .close-button { 23 | position: absolute; /* Position the button relative to the modal */ 24 | top: 10px; /* Position from the top */ 25 | right: 10px; /* Position from the right */ 26 | background: transparent; /* No background */ 27 | border: none; /* No border */ 28 | font-size: 1.5em; /* Make the "X" bigger */ 29 | cursor: pointer; /* Change cursor to pointer on hover */ 30 | } 31 | 32 | .text-height { 33 | height: 50px; 34 | } 35 | 36 | .overlay { 37 | position: fixed; 38 | top: 0; 39 | left: 0; 40 | width: 100%; 41 | height: 100%; 42 | background-color: rgba(0, 0, 0, 0.5); 43 | display: flex; 44 | justify-content: center; 45 | align-items: center; 46 | color: white; 47 | font-size: 2em; 48 | } -------------------------------------------------------------------------------- /ui/src/unstructured-import/App.tsx: -------------------------------------------------------------------------------- 1 | import "@neo4j-ndl/base/lib/neo4j-ds-styles.css"; 2 | import { useState, useEffect } from "react"; 3 | import { runImport } from "./utils/fetch-utils"; 4 | import { Switch } from "../components/switch"; 5 | import { graphSchemaToModelSchema } from "./utils/graph-schema-utils"; 6 | import KeyModal from "../components/keymodal"; 7 | import { ImportResult } from "./types/respons-types"; 8 | import { NeoGraph2D } from "../components/neo-graph-2d"; 9 | import "./App.css"; 10 | import Modal from "react-modal"; 11 | 12 | 13 | const HAS_API_KEY_URI = 14 | import.meta.env.VITE_HAS_API_KEY_ENDPOINT ?? 15 | "http://localhost:7860/hasapikey"; 16 | 17 | function loadKeyFromStorage() { 18 | return localStorage.getItem("api_key"); 19 | } 20 | 21 | function App() { 22 | const [serverAvailable, setServerAvailable] = useState(true); 23 | const [needsApiKeyLoading, setNeedsApiKeyLoading] = useState(true); 24 | const [needsApiKey, setNeedsApiKey] = useState(true); 25 | const [useSchema, setUseSchema] = useState(false); 26 | const [loading, setLoading] = useState(false); 27 | const [result, setResult] = useState(null); 28 | const [schema, setSchema] = useState(""); 29 | const [modalIsOpen, setModalIsOpen] = useState(false); 30 | const [apiKey, setApiKey] = useState(loadKeyFromStorage() || ""); 31 | 32 | const initDone = serverAvailable && !needsApiKeyLoading; 33 | 34 | const [saveModalIsOpen, setSaveModalIsOpen] = useState(false); 35 | const [isSaving, setIsSaving] = useState(false); 36 | const [updateGraphInProgress, setUpdateGraphInProgress] = useState(false); 37 | 38 | 39 | useEffect(() => { 40 | fetch(HAS_API_KEY_URI).then( 41 | (response) => { 42 | response.json().then( 43 | (result) => { 44 | const needsKey = !result.output; 45 | setNeedsApiKey(needsKey); 46 | setNeedsApiKeyLoading(false); 47 | if (needsKey) { 48 | const api_key = loadKeyFromStorage(); 49 | if (api_key) { 50 | setApiKey(api_key); 51 | } else { 52 | setModalIsOpen(true); 53 | } 54 | } 55 | }, 56 | (error) => { 57 | setNeedsApiKeyLoading(false); 58 | setServerAvailable(false); 59 | }, 60 | ); 61 | }, 62 | (error) => { 63 | setNeedsApiKeyLoading(false); 64 | setServerAvailable(false); 65 | }, 66 | ); 67 | }, []); 68 | 69 | const openModal = () => { 70 | setModalIsOpen(true); 71 | }; 72 | 73 | const onCloseModal = () => { 74 | setModalIsOpen(false); 75 | }; 76 | 77 | const onApiKeyChange = (newApiKey: string) => { 78 | setApiKey(newApiKey); 79 | localStorage.setItem("api_key", newApiKey); 80 | }; 81 | 82 | const handleImport = async () => { 83 | if (!serverAvailable || needsApiKeyLoading) { 84 | return; 85 | } 86 | setLoading(true); 87 | setResult(null); 88 | const file = document.querySelector(".file-input") as HTMLInputElement; 89 | const reader = new FileReader(); 90 | reader.onload = async () => { 91 | console.log(reader.result); 92 | try { 93 | console.log("running import"); 94 | console.log("raw schema", schema); 95 | const schemaJson = useSchema 96 | ? graphSchemaToModelSchema(schema) 97 | : undefined; 98 | console.log("schema json", schemaJson); 99 | const importResult = await runImport( 100 | reader.result as string, 101 | schemaJson, 102 | needsApiKey ? apiKey : undefined, 103 | ); 104 | console.log("import result", importResult); 105 | 106 | if (importResult) { 107 | console.log(importResult); 108 | setResult(importResult); 109 | } 110 | } catch (e) { 111 | console.error(e); 112 | } finally { 113 | setLoading(false); 114 | } 115 | }; 116 | reader.readAsText(file.files![0]); 117 | }; 118 | 119 | const handleSaveToNeo4j = async () => { 120 | setIsSaving(true); 121 | // sleep to make sure the modal is shown 122 | await new Promise((resolve) => setTimeout(resolve, 5000)); 123 | const response = await fetch( 124 | `${import.meta.env.VITE_UNSTRUCTURED_IMPORT_BACKEND_ENDPOINT}/save_graph`, 125 | { 126 | method: "POST", 127 | headers: { 128 | "Content-Type": "application/json", 129 | }, 130 | body: JSON.stringify(result), 131 | }, 132 | ); 133 | 134 | const saveResult = await response.json(); 135 | console.log("saved result: ", saveResult); 136 | // Show the modal when the save is done 137 | setSaveModalIsOpen(true); 138 | setIsSaving(false); 139 | }; 140 | 141 | const handleChangeGraphCommand = async (e: any, apiKey?: string) => { 142 | const body = { 143 | user_input: e.target.value, 144 | graph_data: JSON.stringify(result), 145 | }; 146 | if (apiKey) { 147 | // @ts-ignore 148 | body.api_key = apiKey; 149 | } 150 | if (e.key == "Enter") { 151 | setUpdateGraphInProgress(true); 152 | const response = await fetch( 153 | `${import.meta.env.VITE_UNSTRUCTURED_IMPORT_BACKEND_ENDPOINT}/update_graph`, 154 | { 155 | method: "POST", 156 | headers: { 157 | "Content-Type": "application/json", 158 | }, 159 | body: JSON.stringify(body), 160 | }, 161 | ); 162 | if (!response.ok) { 163 | setUpdateGraphInProgress(false); 164 | return Promise.reject( 165 | new Error(`Failed to update: ${response.statusText}`), 166 | ); 167 | } 168 | 169 | const res = await response.json(); 170 | 171 | if (res.status == "ok") { 172 | let graph_data = res.data; 173 | // if "source" not in graph_data add it 174 | if (!graph_data["source"]) { 175 | graph_data["source"] = { 176 | "page_content": "", 177 | "metadata": {}, 178 | "type": "Document", 179 | }; 180 | } 181 | 182 | console.log("updated graph data from llm: ", graph_data); 183 | setResult(graph_data); 184 | } else { 185 | console.log("failed to update graph data from llm: "); 186 | } 187 | 188 | setUpdateGraphInProgress(false); 189 | } 190 | }; 191 | 192 | if (serverAvailable) { 193 | return ( 194 |
195 | {needsApiKey && ( 196 |
197 | 198 |
199 | )} 200 | 206 |
207 |
208 |

Import data

209 |

210 | This tool is used to import unstructured data into Neo4j. It takes 211 | a file as input and optionally a schema in the form of{" "} 212 | 213 | graph data model 214 | {" "} 215 | which is used to limit the data that is extracted from the file. 216 | It's important to give the schema descriptive tokens so the tool 217 | can identify the data that is imported. 218 |

219 | 220 | setUseSchema(!useSchema)} 224 | /> 225 | {useSchema ? ( 226 |
227 | {"Please provide your schema in json format:"} 228 |