# 傳統(tǒng)RAG實(shí)現(xiàn)（基于向量相似度）
from llama_index.core import VectorStoreIndex

vector_index = VectorStoreIndex.from_documents(documents)
response = vector_index.as_query_engine().query("利率變動的影響")

# Graph RAG實(shí)現(xiàn)（基于知識圖譜）
from llama_index.core import KnowledgeGraphIndex

kg_index = KnowledgeGraphIndex.from_documents(
    documents,
    llm=llm,
    embed_model=embed_model
)
response = kg_index.as_query_engine().query("央行加息如何通過產(chǎn)業(yè)鏈影響消費(fèi)市場")

import pandas as pd
from llama_index.core import Document

def load_news_data(data_path):
    """
    加載并預(yù)處理新聞數(shù)據(jù)集
    :param data_path: 新聞數(shù)據(jù)目錄路徑
    :return: 結(jié)構(gòu)化Document列表
    """
    categories = []
    texts = []

    # 遞歸遍歷數(shù)據(jù)目錄
    for root, dirs, files in os.walk(data_path):
        if os.path.basename(root) in ['經(jīng)濟(jì)', '政治', '國際', '社會']:
            for file in files:
                if file.endswith('.txt'):
                    with open(os.path.join(root, file), 'r', encoding='utf-8') as f:
                        text = f.read().strip()
                        if len(text) > 50:  # 過濾空文檔
                            categories.append(os.path.basename(root))
                            texts.append(text)

    # 構(gòu)建結(jié)構(gòu)化DataFrame
    df = pd.DataFrame({'category': categories, 'text': texts})

    # 轉(zhuǎn)換為LlamaIndex文檔對象
    documents = [
        Document(
            text=row['text'],
            metadata={"category": row['category']}
        ) for _, row in df.iterrows()
    ]

    return documents

# 使用示例
news_documents = load_news_data("path/to/news_dataset")
print(f"成功加載 {len(news_documents)} 篇新聞文檔")

from llama_index.core import StorageContext
from llama_index.core.graph_stores import SimpleGraphStore

def build_knowledge_graph(documents, llm, embed_model):
    """
    構(gòu)建知識圖譜索引
    :param documents: 預(yù)處理后的文檔列表
    :param llm: 大語言模型實(shí)例
    :param embed_model: 嵌入模型
    :return: 知識圖譜索引對象
    """
    # 初始化圖存儲
    graph_store = SimpleGraphStore()
    storage_context = StorageContext.from_defaults(
        graph_store=graph_store
    )

    # 圖譜構(gòu)建配置
    kg_config = {
        "max_triplets_per_chunk": 2,  # 每個文本塊提取的關(guān)系數(shù)
        "include_embeddings": True,   # 包含嵌入向量
        "show_progress": True         # 顯示進(jìn)度條
    }

    # 構(gòu)建圖譜索引
    kg_index = KnowledgeGraphIndex.from_documents(
        documents=documents[:5000],  # 首批處理5000篇
        storage_context=storage_context,
        llm=llm,
        embed_model=embed_model,
        **kg_config
    )

    return kg_index

# 使用示例
llm = Anthropic(model="claude-3-opus-20240229")
embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-large-zh-v1.5")
kg_index = build_knowledge_graph(news_documents, llm, embed_model)

from llama_index.core.retrievers import KGRetriever, VectorIndexRetriever
from llama_index.core.query_engine import RetrieverQueryEngine

class HybridRAGSystem:
    def __init__(self, vector_index, kg_index):
        """
        初始化混合檢索系統(tǒng)
        :param vector_index: 向量索引
        :param kg_index: 知識圖譜索引
        """
        self.vector_retriever = VectorIndexRetriever(
            index=vector_index,
            similarity_top_k=3
        )
        self.kg_retriever = KGRetriever(
            index=kg_index,
            similarity_top_k=2,
            include_text=True
        )

    def query(self, question, llm):
        """
        執(zhí)行混合檢索查詢
        :param question: 自然語言問題
        :param llm: 語言模型實(shí)例
        :return: 增強(qiáng)后的回答
        """
        # 并行執(zhí)行兩種檢索
        vector_results = self.vector_retriever.retrieve(question)
        kg_results = self.kg_retriever.retrieve(question)

        # 結(jié)果融合與去重
        unique_results = self._merge_results(vector_results, kg_results)

        # 生成最終回答
        response = llm.generate(
            f"基于以下證據(jù)回答問題:\n{unique_results}\n\n問題:{question}"
        )

        return response

    def _merge_results(self, vec_res, kg_res):
        """合并檢索結(jié)果并去重"""
        merged = []
        seen_ids = set()

        for res in vec_res + kg_res:
            if res.node.node_id not in seen_ids:
                merged.append(res)
                seen_ids.add(res.node.node_id)

        return merged[:5]  # 返回Top5結(jié)果

# 使用示例
hybrid_system = HybridRAGSystem(vector_index, kg_index)
answer = hybrid_system.query(
    "土耳其央行連續(xù)加息對歐元區(qū)進(jìn)出口貿(mào)易的影響路徑",
    llm=llm
)
print(answer)

# 構(gòu)建金融專用圖譜
finance_docs = [d for d in news_documents if d.metadata['category'] == '經(jīng)濟(jì)']
finance_kg = build_knowledge_graph(finance_docs, llm, embed_model)

# 專用查詢引擎
finance_query_engine = finance_kg.as_query_engine(
    llm=llm,
    include_text=True,
    response_mode="tree_summarize"
)

# 執(zhí)行多跳查詢
response = finance_query_engine.query(
    "分析2023年土耳其央行貨幣政策調(diào)整如何影響里拉匯率和進(jìn)口商品價格"
)

# 時間關(guān)系增強(qiáng)處理
def add_temporal_relations(kg_index):
    """為圖譜添加時間關(guān)系邊"""
    temporal_rules = {
        "之后": ["發(fā)布", "回應(yīng)", "導(dǎo)致"],
        "之前": ["準(zhǔn)備", "計劃", "預(yù)測"]
    }

    graph = kg_index.graph_store.get_graph()
    for node in graph.nodes:
        if 'date' in node.metadata:
            # 添加時間關(guān)系邏輯
            pass

    return kg_index

# 時間敏感查詢
timeline_query = """
梳理2024年土耳其地方選舉期間，
主要政黨候選人的言論如何逐步影響選民傾向
"""

from multiprocessing import Pool

def process_document_batch(batch):
    """多進(jìn)程處理文檔批次"""
    local_llm = Anthropic()  # 每個進(jìn)程獨(dú)立實(shí)例
    return KnowledgeGraphIndex.from_documents(
        batch,
        llm=local_llm,
        embed_model=embed_model
    )

# 分批次并行處理
batch_size = 1000
batches = [news_documents[i:i+batch_size] for i in range(0, len(news_documents), batch_size)]

with Pool(4) as pool:  # 使用4個進(jìn)程
    results = pool.map(process_document_batch, batches)

class DynamicGraphUpdater:
    def __init__(self, kg_index):
        self.index = kg_index

    def update_with_new_docs(self, new_documents):
        """增量更新知識圖譜"""
        for doc in new_documents:
            self.index.insert(
                document=doc,
                llm=llm,
                embed_model=embed_model
            )

    def prune_old_nodes(self, max_nodes=10000):
        """圖譜剪枝保持性能"""
        current_nodes = len(self.index.graph_store.get_all_nodes())
        if current_nodes > max_nodes:
            # 實(shí)現(xiàn)LRU剪枝策略
            pass

class ProductionRAGSystem:
 def __init__(self):
     self.vector_index = load_vector_index()
     self.kg_index = load_knowledge_graph()
     self.fast_llm = Anthropic(model="claude-instant")  # 快速響應(yīng)
     self.expert_llm = Anthropic(model="claude-opus")   # 復(fù)雜推理

 def route_query(self, question):
     """智能路由查詢請求"""
     if requires_kg(question):  # 需要關(guān)系推理
         return self.kg_index.query(question, llm=self.expert_llm)
     else:  # 簡單事實(shí)查詢
         return self.vector_index.query(question, llm=self.fast_llm)