base_rag/examples/simple_test.py

154 lines
4.9 KiB
Python

#!/usr/bin/env python3
"""
简单测试示例 - 基础RAG功能验证
"""
import sys
import os
import asyncio
import warnings
from pathlib import Path
import shutil
# 过滤掉PyTorch的FutureWarning
warnings.filterwarnings("ignore", category=FutureWarning, module="torch")
# 添加源码路径
sys.path.append(os.path.join(os.path.dirname(__file__), "..", "src"))
from base_rag.core import BaseRAG
class SimpleTestRAG(BaseRAG):
"""简单测试RAG实现"""
async def ingest(self, file_path: str, **kwargs):
"""文档导入"""
return await self.process_file_to_vector_store(file_path, **kwargs)
async def query(self, question: str) -> str:
"""查询实现"""
docs = await self.similarity_search_with_rerank(question, k=3)
if not docs:
return "抱歉,没有找到相关信息。"
# 整理搜索结果
sources = []
contexts = []
for doc in docs:
source = doc.metadata.get("source_file", "未知来源")
content = doc.page_content.strip()
if source not in sources:
sources.append(source)
contexts.append(content)
context = "\n\n".join(contexts)
sources_str = "".join(sources)
return f"基于文档({sources_str})的信息:\n\n{context}"
async def test_basic_functionality():
"""测试基础RAG功能"""
print("🔧 基础RAG功能测试")
print("=" * 50)
# 清理向量数据库
db_path = Path("/Users/liruwei/Documents/code/project/demo/base_rag/chroma_db/simple_test")
if db_path.exists():
shutil.rmtree(db_path)
print("🧹 已清理向量数据库")
# 创建RAG实例 - 禁用图片处理用于基础测试
rag = SimpleTestRAG(
vector_store_name="simple_test",
retriever_top_k=3,
storage_directory="/Users/liruwei/Documents/code/project/demo/base_rag/test_files",
status_db_path="/Users/liruwei/Documents/code/project/demo/base_rag/simple_test_status.db",
image_config={"enabled": False} # 基础测试禁用图片
)
print("✅ RAG实例创建成功")
print()
# 测试基础文档
test_files = ["test_document.txt", "test_markdown.md", "python_basics.txt", "data_science.txt"]
print("📂 处理基础文档...")
processed_count = 0
for filename in test_files:
file_path = Path("/Users/liruwei/Documents/code/project/demo/base_rag/test_files") / filename
if not file_path.exists():
print(f"⚠️ {filename} - 文件不存在,跳过")
continue
print(f"📄 处理: {filename}")
try:
result = await rag.ingest(str(file_path))
if result and result.get('success'):
print(f" ✅ 成功: {result['chunks_count']} 个片段")
processed_count += 1
else:
message = result.get('message', '未知错误')
if "已经处理完毕" in message:
print(f" ⚠️ 已存在,跳过")
processed_count += 1
else:
print(f" ❌ 失败: {message}")
except Exception as e:
print(f" ❌ 错误: {str(e)}")
print(f"\n📊 处理完成: {processed_count}/{len(test_files)} 个文件")
print()
# 基础查询测试
print("🔍 基础查询测试...")
test_queries = [
"Python编程语言的特点",
"数据科学的核心技术",
"机器学习的应用",
"什么是深度学习"
]
for i, question in enumerate(test_queries, 1):
print(f"\n❓ 查询 {i}: {question}")
try:
answer = await rag.query(question)
if "抱歉" not in answer:
# 显示结果摘要
lines = answer.split('\n')
source_line = lines[0] if lines[0].startswith('基于文档') else "来源未知"
print(f" 📚 {source_line}")
# 显示内容预览
content_start = answer.find('\n\n')
if content_start > 0:
content = answer[content_start+2:]
preview = content[:150] + "..." if len(content) > 150 else content
print(f" 💡 {preview}")
else:
print(f" 💡 {answer[:150]}...")
else:
print(f" 💡 {answer}")
except Exception as e:
print(f" ❌ 查询失败: {str(e)}")
print("\n" + "=" * 50)
print("🎉 基础功能测试完成!")
print("✅ 验证项目:")
print(" 📄 文档加载和切分")
print(" 🔍 文本向量化和存储")
print(" 🔎 相似性搜索")
print(" 📝 查询结果整合")
if __name__ == "__main__":
asyncio.run(test_basic_functionality())