Agent/test_embeddings.py

"""
测试 Embeddings 模块（不依赖 Milvus）
"""

import asyncio
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))

# 加载环境变量
from dotenv import load_dotenv
load_dotenv(Path(__file__).parent / ".env")

from knowhub.embeddings import get_embedding, get_embeddings_batch


async def test_embeddings():
    print("=" * 60)
    print("测试 Embeddings 模块")
    print("=" * 60)

    # 测试单条
    print("\n1. 测试单条 embedding 生成...")
    text = "如何使用 Python 读取 PDF 文件"
    try:
        embedding = await get_embedding(text)
        print(f"✓ 成功生成 embedding")
        print(f"  文本: {text}")
        print(f"  向量维度: {len(embedding)}")
        print(f"  前 5 个值: {embedding[:5]}")
    except Exception as e:
        print(f"✗ 失败: {e}")
        return

    # 测试批量
    print("\n2. 测试批量 embedding 生成...")
    texts = [
        "使用 pymupdf 读取 PDF",
        "使用 selenium 进行网页自动化",
        "使用 pandas 处理数据"
    ]
    try:
        embeddings = await get_embeddings_batch(texts)
        print(f"✓ 成功生成批量 embeddings")
        print(f"  文本数量: {len(texts)}")
        print(f"  向量数量: {len(embeddings)}")
        print(f"  每个向量维度: {len(embeddings[0])}")
    except Exception as e:
        print(f"✗ 失败: {e}")
        return

    print("\n" + "=" * 60)
    print("Embeddings 模块测试通过！")
    print("=" * 60)


if __name__ == "__main__":
    asyncio.run(test_embeddings())