import os
import sys
import pandas as pd
from llama_index.core import SimpleDirectoryReader, VectorStoreIndex, Settings
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
# ========== 3. 生成测试Excel ==========
test_excel = "./docs/product_price.xlsx"
if not os.path.exists(test_excel):
    df = pd.DataFrame({
        "产品ID": ["P001", "P002"],
        "产品名称": ["产品A", "产品B"],
        "年份": [2025, 2025],
        "单价(元)": [999, 1299],
        "库存": [500, 300]
    })
    df.to_excel(test_excel, index=False)
    print(f"✅ 测试Excel已生成：{test_excel}")
# ========== 4. 读取Excel（核心步骤） ==========
docs = SimpleDirectoryReader(input_files=[test_excel]).load_data()
print(f"✅ 表格解析完成，共{len(docs)}条数据")
# ========== 5. 配置Embedding ==========
Settings.embed_model = HuggingFaceEmbedding(
    model_name="D:/modelscope/hub/models/sentence-transformers/paraphrase-MiniLM-L6-v2",
    model_kwargs={"device": "cpu"}
)
# ========== 6. 构建向量索引 ==========
index = VectorStoreIndex.from_documents(docs)
print("✅ 向量索引构建完成")
# ========== 7. 加载本地Qwen模型（可选，不影响核心验证） ==========
try:
    from modelscope.hub.snapshot_download import snapshot_download
    from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
    from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
    MODEL_NAME = "qwen/Qwen1.5-1.8B-Chat"
    local_path = snapshot_download(MODEL_NAME, cache_dir="D:\\modelscope\\hub")
    
    tokenizer = AutoTokenizer.from_pretrained(local_path, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        local_path, trust_remote_code=True, device_map="auto", torch_dtype="auto"
    )
    llm_pipe = pipeline(
        "text-generation", model=model, tokenizer=tokenizer,
        max_new_tokens=100, temperature=0.1
    )
    llm = HuggingFacePipeline(pipeline=llm_pipe)
    print("✅ 本地Qwen模型加载成功")
except Exception as e:
    print(f"⚠️ 模型加载失败（仅影响问答，表格读取已验证）：{e}")
    llm = None
# ========== 8. 问答验证 ==========
if llm:
    query = "2025产品A的单价是多少？"
    retriever = index.as_retriever(similarity_top_k=1)
    context = "\n".join([n.text for n in retriever.retrieve(query)])
    response = llm.invoke(f"仅基于以下数据回答：{context}\n问题：{query}")
    print(f"\n📌 问答结果：")
    print(f"问题：{query}")
    print(f"回答：{response.strip()}")
else:
    print("\n✅ 核心验证完成（表格读取+索引构建），模型加载失败不影响核心功能")

import os
import sys
import pandas as pd
from modelscope.hub.snapshot_download import snapshot_download
# ===================== 0.11.23 纯净版导入（无任何废弃参数） =====================
from llama_index.core import (
    VectorStoreIndex, Settings, StorageContext,
    load_index_from_storage
)
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
# 适配0.11.23：极简Excel读取（完全移除return_single_document）
class ExcelReader:
    """0.11.23专用Excel读取类，无任何废弃参数"""
    def __init__(self, sheet_name=None):  # 仅保留必要参数
        self.sheet_name = sheet_name
    def load_data(self, file):
        from llama_index.core import SimpleDirectoryReader
        # 0.11.23 SimpleDirectoryReader原生用法（仅传input_files）
        reader = SimpleDirectoryReader(input_files=[file])
        docs = reader.load_data()
        
        # 手动拆分表格行（实现按行检索的效果）
        split_docs = []
        for doc in docs:
            # 按换行拆分，跳过表头，仅保留数据行
            lines = doc.text.strip().split("\n")
            if len(lines) <= 1:
                continue
            # 提取表头，拼接每行数据为"表头:值"格式
            header = lines[0].split("\t")
            for line in lines[1:]:
                if not line.strip():
                    continue
                values = line.strip().split("\t")
                if len(values) != len(header):
                    continue
                # 构建结构化文本，便于检索
                row_text = ", ".join([f"{h}:{v}" for h, v in zip(header, values)])
                new_doc = doc.copy()
                new_doc.text = row_text
                split_docs.append(new_doc)
        return split_docs
# LLM加载模块
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
# ===================== 1. 自动生成测试表格 =====================
def generate_test_tables():
    # 产品价格表
    price_data = {
        "产品ID": ["P001", "P002"],
        "产品名称": ["产品A", "产品B"],
        "年份": [2025, 2025],
        "单价（元）": [999, 1299],
        "批量折扣": ["9折", "8.5折"]
    }
    pd.DataFrame(price_data).to_excel("./product_price.xlsx", index=False)
    # 销售报表
    sales_data = {
        "产品ID": ["P001", "P002"],
        "产品名称": ["产品A", "产品B"],
        "年份": [2025, 2025],
        "销售额（万元）": [49.95, 38.97],
        "销量（件）": [500, 300]
    }
    pd.DataFrame(sales_data).to_excel("./sales_report.xlsx", index=False)
    print("✅ 测试表格生成完成")
    return ["./product_price.xlsx", "./sales_report.xlsx"]
# ===================== 2. 初始化模型+Embedding =====================
def init_models():
    MODEL_NAME = "qwen/Qwen1.5-1.8B-Chat"
    CACHE_DIR = "D:\\modelscope\\hub"
    # 下载模型
    try:
        local_path = snapshot_download(MODEL_NAME, cache_dir=CACHE_DIR)
    except Exception as e:
        print(f"❌ 模型下载失败：{e}")
        sys.exit(1)
    # 加载模型
    tokenizer = AutoTokenizer.from_pretrained(local_path, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        local_path, trust_remote_code=True, device_map="auto", torch_dtype="auto"
    )
    # 封装LLM
    llm_pipe = pipeline(
        "text-generation", model=model, tokenizer=tokenizer,
        max_new_tokens=300, temperature=0.1
    )
    llm = HuggingFacePipeline(pipeline=llm_pipe)
    # 配置Embedding
    Settings.embed_model = HuggingFaceEmbedding(
        model_name="D:/modelscope/hub/models/sentence-transformers/paraphrase-MiniLM-L6-v2",
        model_kwargs={"device": "cpu"}
    )
    print("✅ 模型初始化完成")
    return llm
# ===================== 3. 多表格索引构建+持久化 =====================
def build_table_index(table_paths, index_dir="./table_storage"):
    if not os.path.exists(index_dir):
        all_docs = []
        for path in table_paths:
            if not os.path.exists(path):
                print(f"❌ 表格不存在：{path}")
                continue
            # 初始化0.11.23专用ExcelReader（无废弃参数）
            reader = ExcelReader()
            docs = reader.load_data(file=path)
            # 添加元数据
            for doc in docs:
                doc.metadata["table_name"] = os.path.basename(path)
                doc.metadata["年份"] = 2025
            all_docs.extend(docs)
        if not all_docs:
            print("❌ 未解析到表格数据")
            return None
        # 构建索引
        table_index = VectorStoreIndex.from_documents(all_docs)
        table_index.storage_context.persist(persist_dir=index_dir)
        print(f"✅ 索引构建完成，解析{len(all_docs)}行数据")
    else:
        # 加载本地索引
        storage_context = StorageContext.from_defaults(persist_dir=index_dir)
        table_index = load_index_from_storage(storage_context)
        print("✅ 加载本地索引完成")
    return table_index
# ===================== 4. 检索后过滤（替代MetadataFilter） =====================
def filter_nodes_by_metadata(nodes, filter_conditions):
    if not filter_conditions:
        return nodes
    filtered_nodes = []
    for node in nodes:
        match = True
        for k, v in filter_conditions.items():
            if node.metadata.get(k) != v:
                match = False
                break
        if match:
            filtered_nodes.append(node)
    return filtered_nodes
# ===================== 5. 表格问答 =====================
def advanced_table_qa(index, llm, query, filter_conditions=None):
    if not index or not llm:
        return "❌ 索引/模型未初始化"
    
    # 基础检索
    retriever = VectorIndexRetriever(index=index, similarity_top_k=5)
    retrieved_nodes = retriever.retrieve(query)
    
    # 过滤
    filtered_nodes = filter_nodes_by_metadata(retrieved_nodes, filter_conditions)
    if not filtered_nodes:
        return "❌ 未检索到匹配数据"
    
    # 拼接上下文
    context = ""
    for node in filtered_nodes:
        context += f"【{node.metadata['table_name']}】{node.text}\n"
    
    # 生成回答
    prompt = f"""仅基于以下数据回答，格式：结果（数据来源）
{context}
问题：{query}"""
    response = llm.invoke(prompt).strip()
    return response
# ===================== 主函数 =====================
if __name__ == "__main__":
    print("===== LlamaIndex 0.11.23 最终纯净版 =====")
    
    # 生成表格
    table_paths = generate_test_tables()
    # 初始化模型
    llm = init_models()
    # 构建索引（核心修复：无参数错误）
    table_index = build_table_index(table_paths)
    if not table_index:
        sys.exit(1)
    
    # 测试问答
    print("\n===== 问答测试 =====")
    query1 = "2025产品A的销售额？"
    filter_cond = {"table_name": "sales_report.xlsx", "年份": 2025}
    print(f"问题：{query1}\n回答：{advanced_table_qa(table_index, llm, query1, filter_cond)}\n")
    
    query2 = "2025产品A的价格和销售额分别是多少？"
    print(f"问题：{query2}\n回答：{advanced_table_qa(table_index, llm, query2)}")