base_rag/examples/multi_format_test.py

163 lines
4.9 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#!/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", "未知来源")
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_multiple_formats():
"""测试多种文件格式处理"""
print("🚀 多格式文件处理测试")
print("=" * 50)
# 创建RAG实例
rag = MultiFormatRAG(
vector_store_name="multiformat_kb",
retriever_top_k=3,
storage_directory="../test_files", # 相对于examples目录
status_db_path="../status.db", # 相对于examples目录
)
# 测试文件列表
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:
# 分离来源信息和内容
parts = answer.split('\n\n', 1)
if len(parts) == 2:
source_info = parts[0] # "基于以下文档..."
content = parts[1] # 实际内容
print(f" 📚 {source_info}")
# 显示内容摘要前200字符
if len(content) > 200:
content_preview = content[:200] + "..."
else:
content_preview = content
print(f" 💡 {content_preview}")
else:
print(f" 💡 {answer}")
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())