152 lines
4.5 KiB
Python
152 lines
4.5 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
多格式文件测试 - 测试 TXT、MD、DOCX 文件格式
|
||
"""
|
||
|
||
import sys
|
||
import os
|
||
import asyncio
|
||
import warnings
|
||
from pathlib import Path
|
||
|
||
# 过滤掉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 MultiFormatRAG(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", "未知来源")
|
||
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 test_multiple_formats():
|
||
"""测试多种文件格式处理"""
|
||
print("🚀 多格式文件处理测试")
|
||
print("=" * 50)
|
||
|
||
# 创建RAG实例
|
||
rag = MultiFormatRAG(
|
||
vector_store_name="multiformat_kb",
|
||
retriever_top_k=3,
|
||
storage_directory="./test_files",
|
||
status_db_path="./status.db",
|
||
)
|
||
|
||
# 测试文件列表
|
||
test_files = [
|
||
{
|
||
"file": "knowledge.txt",
|
||
"format": "TXT",
|
||
"description": "纯文本文件"
|
||
},
|
||
{
|
||
"file": "python_guide.md",
|
||
"format": "MD",
|
||
"description": "Markdown文件"
|
||
},
|
||
{
|
||
"file": "machine_learning.md",
|
||
"format": "MD",
|
||
"description": "Markdown文件"
|
||
},
|
||
{
|
||
"file": "deep_learning_guide.docx",
|
||
"format": "DOCX",
|
||
"description": "Word文档"
|
||
}
|
||
]
|
||
|
||
print("📂 处理文件...")
|
||
processed_count = 0
|
||
|
||
for file_info in test_files:
|
||
filename = file_info["file"]
|
||
format_type = file_info["format"]
|
||
description = file_info["description"]
|
||
|
||
file_path = Path("./test_files") / filename
|
||
|
||
if not file_path.exists():
|
||
print(f"❌ {format_type}: {filename} - 文件不存在")
|
||
continue
|
||
|
||
print(f"📄 处理 {format_type}: {filename} ({description})")
|
||
|
||
try:
|
||
result = await rag.ingest(str(file_path))
|
||
if result and result.get('success'):
|
||
print(f" ✅ 成功: {result['chunks_count']} 个片段")
|
||
processed_count += 1
|
||
else:
|
||
print(f" ⚠️ 跳过: {result.get('message', '可能已存在')}")
|
||
processed_count += 1 # 已存在也算处理过
|
||
except Exception as e:
|
||
print(f" ❌ 失败: {str(e)}")
|
||
|
||
print(f"\n📊 处理完成: {processed_count}/{len(test_files)} 个文件")
|
||
print()
|
||
|
||
# 测试跨格式查询
|
||
print("💬 跨格式查询测试...")
|
||
|
||
queries = [
|
||
"Python有什么特点?",
|
||
"什么是机器学习?",
|
||
"深度学习的应用领域有哪些?",
|
||
"Web框架有哪些?"
|
||
]
|
||
|
||
for query in queries:
|
||
print(f"\n❓ {query}")
|
||
try:
|
||
answer = await rag.query(query)
|
||
if "抱歉" not in answer:
|
||
# 显示简化的回答
|
||
lines = answer.split('\n')
|
||
source_line = next((line for line in lines if line.startswith('基于')), "")
|
||
content_lines = [line for line in lines if line.strip() and not line.startswith('基于')]
|
||
if content_lines:
|
||
print(f" 💡 {content_lines[0][:100]}...")
|
||
if source_line:
|
||
print(f" 📚 {source_line}")
|
||
else:
|
||
print(f" 💡 {answer}")
|
||
except Exception as e:
|
||
print(f" ❌ 查询失败: {str(e)}")
|
||
|
||
print("\n" + "=" * 50)
|
||
print("✅ 多格式文件测试完成!")
|
||
print("支持的格式: TXT, MD, DOCX")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(test_multiple_formats())
|