test_cache_fix.py

#!/usr/bin/env python3
"""
测试脚本：验证 Hugging Face 缓存目录修复是否有效
"""

import os
import sys
import tempfile

# 添加 src 目录到 Python 路径
sys.path.insert(0, 'src')
print(f"Python path: {sys.path[:3]}")

def test_cache_setup():
    """测试缓存目录设置"""
    print("=== 测试缓存目录设置 ===")
    
    # 测试本地环境
    print("\n1. 测试本地环境（无 SPACE_ID）:")
    if 'SPACE_ID' in os.environ:
        del os.environ['SPACE_ID']
    if 'HF_SPACE_ID' in os.environ:
        del os.environ['HF_SPACE_ID']
    
    try:
        print("   正在导入 setup_hf_cache...")
        from demo.path_utils import setup_hf_cache
        print("   导入成功，正在调用 setup_hf_cache...")
        cache_dir = setup_hf_cache()
        print(f"   缓存目录: {cache_dir}")
        print("   ✅ 本地环境测试通过")
    except Exception as e:
        print(f"   ❌ 本地环境测试失败: {e}")
        import traceback
        traceback.print_exc()
    
    # 测试 Hugging Face Spaces 环境
    print("\n2. 测试 Hugging Face Spaces 环境（有 SPACE_ID）:")
    os.environ['SPACE_ID'] = 'test_space_123'
    
    try:
        cache_dir = setup_hf_cache()
        print(f"   缓存目录: {cache_dir}")
        if cache_dir and os.path.exists(cache_dir):
            print("   ✅ Hugging Face Spaces 环境测试通过")
        else:
            print("   ❌ 缓存目录不存在")
    except Exception as e:
        print(f"   ❌ Hugging Face Spaces 环境测试失败: {e}")
        import traceback
        traceback.print_exc()

def test_embedder_initialization():
    """测试 embedder 初始化"""
    print("\n=== 测试 Embedder 初始化 ===")
    
    # 设置 Hugging Face Spaces 环境
    os.environ['SPACE_ID'] = 'test_space_456'
    
    try:
        from demo.path_utils import setup_hf_cache
        from langchain_community.embeddings import HuggingFaceEmbeddings
        
        cache_dir = setup_hf_cache()
        print(f"使用缓存目录: {cache_dir}")
        
        # 尝试初始化 embedder
        print("正在初始化 HuggingFaceEmbeddings...")
        embedder = HuggingFaceEmbeddings(
            model_name="sentence-transformers/all-MiniLM-L6-v2",
            cache_folder=cache_dir
        )
        print("✅ Embedder 初始化成功")
        
        # 测试简单的嵌入
        test_text = "This is a test sentence."
        print("正在测试嵌入...")
        embedding = embedder.embed_query(test_text)
        print(f"✅ 嵌入测试成功，向量维度: {len(embedding)}")
        
    except Exception as e:
        print(f"❌ Embedder 初始化失败: {e}")
        import traceback
        traceback.print_exc()

def test_sentence_transformer():
    """测试 SentenceTransformer"""
    print("\n=== 测试 SentenceTransformer ===")
    
    # 设置 Hugging Face Spaces 环境
    os.environ['SPACE_ID'] = 'test_space_789'
    
    try:
        from demo.path_utils import setup_hf_cache
        from sentence_transformers import SentenceTransformer
        
        cache_dir = setup_hf_cache()
        print(f"使用缓存目录: {cache_dir}")
        
        # 尝试初始化 SentenceTransformer
        print("正在初始化 SentenceTransformer...")
        model = SentenceTransformer(
            "nomic-ai/nomic-embed-text-v1", 
            trust_remote_code=True, 
            cache_folder=cache_dir
        )
        print("✅ SentenceTransformer 初始化成功")
        
        # 测试简单的嵌入
        test_text = "This is a test sentence."
        print("正在测试嵌入...")
        embedding = model.encode(test_text)
        print(f"✅ 嵌入测试成功，向量维度: {len(embedding)}")
        
    except Exception as e:
        print(f"❌ SentenceTransformer 初始化失败: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    print("开始测试 Hugging Face 缓存目录修复...")
    
    test_cache_setup()
    test_embedder_initialization()
    test_sentence_transformer()
    
    print("\n=== 测试完成 ===")