# 查看 Python 版本
!python --version

import os
from pathlib import Path
from dotenv import load_dotenv

# 加载环境变量，通常用于配置 API 密钥和存储路径
load_dotenv(override=True)

def check_dependencies():
    """
    检查运行多模态检索所需的必要 Python 依赖库是否已安装。
    
    返回:
        bool: 如果所有依赖都已安装返回 True，否则返回 False。
    """
    # 定义需要检查的库：(模块名, pip安装名)
    required_packages = [
        ("llama_index.core", "llama-index"),
        ("llama_index.embeddings.clip", "llama-index-embeddings-clip"),
        ("pymilvus", "pymilvus"),
    ]
    
    missing = []
    # 遍历列表，尝试导入每个模块以验证其是否存在
    for module_name, package_name in required_packages:
        try:
            __import__(module_name)
            print(f"✅ {package_name} 已安装")
        except ImportError:
            missing.append(package_name)
            print(f"❌ {package_name} 未安装")
    
    # 检查 OpenAI 的 CLIP 库，该库通常需要从 GitHub 源码安装
    try:
        import clip
        print("✅ CLIP 已安装")
    except ImportError:
        print("❌ CLIP 未安装")
        print("   请运行: pip install git+https://github.com/openai/CLIP.git")
        missing.append("CLIP")
    
    # 如果有任何缺失的包，提示用户
    if missing:
        print(f"\n⚠️ 请安装缺失的依赖")
        return False
    
    print("\n🎉 所有依赖检查通过!")
    return True

# 执行依赖项检查函数
check_dependencies()

def print_package_versions():
    """打印核心依赖包版本信息"""
    import importlib.metadata
    
    packages = [
        "llama-index",
        "llama-index-embeddings-clip",
        "pymilvus",
        "llama-index-vector-stores-milvus",
        "pillow"
    ]
    
    print("\n📦 核心依赖包版本信息:")
    print("-" * 50)
    
    for package in packages:
        try:
            version = importlib.metadata.version(package)
            print(f"  {package}: {version}")
        except importlib.metadata.PackageNotFoundError:
            print(f"  {package}: 未安装")
    
    # 特殊处理 CLIP
    try:
        import clip
        # CLIP 没有标准版本号，打印安装状态
        print(f"  CLIP (OpenAI): 已安装")
    except ImportError:
        print(f"  CLIP (OpenAI): 未安装")
    
    print("-" * 50)

# 打印版本信息
print_package_versions()

import os
from dotenv import load_dotenv

# 加载环境变量（.env 文件中需要配置 OPENAI_API_KEY
load_dotenv(override=True)

# 验证 API Key
api_key = os.getenv("OPENAI_API_KEY")
base_url = os.getenv("OPENAI_BASE_URL")
if api_key:
    print(f"✅ OPENAI_API_KEY 已配置: {api_key[:20]}...")
else:
    raise ValueError("❌ 请设置 OPENAI_API_KEY 环境变量")

from llama_index.core import Settings
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding

# 配置 LLM（大语言模型）
Settings.llm = OpenAI(
    model="gpt-4o",           # 使用 GPT-4o 模型
    temperature=0,             # 降低随机性，输出更稳定
    api_key=api_key,
    api_base=base_url          # 可选：使用代理 API
)

# 简单文本测试
response = Settings.llm.complete("你好，请用一句话介绍你自己。")
print("LLM 响应:")
print(response.text)

# 在存在docker-compose.yml文件相同的目录下，启动Milvus
docker-compose up -d

def check_milvus_connection():
    """检查 Milvus 连接是否正常"""
    try:
        from pymilvus import connections, utility
        
        MILVUS_HOST = "localhost"
        MILVUS_PORT = 19530
        
        connections.connect(
            alias="default",
            host=MILVUS_HOST,
            port=MILVUS_PORT
        )
        
        # 测试连接
        version = utility.get_server_version()
        print(f"✅ Milvus 连接成功 (版本: {version})")
        
        connections.disconnect("default")
        return True
        
    except Exception as e:
        print(f"❌ Milvus 连接失败: {e}")
        print(f"   请确保 Milvus 正在运行: docker-compose up -d")
        return False

# 执行检查
check_milvus_connection()

from llama_index.embeddings.clip import ClipEmbedding

# 初始化 CLIP 嵌入模型
embed_model = ClipEmbedding()
print("✅ ClipEmbedding 初始化完成")

# 文本向量编码
text_vec = embed_model.get_text_embedding("一只橘色的猫")
print(f"文本向量维度: {len(text_vec)}")
print(f"文本向量前5维: {text_vec[:5]}")

# 图像向量编码（假设有一张猫的图片）
img_vec = embed_model.get_image_embedding("/Users/mac/大模型资料/多模态RAG/七大场景llama-Index进阶实战/测试文档&图片/test_images/Image_cvvlc4cvvlc4cvvl.png")
print(f"图像向量维度: {len(img_vec)}")
print(f"文本向量前5维: {img_vec[:5]}")

from llama_index.core import SimpleDirectoryReader
from pathlib import Path

# 配置图片目录
BASE_DIR = Path.cwd()
image_dir = BASE_DIR / "images"  # 请替换为你的图片目录路径

# 加载图片文档
documents = SimpleDirectoryReader(
    input_dir=str(image_dir),
    required_exts=[".png", ".jpg", ".jpeg", ".gif", ".webp"]
).load_data()

print(f"📷 加载了 {len(documents)} 个图片文档")

from llama_index.core.indices import MultiModalVectorStoreIndex

# 构建多模态索引
index = MultiModalVectorStoreIndex.from_documents(
    documents,
    image_embed_model=embed_model
)
print("✅ MultiModalVectorStoreIndex 构建完成")

# 加载打印图片的方法
from utils_image_display import display_retrieval_results,display_local_image

# 创建检索器
retriever = index.as_retriever(
    similarity_top_k=3,
    image_similarity_top_k=3
)

# 执行文搜图检索
query = "架构图"
results = retriever.text_to_image_retrieve(query)

print(f"\n🔍 查询: \"{query}\"")
print("-" * 50)
for i, result in enumerate(results, 1):
    file_path = result.node.metadata.get("file_path", "未知")
    score = result.score if hasattr(result, 'score') else 0.0
    print(f"  [{i}] {Path(file_path).name}: {score:.4f}")

# 3. 显示结果
display_retrieval_results(results, title="文搜图: 架构图", max_display=6)

test_queries = [
    ("架构图", "architecture diagram"),
    ("流程图", "flowchart"),
]

for cn_query, en_query in test_queries:
    print(f"\n中文查询: \"{cn_query}\" | 英文查询: \"{en_query}\"")
    print("-" * 40)
    
    cn_results = retriever.text_to_image_retrieve(cn_query)
    en_results = retriever.text_to_image_retrieve(en_query)
    
    if cn_results:
        cn_top_score = cn_results[0].score if hasattr(cn_results[0], 'score') else 0.0
        print(f"  中文Top1分数: {cn_top_score:.4f}")
    
    if en_results:
        en_top_score = en_results[0].score if hasattr(en_results[0], 'score') else 0.0
        print(f"  英文Top1分数: {en_top_score:.4f}")
        print(f"  提升幅度: {(en_top_score - cn_top_score):.4f}")

from pathlib import Path

# 配置图片目录
BASE_DIR = Path.cwd()
image_dir = BASE_DIR / "images"  # 请替换为你的图片目录路径

# 选择一张查询图片
image_files = list(image_dir.glob("*.png")) + list(image_dir.glob("*.jpg"))
if image_files:
    query_image = image_files[0]
    
    print(f"\n🔍 图搜图查询: {query_image.name}")
    print("-" * 50)
    
    # 执行图搜图检索
    retriever_img = index.as_retriever(image_similarity_top_k=4)
    results = retriever_img.image_to_image_retrieve(str(query_image))
    
    for i, result in enumerate(results, 1):
        file_path = result.node.metadata.get("file_path", "未知")
        score = result.score if hasattr(result, 'score') else 0.0
        print(f"  [{i}] {Path(file_path).name}: {score:.4f}")

    # 3. 显示结果
    display_retrieval_results(results, title="图搜图", max_display=6)

# 优化后的图搜图：过滤查询图片本身
query_image_path = Path(query_image)
results = retriever_img.image_to_image_retrieve(str(query_image))

filtered_results = []
for result in results:
    file_path = Path(result.node.metadata.get("file_path", ""))
    
    # 跳过查询图片本身
    if file_path.resolve() == query_image_path.resolve():
        print(f"  [跳过] {file_path.name} (查询图片本身, 分数: {result.score:.4f})")
        continue
    
    filtered_results.append(result)
    print(f"  [{len(filtered_results)}] {file_path.name}: {result.score:.4f}")

# 3. 显示结果
display_retrieval_results(filtered_results, title="图搜图", max_display=6)

from llama_index.core import StorageContext
from llama_index.vector_stores.milvus import MilvusVectorStore

# 配置 Milvus 向量存储，链接地址，可换成生产环境地址
MILVUS_URI = "http://localhost:19530"

# 文本集合名称
TEXT_COLLECTION = "multimodal_text_collection"

# 图片集合名称
IMAGE_COLLECTION = "multimodal_image_collection"

# 文本向量存储
text_store = MilvusVectorStore(
    uri=MILVUS_URI,
    collection_name=TEXT_COLLECTION,
    dim=512,  # CLIP 文本向量维度
    overwrite=True,  # 重新创建集合
)

# 图像向量存储
image_store = MilvusVectorStore(
    uri=MILVUS_URI,
    collection_name=IMAGE_COLLECTION,
    dim=512,  # CLIP 图像向量维度
    overwrite=True,
)

# 创建存储上下文
storage_context = StorageContext.from_defaults(
    vector_store=text_store,  # 这里没有用到文本数据，只作为占位符
    image_store=image_store
)
print("✅ Milvus 向量存储配置完成")

# 构建持久化索引
index_milvus = MultiModalVectorStoreIndex.from_documents(
    documents,
    storage_context=storage_context,
    embed_model=embed_model,
    image_embed_model=embed_model,
)
print("✅ MultiModalVectorStoreIndex 已持久化到 Milvus")

# 测试持久化后的检索
retriever_milvus = index_milvus.as_retriever(image_similarity_top_k=3)

# 修正：Standard CLIP 模型对中文支持较弱，改用英文 Prompt 以修复相关性问题
query_text = "architecture diagram"
results = retriever_milvus.text_to_image_retrieve(query_text)

print(f"\n Milvus持久化索引检索结果 (Query: {query_text}):")
for i, result in enumerate(results, 1):
    file_path = result.node.metadata.get("file_path", "未知")
    score = result.score if hasattr(result, 'score') else 0.0
    print(f"  [{i}] {Path(file_path).name}: {score:.4f}")

# 显示结果
display_retrieval_results(results, title=f"Retrieval: {query_text}", max_display=6)

# 加载已有索引（程序重启后）
from llama_index.core.indices import MultiModalVectorStoreIndex

# 注意：这里直接使用上面定义好的 text_store 和 image_store
# 如果是完全重启环境，需要重新定义 MilvusVectorStore(..., overwrite=False)
index_reload = MultiModalVectorStoreIndex.from_vector_store(
    vector_store=text_store,
    image_vector_store=image_store,
    image_embed_model=embed_model
)
print("✅ 已成功从 Milvus 加载 MultiModalVectorStoreIndex")

# 创建检索器
retriever_reload = index_reload.as_retriever()

# 执行文搜图检索
reload_results = retriever_reload.text_to_image_retrieve("architecture diagram")

print(f"\n Milvus持久化索引检索结果 (Query: architecture diagram):")
for i, result in enumerate(reload_results, 1):
    file_path = result.node.metadata.get("file_path", "未知")
    score = result.score if hasattr(result, 'score') else 0.0
    print(f"  [{i}] {Path(file_path).name}: {score:.4f}")

# 显示结果
display_retrieval_results(reload_results, title=f"Retrieval: {query_text}", max_display=6)

# 执行图搜图检索
images_results = retriever_reload.image_to_image_retrieve("images/Image_gkvpywgkvpywgkvp.png")

print(f"\n Milvus持久化索引检索结果 (Query: architecture diagram):")
for i, result in enumerate(images_results, 1):
    file_path = result.node.metadata.get("file_path", "未知")
    score = result.score if hasattr(result, 'score') else 0.0
    print(f"  [{i}] {Path(file_path).name}: {score:.4f}")

# 显示结果
display_retrieval_results(images_results, title=f"Retrieval: {query_text}", max_display=6)

class VLMSelector:
    """VLM 自适应选择器 - 根据图片数量智能选择最优VLM"""
    
    SMALL_SCALE_THRESHOLD = 50  # 小规模阈值
    
    @staticmethod
    def select_vlm(
        num_images: int,
        user_preference: str = "auto"
    ):
        """
        根据图片数量自动选择VLM
        
        Returns:
            (vlm_model, reason)  # 选择的VLM和原因说明
        """
        if user_preference != "auto":
            return user_preference, f"用户指定: {user_preference}"
        
        # 小规模（< 50张）：GPT-4o（效果最好）
        if num_images < VLMSelector.SMALL_SCALE_THRESHOLD:
            cost_estimate = f"约${num_images * 0.003:.2f}"
            return "gpt-4o", f"小规模({num_images}张) → GPT-4o (最佳效果, 成本{cost_estimate})"
        
        # 大规模（>= 50张）：Qwen-VL-Max（性价比）
        else:
            cost_estimate = f"约¥{num_images * 0.01:.2f}"
            return "qwen-vl-max", f"大规模({num_images}张) → Qwen-VL-Max (性价比, 成本{cost_estimate})"

# 使用示例
num_images = 10
vlm_model, reason = VLMSelector.select_vlm(num_images)
print(f"✅ {reason}")

from dotenv import load_dotenv
load_dotenv(override=True)

from llama_index.core.llms import ChatMessage, ImageBlock, TextBlock
from llama_index.llms.openai import OpenAI

def generate_image_caption_gpt4o(image_path: str) -> str:
    """使用 GPT-4o 生成图片描述"""
    
    # Prompt模板
    prompt_text = """请详细描述这张图片，包括以下信息：
    1. 图片的主要内容和主题
    2. 图片中的所有可见文字（如有）
    3. 如果是图表，描述数据趋势和关键数值
    4. 颜色、布局和设计风格
    5. 任何其他值得注意的视觉元素

    请用中文回答，描述要详细具体，便于后续检索。"""
    
    # 构造ChatMessage
    message = ChatMessage(
        role="user",
        blocks=[
            ImageBlock(path=image_path),
            TextBlock(text=prompt_text)
        ]
    )
    
    # 使用Settings.llm设置的全局LLM，Gpt-4o模型
    response = Settings.llm.chat([message])
    return response.message.content

# 测试单张图片
image_path = "images/20251216110324_89_50.png"
caption = generate_image_caption_gpt4o(image_path)
print(f"描述: {caption}")

# 显示结果
display_local_image(image_path, title=f"Retrieval: 20251216110324_89_50.png")

from pathlib import Path

CAPTION_CACHE_DIR = Path("./caption_cache")
CAPTION_CACHE_DIR.mkdir(exist_ok=True)

def batch_generate_captions(
    image_dir: Path,          # 图片目录路径
    max_images: int = 10,     # 最大处理图片数量
    vlm_model: str = "gpt-4o" # 使用的视觉语言模型名称
):
    """批量为图片生成描述（带缓存）"""
    # 获取图片列表
    image_files = list(image_dir.glob("*.png")) + list(image_dir.glob("*.jpg"))
    # 取最大图片数量
    image_files = image_files[:max_images]
    
    results = []
    
    # 迭代处理选定的图片文件
    for i, image_path in enumerate(image_files, 1):
        print(f"\n[{i}/{len(image_files)}] 处理: {image_path.name}")
        
        # 定义缓存文件路径 (文件名_模型名.txt)
        cache_file = CAPTION_CACHE_DIR / f"{image_path.stem}_{vlm_model}.txt"
        
        # 优先从本地缓存读取，避免重复调用 API 产生费用
        if cache_file.exists():
            print(f"  📂 使用缓存描述")
            caption = cache_file.read_text(encoding="utf-8")
        else:
            print(f"  🚀 调用 {vlm_model} 生成描述...")
            # 核心步骤：调用视觉模型生成描述
            caption = generate_image_caption_gpt4o(str(image_path))
            
            # 生成成功后写入缓存文件供下次使用
            if caption:
                cache_file.write_text(caption, encoding="utf-8")
                print(f"  💾 已缓存到 {cache_file.name}")
        
        # 将有效的描述结果及元数据存入结果列表
        if caption:
            results.append({
                "image_path": str(image_path),
                "image_name": image_path.name,
                "caption": caption
            })
            print(f"  ✅ 描述长度: {len(caption)} 字符")
            print(f"  📝 预览: {caption[:100]}...")
    
    return results

# 执行批量生成
captions = batch_generate_captions(
    image_dir=Path("images"),
    max_images=12,
    vlm_model="gpt-4o"
)

from llama_index.core import Settings
from llama_index.embeddings.openai import OpenAIEmbedding

# 显式重置 embed_model 为 OpenAI，防止沿用之前的 CLIP 模型
embed_model = OpenAIEmbedding(model="text-embedding-3-small")

# 设置全局嵌入模型，可以将之前的CLIP模型覆盖掉
Settings.embed_model = embed_model

from llama_index.core.schema import ImageNode
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.core import VectorStoreIndex
from typing import List

def build_caption_based_index(captions: List[dict]):
    """基于图片描述构建向量索引"""
    
    # 1. 创建 ImageNode 列表
    nodes = []
    for item in captions:
        node = ImageNode(
            text=item["caption"],  # 描述文本
            image_path=item["image_path"],  # 图片路径
            metadata={
                "file_name": item["image_name"],  # 图片文件名
                "caption_length": len(item["caption"])  # 描述文本长度
            }
        )
        nodes.append(node)
    
    print(f"  创建了 {len(nodes)} 个 ImageNode")
    
    # 2. 配置 OpenAI 嵌入模型
    embed_model = OpenAIEmbedding(model="text-embedding-3-small")
    print(f"  使用 text-embedding-3-small 嵌入模型（1536维）")
    
    # 3. 构建索引（暂用内存存储）
    index = VectorStoreIndex(
        nodes=nodes,
        embed_model=embed_model
    )
    
    print("✅ 基于描述的向量索引构建完成")
    return index ,nodes

# 构建索引
index,nodes = build_caption_based_index(captions)

from llama_index.core import StorageContext
from llama_index.vector_stores.milvus import MilvusVectorStore

# Milvus 配置
MILVUS_URI = "http://localhost:19530"

# 向量存储集合名称
CAPTION_COLLECTION = "vlm_caption_collection"

# 配置 Milvus 向量存储
vector_store = MilvusVectorStore(
    uri=MILVUS_URI,
    collection_name=CAPTION_COLLECTION,
    dim=1536,  # text-embedding-3-small 维度
    overwrite=True
)

# 创建存储上下文，并将 Milvus 向量存储配置其中
storage_context = StorageContext.from_defaults(
    vector_store=vector_store
)

# 构建持久化索引
index = VectorStoreIndex(
    nodes=nodes,
    storage_context=storage_context,
    embed_model=embed_model
)

print("✅ 向量索引已持久化到 Milvus")

from llama_index.core import VectorStoreIndex

# 加载已有索引
vector_store = MilvusVectorStore(
    uri=MILVUS_URI,
    collection_name=CAPTION_COLLECTION,
    dim=1536,
    overwrite=False  # 不覆盖已有集合
)

index_reload = VectorStoreIndex.from_vector_store(
    vector_store=vector_store,
    embed_model=embed_model
)
print("✅ 从 Milvus 加载已有索引")

# 创建查询引擎
query_engine = index_reload.as_query_engine()

# 查询
response = query_engine.query("我需要了解房屋架构图，给我描述一下，一个房屋的架构图都有什么东西？")

# 显示结果
print(response)

from llama_index.retrievers.bm25 import BM25Retriever
from llama_index.core import Document
from llama_index.core.node_parser import SentenceSplitter

def build_bm25_index(captions: List[dict]):
    """基于图片描述构建BM25索引（关键词检索）"""
    
    print("\n📊 构建BM25索引（关键词检索）...")
    
    # 1. 创建 Document 列表
    documents = []
    for item in captions:
        doc = Document(
            text=item["caption"],
            metadata={
                "image_path": item["image_path"],
                "file_name": item["image_name"]
            }
        )
        documents.append(doc)
    
    print(f"  创建了 {len(documents)} 个 Document")
    
    # 2. 将Documents转换为nodes
    splitter = SentenceSplitter(chunk_size=512)
    nodes = splitter.get_nodes_from_documents(documents)
    
    print(f"  转换为 {len(nodes)} 个 Node")
    
    # 3. 构建 BM25 检索器
    bm25_retriever = BM25Retriever.from_defaults(
        nodes=nodes,
        similarity_top_k=3
    )
    
    print("✅ BM25索引构建完成")
    
    return bm25_retriever

# 构建BM25索引
bm25_retriever = build_bm25_index(captions)

# 纯BM25检索
query = "卧室"
results = bm25_retriever.retrieve(query)

print(f"\n🔍 【纯BM25检索】查询: \"{query}\"")
print("-" * 50)

for i, result in enumerate(results[:3], 1):
    file_name = result.node.metadata.get("file_name", "未知")
    caption = result.node.text[:100] + "..."
    print(f"  [{i}] {file_name}")
    print(f"      描述: {caption}")

from llama_index.core.retrievers import QueryFusionRetriever

def hybrid_search(
    vector_index,
    bm25_retriever,
    query: str,
    top_k: int = 5
):
    """混合检索：向量检索 + BM25检索 + RRF融合"""
    
    print(f"\n🔍 【混合检索】查询: \"{query}\"")
    print("-" * 50)
    
    # 1. 创建向量检索器
    vector_retriever = vector_index.as_retriever(similarity_top_k=10)
    
    # 2. 创建融合检索器（RRF融合）
    print("  路径1: 向量检索（语义理解）")
    print("  路径2: BM25检索（关键词匹配）")
    print("  融合: Reciprocal Rank Fusion (RRF)")
    
    fusion_retriever = QueryFusionRetriever(
        retrievers=[vector_retriever, bm25_retriever],
        similarity_top_k=top_k,
        num_queries=1,  # 不生成query变体
        mode="reciprocal_rerank",  # RRF融合
        use_async=False
    )
    
    # 3. 执行检索
    results = fusion_retriever.retrieve(query)
    
    # 4. 格式化输出
    for i, result in enumerate(results, 1):
        file_name = result.node.metadata.get("file_name", "未知")
        image_path = result.node.metadata.get("image_path", "")
        score = result.score if hasattr(result, 'score') else 0.0
        caption = result.node.text[:100] + "..."
        
        print(f"  [{i}] {file_name} (融合分数: {score:.4f})")
        print(f"      描述: {caption}")
    
    return results

# 对比测试
query = "跟建筑平面图"

# 纯向量检索
vector_retriever = index.as_retriever(similarity_top_k=3)
vector_results = vector_retriever.retrieve(query)
print("\n【纯向量检索】结果:")
for i, r in enumerate(vector_results, 1):
    print(f"  [{i}] {r.node.metadata['file_name']}")

# 显示结果
display_retrieval_results(vector_results, title=f"Retrieval: 纯向量检索", max_display=6)

# 纯BM25检索
bm25_results = bm25_retriever.retrieve(query)
print("\n【纯BM25检索】结果:")
for i, r in enumerate(bm25_results[:3], 1):
    print(f"  [{i}] {r.node.metadata['file_name']}")
    print(f"      描述: {r.node.text[:100]}...")

# 混合检索
hybrid_results = hybrid_search(index, bm25_retriever, query, top_k=3)

from modelscope import snapshot_download
from pathlib import Path

MODEL_CACHE_DIR = Path("./models")
MODEL_CACHE_DIR.mkdir(exist_ok=True)

EMBEDDING_MODEL_MS = "qwen/Qwen3-VL-Embedding-2B"

# 下载模型
print(f"\n📥 从 ModelScope 下载: {EMBEDDING_MODEL_MS}")
model_dir = snapshot_download(
    EMBEDDING_MODEL_MS,
    cache_dir=str(MODEL_CACHE_DIR),
    revision='master'
)
print(f"✅ 下载完成: {model_dir}")

import torch
import numpy as np
from transformers import AutoModel, AutoProcessor
from qwen_vl_utils import process_vision_info
from PIL import Image

class Qwen3VLEmbedding:
    """Qwen3-VL-Embedding-2B 封装"""
    
    def __init__(self, model_path: str = None, output_dim: int = 512, device: str = "auto"):
        self.output_dim = output_dim
        self.process_vision_info = process_vision_info
        
        print(f"\n🚀 初始化 Qwen3-VL-Embedding (输出维度: {output_dim})")
        
        # 自动选择设备
        if device == "auto":
            if torch.cuda.is_available():
                self.device, dtype = "cuda", torch.bfloat16
            elif torch.backends.mps.is_available():
                self.device, dtype = "mps", torch.float16
            else:
                self.device, dtype = "cpu", torch.float32
        else:
            self.device = device
            dtype = torch.float16 if device != "cpu" else torch.float32
        
        print(f"   设备: {self.device}")
        
        # 加载模型
        if model_path is None:
            model_path = model_dir  # 使用下载的模型路径
        
        self.processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
        self.model = AutoModel.from_pretrained(
            model_path,
            trust_remote_code=True,
            torch_dtype=dtype,
            device_map="auto" if self.device == "cuda" else None
        )
        
        if self.device != "cuda":
            self.model = self.model.to(self.device)
        self.model.eval()
        
        print("✅ Qwen3-VL-Embedding 加载完成")
    
    def _get_embedding_from_model(self, messages: list) -> np.ndarray:
        """从模型获取嵌入向量"""
        # 处理消息
        text = self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
        image_inputs, video_inputs = self.process_vision_info(messages)
        
        # 编码
        inputs = self.processor(
            text=[text],
            images=image_inputs,
            videos=video_inputs,
            return_tensors="pt",
            padding=True
        ).to(self.device)
        
        # 推理
        with torch.no_grad():
            outputs = self.model(**inputs, output_hidden_states=True)
            embedding = outputs.hidden_states[-1].mean(dim=1).squeeze()
        
        # 转换为numpy
        embedding = embedding.cpu().numpy().astype(np.float32)
        
        # 截断到指定维度
        if self.output_dim and len(embedding) > self.output_dim:
            embedding = embedding[:self.output_dim]
        
        return embedding
    
    def encode_text(self, text: str) -> np.ndarray:
        """编码文本"""
        messages = [{
            "role": "user",
            "content": [{"type": "text", "text": text}]
        }]
        return self._get_embedding_from_model(messages)
    
    def encode_image(self, image_path: str) -> np.ndarray:
        """编码图像"""
        image = Image.open(image_path).convert("RGB")
        messages = [{
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": "Describe this image."}
            ]
        }]
        return self._get_embedding_from_model(messages)

# 初始化
embedder = Qwen3VLEmbedding(output_dim=512)

# 测试文本向量
text_vec = embedder.encode_text("住宅的平面图，显示了房间的布局和尺寸")
print(f"文本向量维度: {len(text_vec)}")  # 512

# 测试图像向量  
img_vec = embedder.encode_image("images/20251216110526_92_50.png")
print(f"图像向量维度: {len(img_vec)}")  # 512

# 计算余弦相似度
from numpy import dot
from numpy.linalg import norm
similarity = dot(text_vec, img_vec) / (norm(text_vec) * norm(img_vec))
print(f"文本-图像相似度: {similarity:.4f}")  # 约0.20-0.35

from llama_index.core.embeddings import MultiModalEmbedding
from typing import List

class Qwen3VLMultiModalEmbedding(MultiModalEmbedding):
    """
    Qwen3-VL 的 LlamaIndex 多模态适配器
    
    该类将自定义的 Qwen3VLEmbedding 封装为 LlamaIndex 标准的 MultiModalEmbedding 接口，
    使得 Qwen3-VL 模型可以无缝集成到 LlamaIndex 的多模态 RAG 工作流中。
    """
    
    # Pydantic 模型配置：允许任意类型（用于存储非 Pydantic 基本类型的 qwen3_embedder 实例）
    model_config = {"arbitrary_types_allowed": True, "extra": "allow"}
    
    def __init__(self, qwen3_embedder: Qwen3VLEmbedding):
        """
        初始化适配器
        
        Args:
            qwen3_embedder: 底层的 Qwen3VLEmbedding 推理实例
        """
        super().__init__(
            embed_batch_size=1,
            model_name="Qwen3-VL-Embedding-2B"
        )
        # 使用 object.__setattr__ 绕过 Pydantic 的属性校验，确保底层实例被正确存储
        object.__setattr__(self, 'qwen3_embedder', qwen3_embedder)
    
    def _get_text_embedding(self, text: str) -> List[float]:
        """
        获取文本的向量表示
        
        Args:
            text: 输入的文本字符串
        Returns:
            List[float]: 文本嵌入向量
        """
        return self.qwen3_embedder.encode_text(text).tolist()
    
    def _get_query_embedding(self, query: str) -> List[float]:
        """
        获取查询文本的向量表示（通常用于检索阶段）
        
        Args:
            query: 输入的查询字符串
        Returns:
            List[float]: 查询嵌入向量
        """
        return self.qwen3_embedder.encode_text(query).tolist()
    
    def _get_image_embedding(self, img_file_path: str) -> List[float]:
        """
        获取图像的向量表示
        
        Args:
            img_file_path: 本地图像文件的路径
        Returns:
            List[float]: 图像嵌入向量
        """
        return self.qwen3_embedder.encode_image(img_file_path).tolist()
    
    # 异步接口实现：目前通过调用同步方法实现
    async def _aget_text_embedding(self, text: str) -> List[float]:
        """异步获取文本嵌入"""
        return self._get_text_embedding(text)
    
    async def _aget_query_embedding(self, query: str) -> List[float]:
        """异步获取查询嵌入"""
        return self._get_query_embedding(query)
    
    async def _aget_image_embedding(self, img_file_path: str) -> List[float]:
        """异步获取图像嵌入"""
        return self._get_image_embedding(img_file_path)

# 实例化适配器并传入底层的推理模型实例
embed_model = Qwen3VLMultiModalEmbedding(embedder)
print(f"✅ 适配器初始化成功: {embed_model.model_name}")

# 设置全局嵌入模型，需要将之前的模型覆盖掉
Settings.embed_model = embed_model
Settings.image_embed_model = embed_model

# ========================================================
# 构建 Caption 文本索引（qwen3_vl_text）
# ========================================================

from llama_index.core import VectorStoreIndex, Document,StorageContext
from llama_index.core.node_parser import SentenceSplitter
from llama_index.vector_stores.milvus import MilvusVectorStore


MILVUS_URI = "http://localhost:19530"

# 1. 确保已经生成了 captions（需要先运行 3.3 节生成 Caption）
if 'captions' not in globals():
    print("❌ 请先运行 3.3 节生成图片的 Caption 描述")
else:
    print(f"📝 使用已生成的 {len(captions)} 条 Caption")
    
    # 2. 将 Caption 转换为 Document
    caption_documents = []
    for item in captions:
        doc = Document(
            text=item["caption"],
            metadata={
                "image_path": item["image_path"],
                "file_name": item["image_name"]
            }
        )
        caption_documents.append(doc)
    
    print(f"   创建了 {len(caption_documents)} 个 Caption Document")
    
    # 3. 配置 Milvus
    text_store = MilvusVectorStore(
        uri=MILVUS_URI,
        collection_name="qwen3_vl_text",
        dim=512,  # Qwen3-VL 维度
        overwrite=True
    )
    
    # 4. 构建文本索引（使用 Qwen3-VL embed_model）
    caption_index = VectorStoreIndex.from_documents(
        caption_documents,
        storage_context=StorageContext.from_defaults(vector_store=text_store),
        embed_model=embed_model  # ⭐ 使用 Qwen3-VL 文本编码
    )
    
    print("✅ Caption 文本索引构建完成")

from llama_index.core import SimpleDirectoryReader, StorageContext
from llama_index.core.indices import MultiModalVectorStoreIndex
from llama_index.vector_stores.milvus import MilvusVectorStore

# 1. 加载图片
documents = SimpleDirectoryReader(
    input_dir="images",
    required_exts=[".png", ".jpg", ".jpeg"]
).load_data()
print(f"📷 加载了 {len(documents)} 张图片")

# 2. 配置Milvus（双索引）
MILVUS_URI = "http://localhost:19530"

# 初始化图像向量存储
image_store = MilvusVectorStore(
    uri=MILVUS_URI,
    collection_name="qwen3_vl_image",
    dim=512,
    overwrite=True
)

# 创建存储上下文，同时配置文本和图像向量库
storage_context = StorageContext.from_defaults(
    vector_store=text_store,
    image_store=image_store
)

# 3. 构建索引（使用Qwen3-VL）
index = MultiModalVectorStoreIndex.from_documents(
    documents,
    storage_context=storage_context,
    embed_model=embed_model,
    image_embed_model=embed_model  # ⭐ 使用Qwen3-VL
)
print("✅ Qwen3-VL索引构建完成")

# 创建检索器
retriever = index.as_retriever(
    similarity_top_k=3,
    image_similarity_top_k=3
)

# 执行文搜图检索
query = "房屋架构图"
results = retriever.text_to_image_retrieve(query)

print(f"\n🔍 查询: \"{query}\"")
print("-" * 50)
for i, result in enumerate(results, 1):
    file_path = result.node.metadata.get("file_path", "未知")
    score = result.score if hasattr(result, 'score') else 0.0
    print(f"  [{i}] {Path(file_path).name}: {score:.4f}")

# 3. 显示结果
display_retrieval_results(results, title="文搜图: 架构图", max_display=6)

from llama_index.core.postprocessor.types import BaseNodePostprocessor
from llama_index.core.schema import NodeWithScore, QueryBundle
from llama_index.core import QueryBundle
from typing import List, Optional, Any
from pydantic import Field
from pathlib import Path
import sys
import importlib.util

class Qwen3VLNodePostprocessor(BaseNodePostprocessor):
    """Qwen3-VL-Reranker-2B 的 LlamaIndex NodePostprocessor 实现，用于多模态精排"""
    
    top_n: int = Field(default=5, description="精排后返回的前 N 个结果数量")
    model: Any = Field(default=None, description="Reranker 模型实例")
    image_base_dir: str = Field(default="images", 
                                description="图片存储的基础目录，用于解析相对路径")
    
    def __init__(self, model_path: str = None, top_n: int = 5, 
                 image_base_dir: str = "images", **kwargs):
        """初始化 Reranker，加载模型并注入官方脚本"""
        super().__init__(top_n=top_n, image_base_dir=image_base_dir, **kwargs)
        
        print(f"\n🔧 加载 Qwen3-VL-Reranker-2B (NodePostprocessor)")
        
        # 1. 确定模型路径，若未提供则从 ModelScope 下载
        if model_path is None:
            RERANKER_MODEL_MS = "qwen/Qwen3-VL-Reranker-2B"
            model_path = snapshot_download(RERANKER_MODEL_MS, cache_dir=str(MODEL_CACHE_DIR))
        
        # 2. 动态加载模型权重目录下的官方推理脚本
        script_path = Path(model_path) / "scripts" / "qwen3_vl_reranker.py"
        
        if not script_path.exists():
            raise FileNotFoundError(f"未找到官方脚本: {script_path}")
        
        # 使用 importlib 动态导入本地模块
        spec = importlib.util.spec_from_file_location("qwen3_vl_reranker_script", script_path)
        reranker_module = importlib.util.module_from_spec(spec)
        sys.modules["qwen3_vl_reranker_script"] = reranker_module
        spec.loader.exec_module(reranker_module)
        
        # 3. 实例化官方 Reranker 类
        OfficialReranker = reranker_module.Qwen3VLReranker
        # 使用 object.__setattr__ 绕过 Pydantic 对 model 字段的直接赋值限制
        object.__setattr__(self, 'model', OfficialReranker(model_name_or_path=model_path))
        
        print("✅ Reranker 加载完成 (BaseNodePostprocessor)")
    
    @classmethod
    def class_name(cls) -> str:
        return "Qwen3VLNodePostprocessor"
    
    def _get_image_path(self, node: NodeWithScore) -> Optional[str]:
        """从节点中提取图片路径，支持多种 metadata 键名和路径解析逻辑"""
        # 1. 检查节点对象属性
        if hasattr(node.node, 'image_path'):
            return str(Path(node.node.image_path).resolve())
        
        metadata = node.node.metadata
        
        # 2. 检查 metadata 中的 file_path
        if 'file_path' in metadata:
            return str(Path(metadata['file_path']).resolve())
        
        # 3. 检查 metadata 中的 image_path
        if 'image_path' in metadata:
            image_path = metadata['image_path']
            path_obj = Path(image_path)
            if path_obj.is_absolute() or path_obj.exists():
                return str(path_obj.resolve())
            # 尝试从 base_dir 解析相对路径
            full_path = Path(self.image_base_dir) / path_obj
            if full_path.exists():
                return str(full_path.resolve())
            return image_path
        
        # 4. 检查 metadata 中的 file_name，结合 base_dir
        if 'file_name' in metadata:
            file_name = metadata['file_name']
            full_path = Path(self.image_base_dir) / file_name
            return str(full_path.resolve())
        
        return None
    
    def _postprocess_nodes(
        self,
        nodes: List[NodeWithScore],
        query_bundle: Optional[QueryBundle] = None
    ) -> List[NodeWithScore]:
        """
        LlamaIndex 标准后处理接口实现：调用模型并根据分数重新排序
        
        支持两种查询模式：
        - 文搜图：query_str 为文本描述
        - 图搜图：query_str 为图片路径
        """
        
        if not nodes:
            return []
        
        if query_bundle is None:
            return nodes[:self.top_n]
        
        query_str = query_bundle.query_str
        
        # 判断查询类型：图片路径 or 文本
        query_path = Path(query_str)
        is_image_query = (
            query_path.exists() and 
            query_path.suffix.lower() in {'.png', '.jpg', '.jpeg', '.webp', '.gif', '.bmp'}
        )
        
        query_type = "图搜图" if is_image_query else "文搜图"
        print(f"\n✨ VLM 精排 [{query_type}] ({len(nodes)} → {self.top_n})")
        
        # 1. 构造官方 Reranker 的输入文档列表
        documents = []
        valid_indices = []
        
        for i, node in enumerate(nodes):
            image_path = self._get_image_path(node)
            
            if image_path is None:
                print(f"节点 {i+1} 缺少图片路径，跳过")
                continue
            
            documents.append({"image": image_path})
            valid_indices.append(i)
        
        if not documents:
            print("❌ 没有有效的图片节点，跳过精排")
            return nodes[:self.top_n]
        
        print(f"   📊 成功提取 {len(documents)}/{len(nodes)} 个图片路径")
        
        # 2. ⭐ 根据查询类型构造不同的 query 输入
        if is_image_query:
            query_input = {"image": str(query_path.resolve())}
            print(f"    查询图片: {query_path.name}")
        else:
            query_input = {"text": query_str}
            print(f"    查询文本: {query_str[:50]}...")
        
        # 3. 构造推理输入字典
        inputs = {
            "instruction": "Retrieve images relevant to the query.",
            "query": query_input,  # ⭐ 支持图片或文本
            "documents": documents,
            "fps": 1.0
        }
        
        try:
            # 4. 执行模型推理
            scores = self.model.process(inputs)
            
            # 5. 更新有效节点分数
            for i, score in enumerate(scores):
                original_idx = valid_indices[i]
                nodes[original_idx].score = float(score)
            
            # 6. 无效节点分数清零
            for i in range(len(nodes)):
                if i not in valid_indices:
                    nodes[i].score = 0.0
            
            # 7. 根据新分数进行降序排序
            nodes.sort(key=lambda x: x.score or 0.0, reverse=True)
            
            score_list = [f"{n.score:.4f}" for n in nodes[:self.top_n]]
            print(f"\n   📊 Top-{self.top_n} 精排分数: {score_list}")
            print(f"   ✅ VLM精排完成")
            
            return nodes[:self.top_n]
            
        except Exception as e:
            print(f"❌ Reranker 失败: {e}")
            return nodes[:self.top_n]

# 重新初始化 Reranker 实例，设置默认为5个结果
reranker = Qwen3VLNodePostprocessor(top_n=5)