154 lines
4.9 KiB
Python
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())
|