1 line
43 KiB
JSON
1 line
43 KiB
JSON
{"repo": "caicongyang/ai-agent-demo", "files": {"graphrag/graphrag_neo4j_demo.py": {"bc378182259183e6a980bffb9bc815de1c39b907": {"t": 1785746485.0437195, "n": 19446, "h": "f08b730be8477f9f", "c": "\"\"\"\nGraphRAG with Neo4j Implementation\n\nThis module implements a GraphRAG (Graph Retrieval Augmented Generation) system\nusing Neo4j as the graph database backend.\n\nWorkflow:\n1. Load documents\n2. Convert documents to text nodes\n3. Extract entities and relationships from text using LLM\n4. Store the extracted knowledge graph in Neo4j\n5. Query the knowledge graph using natural language\n\nBased on LlamaIndex's GraphRAG approach:\nhttps://docs.llamaindex.ai/en/stable/examples/cookbooks/GraphRAG_v1/\nhttps://docs.llamaindex.ai/en/stable/examples/cookbooks/GraphRAG_v2/\n\"\"\"\n\nimport os\nimport pandas as pd\nfrom dotenv import load_dotenv\nfrom llama_index.core import PropertyGraphIndex\nfrom llama_index.core import Document\nfrom llama_index.graph_stores.neo4j import Neo4jPropertyGraphStore\nimport openai\nfrom neo4j import GraphDatabase\n\n# 自定义 LLM 实现\nclass CustomLLM:\n \"\"\"Custom LLM implementation for DeepSeek model\"\"\"\n \n def __init__(self, api_key, api_base, model):\n self.api_key = api_key\n # 确保 api_base 格式正确\n if not api_base.endswith(\"/v1\"):\n if api_base.endswith(\"/\"):\n api_base = api_base + \"v1\"\n else:\n api_base = api_base + \"/v1\"\n self.api_base = api_base\n self.model = model\n # Initialize OpenAI client\n self.client = openai.OpenAI(\n api_key=self.api_key,\n base_url=self.api_base\n )\n print(f\"Initialized CustomLLM with base_url: {self.api_base}\")\n\n def complete(self, prompt, **kwargs):\n try:\n response = self.client.chat.completions.create(\n model=self.model,\n messages=[{\"role\": \"user\", \"content\": prompt}],\n temperature=0.7,\n )\n class Response:\n def __init__(self, text):\n self.text = text\n return Response(response.choices[0].message.content)\n except Exception as e:\n print(f\"API 调用错误: {e}\")\n # 返回一个空响应而不是抛出异常，以便演示能继续\n class Response:\n def __init__(self, text):\n self.text = f\"{{\\\"entities\\\": [], \\\"relationships\\\": []}}\"\n return Response()\n\n# Neo4j 版本的 GraphRAG 抽取器\nclass GraphRAGExtractor:\n \"\"\"\n GraphRAG 抽取器：使用 LLM 从文本中抽取三元组（实体-关系-实体），适配 Neo4j。\n \"\"\"\n def __init__(self, llm):\n self.llm = llm\n \n async def acall(self, *args, **kwargs):\n \"\"\"\n 异步调用，对于我们的简单实现，直接调用同步方法。\n 对于实际生产环境，应该实现真正的异步处理。\n \"\"\"\n return self.extract(*args, **kwargs)\n \n def extract(self, nodes, show_progress=False, **kwargs):\n \"\"\"\n 从节点列表中抽取实体和关系。\n 返回 EntityNode 和 Relation 的列表。\n \n 参数:\n nodes: 要处理的节点列表\n show_progress: 是否显示进度，兼容PropertyGraphIndex的调用\n **kwargs: 其他可能的参数\n \"\"\"\n all_entities = []\n all_relationships = []\n \n # 输出调试信息\n print(f\"Extract method called with {len(nodes)} nodes\")\n print(f\"Node types: {[type(node).__name__ for node in nodes]}\")\n \n # 检查nodes是否是列表\n if not isinstance(nodes, list):\n print(f\"WARNING: nodes is not a list but a {type(nodes)}\")\n # 尝试转换为列表\n nodes = list(nodes)\n \n for i, node in enumerate(nodes):\n print(f\"Processing node {i} of type {type(node).__name__}\")\n # 确保我们可以访问节点的node_id和text\n try:\n node_id = getattr(node, \"node_id\", None) or str(i)\n node_text = getattr(node, \"text\", None)\n if node_text is None and hasattr(node, \"get_content\"):\n node_text = node.get_content()\n \n if node_text is None:\n print(f\"WARNING: Node {i} has no text attribute\")\n # 尝试其他可能的属性\n for attr in [\"content\", \"get_text\", \"__str__\"]:\n if hasattr(node, attr):\n if callable(getattr(node, attr)):\n node_text = getattr(node, attr)()\n else:\n node_text = getattr(node, attr)\n print(f\"Found text in {attr}: {node_text[:50]}...\")\n break\n \n if node_text is None:\n print(f\"SKIPPING node {i} as no text could be found\")\n continue\n \n prompt = f\"\"\"\n Extract key entities and their relationships from the following text.\n Format your response as a JSON with 'entities' and 'relationships' keys.\n \n For each entity, include:\n - 'entity_id': A unique identifier for the entity\n - 'entity_type': The type of entity (e.g., Person, Organization, Product)\n - 'entity_description': A brief description of the entity\n \n For each relationship, include:\n - 'source': The source entity ID\n - 'target': The target entity ID\n - 'relation': The type of relation\n - 'relationship_description': A description of the relationship\n \n YOUR RESPONSE MUST BE VALID JSON. Make sure to return only a JSON object with no additional text or explanation.\n \n TEXT: {node_text}\n \"\"\"\n \n response = self.llm.complete(prompt)\n \n # 实际应用中应解析 LLM 返回的 JSON 并结构化数据\n try:\n import json\n import re\n \n # Try to extract JSON from the response if there's other text\n response_text = response.text.strip()\n json_match = re.search(r'\\{.*\\}', response_text, re.DOTALL)\n if json_match:\n json_str = json_match.group(0)\n else:\n json_str = response_text\n \n # Parse the JSON\n data = json.loads(json_str)\n \n from llama_index.core.graph_stores.types import EntityNode, Relation\n \n # 构建实体节点\n entities = [\n EntityNode(\n label=entity.get(\"entity_type\", \"Entity\"),\n properties={\n \"id\": entity.get(\"entity_id\", \"\"),\n \"entity_description\": entity.get(\"entity_description\", \"\"),\n \"triplet_source_id\": node_id\n },\n name=entity.get(\"entity_id\", \"\")\n )\n for entity in data.get(\"entities\", [])\n ]\n \n # 构建关系\n relationships = [\n Relation(\n label=relation.get(\"relation\", \"Related\"),\n source_id=relation.get(\"source\", \"\"),\n target_id=relation.get(\"target\", \"\"),\n properties={\n \"relationship_description\": relation.get(\"relationship_description\", \"\"),\n \"triplet_source_id\": node_id\n }\n )\n for relation in data.get(\"relationships\", [])\n ]\n \n all_entities.extend(entities)\n all_relationships.extend(relationships)\n \n except Exception as e:\n print(f\"Error parsing response for node {i}: {e}\")\n print(f\"Response text: {response.text[:100]}...\")\n \n except Exception as e:\n print(f\"Error processing node {i}: {e}\")\n \n print(f\"Extraction completed: {len(all_entities)} entities, {len(all_relationships)} relationships\")\n return all_entities, all_relationships\n\n\nclass GraphRAGStore:\n \"\"\"\n Neo4j 图数据库存储的包装类，提供 GraphRAG 相关接口。\n \"\"\"\n def __init__(self, username, password, url):\n # 初始化 Neo4j 图存储\n self.graph_store = Neo4jPropertyGraphStore(\n username=username,\n password=password,\n url=url\n )\n \n def get_triplets(self):\n \"\"\"\n 获取所有三元组（关系映射）。\n \"\"\"\n return self.graph_store.get_rel_map().items()\n \n def build_communities(self):\n \"\"\"\n 社区发现（本 demo 仅打印提示，未做实际聚类）。\n \"\"\"\n print(\"Building communities... (Demo version - no actual clustering)\")\n # 实际应用中可用社区发现算法并在 Neo4j 中创建社区节点\n\n\nclass GraphRAGQueryEngine:\n \"\"\"\n 基于 Neo4j 图存储的 GraphRAG 查询引擎。\n 使用 Cypher 直接查询 Neo4j。\n \"\"\"\n def __init__(self, graph_store, llm, index=None, similarity_top_k=5):\n self.graph_store = graph_store\n self.llm = llm\n self.index = index\n self.similarity_top_k = similarity_top_k\n \n # Neo4j连接信息\n self.uri = os.getenv(\"NEO4J_URL\")\n self.username = os.getenv(\"NEO4J_USERNAME\")\n self.password = os.getenv(\"NEO4J_PASSWORD\")\n \n def query(self, query_text):\n \"\"\"\n 使用 Cypher 查询图谱，拼接上下文，调用 LLM 进行问答。\n \"\"\"\n # 获取图谱数据作为上下文\n context = self._get_context_from_neo4j()\n \n # 拼接最终 prompt，交给 LLM 回答\n prompt = f\"\"\"\n Use the following knowledge graph information to answer the user's question.\n If the information is not in the knowledge graph, say that you don't know.\n \n {context}\n \n User Question: {query_text}\n \n Answer:\n \"\"\"\n \n response = self.llm.complete(prompt)\n \n class Response:\n def __init__(self, text):\n self.response = text\n \n return Response(response.text)\n \n def _get_context_from_neo4j(self):\n \"\"\"\n 从 Neo4j 获取图谱数据作为上下文。\n \"\"\"\n try:\n # 打开一个Neo4j会话\n driver = GraphDatabase.driver(self.uri, auth=(self.username, self.password))\n \n context = \"Knowledge Graph Information:\\n\\n\"\n \n with driver.session() as session:\n # 获取所有实体\n entities_result = session.run(\"\"\"\n MATCH (entity) \n RETURN entity.name as name, labels(entity) as labels, entity.description as description\n LIMIT 20\n \"\"\")\n \n for record in entities_result:\n entity_name = record[\"name\"]\n entity_label = record[\"labels\"][0] if record[\"labels\"] else \"Entity\"\n entity_desc = record[\"description\"] or \"\"\n context += f\"- {entity_name} ({entity_label}): {entity_desc}\\n\"\n \n # 获取所有关系\n relations_result = session.run(\"\"\"\n MATCH (source)-[rel]->(target)\n RETURN source.name as source, target.name as target, type(rel) as type, rel.description as description\n LIMIT 20\n \"\"\")\n \n for record in relations_result:\n source = record[\"source\"]\n target = record[\"target\"]\n rel_type = record[\"type\"]\n rel_desc = record[\"description\"] or \"\"\n context += f\"- {source} --{rel_type}--> {target}: {rel_desc}\\n\"\n \n driver.close()\n return context\n \n except Exception as e:\n print(f\"Error querying Neo4j: {e}\")\n return \"Unable to retrieve knowledge graph information due to an error.\"\n\n\n# 主演示函数\ndef run_neo4j_demo():\n \"\"\"\n 运行Neo4j GraphRAG演示\n \n 流程:\n 1. 加载环境变量和初始化LLM\n 2. 创建示例文档\n 3. 从文档中提取文本节点\n 4. 使用LLM从文本中提取实体和关系\n 5. 将实体和关系存储到Neo4j中\n 6. 使用自然语言查询知识图谱\n \"\"\"\n # 加载 .env 文件中的环境变量\n load_dotenv()\n \n # 检查 LLM API Key 是否存在 (使用与其他模块一致的变量名)\n api_key = os.getenv(\"LLM_API_KEY\")\n api_base = os.getenv(\"LLM_BASE_URL\", \"https://api.deepseek.com/v1\")\n if not api_key:\n print(\"ERROR: LLM_API_KEY not found in environment variables or .env file\")\n print(\"Please set your LLM API key before running this demo.\")\n return\n \n # 检查 Neo4j 相关配置 \n neo4j_user = os.getenv(\"NEO4J_USERNAME\") \n neo4j_password = os.getenv(\"NEO4J_PASSWORD\") \n neo4j_url = os.getenv(\"NEO4J_URL\")\n \n print(f\"Using API Key: {api_key[:5]}...\")\n print(f\"Using API Base: {api_base}\")\n print(f\"Using Neo4j config: {neo4j_user}@{neo4j_url}\")\n \n # 初始化自定义 LLM\n llm = CustomLLM(\n api_key=api_key,\n api_base=api_base,\n model=\"deepseek-chat\"\n )\n \n # 示例文档\n documents = [\n Document(text=\"Apple Inc. is a technology company founded by Steve Jobs, Steve Wozniak, and Ronald Wayne in 1976. \"\n \"The company is known for its iPhone products and is headquartered in Cupertino, California.\"),\n Document(text=\"Microsoft Corporation was founded by Bill Gates and Paul Allen in 1975. \"\n \"The company develops Windows operating system and Office software suite.\"),\n Document(text=\"Amazon was founded by Jeff Bezos in 1994 as an online bookstore. \"\n \"It has since expanded to sell various products and services including cloud computing through AWS.\")\n ]\n \n try:\n print(\"Step 1: Initializing Neo4j GraphRAG components...\")\n extractor = GraphRAGExtractor(llm)\n \n # 不需要再次验证URL，因为已在前面处理\n try:\n graph_store = GraphRAGStore(\n username=neo4j_user,\n password=neo4j_password,\n url=neo4j_url\n )\n print(\"Neo4j连接成功!\")\n except Exception as neo4j_error:\n print(f\"Neo4j连接错误: {neo4j_error}\")\n print(\"继续尝试处理，但可能会在后续步骤失败...\")\n raise # 由于图存储是必须的，这里还是需要终止执行\n \n print(\"\\nStep 2: Creating nodes from documents...\")\n # 按照 GraphRAG v1 文档，先创建节点/chunks\n from llama_index.core.node_parser import SentenceSplitter\n \n # 创建文本分割器\n splitter = SentenceSplitter(chunk_size=1024, chunk_overlap=20)\n \n # 从文档创建节点\n nodes = []\n for doc in documents:\n # 确保文档有metadata\n if not hasattr(doc, \"metadata\") or doc.metadata is None:\n doc.metadata = {}\n \n # 分割文档成节点\n doc_nodes = splitter.get_nodes_from_documents([doc])\n nodes.extend(doc_nodes)\n \n print(f\"Created {len(nodes)} nodes from {len(documents)} documents\")\n \n print(\"\\nStep 3: Extracting entities and relations directly...\")\n try:\n # 直接使用抽取器提取三元组\n entities, relations = extractor.extract(nodes)\n print(f\"Extracted {len(entities)} entities and {len(relations)} relationships\")\n \n # 直接加载到Neo4j\n print(\"Loading entities and relations into Neo4j...\")\n \n # 获取Neo4j图存储\n neo4j_store = graph_store.graph_store\n \n # 使用更简单的方法 - 使用Neo4j客户端直接执行Cypher\n \n # 获取Neo4j连接信息\n uri = neo4j_url\n username = neo4j_user\n password = neo4j_password\n \n # 打开一个Neo4j会话\n driver = GraphDatabase.driver(uri, auth=(username, password))\n print(\"Connected to Neo4j directly\")\n \n with driver.session() as session:\n # 清除现有数据\n session.run(\"MATCH (n) DETACH DELETE n\")\n print(\"Cleared existing data\")\n \n # 创建实体\n for entity in entities:\n cypher = \"\"\"\n CREATE (e:`{label}` {{name: $name, description: $desc, id: $id}})\n \"\"\".format(label=entity.label)\n \n session.run(cypher, {\n \"name\": entity.name,\n \"desc\": entity.properties.get(\"entity_description\", \"\"),\n \"id\": entity.properties.get(\"id\", entity.name)\n })\n print(f\"Added entity {entity.name} ({entity.label})\")\n \n # 创建关系\n for relation in relations:\n cypher = \"\"\"\n MATCH (source) WHERE source.name = $source_name\n MATCH (target) WHERE target.name = $target_name\n CREATE (source)-[r:`{label}` {{description: $desc}}]->(target)\n \"\"\".format(label=relation.label)\n \n session.run(cypher, {\n \"source_name\": relation.source_id,\n \"target_name\": relation.target_id,\n \"desc\": relation.properties.get(\"relationship_description\", \"\")\n })\n print(f\"Added relation {relation.source_id}--{relation.label}-->{relation.target_id}\")\n \n driver.close()\n print(\"Neo4j connection closed\")\n \n # 创建一个简单的自定义索引，不需要使用PropertyGraphIndex\n class SimpleGraphIndex:\n def __init__(self, graph_store):\n self.property_graph_store = graph_store\n \n # 使用自定义索引\n index = SimpleGraphIndex(neo4j_store)\n print(\"实体和关系加载成功!\")\n \n except Exception as extract_error:\n print(f\"Error extracting or loading entities: {extract_error}\")\n raise\n \n print(\"\\nStep 4: Building communities...\")\n graph_store.build_communities()\n \n print(\"\\nStep 5: Creating query engine...\")\n query_engine = GraphRAGQueryEngine(\n graph_store=None, # 我们不再需要传递graph_store，因为我们直接连接Neo4j\n llm=llm,\n index=None, # 我们不再需要index\n similarity_top_k=5\n )\n \n print(\"\\nStep 6: Ready to answer queries!\")\n \n print(\"\\nExample queries:\")\n queries = [\n \"What companies are mentioned in the documents?\",\n \"Who founded Apple?\",\n \"What is the relationship between Jeff Bezos and Amazon?\"\n ]\n \n for query in queries:\n print(f\"\\nQuery: {query}\")\n response = query_engine.query(query)\n print(f\"Response: {response.response}\")\n \n except Exception as e:\n print(f\"Error: {e}\")\n print(\"\\nNOTE: This demo requires a running Neo4j instance. If you don't have Neo4j set up,\")\n print(\"try running the simpler in-memory demo with: python graphrag/graphrag_demo.py\")\n print(\"\\nTo set up Neo4j:\")\n print(\"1. Install Neo4j Desktop from https://neo4j.com/download/\")\n print(\"2. Create a new database and start it\")\n print(\"3. Add the following to your .env file:\")\n print(\" NEO4J_USERNAME=neo4j\")\n print(\" NEO4J_PASSWORD=your_password\")\n print(\" NEO4J_URL=localhost\")\n\n\nif __name__ == \"__main__\":\n run_neo4j_demo() "}, "62af707b1a2497604e8b85c557d8a28916b9e07b": {"t": 1785764818.0876498, "n": 19446, "h": "f08b730be8477f9f", "c": "\"\"\"\nGraphRAG with Neo4j Implementation\n\nThis module implements a GraphRAG (Graph Retrieval Augmented Generation) system\nusing Neo4j as the graph database backend.\n\nWorkflow:\n1. Load documents\n2. Convert documents to text nodes\n3. Extract entities and relationships from text using LLM\n4. Store the extracted knowledge graph in Neo4j\n5. Query the knowledge graph using natural language\n\nBased on LlamaIndex's GraphRAG approach:\nhttps://docs.llamaindex.ai/en/stable/examples/cookbooks/GraphRAG_v1/\nhttps://docs.llamaindex.ai/en/stable/examples/cookbooks/GraphRAG_v2/\n\"\"\"\n\nimport os\nimport pandas as pd\nfrom dotenv import load_dotenv\nfrom llama_index.core import PropertyGraphIndex\nfrom llama_index.core import Document\nfrom llama_index.graph_stores.neo4j import Neo4jPropertyGraphStore\nimport openai\nfrom neo4j import GraphDatabase\n\n# 自定义 LLM 实现\nclass CustomLLM:\n \"\"\"Custom LLM implementation for DeepSeek model\"\"\"\n \n def __init__(self, api_key, api_base, model):\n self.api_key = api_key\n # 确保 api_base 格式正确\n if not api_base.endswith(\"/v1\"):\n if api_base.endswith(\"/\"):\n api_base = api_base + \"v1\"\n else:\n api_base = api_base + \"/v1\"\n self.api_base = api_base\n self.model = model\n # Initialize OpenAI client\n self.client = openai.OpenAI(\n api_key=self.api_key,\n base_url=self.api_base\n )\n print(f\"Initialized CustomLLM with base_url: {self.api_base}\")\n\n def complete(self, prompt, **kwargs):\n try:\n response = self.client.chat.completions.create(\n model=self.model,\n messages=[{\"role\": \"user\", \"content\": prompt}],\n temperature=0.7,\n )\n class Response:\n def __init__(self, text):\n self.text = text\n return Response(response.choices[0].message.content)\n except Exception as e:\n print(f\"API 调用错误: {e}\")\n # 返回一个空响应而不是抛出异常，以便演示能继续\n class Response:\n def __init__(self, text):\n self.text = f\"{{\\\"entities\\\": [], \\\"relationships\\\": []}}\"\n return Response()\n\n# Neo4j 版本的 GraphRAG 抽取器\nclass GraphRAGExtractor:\n \"\"\"\n GraphRAG 抽取器：使用 LLM 从文本中抽取三元组（实体-关系-实体），适配 Neo4j。\n \"\"\"\n def __init__(self, llm):\n self.llm = llm\n \n async def acall(self, *args, **kwargs):\n \"\"\"\n 异步调用，对于我们的简单实现，直接调用同步方法。\n 对于实际生产环境，应该实现真正的异步处理。\n \"\"\"\n return self.extract(*args, **kwargs)\n \n def extract(self, nodes, show_progress=False, **kwargs):\n \"\"\"\n 从节点列表中抽取实体和关系。\n 返回 EntityNode 和 Relation 的列表。\n \n 参数:\n nodes: 要处理的节点列表\n show_progress: 是否显示进度，兼容PropertyGraphIndex的调用\n **kwargs: 其他可能的参数\n \"\"\"\n all_entities = []\n all_relationships = []\n \n # 输出调试信息\n print(f\"Extract method called with {len(nodes)} nodes\")\n print(f\"Node types: {[type(node).__name__ for node in nodes]}\")\n \n # 检查nodes是否是列表\n if not isinstance(nodes, list):\n print(f\"WARNING: nodes is not a list but a {type(nodes)}\")\n # 尝试转换为列表\n nodes = list(nodes)\n \n for i, node in enumerate(nodes):\n print(f\"Processing node {i} of type {type(node).__name__}\")\n # 确保我们可以访问节点的node_id和text\n try:\n node_id = getattr(node, \"node_id\", None) or str(i)\n node_text = getattr(node, \"text\", None)\n if node_text is None and hasattr(node, \"get_content\"):\n node_text = node.get_content()\n \n if node_text is None:\n print(f\"WARNING: Node {i} has no text attribute\")\n # 尝试其他可能的属性\n for attr in [\"content\", \"get_text\", \"__str__\"]:\n if hasattr(node, attr):\n if callable(getattr(node, attr)):\n node_text = getattr(node, attr)()\n else:\n node_text = getattr(node, attr)\n print(f\"Found text in {attr}: {node_text[:50]}...\")\n break\n \n if node_text is None:\n print(f\"SKIPPING node {i} as no text could be found\")\n continue\n \n prompt = f\"\"\"\n Extract key entities and their relationships from the following text.\n Format your response as a JSON with 'entities' and 'relationships' keys.\n \n For each entity, include:\n - 'entity_id': A unique identifier for the entity\n - 'entity_type': The type of entity (e.g., Person, Organization, Product)\n - 'entity_description': A brief description of the entity\n \n For each relationship, include:\n - 'source': The source entity ID\n - 'target': The target entity ID\n - 'relation': The type of relation\n - 'relationship_description': A description of the relationship\n \n YOUR RESPONSE MUST BE VALID JSON. Make sure to return only a JSON object with no additional text or explanation.\n \n TEXT: {node_text}\n \"\"\"\n \n response = self.llm.complete(prompt)\n \n # 实际应用中应解析 LLM 返回的 JSON 并结构化数据\n try:\n import json\n import re\n \n # Try to extract JSON from the response if there's other text\n response_text = response.text.strip()\n json_match = re.search(r'\\{.*\\}', response_text, re.DOTALL)\n if json_match:\n json_str = json_match.group(0)\n else:\n json_str = response_text\n \n # Parse the JSON\n data = json.loads(json_str)\n \n from llama_index.core.graph_stores.types import EntityNode, Relation\n \n # 构建实体节点\n entities = [\n EntityNode(\n label=entity.get(\"entity_type\", \"Entity\"),\n properties={\n \"id\": entity.get(\"entity_id\", \"\"),\n \"entity_description\": entity.get(\"entity_description\", \"\"),\n \"triplet_source_id\": node_id\n },\n name=entity.get(\"entity_id\", \"\")\n )\n for entity in data.get(\"entities\", [])\n ]\n \n # 构建关系\n relationships = [\n Relation(\n label=relation.get(\"relation\", \"Related\"),\n source_id=relation.get(\"source\", \"\"),\n target_id=relation.get(\"target\", \"\"),\n properties={\n \"relationship_description\": relation.get(\"relationship_description\", \"\"),\n \"triplet_source_id\": node_id\n }\n )\n for relation in data.get(\"relationships\", [])\n ]\n \n all_entities.extend(entities)\n all_relationships.extend(relationships)\n \n except Exception as e:\n print(f\"Error parsing response for node {i}: {e}\")\n print(f\"Response text: {response.text[:100]}...\")\n \n except Exception as e:\n print(f\"Error processing node {i}: {e}\")\n \n print(f\"Extraction completed: {len(all_entities)} entities, {len(all_relationships)} relationships\")\n return all_entities, all_relationships\n\n\nclass GraphRAGStore:\n \"\"\"\n Neo4j 图数据库存储的包装类，提供 GraphRAG 相关接口。\n \"\"\"\n def __init__(self, username, password, url):\n # 初始化 Neo4j 图存储\n self.graph_store = Neo4jPropertyGraphStore(\n username=username,\n password=password,\n url=url\n )\n \n def get_triplets(self):\n \"\"\"\n 获取所有三元组（关系映射）。\n \"\"\"\n return self.graph_store.get_rel_map().items()\n \n def build_communities(self):\n \"\"\"\n 社区发现（本 demo 仅打印提示，未做实际聚类）。\n \"\"\"\n print(\"Building communities... (Demo version - no actual clustering)\")\n # 实际应用中可用社区发现算法并在 Neo4j 中创建社区节点\n\n\nclass GraphRAGQueryEngine:\n \"\"\"\n 基于 Neo4j 图存储的 GraphRAG 查询引擎。\n 使用 Cypher 直接查询 Neo4j。\n \"\"\"\n def __init__(self, graph_store, llm, index=None, similarity_top_k=5):\n self.graph_store = graph_store\n self.llm = llm\n self.index = index\n self.similarity_top_k = similarity_top_k\n \n # Neo4j连接信息\n self.uri = os.getenv(\"NEO4J_URL\")\n self.username = os.getenv(\"NEO4J_USERNAME\")\n self.password = os.getenv(\"NEO4J_PASSWORD\")\n \n def query(self, query_text):\n \"\"\"\n 使用 Cypher 查询图谱，拼接上下文，调用 LLM 进行问答。\n \"\"\"\n # 获取图谱数据作为上下文\n context = self._get_context_from_neo4j()\n \n # 拼接最终 prompt，交给 LLM 回答\n prompt = f\"\"\"\n Use the following knowledge graph information to answer the user's question.\n If the information is not in the knowledge graph, say that you don't know.\n \n {context}\n \n User Question: {query_text}\n \n Answer:\n \"\"\"\n \n response = self.llm.complete(prompt)\n \n class Response:\n def __init__(self, text):\n self.response = text\n \n return Response(response.text)\n \n def _get_context_from_neo4j(self):\n \"\"\"\n 从 Neo4j 获取图谱数据作为上下文。\n \"\"\"\n try:\n # 打开一个Neo4j会话\n driver = GraphDatabase.driver(self.uri, auth=(self.username, self.password))\n \n context = \"Knowledge Graph Information:\\n\\n\"\n \n with driver.session() as session:\n # 获取所有实体\n entities_result = session.run(\"\"\"\n MATCH (entity) \n RETURN entity.name as name, labels(entity) as labels, entity.description as description\n LIMIT 20\n \"\"\")\n \n for record in entities_result:\n entity_name = record[\"name\"]\n entity_label = record[\"labels\"][0] if record[\"labels\"] else \"Entity\"\n entity_desc = record[\"description\"] or \"\"\n context += f\"- {entity_name} ({entity_label}): {entity_desc}\\n\"\n \n # 获取所有关系\n relations_result = session.run(\"\"\"\n MATCH (source)-[rel]->(target)\n RETURN source.name as source, target.name as target, type(rel) as type, rel.description as description\n LIMIT 20\n \"\"\")\n \n for record in relations_result:\n source = record[\"source\"]\n target = record[\"target\"]\n rel_type = record[\"type\"]\n rel_desc = record[\"description\"] or \"\"\n context += f\"- {source} --{rel_type}--> {target}: {rel_desc}\\n\"\n \n driver.close()\n return context\n \n except Exception as e:\n print(f\"Error querying Neo4j: {e}\")\n return \"Unable to retrieve knowledge graph information due to an error.\"\n\n\n# 主演示函数\ndef run_neo4j_demo():\n \"\"\"\n 运行Neo4j GraphRAG演示\n \n 流程:\n 1. 加载环境变量和初始化LLM\n 2. 创建示例文档\n 3. 从文档中提取文本节点\n 4. 使用LLM从文本中提取实体和关系\n 5. 将实体和关系存储到Neo4j中\n 6. 使用自然语言查询知识图谱\n \"\"\"\n # 加载 .env 文件中的环境变量\n load_dotenv()\n \n # 检查 LLM API Key 是否存在 (使用与其他模块一致的变量名)\n api_key = os.getenv(\"LLM_API_KEY\")\n api_base = os.getenv(\"LLM_BASE_URL\", \"https://api.deepseek.com/v1\")\n if not api_key:\n print(\"ERROR: LLM_API_KEY not found in environment variables or .env file\")\n print(\"Please set your LLM API key before running this demo.\")\n return\n \n # 检查 Neo4j 相关配置 \n neo4j_user = os.getenv(\"NEO4J_USERNAME\") \n neo4j_password = os.getenv(\"NEO4J_PASSWORD\") \n neo4j_url = os.getenv(\"NEO4J_URL\")\n \n print(f\"Using API Key: {api_key[:5]}...\")\n print(f\"Using API Base: {api_base}\")\n print(f\"Using Neo4j config: {neo4j_user}@{neo4j_url}\")\n \n # 初始化自定义 LLM\n llm = CustomLLM(\n api_key=api_key,\n api_base=api_base,\n model=\"deepseek-chat\"\n )\n \n # 示例文档\n documents = [\n Document(text=\"Apple Inc. is a technology company founded by Steve Jobs, Steve Wozniak, and Ronald Wayne in 1976. \"\n \"The company is known for its iPhone products and is headquartered in Cupertino, California.\"),\n Document(text=\"Microsoft Corporation was founded by Bill Gates and Paul Allen in 1975. \"\n \"The company develops Windows operating system and Office software suite.\"),\n Document(text=\"Amazon was founded by Jeff Bezos in 1994 as an online bookstore. \"\n \"It has since expanded to sell various products and services including cloud computing through AWS.\")\n ]\n \n try:\n print(\"Step 1: Initializing Neo4j GraphRAG components...\")\n extractor = GraphRAGExtractor(llm)\n \n # 不需要再次验证URL，因为已在前面处理\n try:\n graph_store = GraphRAGStore(\n username=neo4j_user,\n password=neo4j_password,\n url=neo4j_url\n )\n print(\"Neo4j连接成功!\")\n except Exception as neo4j_error:\n print(f\"Neo4j连接错误: {neo4j_error}\")\n print(\"继续尝试处理，但可能会在后续步骤失败...\")\n raise # 由于图存储是必须的，这里还是需要终止执行\n \n print(\"\\nStep 2: Creating nodes from documents...\")\n # 按照 GraphRAG v1 文档，先创建节点/chunks\n from llama_index.core.node_parser import SentenceSplitter\n \n # 创建文本分割器\n splitter = SentenceSplitter(chunk_size=1024, chunk_overlap=20)\n \n # 从文档创建节点\n nodes = []\n for doc in documents:\n # 确保文档有metadata\n if not hasattr(doc, \"metadata\") or doc.metadata is None:\n doc.metadata = {}\n \n # 分割文档成节点\n doc_nodes = splitter.get_nodes_from_documents([doc])\n nodes.extend(doc_nodes)\n \n print(f\"Created {len(nodes)} nodes from {len(documents)} documents\")\n \n print(\"\\nStep 3: Extracting entities and relations directly...\")\n try:\n # 直接使用抽取器提取三元组\n entities, relations = extractor.extract(nodes)\n print(f\"Extracted {len(entities)} entities and {len(relations)} relationships\")\n \n # 直接加载到Neo4j\n print(\"Loading entities and relations into Neo4j...\")\n \n # 获取Neo4j图存储\n neo4j_store = graph_store.graph_store\n \n # 使用更简单的方法 - 使用Neo4j客户端直接执行Cypher\n \n # 获取Neo4j连接信息\n uri = neo4j_url\n username = neo4j_user\n password = neo4j_password\n \n # 打开一个Neo4j会话\n driver = GraphDatabase.driver(uri, auth=(username, password))\n print(\"Connected to Neo4j directly\")\n \n with driver.session() as session:\n # 清除现有数据\n session.run(\"MATCH (n) DETACH DELETE n\")\n print(\"Cleared existing data\")\n \n # 创建实体\n for entity in entities:\n cypher = \"\"\"\n CREATE (e:`{label}` {{name: $name, description: $desc, id: $id}})\n \"\"\".format(label=entity.label)\n \n session.run(cypher, {\n \"name\": entity.name,\n \"desc\": entity.properties.get(\"entity_description\", \"\"),\n \"id\": entity.properties.get(\"id\", entity.name)\n })\n print(f\"Added entity {entity.name} ({entity.label})\")\n \n # 创建关系\n for relation in relations:\n cypher = \"\"\"\n MATCH (source) WHERE source.name = $source_name\n MATCH (target) WHERE target.name = $target_name\n CREATE (source)-[r:`{label}` {{description: $desc}}]->(target)\n \"\"\".format(label=relation.label)\n \n session.run(cypher, {\n \"source_name\": relation.source_id,\n \"target_name\": relation.target_id,\n \"desc\": relation.properties.get(\"relationship_description\", \"\")\n })\n print(f\"Added relation {relation.source_id}--{relation.label}-->{relation.target_id}\")\n \n driver.close()\n print(\"Neo4j connection closed\")\n \n # 创建一个简单的自定义索引，不需要使用PropertyGraphIndex\n class SimpleGraphIndex:\n def __init__(self, graph_store):\n self.property_graph_store = graph_store\n \n # 使用自定义索引\n index = SimpleGraphIndex(neo4j_store)\n print(\"实体和关系加载成功!\")\n \n except Exception as extract_error:\n print(f\"Error extracting or loading entities: {extract_error}\")\n raise\n \n print(\"\\nStep 4: Building communities...\")\n graph_store.build_communities()\n \n print(\"\\nStep 5: Creating query engine...\")\n query_engine = GraphRAGQueryEngine(\n graph_store=None, # 我们不再需要传递graph_store，因为我们直接连接Neo4j\n llm=llm,\n index=None, # 我们不再需要index\n similarity_top_k=5\n )\n \n print(\"\\nStep 6: Ready to answer queries!\")\n \n print(\"\\nExample queries:\")\n queries = [\n \"What companies are mentioned in the documents?\",\n \"Who founded Apple?\",\n \"What is the relationship between Jeff Bezos and Amazon?\"\n ]\n \n for query in queries:\n print(f\"\\nQuery: {query}\")\n response = query_engine.query(query)\n print(f\"Response: {response.response}\")\n \n except Exception as e:\n print(f\"Error: {e}\")\n print(\"\\nNOTE: This demo requires a running Neo4j instance. If you don't have Neo4j set up,\")\n print(\"try running the simpler in-memory demo with: python graphrag/graphrag_demo.py\")\n print(\"\\nTo set up Neo4j:\")\n print(\"1. Install Neo4j Desktop from https://neo4j.com/download/\")\n print(\"2. Create a new database and start it\")\n print(\"3. Add the following to your .env file:\")\n print(\" NEO4J_USERNAME=neo4j\")\n print(\" NEO4J_PASSWORD=your_password\")\n print(\" NEO4J_URL=localhost\")\n\n\nif __name__ == \"__main__\":\n run_neo4j_demo() "}}}} |