# 查看 Python 版本
!python --version

# 查看 Llama-Index 版本
!pip list | grep llama-index

# 安装依赖
!pip install -q llama-index-core llama-index-llms-openai llama-index-embeddings-openai llama-index-llms-dashscope llama-index-embeddings-dashscope
!pip install -q llama-index-vector-stores-milvus llama-index-retrievers-bm25 rank-bm25 llama-parse
!pip install -q sqlalchemy pandas python-dotenv 

print("✅ 依赖安装完成")

# 全局配置
import os
import warnings
from dotenv import load_dotenv
from llama_index.llms.dashscope import DashScope
from llama_index.embeddings.dashscope import DashScopeEmbedding
from llama_index.core.settings import Settings

warnings.filterwarnings('ignore')
# 加载环境变量
load_dotenv(override=True)

# 设置 LLM
Settings.llm = DashScope(
    model_name="qwen-max",
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    api_base=os.getenv("DASHSCOPE_BASE_URL", "https://api.dashscope.aliyuncs.com")
)

# 设置 Embedding模型
Settings.embed_model = DashScopeEmbedding(
    model_name="text-embedding-v4",
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    api_base=os.getenv("DASHSCOPE_BASE_URL", "https://api.dashscope.aliyuncs.com")
)

print("✅ LLM配置完成")
print(f"   模型: {Settings.llm.model_name}")
print(f"   Embedding: {Settings.embed_model.model_name}")

# 基本使用
response = Settings.llm.complete("你好，请介绍一下通义千问。")
print(response)

response_embed = Settings.embed_model.get_query_embedding("你好，请介绍一下通义千问。")
print(response_embed)

import os
import pandas as pd
from llama_index.core import Document, VectorStoreIndex
from llama_index.core.node_parser import MarkdownElementNodeParser
from llama_index.core.retrievers import RecursiveRetriever
from llama_index.llms.openai import OpenAI
from llama_index.core.schema import TextNode, IndexNode # 确保导入这两个类

# 0. 准备测试数据 (确保表格格式标准)
markdown_content = """
# 公司季度财报
以下是我们在2023年第三季度的详细财务表现:
| 部门 | 收入(万元) | 成本(万元) | 利润(万元) |
|------|------------|------------|------------|
| 研发部 | 0 | 500 | -500 |
| 销售部 | 2000 | 800 | 1200 |
| 产品部 | 500 | 200 | 300 |
总结来看,销售部表现最为亮眼。
"""
doc = Document(text=markdown_content)

# 1. 创建解析器
node_parser = MarkdownElementNodeParser(
    llm=Settings.llm, 
    num_workers=1   # 并行解析数
)

# 2. 正确的解析流程
# 第一步: 从文档获取原始节点
raw_nodes = node_parser.get_nodes_from_documents([doc])

# 第二步: 从节点中提取基座节点和表格对象
base_nodes, objects = node_parser.get_nodes_and_objects(raw_nodes)
print(f"解析完成: 生成了 {len(base_nodes)} 个基座节点, {len(objects)} 个表格对象")

# 3. 验证表格对象
print(f"\nobjects 类型: {type(objects)}")

if objects:
    print("\n✅ 表格对象提取成功!")

    for i, obj in enumerate(objects):
        print(f"  对象 {i}:")
        print(f"    类型: {type(obj).__name__}")
        print(f"    ID: {obj.node_id if hasattr(obj, 'node_id') else 'N/A'}")
        # 如果是表格节点,尝试显示内容
        if hasattr(obj, 'text'):
            print(f"    内容预览: {obj.text[:100]}...")
        if hasattr(obj, 'metadata'):
            print(f"    元数据: {obj.metadata}")
else:
    print("\n⚠️ 未提取到表格,调试信息:")
    for i, node in enumerate(raw_nodes):
        print(f"  节点 {i}: {type(node).__name__} - {node.text[:60]}...")

# 4. 过滤 base_nodes (只保留 TextNode 和 IndexNode)
base_nodes_filtered = [n for n in raw_nodes if isinstance(n, (TextNode, IndexNode))]
print(f"\n过滤后基座节点数: {len(base_nodes_filtered)}")

# 5. 建立向量索引 (只对摘要建索引)
print("\n开始建立向量索引...")
vector_index = VectorStoreIndex(nodes=base_nodes_filtered)
vector_retriever = vector_index.as_retriever(similarity_top_k=2)

# 6. 构建递归检索器
# 将 objects 列表转换为字典 {node_id: node}
objects_dict = {obj.node_id: obj for obj in objects}

recursive_retriever = RecursiveRetriever(
    "vector",
    retriever_dict={"vector": vector_retriever},
    node_dict=objects_dict,  # 传入表格对象字典
    verbose=True  # 开启日志，可以看到递归过程
)

print("\n✅ 递归检索器构建完成!")

# 7. 执行检索测试
print("\n" + "="*50)
print("开始检索测试")
print("="*50)

query = "销售部的利润是多少？"
print(f"\n查询: {query}")

retrieved_nodes = recursive_retriever.retrieve(query)

print(f"\n检索到 {len(retrieved_nodes)} 个结果:")
for i, node_with_score in enumerate(retrieved_nodes):
    print(f"\n--- 结果 {i+1} (相关度: {node_with_score.score:.3f}) ---")
    print(f"节点类型: {type(node_with_score.node).__name__}")
    print(f"内容:\n{node_with_score.node.text[:300]}...")
    
# 8. 结合 LLM 生成最终答案
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core import Settings

query_engine = RetrieverQueryEngine.from_args(
    recursive_retriever,
    llm=Settings.llm
)

print("\n" + "="*50)
print("生成最终答案")
print("="*50)

response = query_engine.query(query)
print(f"\n问题: {query}")
print(f"\n答案: {response}")
print(f"\n来源节点数: {len(response.source_nodes)}")

import os
from llama_index.core import Document, VectorStoreIndex, Settings
from llama_index.core.node_parser import MarkdownElementNodeParser
from llama_index.core.retrievers import RecursiveRetriever
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.schema import TextNode, IndexNode
from llama_parse import LlamaParse
from dotenv import load_dotenv

load_dotenv(override=True)


# ========== 步骤 0: 环境配置 ==========
# 1. 安装依赖: pip install llama-parse llama-index
# 2. 获取 API Key: https://cloud.llamaindex.ai (需要注册 LlamaCloud 账号)
# 3. 设置环境变量LLAMA_CLOUD_API_KEY
# os.environ["LLAMA_CLOUD_API_KEY"] = "llx-..."  

# ========== 步骤 1: 使用 LlamaParse 解析 PDF ==========
print("📄 开始解析 PDF...")
parser = LlamaParse(
    result_type="markdown",  # 输出格式为 Markdown (保留表格结构)
    verbose=True,
    language="ch_sim",  # 简体中文
    # 可选: 提供自定义的解析指令
    # parsing_instruction="重点提取财务数据表格，保留表头和数值"
)

# 解析 PDF 文件
pdf_path = "dataset/招股意向书.pdf"
# 添加错误处理
try:
    documents = parser.load_data(pdf_path)
    print(f"✅ PDF 解析完成，共生成 {len(documents)} 个文档片段")
    
    # 检查是否为空
    if not documents:
        print("⚠️ 警告: 解析结果为空!")
        print("可能的原因:")
        print("1. PDF 文件路径不正确")
        print("2. LlamaParse API 调用失败")
        print("3. API Key 未正确配置")
        print(f"\n请检查:")
        print(f"- 文件是否存在: {os.path.exists(pdf_path)}")
        print(f"- API Key 是否已设置: {bool(os.getenv('LLAMA_CLOUD_API_KEY'))}")
    else:
        print(f"预览第一个片段:\n{documents[0].text[:300]}...\n")
        
except Exception as e:
    print(f"❌ PDF 解析失败: {e}")
    print("\n调试信息:")
    print(f"- 文件路径: {pdf_path}")
    print(f"- 文件存在: {os.path.exists(pdf_path)}")
    print(f"- API Key 已设置: {bool(os.getenv('LLAMA_CLOUD_API_KEY'))}")
    raise

# ========== 步骤 2: 使用 MarkdownElementNodeParser 提取表格 ==========
print("🔍 开始提取表格...")
node_parser = MarkdownElementNodeParser(
    llm=Settings.llm,  # 用于生成表格摘要
    num_workers=4      # 并行处理
)

# 从文档获取原始节点
raw_nodes = node_parser.get_nodes_from_documents(documents)

# 提取基座节点和表格对象
base_nodes, objects = node_parser.get_nodes_and_objects(raw_nodes)
print(f"✅ 提取完成:")
print(f"  - 基座节点(文本+摘要): {len(base_nodes)} 个")
print(f"  - 表格对象(完整表格): {len(objects)} 个\n")

# 调试:查看提取到的表格
if objects:
    print("📊 表格对象示例:")
    for i, obj in enumerate(objects[:3]):  # 只显示前3个
        print(f"  表格 {i+1}:")
        print(f"    ID: {obj.node_id}")
        print(f"    内容预览: {obj.text[:100]}...\n")


# ========== 步骤 3: 构建索引和递归检索器 ==========
print("🔧 构建检索系统...")
# 过滤节点(只保留 TextNode 和 IndexNode)
base_nodes_filtered = [n for n in base_nodes if isinstance(n, (TextNode, IndexNode))]

# 建立向量索引
vector_index = VectorStoreIndex(nodes=base_nodes_filtered)
vector_retriever = vector_index.as_retriever(similarity_top_k=5)

# 构建 objects 字典
objects_dict = {obj.node_id: obj for obj in objects}

# 构建递归检索器
recursive_retriever = RecursiveRetriever(
    "vector",
    retriever_dict={"vector": vector_retriever},
    node_dict=objects_dict,
    verbose=True
)

# 构建查询引擎
query_engine = RetrieverQueryEngine.from_args(
    recursive_retriever,
    llm=Settings.llm
)
print("✅ 检索系统构建完成!\n")

# ========== 步骤 4: 测试查询 ==========
print("="*60)
print("开始测试查询")
print("="*60)
# 示例查询(根据招股意向书的实际内容调整)
test_queries = [
    "体外诊断项目中，“输血项目”的检测范围和用途是什么?",
]
for query in test_queries:
    print(f"\n问题: {query}")
    response = query_engine.query(query)
    print(f"答案: {response}\n")
    print("-" * 60)

import pandas as pd
import time
from typing import Dict, Any, List
from llama_index.experimental.query_engine import PandasQueryEngine
from llama_index.llms.openai import OpenAI
from llama_index.core import Settings

import logging
# 禁用 Phoenix/OpenInference 的错误日志，在后续有提到的监控工具
logging.getLogger("openinference.instrumentation.llama_index").setLevel(logging.CRITICAL)

# 加载产品数据
df = pd.read_csv("dataset/products.csv")
print(f"📊 数据加载成功: {len(df)} 行 x {len(df.columns)} 列")
print(f"\n列名: {list(df.columns)}")
print("\n数据预览:")
df.head()

# 优化的指令字符串 - 关键配置
instruction_str = (
    "1. 将查询转换为可执行的 Pandas 代码\n"
    "2. 代码必须是单行表达式，可以被 eval() 执行\n"
    "3. DataFrame 变量名为 'df'\n"
    "4. 只输出表达式本身，不要有其他文字\n"
    "5. 不要使用引号包裹表达式\n"
    "\n示例:\n"
    "问题: 找出价格最高的产品\n"
    "代码: df.loc[df['price'].idxmax()]['name']\n"
)

# 创建查询引擎
query_engine = PandasQueryEngine(
    df=df,   # DataFrame读取的数据
    verbose=True,  # 显示生成的代码
    instruction_str=instruction_str,
    synthesize_response=True  # 综合回答
)

print("✅ PandasQueryEngine 初始化完成")

# 场景 1: 简单聚合查询
query = "电子产品类别有多少个产品？"
print(f"📝 查询: {query}")

response = query_engine.query(query)
print(f"💬 回答: {response}")

# 场景 2: 排序和筛选
query = "列出价格(price)最高的5个产品，显示名称和价格"
print(f"\n📝 查询: {query}")

response = query_engine.query(query)
print(f"💬 回答: {response}")

# 展示Pandas查询结果
df[['name', 'price']].sort_values(by='price', ascending=False).head(5)

class AdvancedPandasQueryEngine:
    """增强的 Pandas 查询引擎 - 带错误处理和性能监控"""
    
    def __init__(self, df: pd.DataFrame, llm_model: str = "gpt-4o-mini"):
        self.df = df
        self.llm = OpenAI(model=llm_model, temperature=0)
        
        # 优化的指令（中文提示效果更好）
        instruction_str = (
            "你是数据分析专家，将自然语言转换为 Pandas 代码。\n"
            "要求:\n"
            "1. 生成单行可 eval() 执行的表达式\n"
            "2. DataFrame 变量名为 'df'\n"
            "3. 只输出代码，无其他文字\n"
            "4. 优先使用向量化操作\n"
        )
        
        self.query_engine = PandasQueryEngine(
            df=df, verbose=True,
            instruction_str=instruction_str,
            synthesize_response=True
        )
        self.query_history = []
    
    def query(self, query_str: str, max_retries: int = 2) -> Dict[str, Any]:
        """执行查询，带重试和性能记录"""
        start_time = time.time()
        
        for attempt in range(max_retries + 1):
            try:
                logger.info(f"执行查询 (尝试 {attempt + 1}): {query_str}")
                response = self.query_engine.query(query_str)
                execution_time = time.time() - start_time
                
                self.query_history.append({
                    "query": query_str,
                    "time": execution_time,
                    "success": True
                })
                
                return {"success": True, "result": str(response), "time": execution_time}
                
            except Exception as e:
                if attempt == max_retries:
                    self.query_history.append({
                        "query": query_str,
                        "time": time.time() - start_time,
                        "success": False,
                        "error": str(e)
                    })
                    return {"success": False, "error": str(e)}
    
    def get_stats(self) -> Dict[str, Any]:
        """获取性能统计"""
        if not self.query_history:
            return {"message": "没有查询历史"}
        
        successful = [q for q in self.query_history if q["success"]]
        return {
            "total": len(self.query_history),
            "success": len(successful),
            "success_rate": f"{len(successful)/len(self.query_history)*100:.1f}%",
            "avg_time": f"{sum(q['time'] for q in successful)/len(successful):.2f}s" if successful else "N/A"
        }

# 初始化高级引擎
advanced_engine = AdvancedPandasQueryEngine(df)
print("✅ AdvancedPandasQueryEngine 初始化完成")

# 进阶查询示例
advanced_queries = [
    "计算每个品牌(brand)的产品数量，按数量降序排列",
    "找出评分(rating)高于4.5的产品名称和价格",
    "显示各类别(category)的平均价格"
]

print("=" * 60)
print("📊 进阶查询测试")
print("=" * 60)

for query in advanced_queries:
    print(f"\n📝 查询: {query}")
    result = advanced_engine.query(query)
    if result["success"]:
        print(f"💬 回答: {result['result']}")
        print(f"⏱️  耗时: {result['time']:.2f}s")
    else:
        print(f"❌ 错误: {result['error']}")

# 显示统计
print("\n" + "=" * 60)
print("📈 性能统计:")
stats = advanced_engine.get_stats()
for k, v in stats.items():
    print(f"   {k}: {v}")

class ChainOfTablesEngine:
    """表格思维链引擎 - 多步骤推理（修正版）"""
    
    def __init__(self, df: pd.DataFrame):
        self.original_df = df.copy()    # 保存原始数据
        self.current_df = df.copy()    # 当前工作表
        self.operation_history = []    # 操作历史
        self.llm = Settings.llm        # LLM
    
    def filter(self, condition: str) -> pd.DataFrame:
        """筛选操作 - 真正更新 DataFrame"""
        print(f"[筛选] {condition}")
        rows_before = len(self.current_df)
        
        # 构造prompt，让LLM生成筛选代码
        prompt = f"""
        你是一位 Pandas 专家。用户想要筛选数据。
        当前 DataFrame 的列名是: {list(self.current_df.columns)}
        用户的筛选条件是: {condition}

        请生成一行 Python 代码来完成筛选，要求：
        1. 变量名必须是 'df'
        2. 返回筛选后的 DataFrame
        3. 只输出代码，不要有任何解释

        示例: df[df['price'] > 5000]
        """
        
        # 调用LLM生成代码
        response = self.llm.complete(prompt)
        generated_code = response.text.strip()
        
        # 清理代码（去掉可能的markdown标记）
        if generated_code.startswith("```"):
            generated_code = generated_code.split("\n")[1]
        if generated_code.endswith("```"):
            generated_code = generated_code.rsplit("\n", 1)[0]
        generated_code = generated_code.strip()
        
        print(f"生成的代码: {generated_code}")
        
        # 执行代码并更新 current_df
        try:
            df = self.current_df  # 为了代码能找到 df 变量
            filtered_df = eval(generated_code)   # 执行代码
            self.current_df = filtered_df    # 更新当前工作表
            print(f"✅ 筛选成功: {rows_before} 行 → {len(self.current_df)} 行")
        except Exception as e:
            print(f"❌ 代码执行失败: {e}")
            print(f"保持原数据不变")
        
        # 记录操作
        self.operation_history.append({
            'step': len(self.operation_history) + 1,
            'operation': 'FILTER',
            'description': condition,
            'code': generated_code,
            'rows': f"{rows_before} → {len(self.current_df)}"
        })
        
        return self.current_df
    
    def aggregate(self, aggregation: str) -> str:
        """聚合操作 - 基于当前表计算"""
        print(f"[聚合] {aggregation}")
        
        # 构造prompt
        prompt = f"""
        你是一位 Pandas 专家。
        当前 DataFrame 有 {len(self.current_df)} 行数据，列名是: {list(self.current_df.columns)}
        用户想要: {aggregation}

        请生成一行 Python 代码来完成计算，要求：
        1. 变量名必须是 'df'
        2. 返回计算结果（数值或Series）
        3. 只输出代码，不要解释

        示例: df['rating'].mean()
        """
        
        response = self.llm.complete(prompt)   # 调用LLM生成代码
        generated_code = response.text.strip()  # 去除前后的空格
        
        # 清理代码
        if generated_code.startswith("```"):
            generated_code = generated_code.split("\n")[1]
        if generated_code.endswith("```"):
            generated_code = generated_code.rsplit("\n", 1)[0]
        generated_code = generated_code.strip()
        
        print(f"生成的代码: {generated_code}")
        
        # 执行计算
        try:
            df = self.current_df
            result = eval(generated_code)
            print(f"✅ 计算结果: {result}")
        except Exception as e:
            result = f"计算失败: {e}"
            print(f"❌ {result}")
        
        # 记录操作
        self.operation_history.append({
            'step': len(self.operation_history) + 1,
            'operation': 'AGGREGATE',
            'description': aggregation,
            'code': generated_code,
            'result': str(result)[:100]
        })
        
        return str(result)
    
    def reset(self):
        """重置到原始表"""
        print("🔄 [重置] 恢复原始表格")
        self.current_df = self.original_df.copy()
        self.operation_history = []
    
    def show_history(self):
        """显示操作历史"""
        print("\n📜 操作历史:")
        print("-" * 60)
        for op in self.operation_history:
            print(f"  步骤 {op['step']}: {op['operation']}")
            print(f"    描述: {op['description']}")
            print(f"    代码: {op.get('code', 'N/A')}")
            if 'rows' in op:
                print(f"    行数变化: {op['rows']}")
            if 'result' in op:
                print(f"    结果: {op['result']}")
            print()

# 重新初始化
cot_engine = ChainOfTablesEngine(df)
print("✅ ChainOfTablesEngine 初始化完成（修正版）")

print("=" * 60)
print("🔗 Chain of Tables 实战示例")
print("=" * 60)

# 场景: 分析高端产品
print("\n📌 任务: 找出高价产品的平均评分")
print("思维链: 筛选(价格>5000) → 聚合(平均评分)")

# 步骤1: 筛选
print("\n" + "-" * 40)
cot_engine.filter("价格(price)大于5000")

# 步骤2: 聚合
print("\n" + "-" * 40)
result = cot_engine.aggregate("计算评分(rating)的平均值")
print(f"\n🎯 最终结果: {result}")

# 显示完整历史
cot_engine.show_history()

# 重置并执行另一个分析
print("\n" + "=" * 60)
print("📌 任务: 各类别产品数量排名")
print("思维链: 分组 → 计数 → 排序")
print("=" * 60)

cot_engine.reset()
result = cot_engine.aggregate("按类别(category)分组，计算每个类别的产品数量，按降序排列")
print(f"\n🎯 结果: {result}")
cot_engine.show_history()

# 数据验证
df1 = df[df['price'] > 5000]
df1

df1['rating'].mean()

# 将按类别(category)分组，计算每个类别的产品数量，按降序排列，第二个任务的结果获取出
query_result = df.groupby('category').size().sort_values(ascending=False)
query_result

#可以通过matplotlib来对生成的Pandas数据框进行可视化
import matplotlib.pyplot as plt
import platform

# 1. 解决中文显示问题
system_name = platform.system()
if system_name == 'Darwin':  # macOS
    plt.rcParams['font.sans-serif'] = ['Arial Unicode MS']
elif system_name == 'Windows': # Windows
    plt.rcParams['font.sans-serif'] = ['SimHei']
else: # Linux
    plt.rcParams['font.sans-serif'] = ['WenQuanYi Micro Hei']
plt.rcParams['axes.unicode_minus'] = False 

# 2. 定义标题
title = '各类别产品数量分布'

# 3. 绘图
plt.figure(figsize=(10, 6))
ax = query_result.plot(kind='bar', color='steelblue')
plt.title(title)
plt.xlabel('类别')
plt.ylabel('数量')
plt.xticks(rotation=45, ha='right')

# 4. 添加数值标签
for p in ax.patches:
    ax.annotate(str(p.get_height()), (p.get_x() + p.get_width() / 2., p.get_height()),
                ha='center', va='center', xytext=(0, 5), textcoords='offset points')

plt.tight_layout()
plt.show()

# 示例：向量检索实现
from llama_index.core import VectorStoreIndex, Document
import pandas as pd
from sqlalchemy import create_engine

# 数据库连接
db_url = "sqlite:///dataset/ecommerce.db"
engine = create_engine(db_url)

# 加载数据
df = pd.read_sql("SELECT * FROM products LIMIT 15", engine)
print(df.head(5))

# 将每行转为文本
documents = []
for _, row in df.iterrows():
    text = f"""
    产品名称: {row.get('name', row.get('product_name', 'Unknown'))}
    类别: {row.get('category', 'N/A')}
    价格: {row.get('price', row.get('unit_price', 0))}元
    """
    documents.append(Document(text=text, metadata=row.to_dict()))

# 创建向量索引
vector_index = VectorStoreIndex.from_documents(
    documents,
    insert_batch_size=10  # ✅ DashScope同步接口限制：每批最多10条
)
vector_retriever = vector_index.as_retriever(similarity_top_k=3)

print(f"✅ 向量索引创建完成: {len(documents)}个文档")

documents

# 示例：分块向量化实现
from llama_index.core import VectorStoreIndex, Document
import pandas as pd
from sqlalchemy import create_engine

# 数据库连接
db_url = "sqlite:///dataset/ecommerce.db"
engine = create_engine(db_url)

# 加载数据
df = pd.read_sql("SELECT * FROM orders", engine)
print(f"总数据量: {len(df)}行")

# ========== 分块参数 ==========
CHUNK_SIZE = 10  # 每个分块包含的行数

# 将多行合并为一个文档（分块向量化）
documents = []
for chunk_start in range(0, len(df), CHUNK_SIZE):
    # 取出当前分块的数据
    chunk_df = df.iloc[chunk_start:chunk_start + CHUNK_SIZE]
    
    # 将分块内所有行拼接为文本
    text_parts = []
    for _, row in chunk_df.iterrows():
        item = f"产品id: {row.get('product_id')} | 用户id: {row.get('user_id')} | 数量: {row.get('quantity')} | 订单日期: {row.get('order_date')} | 订单金额: {row.get('amount')}元 | 订单状态: {row.get('status')} "
        text_parts.append(item)
    
    # 合并为单个文档
    chunk_text = "\n".join(text_parts)
    
    # 元数据记录分块信息
    metadata = {
        "chunk_id": chunk_start // CHUNK_SIZE,
        "row_range": f"{chunk_start}-{min(chunk_start + CHUNK_SIZE, len(df))}",
        "status": list(chunk_df['status'].unique())  # 分块包含的状态
    }
    
    documents.append(Document(text=chunk_text, metadata=metadata))

print(f"生成 {len(documents)} 个分块文档 (每块 {CHUNK_SIZE} 行)")

# 创建向量索引
vector_index = VectorStoreIndex.from_documents(
    documents,
    insert_batch_size=10  # DashScope限制
)
vector_retriever = vector_index.as_retriever(similarity_top_k=3)

print(f"✅ 分块向量索引创建完成")

print(documents[0])

"""
Milvus混合向量化示例

核心优势:
Milvus原生支持元数据(标量字段)过滤
性能优异,适合大规模生产环境
支持本地Lite模式(./milvus.db)和服务器模式
"""

from typing import Optional, Dict, Any, List, Tuple, Union
import pandas as pd
import numpy as np
from sqlalchemy import create_engine
from tqdm import tqdm

from llama_index.core import (
    VectorStoreIndex,
    Document,
    StorageContext,
    Settings
)
from llama_index.core.vector_stores import (
    MetadataFilters,
    MetadataFilter,
    FilterOperator
)
from llama_index.vector_stores.milvus import MilvusVectorStore
from llama_index.embeddings.dashscope import DashScopeEmbedding
import os
from dotenv import load_dotenv
from sqlalchemy import inspect


# 加载环境变量
load_dotenv()

class UniversalHybridVectorizer:
    """
    通用混合索引 - 使用Milvus支持元数据过滤
    
    核心功能:
    1. 支持CSV/Excel/数据库多种数据源
    2. 数据库自动推断字段：利用 sqlalchemy.inspect 扫描表结构，自动将字段归类：
            长文本（Text/LongText）：自动划为 text_fields，用于生成向量，提供语义搜索能力。
            数值/日期（Int/Float/Date）：自动划为 metadata_fields，用于范围查询（如“价格>5000”）。
            短文本（Tags/Category）：自动划为 metadata_fields，用于精确匹配（如“品牌=Apple”）。
    3. 使用Milvus实现高性能向量检索 + 标量过滤
    """
    
    def __init__(
        self, 
        db_url: Optional[str] = None,
        table: Optional[str] = None,
        # ===== 可选：手动覆盖自动推断 =====
        text_fields: Optional[List[str]] = None,
        metadata_fields: Optional[List[str]] = None,
        id_field: Optional[str] = None,
        numeric_fields: Optional[List[str]] = None,
        categorical_fields: Optional[List[str]] = None
    ):
        """
        初始化向量化器
        
        方式1（推荐）：传入数据库连接，自动推断字段
            vectorizer = UniversalHybridVectorizer(
                db_url="sqlite:///ecommerce.db",
                table="products"
            )
        
        方式2：手动指定所有字段（用于CSV/Excel或需要精确控制）
            vectorizer = UniversalHybridVectorizer(
                text_fields=["name", "description"],
                metadata_fields=["price", "category"]
            )
        
        Args:
            db_url: 数据库连接字符串（可选）
            table: 表名，用于自动推断字段（可选）
            text_fields: 需要向量化的字段列表
            metadata_fields: 作为元数据的字段列表
            id_field: 唯一标识符字段名
            numeric_fields: 数值字段列表（用于范围查询）
            categorical_fields: 类别字段列表（用于精确/IN查询）
        """
        self.engine = None
        self.table = table
        
        # 如果提供了数据库连接，创建引擎
        if db_url:
            self.engine = create_engine(db_url)
            print(f"🔌 连接数据库: {db_url}")
        
        # 自动推断字段（如果提供了db_url和table且未手动指定字段）
        if self.engine and table and not text_fields:
            auto_fields = self._auto_detect_fields(table)
            self.text_fields = text_fields or auto_fields.get('text', [])
            self.metadata_fields = metadata_fields or auto_fields.get('metadata', [])
            self.id_field = id_field or auto_fields.get('id')
            self.numeric_fields = numeric_fields or auto_fields.get('numeric', [])
            self.categorical_fields = categorical_fields or auto_fields.get('categorical', [])
            
            print(f"📊 自动推断字段配置:")
            print(f"   文本字段: {self.text_fields}")
            print(f"   数值字段: {self.numeric_fields}")
            print(f"   分类字段: {self.categorical_fields}")
            print(f"   ID字段: {self.id_field}")
        else:
            # 使用手动指定的字段
            self.text_fields = text_fields or []
            self.metadata_fields = metadata_fields or []
            self.id_field = id_field
            self.numeric_fields = numeric_fields or []
            self.categorical_fields = categorical_fields or []
        
        self.index = None
    
    def _auto_detect_fields(self, table: str) -> Dict[str, Any]:
        """
        从数据库表结构自动推断字段分类
        
        分类规则：
        - TEXT, LONGTEXT, VARCHAR(255+) → text_fields
        - INTEGER, FLOAT, DECIMAL → numeric_fields + metadata_fields
        - VARCHAR(<255) → categorical_fields + metadata_fields
        - PRIMARY KEY → id_field
        """
        
        # 获取表结构
        inspector = inspect(self.engine)

        # 获取表的列信息
        columns = inspector.get_columns(table)

        # 获取主键约束
        pk_constraint = inspector.get_pk_constraint(table)
        
        result = {
            'text': [],
            'metadata': [],
            'numeric': [],
            'categorical': [],
            'id': None
        }
        
        # 获取主键
        if pk_constraint and pk_constraint.get('constrained_columns'):
            result['id'] = pk_constraint['constrained_columns'][0]
        
        for col in columns:
            col_name = col['name']
            col_type = str(col['type']).upper()
            
            # 跳过主键（已单独处理）
            if col_name == result['id']:
                continue
            
            # 根据类型分类
            if 'TEXT' in col_type or 'LONGTEXT' in col_type:
                result['text'].append(col_name)
            elif 'VARCHAR' in col_type:
                # 解析 VARCHAR 长度
                import re
                match = re.search(r'VARCHAR\((\d+)\)', col_type)
                length = int(match.group(1)) if match else 50
                
                if length >= 255:
                    result['text'].append(col_name)
                else:
                    result['categorical'].append(col_name)
                    result['metadata'].append(col_name)
            elif any(t in col_type for t in ['INT', 'FLOAT', 'DECIMAL', 'NUMERIC', 'REAL']):
                result['numeric'].append(col_name)
                result['metadata'].append(col_name)
            elif any(t in col_type for t in ['DATE', 'TIME', 'DATETIME', 'TIMESTAMP']):
                result['metadata'].append(col_name)
            else:
                # 其他类型默认加入metadata
                result['metadata'].append(col_name)
        
        return result
    
    def load_from_database(
        self,
        table: Optional[str] = None,
        query: Optional[str] = None,
        limit: Optional[int] = None,
        chunk_size: int = 1000
    ) -> List[Document]:
        """
        从数据库加载数据
        
        Args:
            table: 表名（默认使用初始化时的表）
            query: 自定义SQL查询（支持JOIN，优先于table）
            limit: 限制返回行数
            chunk_size: 分批读取的行数
            
        Examples:
            # 简单用法：加载初始化时指定的表
            docs = vectorizer.load_from_database()
            
            # 指定表和限制
            docs = vectorizer.load_from_database(table="products", limit=1000)
            
            # 自定义SQL（支持JOIN）
            docs = vectorizer.load_from_database(
                query=\"\"\"
                    SELECT p.*, r.content as review_content 
                    FROM products p 
                    LEFT JOIN reviews r ON p.id = r.product_id
                \"\"\"
            )
        """
        if not self.engine:
            raise ValueError("未配置数据库连接，请在初始化时传入 db_url")
        
        # 确定要使用的表
        target_table = table or self.table
        
        # 构建SQL
        if query:
            sql = query
            print(f"🗄️  执行自定义查询...")
            # 检测查询中的新列并尝试自动推断
            self._detect_new_columns_from_query(query)
        elif target_table:
            sql = f"SELECT * FROM {target_table}"
            if limit:
                sql += f" LIMIT {limit}"
            print(f"🗄️  加载表: {target_table}" + (f" (LIMIT {limit})" if limit else ""))
        else:
            raise ValueError("必须指定 table 或 query")
        
        # 分批读取
        all_documents = []
        chunks = pd.read_sql(sql, self.engine, chunksize=chunk_size)
        
        for df_chunk in tqdm(chunks, desc="处理数据库数据"):
            documents = self._convert_to_hybrid_documents(df_chunk)
            all_documents.extend(documents)
        
        print(f"✅ 生成 {len(all_documents)} 个文档")
        return all_documents
    
    def _convert_to_hybrid_documents(self, df: pd.DataFrame) -> List[Document]:
        """
        “结构化”转“非结构化”的关键中间层，转换为混合文档
        
        策略：
        1. 文本化（Text Parts）：将所有 text_fields 的内容拼接成一段自然语言描述（例如："name: iPhone 15\ndescription: 最新款苹果手机..."），
            这部分会被 Embedding 模型转化为向量。
        2. 元数据化（Metadata）：将所有 metadata_fields（如价格、类别、库存）保留为原始键值对（Key-Value），这部分会被存入 Milvus 的标量字段中，
            用于后续的 Filter 操作。
        这实现了**“语义在向量里，属性在字段里”**的双重存储结构。
        """
        documents = []
        
        for idx, row in df.iterrows():
            # 1. 构建向量化文本
            text_parts = []
            for field in self.text_fields:
                if field in row and pd.notna(row[field]):
                    text_parts.append(f"{field}: {row[field]}")
            
            text = "\n".join(text_parts) if text_parts else "N/A"
            
            # 2. 构建元数据
            metadata = {}
            for field in self.metadata_fields:
                if field in row and pd.notna(row[field]):
                    value = row[field]
                    # 处理pandas特殊类型
                    if isinstance(value, (pd.Timestamp, pd.Timedelta)):
                        value = str(value)
                    metadata[field] = value
            
            # 3. 确定文档ID
            if self.id_field and self.id_field in row:
                doc_id = str(row[self.id_field])
            else:
                doc_id = str(idx)
            
            # 4. 创建文档
            doc = Document(
                text=text,
                metadata=metadata,
                doc_id=doc_id
            )
            documents.append(doc)
        
        return documents
    
    def build_index(self, documents: List[Document], 
                      save_path: str = "./hybrid_index",
                      collection_name: str = "hybrid_collection",
                      milvus_uri: str = "./milvus_demo.db"):
        """
        构建混合索引（使用Milvus支持元数据过滤）
        
        Args:
            documents: 文档列表
            save_path: 索引保存路径（兼容参数，实际数据存储在Milvus）
            collection_name: Milvus集合名称
            milvus_uri: Milvus连接URI
                - 本地模式: "./milvus_demo.db" (使用Milvus Lite)
                - 服务器模式: "http://localhost:19530"
        """
        print(f"🔨 构建混合索引，文档数: {len(documents)}")
        
        from llama_index.vector_stores.milvus import MilvusVectorStore
        
        # 使用Milvus构建向量存储（支持元数据过滤！）
        # dim=1024 对应 DashScope text-embedding-v4 的向量维度
        vector_store = MilvusVectorStore(
            uri=milvus_uri,            # Milvus连接URI
            collection_name=collection_name,  # Milvus集合名称
            dim=1024,          # 向量维度，最好跟Embedding模型长度一致
            overwrite=True  # 如果存在同名collection则覆盖
        )
        
        # 使用MilvusVectorStore构建存储上下文
        storage_context = StorageContext.from_defaults(
            vector_store=vector_store
        )
        
        # 使用VectorStoreIndex构建向量索引
        self.index = VectorStoreIndex.from_documents(
            documents,
            storage_context=storage_context,  # 使用MilvusVectorStore
            show_progress=True,   # 显示进度
            insert_batch_size=10  # ✅ DashScope同步接口限制：每批最多10条
        )
        
        # Milvus数据是自托管的,不需要persist到文件系统
        print(f"✅ 索引已创建: collection={collection_name}, uri={milvus_uri}")
        
        vector_retriever = self.index.as_retriever(similarity_top_k=5)
        return vector_retriever
    
    

print("\n" + "="*60)
print("🚀 Milvus混合索引示例 - 电商产品数据")
print("="*60)

# 1. 创建向量化器
vectorizer = UniversalHybridVectorizer(
        db_url="sqlite:///dataset/ecommerce.db",
        table="products"
    )
# 自动打印：文本字段: ['description'], 数值字段: ['price', 'stock']...

# 2. 加载数据
docs = vectorizer.load_from_database()  # 使用初始化时的表

# 3. 构建Milvus检索器
vector_retriever = vectorizer.build_index(
    docs, 
    collection_name="products",
    milvus_uri="http://localhost:19530"
)

from llama_index.core.schema import TextNode
from llama_index.retrievers.bm25 import BM25Retriever

import pandas as pd
from sqlalchemy import create_engine

# 数据库连接
db_url = "sqlite:///dataset/ecommerce.db"
engine = create_engine(db_url)

# 加载 products 表数据
df = pd.read_sql("SELECT * FROM products", engine)
print(f"总数据量: {len(df)}行")

# 直接创建 Nodes（跳过 Documents）
nodes = []
for idx, row in df.iterrows():
    # 将所有字段都添加到 text 中
    text = f"""产品ID: {row.get('product_id')}
            产品名称: {row.get('name')}
            类别: {row.get('category')}
            价格: {row.get('price')}元
            库存: {row.get('stock')}件
            评分: {row.get('rating')}
            描述: {row.get('description')}"""

    node = TextNode(
        text=text,
        metadata=row.to_dict(),
        id_=f"product_{idx}"
    )
    nodes.append(node)

# 创建BM25 检索器
bm25_retriever = BM25Retriever.from_defaults(
    nodes=nodes,
    similarity_top_k=5
)

print(f"✅ 已创建 {len(nodes)} 个产品节点")