# 一个简化的公司投资领域本体概念示例（使用RDF/OWL风格描述）
class Person:
    properties: [name, birthDate]

class Company:
    properties: [name, foundingDate, industry]

class Investment(Relation):
    domain: Person  # 定义域：关系的主体来自Person类
    range: Company  # 值域：关系的客体来自Company类
    properties: [amount, date]

class ServesAs(Relation):
    domain: Person
    range: Company
    properties: [title, startDate, endDate]

import os
from langchain.graphs import Neo4jGraph
from dotenv import load_dotenv

load_dotenv()

# 从环境变量读取配置
NEO4J_URI = os.getenv("NEO4J_URI")
NEO4J_USERNAME = os.getenv("NEO4J_USERNAME")
NEO4J_PASSWORD = os.getenv("NEO4J_PASSWORD")

# 初始化Neo4j图对象
graph = Neo4jGraph(
    url=NEO4J_URI,
    username=NEO4J_USERNAME,
    password=NEO4J_PASSWORD
)

# 假设我们已通过其他方式将数据导入Neo4j，这里是一个简单的示例数据插入
init_cypher = """
// 创建一些示例节点和关系
MERGE (p:Person {name: 'Elon Musk', birthYear: 1971})
MERGE (c1:Company {name: 'SpaceX', industry: 'Aerospace', founded: 2002})
MERGE (c2:Company {name: 'Tesla', industry: 'Automotive', founded: 2003})
MERGE (c3:Company {name: 'Neuralink', industry: 'Neurotechnology', founded: 2016})
MERGE (p)-[:FOUNDED {year: 2002}]->(c1)
MERGE (p)-[:CEO_OF]->(c2)
MERGE (p)-[:FOUNDED {year: 2016}]->(c3)
"""
graph.query(init_cypher)
print("示例知识图谱数据已初始化。")

from langchain.chains import GraphCypherQAChain
from langchain_openai import ChatOpenAI
from langchain.prompts import PromptTemplate

# 初始化LLM (请替换为你的OpenAI API Key)
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
llm = ChatOpenAI(model="gpt-4", temperature=0, openai_api_key=OPENAI_API_KEY)

# 自定义提示模板，以更好地控制LLM的行为
CYPHER_GENERATION_TEMPLATE = """
你是一个专业的Neo4j Cypher查询生成器。
根据用户的问题，生成一个Cypher查询来从知识图谱中检索相关信息。
仅使用图中存在的标签、关系和属性。
图谱Schema如下：
{schema}

问题：{question}
"""
cypher_prompt = PromptTemplate.from_template(CYPHER_GENERATION_TEMPLATE)

# 创建GraphCypherQAChain
chain = GraphCypherQAChain.from_llm(
    llm=llm,
    graph=graph,
    verbose=True,  # 打印中间步骤，便于调试
    cypher_prompt=cypher_prompt,
    return_direct=False,  # 设为True可跳过LLM，直接返回Cypher查询结果
    top_k=10  # 限制返回的子图大小
)

# 进行问答测试
question = "埃隆·马斯克创立了哪些公司？"
result = chain.invoke({"query": question})
print(f"\n问题：{question}")
print(f"答案：{result['result']}")

# 一个更复杂的多跳查询示例
complex_question = "马斯克创立的公司里，有哪些属于高科技行业？"
complex_result = chain.invoke({"query": complex_question})
print(f"\n问题：{complex_question}")
print(f"答案：{complex_result['result']}")

from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import LLMChainExtractor
from langchain_community.vectorstores import Neo4jVector  # Neo4j 5+ 支持向量索引
from langchain_openai import OpenAIEmbeddings
from langchain.schema import Document
from typing import List

# 1. 图检索函数
def retrieve_from_graph(query: str) -> List[Document]:
    """将图查询结果转换为LangChain Document对象"""
    try:
        # 这里简化处理，实际应使用更鲁棒的查询生成
        cypher_query = f"""
        MATCH (n)
        WHERE toLower(n.name) CONTAINS toLower('{query}') OR toLower(n.industry) CONTAINS toLower('{query}')
        RETURN n.name as name, n.industry as industry, labels(n)[0] as type LIMIT 5
        """
        data = graph.query(cypher_query)
        docs = []
        for record in data:
            text = f"实体：{record['name']}， 类型：{record['type']}， 行业：{record.get('industry', 'N/A')}"
            doc = Document(page_content=text, metadata={"source": "graph", "entity": record['name']})
            docs.append(doc)
        return docs
    except Exception as e:
        print(f"图检索失败: {e}")
        return []

# 2. 向量检索 (假设我们已将公司描述文本向量化并存入Neo4j Vector)
# 初始化Neo4j Vector索引（需提前创建）
# vector_index = Neo4jVector.from_existing_graph(
#     OpenAIEmbeddings(),
#     url=NEO4J_URI,
#     username=NEO4J_USERNAME,
#     password=NEO4J_PASSWORD,
#     index_name='company_descriptions',
#     node_label="Company",
#     text_node_properties=['description'],
#     embedding_node_property='embedding'
# )
# vector_retriever = vector_index.as_retriever(search_kwargs={"k": 3})

# 3. 简单的融合检索器 (轮询或加权平均)
class HybridRetriever:
    def __init__(self, graph_retriever_func):
        self.graph_retriever = graph_retriever_func
        # self.vector_retriever = vector_retriever

    def get_relevant_documents(self, query: str) -> List[Document]:
        graph_docs = self.graph_retriever(query)
        # vector_docs = self.vector_retriever.get_relevant_documents(query)
        # 简单的去重和融合逻辑
        all_docs = graph_docs # + vector_docs
        seen = set()
        unique_docs = []
        for doc in all_docs:
            if doc.page_content not in seen:
                seen.add(doc.page_content)
                unique_docs.append(doc)
        return unique_docs[:5]  # 返回Top 5

hybrid_retriever = HybridRetriever(retrieve_from_graph)

# 4. 使用检索器增强生成
from langchain.chains import RetrievalQA

qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever=hybrid_retriever,
    return_source_documents=True
)

hybrid_result = qa_chain.invoke({"query": "请介绍下马斯克和航空航天相关的公司"})
print(f"\n混合检索结果 - 答案：{hybrid_result['result']}")
print("来源文档：")
for doc in hybrid_result['source_documents']:
    print(f"  - {doc.page_content} (来源：{doc.metadata.get('source', 'unknown')})")