query_chunk.py

"""
按甜品名称从 ChromaDB 向量数据库查询对应的文档块

用法：
    # 精确匹配（与元数据 name 完全一致）
    python query_chunk.py "抹茶白豆腐"
    python query_chunk.py "煤球冰淇淋"

    # 模糊匹配（名称包含关键字，忽略大小写）
    python query_chunk.py --like "煤球"
    python query_chunk.py -l "豆腐"
"""

import sys
import os
import argparse
import chromadb
from chromadb.config import Settings
from dotenv import load_dotenv
from typing import Optional, Dict, Any, List


# 加载环境变量（.env）
load_dotenv()


def _safe_print_results(results: Any) -> None:
    """
    统一的结果打印函数，避免对 numpy 数组做布尔判断。
    接受 ChromaDB GetResult 对象或 dict 类型的查询结果。
    """
    embeddings = results.get("embeddings", None)
    documents: List[str] = results.get("documents") or []
    metadatas: List[Dict[str, Any]] = results.get("metadatas") or []

    emb_len = len(embeddings) if embeddings is not None else 0
    count = min(len(documents), len(metadatas))

    for idx in range(count):
        doc = documents[idx]
        meta = metadatas[idx]
        embedding = embeddings[idx] if (embeddings is not None and idx < emb_len) else None

        print(f"\n✅ 找到文档块 #{idx + 1}：")
        print("=" * 80)

        print("\n📦 文档块数据：")
        print("-" * 80)
        print("{")

        # 文档内容（转义并截断）
        doc_display = (
            doc.replace("\\", "\\\\")
            .replace('"', '\\"')
            .replace("\n", "\\n")
            .replace("\r", "\\r")
            .replace("\t", "\\t")
        )
        if len(doc_display) > 1000:
            doc_display = doc_display[:1000] + "..."
        print(f'  "document": "{doc_display}",')

        # 元数据
        print('  "metadata": {')
        meta_items = list(sorted((meta or {}).items()))
        for i, (key, value) in enumerate(meta_items):
            comma = "," if i < len(meta_items) - 1 else ""
            print(f'    "{key}": "{value}"{comma}')
        print("  },")

        # 向量信息（只展示前 30 项）
        if embedding is not None:
            first_30 = [round(float(v), 9) for v in embedding[:30]]
            print('  "embedding": [')
            print(f"    {first_30},")
            print(f"    ... // ({len(embedding)} total values, generated by text-embedding-v4)")
            print("  ]")

        print("}")
        print("-" * 80)

    print(f"✅ 共找到文档块数量：{len(documents)}")


def query_chunk_by_name(name: str, like: bool = False, max_results: int = 50) -> None:
    """
    根据甜品名称查询文档块内容

    参数：
        name: 要查询的甜品名称或关键字
        like: 是否启用模糊匹配（包含关系，忽略大小写）
        max_results: 最多返回的结果数量（用于本地过滤的回退方案）
    """
    # 从环境变量获取数据库路径
    db_path = os.getenv("CHROMA_DB_PATH", "./vector_db")

    # 初始化 ChromaDB 客户端
    client = chromadb.PersistentClient(
        path=db_path,
        settings=Settings(anonymized_telemetry=False),
    )

    # 获取集合（若不存在会报错，这里给出友好提示）
    try:
        collection = client.get_collection("knowledge_base")
    except Exception:
        print("❌ 未找到集合 'knowledge_base'，请先使用 data_loader.py 导入数据。")
        return

    print("=" * 80)
    if like:
        print(f"🔍 正在按名称模糊查询：'{name}'（包含关系，忽略大小写）")
    else:
        print(f"🔍 正在按名称精确查询：'{name}'")
    print("=" * 80)

    try:
        if not like:
            # 精确匹配
            results = collection.get(
                where={"name": name},
                include=["documents", "metadatas", "embeddings"],
            )
            if not results.get("documents"):
                print(f"\n❌ 未找到名称为 '{name}' 的文档块")
                print("\n💡 提示：")
                print("   - 请确认已通过 data_loader.py 导入了 data/ 目录中的 JSON/TXT/MD 文档")
                print("   - 名称需与 JSON 中的 ‘名称/name’ 字段一致（区分大小写）")
                print("   - 参阅 README 的‘数据目录说明’了解示例数据与命名")
                return
            _safe_print_results(results)
            return

        # 模糊匹配：本地过滤（拉取一定数量元数据后在内存中筛选）
        approx_limit = max(max_results * 20, 200)
        snapshot = collection.get(
            include=["metadatas"],
            limit=approx_limit,
        )
        ids = snapshot.get("ids") or []
        metas = snapshot.get("metadatas") or []

        needle = name.casefold()
        matched_ids: List[str] = []
        for i in range(min(len(ids), len(metas))):
            nm = (metas[i] or {}).get("name")
            if isinstance(nm, str) and needle in nm.casefold():
                matched_ids.append(ids[i])
                if len(matched_ids) >= max_results:
                    break

        if not matched_ids:
            print(f"\n❌ 未找到包含关键字 '{name}' 的文档块")
            print("\n💡 提示：考虑更换关键词，或先运行精确查询确认名称")
            return

        final_results = collection.get(
            ids=matched_ids,
            include=["documents", "metadatas", "embeddings"],
        )
        _safe_print_results(final_results)

    except Exception as e:
        print(f"\n❌ 查询数据库出错：{str(e)}")
        raise


def main() -> None:
    """命令行入口"""
    parser = argparse.ArgumentParser(
        description="按甜品名称查询文档块（支持精确与模糊）",
    )
    parser.add_argument("name", help="甜品名称或关键字")
    parser.add_argument(
        "-l", "--like",
        action="store_true",
        help="启用模糊匹配（名称包含关键字，忽略大小写）",
    )
    parser.add_argument(
        "--max",
        type=int,
        default=50,
        help="最多返回多少条（用于本地过滤回退，默认 50）",
    )
    args = parser.parse_args()

    query_chunk_by_name(args.name, like=args.like, max_results=args.max)


if __name__ == "__main__":
    main()