第二十四章 多模态RAG与向量数据库进阶¶

"""CLIP模型基础用法"""

from transformers import CLIPProcessor, CLIPModel
from PIL import Image
import torch
import numpy as np

# 加载模型
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

# --- 图像编码 ---
def encode_image(image_path: str) -> np.ndarray:
    image = Image.open(image_path).convert("RGB")
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():  # 禁用梯度计算，节省内存（推理时使用）
        features = model.get_image_features(**inputs)
    # L2归一化
    features = features / features.norm(dim=-1, keepdim=True)
    return features.squeeze().numpy()

# --- 文本编码 ---
def encode_text(text: str) -> np.ndarray:
    inputs = processor(text=[text], return_tensors="pt", padding=True)
    with torch.no_grad():
        features = model.get_text_features(**inputs)
    features = features / features.norm(dim=-1, keepdim=True)
    return features.squeeze().numpy()

# --- 文本-图像相似度 ---
def text_image_similarity(text: str, image_path: str) -> float:
    text_emb = encode_text(text)
    image_emb = encode_image(image_path)
    return float(np.dot(text_emb, image_emb))

# 示例
# score = text_image_similarity("a cat sitting on a couch", "cat.jpg")
# print(f"相似度: {score:.4f}")

"""BGE-M3嵌入 —— 支持Dense + Sparse + ColBERT多种检索方式"""

from sentence_transformers import SentenceTransformer
import numpy as np

# 加载BGE-M3模型
model = SentenceTransformer("BAAI/bge-m3")

def get_dense_embeddings(texts: list[str]) -> np.ndarray:
    """获取Dense Embedding（标准向量检索）"""
    embeddings = model.encode(
        texts,
        normalize_embeddings=True,  # L2归一化
        show_progress_bar=True,
    )
    return embeddings

# 示例
texts = [
    "什么是Transformer架构？",
    "注意力机制的计算过程",
    "如何训练大语言模型",
]
embeddings = get_dense_embeddings(texts)
print(f"Embedding维度: {embeddings.shape}")  # (3, 1024)

# 计算相似度矩阵
similarity_matrix = np.dot(embeddings, embeddings.T)
print("相似度矩阵:")
print(similarity_matrix.round(3))

"""方案A：Vision LLM生成描述 + 文本Embedding（实用方案）"""

import base64
from openai import OpenAI
from pathlib import Path

client = OpenAI()

def image_to_description(image_path: str) -> str:
    """用GPT-4o-mini为图像生成详细描述"""
    with open(image_path, "rb") as f:  # with自动管理文件关闭
        b64 = base64.b64encode(f.read()).decode("utf-8")

    suffix = Path(image_path).suffix.lower()
    mime_map = {".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg"}
    mime_type = mime_map.get(suffix, "image/png")

    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": (
                            "请详细描述这张图片的内容，包括：\n"
                            "1. 图片类型（图表/照片/架构图/截图等）\n"
                            "2. 主要内容和关键信息\n"
                            "3. 如果是图表，描述数据趋势和关键数值\n"
                            "4. 如果是架构图，描述组件和关系\n"
                            "用中文回答，尽量详细。"
                        ),
                    },
                    {
                        "type": "image_url",
                        "image_url": {"url": f"data:{mime_type};base64,{b64}"},
                    },
                ],
            }
        ],
        max_tokens=500,
    )
    return response.choices[0].message.content

def get_text_embedding(text: str) -> list[float]:
    """获取文本Embedding"""
    resp = client.embeddings.create(model="text-embedding-3-small", input=[text])
    return resp.data[0].embedding

# 图像处理流水线
def process_image_for_rag(image_path: str) -> dict:
    """将图像转换为可检索的文档"""
    description = image_to_description(image_path)
    embedding = get_text_embedding(description)
    return {
        "type": "image",
        "source": image_path,
        "description": description,
        "embedding": embedding,
    }

"""统一检索实现：文本和图像在同一向量空间中检索"""

import numpy as np
from dataclasses import dataclass

@dataclass  # @dataclass自动生成__init__等方法
class Document:
    doc_id: str
    doc_type: str       # "text" | "image" | "table"
    content: str        # 原始内容或图像路径
    description: str    # 文本描述（图像为LLM生成的描述）
    embedding: list[float]
    metadata: dict = None

class MultimodalIndex:
    """多模态统一索引"""

    def __init__(self):
        self.documents: list[Document] = []

    def add_document(self, doc: Document):
        self.documents.append(doc)

    def search(self, query_embedding: list[float], top_k: int = 5) -> list[tuple[Document, float]]:
        """向量相似度搜索"""
        q = np.array(query_embedding)
        results = []
        for doc in self.documents:
            d = np.array(doc.embedding)
            score = float(np.dot(q, d) / (np.linalg.norm(q) * np.linalg.norm(d)))
            results.append((doc, score))
        results.sort(key=lambda x: x[1], reverse=True)  # lambda匿名函数
        return results[:top_k]

    def hybrid_search(
        self,
        query_embedding: list[float],
        doc_type_filter: str | None = None,
        top_k: int = 5,
    ) -> list[tuple[Document, float]]:
        """混合搜索：向量检索 + 标量过滤"""
        q = np.array(query_embedding)
        results = []
        for doc in self.documents:
            if doc_type_filter and doc.doc_type != doc_type_filter:
                continue
            d = np.array(doc.embedding)
            score = float(np.dot(q, d) / (np.linalg.norm(q) * np.linalg.norm(d)))
            results.append((doc, score))
        results.sort(key=lambda x: x[1], reverse=True)
        return results[:top_k]

"""Milvus Lite快速上手"""

from pymilvus import MilvusClient

# 创建客户端（使用本地文件存储）
client = MilvusClient("./milvus_demo.db")

# 创建Collection
client.create_collection(
    collection_name="docs",
    dimension=1536,  # text-embedding-3-small的维度
)

# 插入数据
data = [
    {"id": 1, "vector": [0.1] * 1536, "text": "什么是RAG？", "source": "doc1.pdf"},
    {"id": 2, "vector": [0.2] * 1536, "text": "向量数据库原理", "source": "doc2.pdf"},
    {"id": 3, "vector": [0.3] * 1536, "text": "多模态嵌入技术", "source": "doc3.pdf"},
]
client.insert(collection_name="docs", data=data)

# 搜索
results = client.search(
    collection_name="docs",
    data=[[0.15] * 1536],  # 查询向量
    limit=2,               # 返回top-2
    output_fields=["text", "source"],  # 返回的字段
)

for result in results[0]:
    print(f"  ID: {result['id']}, 距离: {result['distance']:.4f}")
    print(f"  文本: {result['entity']['text']}")

"""高级Collection配置：Schema定义 + 自定义索引"""

from pymilvus import (
    MilvusClient,
    CollectionSchema,
    FieldSchema,
    DataType,
)

client = MilvusClient("./milvus_advanced.db")

# --- 方法1：使用Schema精确定义 ---

fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
    FieldSchema(name="doc_type", dtype=DataType.VARCHAR, max_length=20),
    FieldSchema(name="content", dtype=DataType.VARCHAR, max_length=10000),
    FieldSchema(name="source", dtype=DataType.VARCHAR, max_length=500),
    FieldSchema(name="page_num", dtype=DataType.INT32),
    FieldSchema(name="dense_vector", dtype=DataType.FLOAT_VECTOR, dim=1536),
]

schema = CollectionSchema(fields=fields, description="多模态文档集合")

# 创建Collection
client.create_collection(
    collection_name="multimodal_docs",
    schema=schema,
)

# 创建索引（HNSW索引适合大多数场景）
index_params = client.prepare_index_params()
index_params.add_index(
    field_name="dense_vector",
    index_type="HNSW",
    metric_type="COSINE",
    params={
        "M": 16,              # 每层连接数（越大越精确，越慢）
        "efConstruction": 256, # 构建时搜索宽度
    },
)

client.create_index(
    collection_name="multimodal_docs",
    index_params=index_params,
)

print("✅ Collection和索引创建完成")

"""完整的插入与检索流程"""

from pymilvus import MilvusClient
from openai import OpenAI

openai_client = OpenAI()
milvus_client = MilvusClient("./milvus_rag.db")

# 创建collection
milvus_client.create_collection(
    collection_name="knowledge_base",
    dimension=1536,
)

def embed(texts: list[str]) -> list[list[float]]:
    resp = openai_client.embeddings.create(
        model="text-embedding-3-small", input=texts
    )
    return [d.embedding for d in resp.data]

# ---- 批量插入 ----

documents = [
    {"text": "RAG通过检索增强生成，提高LLM的准确性和时效性。", "source": "rag_intro.md"},
    {"text": "向量数据库存储高维向量，支持近似最近邻搜索。", "source": "vector_db.md"},
    {"text": "HNSW是一种基于图的近似最近邻搜索算法。", "source": "hnsw.md"},
    {"text": "多模态RAG可以处理文本、图像、表格等多种数据类型。", "source": "multimodal.md"},
    {"text": "Text2SQL将自然语言查询转换为SQL语句。", "source": "text2sql.md"},
]

texts = [d["text"] for d in documents]
vectors = embed(texts)

data = [
    {
        "id": i,
        "vector": vectors[i],
        "text": documents[i]["text"],
        "source": documents[i]["source"],
    }
    for i in range(len(documents))
]

milvus_client.insert(collection_name="knowledge_base", data=data)
print(f"✅ 插入 {len(data)} 条数据")

# ---- 检索 ----

def search(query: str, top_k: int = 3) -> list[dict]:
    query_vector = embed([query])[0]
    results = milvus_client.search(
        collection_name="knowledge_base",
        data=[query_vector],
        limit=top_k,
        output_fields=["text", "source"],
    )
    return [
        {
            "text": r["entity"]["text"],
            "source": r["entity"]["source"],
            "score": r["distance"],
        }
        for r in results[0]
    ]

# 测试检索
results = search("什么是向量搜索？")
for r in results:
    print(f"  [{r['score']:.4f}] {r['text']} ({r['source']})")

"""混合检索：向量相似度 + 标量过滤条件"""

from pymilvus import MilvusClient

client = MilvusClient("./milvus_hybrid.db")

# 创建collection
client.create_collection(
    collection_name="articles",
    dimension=1536,
)

# 假设已插入数据，每条包含: id, vector, text, category, date

# --- 带过滤条件的检索 ---

def hybrid_search(
    query_vector: list[float],
    category: str | None = None,
    date_after: str | None = None,
    top_k: int = 5,
) -> list:
    """混合检索：向量相似度 + 标量过滤"""
    # 构建过滤表达式
    filters = []
    if category:
        filters.append(f'category == "{category}"')
    if date_after:
        filters.append(f'date >= "{date_after}"')

    filter_expr = " and ".join(filters) if filters else ""

    results = client.search(
        collection_name="articles",
        data=[query_vector],
        limit=top_k,
        filter=filter_expr,  # 标量过滤
        output_fields=["text", "category", "date"],
    )
    return results[0]

# 使用示例
# results = hybrid_search(
#     query_vector=embed(["RAG技术"])[0],
#     category="AI",
#     date_after="2024-01-01",
#     top_k=5,
# )

"""Milvus性能调优要点"""

# 1. 索引类型选择
index_configs = {
    # 小规模（<100万向量）：精确且快速
    "small_scale": {
        "index_type": "HNSW",
        "metric_type": "COSINE",
        "params": {"M": 16, "efConstruction": 256},
    },
    # 中规模（100万-1000万）：平衡精度和速度
    "medium_scale": {
        "index_type": "IVF_FLAT",
        "metric_type": "COSINE",
        "params": {"nlist": 1024},
    },
    # 大规模（>1000万）：用量化压缩内存
    "large_scale": {
        "index_type": "IVF_PQ",
        "metric_type": "L2",
        "params": {"nlist": 2048, "m": 16, "nbits": 8},
    },
}

# 2. 搜索参数调优
search_params = {
    "HNSW": {"ef": 128},           # ef越大越精确，越慢
    "IVF_FLAT": {"nprobe": 32},    # nprobe越大越精确，越慢
    "IVF_PQ": {"nprobe": 64},
}

# 3. 批量操作（比逐条插入快10-100倍）
# ✅ 推荐：批量插入
# client.insert(collection_name="docs", data=batch_data)

# ❌ 避免：逐条插入
# for item in data:
#     client.insert(collection_name="docs", data=[item])

"""完整Milvus RAG实战：从文档到检索到生成"""

from pymilvus import MilvusClient
from openai import OpenAI
import json

openai_client = OpenAI()
milvus_client = MilvusClient("./rag_production.db")

COLLECTION_NAME = "production_kb"
EMBEDDING_DIM = 1536

# ---- 初始化 ----

def init_collection():
    """初始化Collection"""
    if milvus_client.has_collection(COLLECTION_NAME):
        milvus_client.drop_collection(COLLECTION_NAME)

    milvus_client.create_collection(
        collection_name=COLLECTION_NAME,
        dimension=EMBEDDING_DIM,
    )
    print(f"✅ Collection '{COLLECTION_NAME}' 已创建")

# ---- 数据预处理 ----

def chunk_documents(texts: list[str], chunk_size: int = 500, overlap: int = 50) -> list[dict]:
    """切分文档"""
    chunks = []
    for doc_idx, text in enumerate(texts):  # enumerate同时获取索引和元素
        start = 0
        chunk_idx = 0
        while start < len(text):
            end = min(start + chunk_size, len(text))
            chunks.append({
                "text": text[start:end],
                "doc_id": doc_idx,
                "chunk_id": chunk_idx,
            })
            start += chunk_size - overlap
            chunk_idx += 1
    return chunks

def embed_batch(texts: list[str], batch_size: int = 100) -> list[list[float]]:
    """批量Embedding"""
    all_embeddings = []
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i + batch_size]
        resp = openai_client.embeddings.create(
            model="text-embedding-3-small", input=batch
        )
        all_embeddings.extend([d.embedding for d in resp.data])
    return all_embeddings

# ---- 索引构建 ----

def build_index(documents: list[str]):
    """构建索引"""
    init_collection()

    # 切分
    chunks = chunk_documents(documents)
    texts = [c["text"] for c in chunks]

    # Embedding
    print(f"正在计算 {len(texts)} 个文本块的Embedding...")
    embeddings = embed_batch(texts)

    # 插入Milvus
    data = [
        {
            "id": i,
            "vector": embeddings[i],
            "text": chunks[i]["text"],
            "doc_id": chunks[i]["doc_id"],
            "chunk_id": chunks[i]["chunk_id"],
        }
        for i in range(len(chunks))
    ]

    milvus_client.insert(collection_name=COLLECTION_NAME, data=data)
    print(f"✅ 已索引 {len(data)} 个文本块")

# ---- 检索与生成 ----

def retrieve_and_generate(query: str, top_k: int = 3) -> str:
    """RAG：检索 + 生成"""
    # 检索
    query_emb = embed_batch([query])[0]
    results = milvus_client.search(
        collection_name=COLLECTION_NAME,
        data=[query_emb],
        limit=top_k,
        output_fields=["text", "doc_id", "chunk_id"],
    )

    contexts = []
    for r in results[0]:
        contexts.append(r["entity"]["text"])

    context_text = "\n---\n".join(contexts)

    # 生成
    response = openai_client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {
                "role": "system",
                "content": (
                    "根据检索到的文档回答问题。引用相关内容。"
                    f"\n\n【检索结果】\n{context_text}"
                ),
            },
            {"role": "user", "content": query},
        ],
    )
    return response.choices[0].message.content

# ---- 使用示例 ----
if __name__ == "__main__":
    # 构建索引
    sample_docs = [
        "RAG（Retrieval-Augmented Generation）是一种将检索和生成结合的技术...",
        "Milvus是一个开源的向量数据库，专为相似性搜索和AI应用设计...",
    ]
    build_index(sample_docs)

    # 查询
    answer = retrieve_and_generate("RAG的核心原理是什么？")
    print(f"\n回答：{answer}")

"""从PDF中提取并处理图像"""

import pymupdf
from pathlib import Path
import base64
from openai import OpenAI

client = OpenAI()

def extract_pdf_content(pdf_path: str) -> list[dict]:
    """从PDF提取文本和图像"""
    doc = pymupdf.open(pdf_path)
    contents = []

    for page_num, page in enumerate(doc):
        # 提取文本
        text = page.get_text()
        if text.strip():  # 链式调用：strip去除空白
            contents.append({
                "type": "text",
                "content": text.strip(),
                "page": page_num + 1,
                "source": pdf_path,
            })

        # 提取图像
        images = page.get_images(full=True)
        for img_idx, img in enumerate(images):
            xref = img[0]
            pix = pymupdf.Pixmap(doc, xref)
            if pix.n > 4:  # CMYK转RGB
                pix = pymupdf.Pixmap(pymupdf.csRGB, pix)

            img_bytes = pix.tobytes("png")
            b64 = base64.b64encode(img_bytes).decode("utf-8")

            # 用Vision LLM生成图像描述
            description = describe_image(b64)

            contents.append({
                "type": "image",
                "content": b64,  # base64编码的图像
                "description": description,
                "page": page_num + 1,
                "source": pdf_path,
            })

    return contents

def describe_image(image_b64: str) -> str:
    """用GPT-4o生成图像描述"""
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "请详细描述这张图片的内容，用中文回答。"},
                    {
                        "type": "image_url",
                        "image_url": {"url": f"data:image/png;base64,{image_b64}"},
                    },
                ],
            }
        ],
        max_tokens=300,
    )
    return response.choices[0].message.content

"""表格数据处理：保留结构信息"""

import pandas as pd
from openai import OpenAI

client = OpenAI()

def table_to_searchable_text(df: pd.DataFrame, table_name: str = "") -> list[str]:
    """将表格转换为可检索的文本块"""
    texts = []

    # 1. 表格整体描述
    overview = f"表格: {table_name}\n"
    overview += f"列: {', '.join(df.columns)}\n"
    overview += f"行数: {len(df)}\n"
    overview += f"数据类型: {dict(df.dtypes)}\n"
    texts.append(overview)

    # 2. 表格Markdown格式（适合小表格）
    if len(df) <= 50:
        md = df.to_markdown(index=False)
        texts.append(f"表格 {table_name} 的完整数据:\n{md}")

    # 3. 按行转换（适合大表格）
    for idx, row in df.iterrows():
        row_text = f"表格 {table_name} 第{idx+1}行: "
        row_text += "; ".join(f"{col}={val}" for col, val in row.items())
        texts.append(row_text)

    # 4. 统计摘要
    numeric_cols = df.select_dtypes(include="number").columns
    if len(numeric_cols) > 0:
        summary = f"表格 {table_name} 统计摘要:\n"
        summary += df[numeric_cols].describe().to_markdown()
        texts.append(summary)

    return texts

# 示例
# df = pd.read_csv("sales_data.csv")
# chunks = table_to_searchable_text(df, "销售数据")

"""多模态统一索引：文本 + 图像 + 表格"""

from pymilvus import MilvusClient
from openai import OpenAI
import json

openai_client = OpenAI()
milvus_client = MilvusClient("./multimodal_rag.db")

COLLECTION = "multimodal_kb"

def init_multimodal_collection():
    if milvus_client.has_collection(COLLECTION):
        milvus_client.drop_collection(COLLECTION)
    milvus_client.create_collection(
        collection_name=COLLECTION,
        dimension=1536,
    )

def embed(texts: list[str]) -> list[list[float]]:
    resp = openai_client.embeddings.create(
        model="text-embedding-3-small", input=texts
    )
    return [d.embedding for d in resp.data]

def build_multimodal_index(contents: list[dict]):
    """构建多模态统一索引"""
    init_multimodal_collection()

    data = []
    texts_to_embed = []

    for item in contents:
        if item["type"] == "text":
            texts_to_embed.append(item["content"])
        elif item["type"] == "image":
            texts_to_embed.append(item["description"])
        elif item["type"] == "table":
            texts_to_embed.append(item["content"])

    embeddings = embed(texts_to_embed)

    for i, item in enumerate(contents):
        data.append({
            "id": i,
            "vector": embeddings[i],
            "doc_type": item["type"],
            "content": item.get("description", item["content"])[:5000],
            "page": item.get("page", 0),
            "source": item.get("source", ""),
        })

    milvus_client.insert(collection_name=COLLECTION, data=data)
    print(f"✅ 已索引 {len(data)} 个多模态文档块")

"""多模态检索 + Cross-Encoder重排序"""

from sentence_transformers import CrossEncoder
import numpy as np

# 加载重排序模型
reranker = CrossEncoder("BAAI/bge-reranker-v2-m3", max_length=512)

def multimodal_search_with_rerank(
    query: str,
    top_k: int = 5,
    rerank_top_n: int = 20,
) -> list[dict]:
    """多模态检索 + 重排序"""
    # Step 1: 向量检索（召回较多候选）
    query_emb = embed([query])[0]
    candidates = milvus_client.search(
        collection_name=COLLECTION,
        data=[query_emb],
        limit=rerank_top_n,
        output_fields=["doc_type", "content", "page", "source"],
    )[0]

    # Step 2: Cross-Encoder重排序
    pairs = [(query, c["entity"]["content"]) for c in candidates]
    rerank_scores = reranker.predict(pairs)

    # Step 3: 合并分数并排序
    results = []
    for i, candidate in enumerate(candidates):
        results.append({
            "content": candidate["entity"]["content"],
            "doc_type": candidate["entity"]["doc_type"],
            "page": candidate["entity"]["page"],
            "source": candidate["entity"]["source"],
            "vector_score": candidate["distance"],
            "rerank_score": float(rerank_scores[i]),
        })

    results.sort(key=lambda x: x["rerank_score"], reverse=True)
    return results[:top_k]

"""多模态RAG生成：将检索到的文本、图像描述、表格整合为上下文"""

from openai import OpenAI

client = OpenAI()

def multimodal_generate(query: str, retrieved_docs: list[dict]) -> str:
    """基于多模态检索结果生成回答"""
    # 按类型组织上下文
    text_contexts = []
    image_contexts = []
    table_contexts = []

    for doc in retrieved_docs:
        if doc["doc_type"] == "text":
            text_contexts.append(f"[文本, 第{doc['page']}页] {doc['content']}")
        elif doc["doc_type"] == "image":
            image_contexts.append(f"[图像描述, 第{doc['page']}页] {doc['content']}")
        elif doc["doc_type"] == "table":
            table_contexts.append(f"[表格, 第{doc['page']}页] {doc['content']}")

    context_parts = []
    if text_contexts:
        context_parts.append("## 文本内容\n" + "\n\n".join(text_contexts))
    if image_contexts:
        context_parts.append("## 图像内容\n" + "\n\n".join(image_contexts))
    if table_contexts:
        context_parts.append("## 表格数据\n" + "\n\n".join(table_contexts))

    context = "\n\n".join(context_parts)

    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {
                "role": "system",
                "content": (
                    "你是一个知识库问答助手。根据检索到的多模态内容回答问题。\n"
                    "内容可能包含文本、图像描述和表格数据，请综合利用。\n"
                    "回答要准确，引用来源（页码）。\n\n"
                    f"【检索内容】\n{context}"
                ),
            },
            {"role": "user", "content": query},
        ],
    )
    return response.choices[0].message.content

"""Text2SQL：自然语言转SQL查询"""

import sqlite3
from openai import OpenAI

client = OpenAI()

# ---- 创建示例数据库 ----

def create_sample_db(db_path: str = "sample.db"):
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()

    cursor.execute("""
        CREATE TABLE IF NOT EXISTS sales (
            id INTEGER PRIMARY KEY,
            product TEXT NOT NULL,
            category TEXT NOT NULL,
            amount REAL NOT NULL,
            quantity INTEGER NOT NULL,
            date TEXT NOT NULL,
            region TEXT NOT NULL
        )
    """)

    sample_data = [
        ("iPhone 15", "手机", 5999, 120, "2024-01-15", "华东"),
        ("MacBook Pro", "笔记本", 14999, 45, "2024-01-20", "华北"),
        ("iPad Air", "平板", 4799, 80, "2024-02-10", "华东"),
        ("AirPods Pro", "耳机", 1899, 200, "2024-02-15", "华南"),
        ("iPhone 15", "手机", 5999, 150, "2024-03-01", "华北"),
        ("MacBook Air", "笔记本", 8999, 60, "2024-03-10", "华南"),
        ("Apple Watch", "手表", 2999, 90, "2024-03-20", "华东"),
    ]

    cursor.executemany(
        "INSERT INTO sales (product, category, amount, quantity, date, region) VALUES (?, ?, ?, ?, ?, ?)",
        sample_data,
    )
    conn.commit()
    conn.close()

# ---- Schema提取 ----

def get_schema(db_path: str) -> str:
    """提取数据库Schema信息"""
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()

    cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
    tables = cursor.fetchall()

    schema_parts = []
    for (table_name,) in tables:
        cursor.execute(f"PRAGMA table_info({table_name})")
        columns = cursor.fetchall()

        col_defs = []
        for col in columns:
            col_defs.append(f"  {col[1]} {col[2]}")
        schema_parts.append(f"CREATE TABLE {table_name} (\n" + ",\n".join(col_defs) + "\n)")

        # 获取样本数据
        cursor.execute(f"SELECT * FROM {table_name} LIMIT 3")
        samples = cursor.fetchall()
        if samples:
            col_names = [c[1] for c in columns]
            schema_parts.append(f"-- 样本数据 ({table_name}):")
            for sample in samples:
                row = dict(zip(col_names, sample))
                schema_parts.append(f"--   {row}")

    conn.close()
    return "\n\n".join(schema_parts)

"""Schema-aware Text2SQL Prompt"""

TEXT2SQL_PROMPT = """你是一个SQL专家。根据用户的自然语言问题，生成对应的SQLite SQL查询。

### 数据库Schema:
{schema}

### 规则:
1. 只生成SELECT查询（不允许INSERT/UPDATE/DELETE）
2. 使用SQLite语法
3. 表名和列名与Schema完全一致
4. 如果问题无法用SQL回答，返回"CANNOT_ANSWER"
5. 只返回SQL语句，不要解释

### 用户问题: {question}

### SQL查询:"""

def text_to_sql(question: str, db_path: str = "sample.db") -> str:
    """自然语言转SQL"""
    schema = get_schema(db_path)

    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {
                "role": "user",
                "content": TEXT2SQL_PROMPT.format(schema=schema, question=question),
            }
        ],
        temperature=0,  # SQL生成用低温度确保确定性
    )
    sql = response.choices[0].message.content.strip()
    # 清理可能的markdown代码块
    sql = sql.replace("```sql", "").replace("```", "").strip()
    return sql

"""SQL安全验证"""

import re
import sqlite3

FORBIDDEN_KEYWORDS = [
    "DROP", "DELETE", "INSERT", "UPDATE", "ALTER",
    "CREATE", "TRUNCATE", "EXEC", "EXECUTE",
    "--",  # SQL注释（可能用于注入）
]

def validate_sql(sql: str) -> tuple[bool, str]:
    """验证SQL安全性"""
    upper_sql = sql.upper().strip()

    # 检查是否为SELECT语句
    if not upper_sql.startswith("SELECT"):
        return False, "只允许SELECT查询"

    # 检查禁止的关键字
    for keyword in FORBIDDEN_KEYWORDS:
        if keyword in upper_sql:
            return False, f"包含禁止的关键字: {keyword}"

    # 检查分号（防止多语句注入）
    if sql.count(";") > 1:
        return False, "不允许多条SQL语句"

    return True, "OK"

def safe_execute_sql(sql: str, db_path: str = "sample.db") -> dict:
    """安全执行SQL并返回结果"""
    # 验证
    is_valid, message = validate_sql(sql)
    if not is_valid:
        return {"success": False, "error": message, "data": None}

    try:  # try/except捕获异常，防止程序崩溃
        conn = sqlite3.connect(db_path)
        conn.execute("PRAGMA query_only = ON")  # 只读模式
        cursor = conn.cursor()

        cursor.execute(sql)
        columns = [desc[0] for desc in cursor.description]
        rows = cursor.fetchall()

        conn.close()

        return {
            "success": True,
            "columns": columns,
            "data": [dict(zip(columns, row)) for row in rows],
            "row_count": len(rows),
        }
    except Exception as e:
        return {"success": False, "error": str(e), "data": None}

"""混合查询路由：自动判断使用RAG还是Text2SQL"""

from openai import OpenAI

client = OpenAI()

ROUTER_PROMPT = """判断用户的问题应该使用哪种方式回答：

1. "SQL" - 如果问题涉及结构化数据查询、统计、聚合、排序、过滤等
2. "RAG" - 如果问题涉及概念解释、原理说明、非结构化知识

只回答 "SQL" 或 "RAG"，不要解释。

用户问题: {question}"""

def route_query(question: str) -> str:
    """路由查询到合适的处理方式"""
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "user", "content": ROUTER_PROMPT.format(question=question)},
        ],
        temperature=0,
        max_tokens=5,
    )
    route = response.choices[0].message.content.strip().upper()
    return "SQL" if "SQL" in route else "RAG"

def hybrid_answer(question: str, db_path: str = "sample.db") -> str:
    """混合问答：自动选择RAG或Text2SQL"""
    route = route_query(question)

    if route == "SQL":
        # Text2SQL路径
        sql = text_to_sql(question, db_path)
        result = safe_execute_sql(sql, db_path)

        if result["success"]:
            # 将SQL结果交给LLM生成自然语言回答
            response = client.chat.completions.create(
                model="gpt-4o-mini",
                messages=[
                    {
                        "role": "system",
                        "content": "根据SQL查询结果，用自然语言回答用户问题。",
                    },
                    {
                        "role": "user",
                        "content": (
                            f"问题: {question}\n"
                            f"SQL: {sql}\n"
                            f"结果: {result['data']}"
                        ),
                    },
                ],
            )
            return f"📊 [SQL查询]\n{response.choices[0].message.content}"
        else:
            return f"❌ SQL执行失败: {result['error']}"
    else:
        # RAG路径（使用之前实现的retrieve_and_generate）
        answer = retrieve_and_generate(question)
        return f"📚 [知识检索]\n{answer}"

# 示例
# print(hybrid_answer("2024年1月总销售额是多少？"))  → SQL
# print(hybrid_answer("什么是RAG技术？"))              → RAG

"""完整的Text2SQL问答系统"""

import sqlite3
import gradio as gr
from openai import OpenAI

client = OpenAI()

def text2sql_app(question: str, db_path: str = "sample.db") -> tuple[str, str, str]:
    """Text2SQL完整流程，返回(SQL, 查询结果, 自然语言回答)"""
    # 1. 生成SQL
    schema = get_schema(db_path)
    sql = text_to_sql(question, db_path)

    if sql == "CANNOT_ANSWER":
        return "无法生成SQL", "", "该问题无法用SQL回答"

    # 2. 验证并执行
    result = safe_execute_sql(sql, db_path)

    if not result["success"]:
        return sql, f"执行失败: {result['error']}", ""

    # 3. 格式化结果
    import pandas as pd
    df = pd.DataFrame(result["data"])
    result_str = df.to_markdown(index=False)

    # 4. 生成自然语言回答
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {
                "role": "system",
                "content": "根据SQL查询结果回答用户的问题，用中文回答，简洁准确。",
            },
            {
                "role": "user",
                "content": f"问题: {question}\nSQL: {sql}\n结果:\n{result_str}",
            },
        ],
    )
    answer = response.choices[0].message.content

    return sql, result_str, answer

# Gradio界面
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 📊 Text2SQL智能问答")
    question = gr.Textbox(label="输入问题", placeholder="例如：哪个产品销量最高？")
    btn = gr.Button("查询", variant="primary")
    sql_output = gr.Code(label="生成的SQL", language="sql")
    table_output = gr.Markdown(label="查询结果")
    answer_output = gr.Textbox(label="回答", lines=3)

    btn.click(fn=text2sql_app, inputs=question, outputs=[sql_output, table_output, answer_output])

    gr.Examples(
        examples=[
            "2024年总销售额是多少？",
            "哪个区域的销量最高？",
            "按类别统计平均单价",
            "列出所有手机类产品",
        ],
        inputs=question,
    )

"""RAG评估实现（基于LLM-as-Judge）"""

from openai import OpenAI
import json

client = OpenAI()

# ---- 评估指标 ----

def evaluate_faithfulness(answer: str, contexts: list[str]) -> float:
    """评估回答对上下文的忠实度"""
    context_text = "\n---\n".join(contexts)
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {
                "role": "system",
                "content": (
                    "你是一个评估专家。判断以下回答是否完全基于给定的上下文。\n"
                    "评分标准：\n"
                    "1.0 = 回答完全基于上下文，没有编造任何信息\n"
                    "0.5 = 回答部分基于上下文，但有些内容无法从上下文中验证\n"
                    "0.0 = 回答主要基于编造的信息\n"
                    "只返回一个数字（0.0/0.5/1.0）。"
                ),
            },
            {
                "role": "user",
                "content": f"上下文:\n{context_text}\n\n回答:\n{answer}",
            },
        ],
        temperature=0,
    )
    try:
        return float(response.choices[0].message.content.strip())
    except ValueError:
        return 0.5

def evaluate_answer_relevancy(question: str, answer: str) -> float:
    """评估回答对问题的相关性"""
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {
                "role": "system",
                "content": (
                    "你是一个评估专家。判断以下回答是否充分回答了用户的问题。\n"
                    "评分：1.0=完全回答, 0.7=大部分回答, 0.4=部分回答, 0.0=完全无关\n"
                    "只返回一个数字。"
                ),
            },
            {
                "role": "user",
                "content": f"问题:\n{question}\n\n回答:\n{answer}",
            },
        ],
        temperature=0,
    )
    try:
        return float(response.choices[0].message.content.strip())
    except ValueError:
        return 0.5

def evaluate_context_precision(question: str, contexts: list[str]) -> float:
    """评估检索上下文的精确度（相关内容是否排在前面）"""
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {
                "role": "system",
                "content": (
                    "你是一个评估专家。判断以下检索结果对回答问题的有用程度。\n"
                    "对每个检索结果标记 1（相关）或 0（不相关）。\n"
                    "返回JSON数组，例如: [1, 0, 1]"
                ),
            },
            {
                "role": "user",
                "content": (
                    f"问题: {question}\n\n"
                    + "\n\n".join(f"检索结果{i+1}: {c}" for i, c in enumerate(contexts))
                ),
            },
        ],
        temperature=0,
    )
    try:
        relevance = json.loads(response.choices[0].message.content.strip())  # json.loads将JSON字符串→Python对象
        if not relevance:
            return 0.0
        # Average Precision计算
        hits = 0
        precision_sum = 0.0
        for i, rel in enumerate(relevance):
            if rel == 1:
                hits += 1
                precision_sum += hits / (i + 1)
        return precision_sum / max(sum(relevance), 1)
    except (json.JSONDecodeError, ValueError):
        return 0.5

# ---- 综合评估 ----

def evaluate_rag_sample(
    question: str,
    answer: str,
    contexts: list[str],
) -> dict:
    """综合评估一个RAG样本"""
    scores = {
        "faithfulness": evaluate_faithfulness(answer, contexts),
        "answer_relevancy": evaluate_answer_relevancy(question, answer),
        "context_precision": evaluate_context_precision(question, contexts),
    }
    scores["overall"] = sum(scores.values()) / len(scores)
    return scores

def evaluate_rag_dataset(dataset: list[dict]) -> dict:
    """评估整个数据集"""
    all_scores = []
    for sample in dataset:
        scores = evaluate_rag_sample(
            question=sample["question"],
            answer=sample["answer"],
            contexts=sample["contexts"],
        )
        all_scores.append(scores)
        print(f"  Q: {sample['question'][:50]}... → 综合: {scores['overall']:.2f}")

    # 汇总
    avg_scores = {}
    for key in all_scores[0]:
        avg_scores[key] = sum(s[key] for s in all_scores) / len(all_scores)
    return avg_scores

# 使用示例
eval_dataset = [
    {
        "question": "什么是RAG？",
        "answer": "RAG是检索增强生成技术，通过检索相关文档来增强LLM的生成质量。",
        "contexts": ["RAG(Retrieval-Augmented Generation)通过检索增强生成..."],
    },
    # ... 更多评估样本
]
# results = evaluate_rag_dataset(eval_dataset)
# print(f"评估结果: {results}")