129 lines
4.0 KiB
Python
129 lines
4.0 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
最终文件格式测试 - 清理后重新测试
|
|
"""
|
|
import asyncio
|
|
import sys
|
|
import os
|
|
import shutil
|
|
from pathlib import Path
|
|
|
|
# 添加项目路径
|
|
sys.path.append('/Users/liruwei/Documents/code/project/demo/base_rag/src')
|
|
|
|
from base_rag.core import BaseRAG
|
|
|
|
|
|
class SimpleRAG(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=2)
|
|
|
|
if not docs:
|
|
return "抱歉,没有找到相关信息。"
|
|
|
|
# 显示搜索到的文档来源
|
|
sources = []
|
|
contexts = []
|
|
for doc in docs:
|
|
source = doc.metadata.get("source_file", "未知来源")
|
|
if source not in sources:
|
|
sources.append(source)
|
|
contexts.append(doc.page_content.strip())
|
|
|
|
context = "\n\n".join(contexts)
|
|
sources_str = "、".join(sources)
|
|
|
|
return f"基于以下文档({sources_str})的信息:\n\n{context}"
|
|
|
|
|
|
async def final_format_test():
|
|
"""最终文件格式测试"""
|
|
print("🧹 清理测试环境...")
|
|
|
|
# 删除旧的向量数据库目录
|
|
test_db_path = Path("/Users/liruwei/Documents/code/project/demo/base_rag/chroma_db/final_test")
|
|
if test_db_path.exists():
|
|
shutil.rmtree(test_db_path)
|
|
|
|
print("✅ 环境清理完成\n")
|
|
|
|
print("🚀 文件格式支持最终测试")
|
|
print("=" * 50)
|
|
|
|
# 初始化新的RAG系统
|
|
rag = SimpleRAG(
|
|
vector_store_name="final_test",
|
|
retriever_top_k=2,
|
|
persist_directory="/Users/liruwei/Documents/code/project/demo/base_rag/chroma_db",
|
|
status_db_path="/Users/liruwei/Documents/code/project/demo/base_rag/final_test.db"
|
|
)
|
|
|
|
# 测试文件
|
|
test_files = [
|
|
{
|
|
"name": "machine_learning.md",
|
|
"path": "/Users/liruwei/Documents/code/project/demo/base_rag/test_files/machine_learning.md",
|
|
"type": "Markdown"
|
|
},
|
|
{
|
|
"name": "deep_learning_guide.docx",
|
|
"path": "/Users/liruwei/Documents/code/project/demo/base_rag/test_files/deep_learning_guide.docx",
|
|
"type": "Word文档"
|
|
}
|
|
]
|
|
|
|
print("📄 处理文件...")
|
|
for file_info in test_files:
|
|
name = file_info["name"]
|
|
path = file_info["path"]
|
|
file_type = file_info["type"]
|
|
|
|
print(f" {file_type}: {name}")
|
|
|
|
try:
|
|
result = await rag.process_file_to_vector_store(path)
|
|
if result and result.get('success'):
|
|
print(f" ✅ 成功: {result['chunks_count']} 个片段")
|
|
else:
|
|
print(f" ⚠️ {result.get('message', '处理失败')}")
|
|
except Exception as e:
|
|
print(f" ❌ 错误: {str(e)}")
|
|
|
|
print("\n💬 测试查询...")
|
|
|
|
queries = [
|
|
"什么是机器学习?",
|
|
"深度学习的应用领域有哪些?",
|
|
"神经网络的架构类型"
|
|
]
|
|
|
|
for query in queries:
|
|
print(f"\n❓ {query}")
|
|
try:
|
|
answer = await rag.query(query)
|
|
# 显示简化的回答
|
|
if "抱歉" not in answer:
|
|
lines = answer.split('\n')
|
|
first_content = next((line for line in lines if line.strip() and not line.startswith('基于')), "")
|
|
print(f" 💡 {first_content[:100]}...")
|
|
else:
|
|
print(f" 💡 {answer}")
|
|
except Exception as e:
|
|
print(f" ❌ {str(e)}")
|
|
|
|
print("\n" + "=" * 50)
|
|
print("🎉 测试完成!")
|
|
print("✅ 支持格式: TXT, MD, DOCX")
|
|
print("✅ 异步处理: 完全支持")
|
|
print("✅ 跨格式查询: 完全支持")
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(final_format_test())
|