大模型应用：向量与元数据联动：解锁向量数据库复合查询的核心能力.30

import chromadb
import datetime
from chromadb.utils import embedding_functions

# 工具函数：日期字符串转Unix时间戳（int）
def date_to_timestamp(date_str):
    dt = datetime.datetime.strptime(date_str, "%Y-%m-%d")
    return int(dt.timestamp())

# 1. 初始化环境
ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="all-MiniLM-L6-v2")
client = chromadb.Client()
collection = client.create_collection("ecommerce_products", embedding_function=ef)

# 2. 插入测试数据（日期转时间戳，统一为int类型）
products = [
    "新款运动鞋 轻便透气",
    "复古帆布鞋 百搭休闲",
    "专业跑步鞋 减震防滑",
    "篮球鞋 高帮护踝"
]
# 替换shelf_date为时间戳（int），保留原日期字符串作为冗余字段（方便展示）
metadatas = [
    {"price": 199, "shelf_date": "2024-05-10", "shelf_timestamp": date_to_timestamp("2024-05-10"), "sales": 500},
    {"price": 89, "shelf_date": "2024-03-15", "shelf_timestamp": date_to_timestamp("2024-03-15"), "sales": 1200},
    {"price": 299, "shelf_date": "2024-06-20", "shelf_timestamp": date_to_timestamp("2024-06-20"), "sales": 800},
    {"price": 399, "shelf_date": "2024-07-05", "shelf_timestamp": date_to_timestamp("2024-07-05"), "sales": 300}
]
ids = ["p1", "p2", "p3", "p4"]
collection.add(documents=products, metadatas=metadatas, ids=ids)

# 3. $gte 多条件过滤（时间戳数字做数值过滤）
query_text = "运动鞋 舒适"
# 目标日期转时间戳
target_timestamp = date_to_timestamp("2024-05-01")
results = collection.query(
    query_texts=[query_text],
    where={
        "$and": [
            {"price": {"$gte": 100}},          # 价格≥100（int）
            {"shelf_timestamp": {"$gte": target_timestamp}},  # 上架时间≥2024-05-01（时间戳int）
            {"sales": {"$gte": 500}}           # 销量≥500（int）
        ]
    },
    n_results=2
)

# 4. 输出结果
print("===== 过滤结果 =====")
for i, (doc, meta) in enumerate(zip(results["documents"][0], results["metadatas"][0])):
    print(f"商品{i+1}：{doc}")
    print(f"元数据：价格={meta['price']} | 上架时间={meta['shelf_date']} | 销量={meta['sales']}\n")

import chromadb
import pandas as pd
import matplotlib.pyplot as plt
from chromadb.utils import embedding_functions
from datetime import datetime

# 1. 初始化Embedding模型（使用Sentence-BERT）
sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name="D:/modelscope/hub/models/sentence-transformers/all-MiniLM-L6-v2"
)

# 2. 初始化Chroma客户端（内存模式）
client = chromadb.Client()

# 3. 创建集合：指定Embedding函数，支持元数据过滤
collection = client.create_collection(
    name="tech_docs",
    embedding_function=sentence_transformer_ef,
    metadata={"hnsw:space": "cosine"}  # 余弦相似度计算
)

# 4. 准备测试数据：包含文档内容、元数据（标签、发布时间、来源）
docs = [
    "RAG技术核心是将检索到的知识融入大模型生成过程，提升回答准确性",
    "向量数据库的索引算法HNSW能够显著提升高维向量的检索效率",
    "大模型的微调技术包括LoRA和QLoRA，可在低资源设备上实现模型优化",
    "向量与元数据联动是实现复合查询的关键，广泛应用于RAG场景",
    "金融领域RAG应用需解决数据隐私问题，常用联邦学习方案",
    "2024年最新HNSW优化算法可将向量检索速度提升300%"
]

# 元数据设计：标签（tag）、发布时间（publish_date）、来源（source）
metadatas = [
    {"tag": "RAG", "publish_date": "2024-05-10", "source": "技术博客"},
    {"tag": "向量数据库", "publish_date": "2024-06-15", "source": "学术论文"},
    {"tag": "大模型微调", "publish_date": "2023-12-20", "source": "教程"},
    {"tag": "向量数据库", "publish_date": "2024-07-02", "source": "技术博客"},
    {"tag": "RAG", "publish_date": "2024-08-18", "source": "行业报告"},
    {"tag": "向量数据库", "publish_date": "2024-09-05", "source": "学术论文"}
]

# 文档ID
ids = ["doc1", "doc2", "doc3", "doc4", "doc5", "doc6"]

# 5. 将数据添加到集合（自动转化为向量+存储元数据）
collection.add(
    documents=docs,
    metadatas=metadatas,
    ids=ids
)

print("数据入库完成！共存储{}篇文档".format(len(ids)))

import chromadb
import pandas as pd
import matplotlib.pyplot as plt
from chromadb.utils import embedding_functions
from datetime import datetime

# ===================== 新增：日期兼容处理工具函数 =====================
def date_to_timestamp(date_str):
    """将YYYY-MM-DD格式的日期字符串转为Unix时间戳（int），适配低版本Chroma的数值过滤"""
    dt = datetime.strptime(date_str, "%Y-%m-%d")
    return int(dt.timestamp())

def timestamp_to_date(timestamp):
    """将时间戳转回日期字符串，用于结果展示"""
    return datetime.fromtimestamp(timestamp).strftime("%Y-%m-%d")

# ===================== 初始化环境 + 数据入库（兼容日期格式） =====================
# 1. 初始化Embedding模型
sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name="all-MiniLM-L6-v2"
)

# 2. 初始化Chroma客户端
client = chromadb.Client()

# 3. 创建集合
collection = client.create_collection(
    name="tech_docs",
    embedding_function=sentence_transformer_ef,
    metadata={"hnsw:space": "cosine"}  # 余弦相似度计算
)

# 4. 准备测试数据（新增publish_timestamp字段，用于数值过滤）
docs = [
    "RAG技术核心是将检索到的知识融入大模型生成过程，提升回答准确性",
    "向量数据库的索引算法HNSW能够显著提升高维向量的检索效率",
    "大模型的微调技术包括LoRA和QLoRA，可在低资源设备上实现模型优化",
    "向量与元数据联动是实现复合查询的关键，广泛应用于RAG场景",
    "金融领域RAG应用需解决数据隐私问题，常用联邦学习方案",
    "2024年最新HNSW优化算法可将向量检索速度提升300%"
]

# 元数据：保留publish_date（字符串，展示用）+ 新增publish_timestamp（时间戳int，过滤用）
metadatas = [
    {"tag": "RAG", "publish_date": "2024-05-10", "publish_timestamp": date_to_timestamp("2024-05-10"), "source": "技术博客"},
    {"tag": "向量数据库", "publish_date": "2024-06-15", "publish_timestamp": date_to_timestamp("2024-06-15"), "source": "学术论文"},
    {"tag": "大模型微调", "publish_date": "2023-12-20", "publish_timestamp": date_to_timestamp("2023-12-20"), "source": "教程"},
    {"tag": "向量数据库", "publish_date": "2024-07-02", "publish_timestamp": date_to_timestamp("2024-07-02"), "source": "技术博客"},
    {"tag": "RAG", "publish_date": "2024-08-18", "publish_timestamp": date_to_timestamp("2024-08-18"), "source": "行业报告"},
    {"tag": "向量数据库", "publish_date": "2024-09-05", "publish_timestamp": date_to_timestamp("2024-09-05"), "source": "学术论文"}
]

ids = ["doc1", "doc2", "doc3", "doc4", "doc5", "doc6"]

# 5. 数据入库
collection.add(
    documents=docs,
    metadatas=metadatas,
    ids=ids
)

# ===================== 工具函数（适配兼容后的元数据） =====================
def plot_similarity(results, title):
    """绘制相似度（距离）分布柱状图"""
    distances = results["distances"][0]
    doc_ids = [f"文档{idx+1}" for idx in range(len(distances))]
    
    plt.rcParams["font.sans-serif"] = ["SimHei"]  # 支持中文
    plt.figure(figsize=(8, 4))
    plt.bar(doc_ids, distances, color="#4285F4")
    plt.title(title)
    plt.xlabel("文档")
    plt.ylabel("余弦距离（越小越相似）")
    plt.grid(axis="y", linestyle="--", alpha=0.7)
    plt.tight_layout()
    plt.show()

def print_formatted_results(results, scene_name):
    """格式化输出查询结果（展示原日期字符串）"""
    print(f"\n===== {scene_name} 查询结果 =====")
    for i, (doc, meta, dist) in enumerate(zip(
        results["documents"][0], 
        results["metadatas"][0], 
        results["distances"][0]
    )):
        print(f"\n【结果{i+1}】")
        print(f"文档内容：{doc}")
        # 展示原日期字符串，而非时间戳
        print(f"元数据：标签={meta['tag']} | 发布时间={meta['publish_date']} | 来源={meta['source']}")
        print(f"余弦距离：{dist:.4f}（越小越相似）")

# ===================== 场景A：过滤优先（日期用时间戳过滤） =====================
query_text_a = "向量数据库优化算法"
# 日期条件转为时间戳数值过滤
target_timestamp_a = date_to_timestamp("2024-06-01")
where_clause_a = {
    "$and": [
        {"publish_timestamp": {"$gte": target_timestamp_a}},  # 用时间戳int过滤
        {"source": "学术论文"},
        {"tag": "向量数据库"}
    ]
}

results_a = collection.query(
    query_texts=[query_text_a],
    where=where_clause_a,
    n_results=2
)
print_formatted_results(results_a, "场景A（过滤优先）")
plot_similarity(results_a, "场景A：向量数据库学术论文相似度分布")

# ===================== 场景B：检索优先（过滤逻辑不变，展示兼容） =====================
query_text_b = "大模型落地应用"
# 先检索Top-5
raw_results_b = collection.query(
    query_texts=[query_text_b],
    n_results=5
)
# 过滤标签为RAG的文档
filtered_docs = []
filtered_metas = []
filtered_dists = []
for doc, meta, dist in zip(raw_results_b["documents"][0], raw_results_b["metadatas"][0], raw_results_b["distances"][0]):
    if meta["tag"] == "RAG":
        filtered_docs.append(doc)
        filtered_metas.append(meta)
        filtered_dists.append(dist)
# 构造过滤后的结果格式
results_b = {
    "documents": [filtered_docs],
    "metadatas": [filtered_metas],
    "distances": [filtered_dists]
}
print_formatted_results(results_b, "场景B（检索优先）")
plot_similarity(results_b, "场景B：大模型应用-RAG文档相似度分布")

# ===================== 场景C：多条件嵌套（日期范围用时间戳过滤） =====================
query_text_c = "AI技术落地实践"
# 日期范围转为时间戳数值
start_timestamp = date_to_timestamp("2024-05-01")
end_timestamp = date_to_timestamp("2024-08-31")
where_clause_c = {
    "$and": [
        {"publish_timestamp": {"$gte": start_timestamp}},
        {"publish_timestamp": {"$lte": end_timestamp}},
        {"source": "技术博客"},
        {"tag": {"$in": ["RAG", "向量数据库"]}}
    ]
}
results_c = collection.query(
    query_texts=[query_text_c],
    where=where_clause_c,
    n_results=3
)
print_formatted_results(results_c, "场景C（多条件嵌套）")
plot_similarity(results_c, "场景C：AI技术落地博客相似度分布")