# 安装: pip install neo4j langchain-neo4j
from neo4j import GraphDatabase
import os

# 1. 连接 Neo4j（社区版免费）
driver = GraphDatabase.driver(
    "bolt://localhost:7687",
    auth=("neo4j", "your_password")
)

# 2. 创建知识图谱
def build_tech_kg(driver):
    """构建一个科技公司知识图谱"""
    with driver.session() as session:
        # 创建公司节点
        companies = [
            ("Apple", "AAPL", 1976, "消费电子"),
            ("Google", "GOOGL", 1998, "互联网"),
            ("Microsoft", "MSFT", 1975, "软件"),
            ("Tesla", "TSLA", 2003, "汽车")
        ]
        for name, ticker, year, industry in companies:
            session.run("""
                MERGE (c:Company {name: $name})
                SET c.ticker = $ticker,
                    c.founded = $year,
                    c.industry = $industry
            """, name=name, ticker=ticker, year=year, industry=industry)

        # 创建创始人关系
        founders = [
            ("Steve Jobs", "Apple", "联合创始人"),
            ("Elon Musk", "Tesla", "CEO"),
            ("Bill Gates", "Microsoft", "联合创始人"),
            ("Larry Page", "Google", "联合创始人")
        ]
        for person, company, role in founders:
            session.run("""
                MERGE (p:Person {name: $person})
                MERGE (c:Company {name: $company})
                MERGE (p)-[:FOUNDED {role: $role}]->(c)
            """, person=person, company=company, role=role)

        # 创建产品关系
        products = [
            ("Apple", "iPhone", "智能手机"),
            ("Google", "Gemini", "AI模型"),
            ("Microsoft", "Copilot", "AI助手"),
            ("Tesla", "Model 3", "电动汽车")
        ]
        for company, product, ptype in products:
            session.run("""
                MERGE (c:Company {name: $company})
                MERGE (p:Product {name: $product})
                MERGE (c)-[:PRODUCES]->(p)
                SET p.type = $type
            """, company=company, product=product, type=ptype)

    print("知识图谱构建完成！")

build_tech_kg(driver)

# 3. 查询：多跳推理
def multi_hop_query(driver, question: str):
    """知识图谱多跳推理"""
    with driver.session() as session:
        # 示例：「Tesla 的 CEO 创建了哪些产品？」
        result = session.run("""
            MATCH (p:Person)-[:FOUNDED]->(c:Company {name: "Tesla"})
            MATCH (c)-[:PRODUCES]->(prod:Product)
            RETURN p.name AS founder, prod.name AS product, prod.type AS type
        """)
        return [(r["founder"], r["product"], r["type"]) for r in result]

results = multi_hop_query(driver, "")
for founder, product, ptype in results:
    print(f"{founder} 的 Tesla 生产了 {product}（{ptype}）")

# 用 LLM 从文本中提取知识图谱三元组
from openai import OpenAI
import json

client = OpenAI()

def extract_triples(text: str) -> list[tuple]:
    """用 LLM 从文本中提取三元组"""
    prompt = f"""
    从以下文本中提取所有「实体-关系-实体」三元组：
    
    文本：{text}
    
    输出 JSON 格式：
    triples
    
    规则：
    - 实体用全名（如「埃隆·马斯克」而非「马斯克」）
    - 关系用简洁动词（如「创立」「出生于」「位于」）
    """
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": prompt}],
        response_format={"type": "json_object"}
    )
    return json.loads(response.choices[0].message.content)["triples"]

text = "埃隆·马斯克于 2004 年加入特斯拉，成为董事长。"
triples = extract_triples(text)
print(triples)
# [["埃隆·马斯克", "加入", "特斯拉"], ...]